Add the ability to specify a local or S3 HF mirror for more guaranteed loading of pre-trained HF models.

[justinxzhao]

Overview

We've seen that HF servers can go down sometimes. This PR adds the ability to specify a local or S3 HF mirror for more guaranteed loading of pre-trained HF models by setting a LUDWIG_PRETRAINED_MODELS_DIR env variable.

Implementation

If the LUDWIG_PRETRAINED_MODELS_DIR environment variable is set, we attempt to load the HF model from this directory, falling back to downloading from the HF hub if the model is not found, downloading fails, or if model initialization fails.

LUDWIG_PRETRAINED_MODELS_DIR can be an s3 path. Weights are copied to a local temporary directory, and the model is loaded from there.

The expected structure of the LUDWIG_PRETRAINED_MODELS_DIR directory is:

{LUDWIG_PRETRAINED_MODELS_DIR}/{pretrained_model_name_or_path}/pytorch_model.bin
{LUDWIG_PRETRAINED_MODELS_DIR}/{pretrained_model_name_or_path}/config.json

For example, if LUDWIG_PRETRAINED_MODELS_DIR is set to s3://my-bucket/pretrained-models, and pretrained_model_name_or_path is set to bert-base-uncased, we expect to find the following files:

s3://my-bucket/bert-base-uncased/
- pytorch_model.bin
- config.json

If the LUDWIG_PRETRAINED_MODELS_DIR environment variable is not set, we download the model from the HF hub.

[github-actions]

Unit Test Results

         6 files           6 suites   6h 45m 55s ⏱️
  4 046 tests   4 003 ✔️   43 💤 0 ❌
12 107 runs  11 975 ✔️ 132 💤 0 ❌

Results for commit 5f815ad.

♻️ This comment has been updated with latest results.

[tgaddair]

Looks good, just a few nits.

[arnavgarg1]

LGTM! Agree with @tgaddair's comments around the tests :) Great addition!