cog-hf-template 0.0.1

Cog template for Hugging Face.

cog-huggingface-template

Tools to help deploy Hugging Face models to replicate. For now, just text generation models.

Requirements

• GPU environment
• Cog

Setup

Clone the repo:

git clone https://github.com/replicate/cog-huggingface-template
cd cog-huggingface-template

Local install:

pip install -e .

Usage

For now, see predict.py for a working example.

Optional - Use GCP for Weights

0. install pget

sudo curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.0.2/pget" && sudo chmod +x /usr/local/bin/pget

1. Run predict.py to download local weights cache

python predict.py

Inside it should be something like:

from cog_hf_template import get_predictor

Predictor = get_predictor(
    task="image-classification",
    model_name_or_path="nateraw/rare-puppers",
    cache_dir="./weights-cache",
)

if __name__ == "__main__":
    predictor = Predictor()
    predictor.setup()

.setup() will download the weights from huggingface to ./weights-cache.

2. Tar up weights cache

Now that you have the weights locally, you need to tar them up before pushing to GCP. Note that you have to tar from the weights-cache directory itself to avoid the weights-cache directory being included in the tar.

cd weights-cache
tar cvf weights-cache.tar *

3. Push to GCP bucket

Take the weights-cache.tar you just created and push to GCP bucket.

gsutil cp weights-cache.tar gs://your-bucket/nateraw/rare-puppers/

4. Specify bucket in predictor to download/extract to cache with pget

Before cog push-ing, make sure you have removed the weights-cache directory from your local or have ignored it in .cogignore so the weights-cache dir doesn't get included in the cog build.

# Under the hood it runs this to download and extract weights cache in same structure as you had it locally
# pget -x https://storage.googleapis.com/your-bucket/nateraw/rare-puppers/weights-cache.tar ./weights-cache
from cog_hf_template import get_predictor


Predictor = get_predictor(
    task="image-classification",
    model_name_or_path="nateraw/rare-puppers",
    cache_dir="./weights-cache",
    gcp_bucket_weights="gs://your-bucket/nateraw/rare-puppers/weights-cache.tar",
)

Notes

Optionally, if you run in interactive mode, you can play with the resulting predictor. In this example we are using a text-generation model...

generator = predictor.predict("Write me a recipe for chocolate chip cookies", 1024, 1.0, 1.0, 50)
for out in generator:
    print(out, end='')