Push a Transformers model
Learn how to push a Hugging Face Transformers language model to Replicate.
Transformers is an open-source Python library that provides a consistent interface for using language models. The library contains multiple open-source generative language models like FLAN, GPT-J, GPT Neo, LLaMA, BLOOM, and others, which have been pre-trained on large text corpora and can be fine-tuned for specific tasks with relatively small amounts of training data.
Transformers also contains models like Longformer, BERT, and RoBERTa, which are generally used for more traditional natural language processing tasks like classification, named entity recognition, and so on. The process we’re walking through here will work for both kinds of models; in fact, it should work for every model on Transformers.
In this guide we’ll walk you through the process of taking an existing Transformers model and pushing it to Replicate as your own public or private model with a stable API.
To follow this guide, you’ll need:
- An account on Replicate.
- Docker. You’ll be using the Cog command-line tool to build and push your model. Cog uses Docker to create a container for your model. You’ll need to install and start Docker before you can run Cog. You can confirm Docker is running by typing
docker infoin your terminal.
- Linux machine with an NVIDIA GPU. If you don’t already have access to a machine with a GPU, check out our guide to getting a GPU machine.
Step 1: Create a model
First, create a model on Replicate at replicate.com/create. If you haven’t used Replicate before, you’ll need to sign in with your GitHub account. You can configure the model to be private so that only you can use it, or public so anyone can use it.
Step 2: Set up your GPU machine
To build and push your new model, you’ll need a Linux machine with an NVIDIA GPU. If you don’t already have access to a machine with a GPU, check out our guide to getting a GPU machine.
Once you’ve SSHed into your GPU machine, install Cog:
curl https://replicate.github.io/codespaces/scripts/install-cog.sh | bash
Confirm that Cog is installed by running
cog --version # cog version 0.6.1 (built 2022-12-18T04:42:41+0000)
Step 3: Initialize your project
From your GPU machine, create a new directory and initialize a new Cog project:
mkdir my-cool-model cd my-cool-model cog init
This will create two files,
predict.py, which you’ll use to configure your dependencies and define the inputs and outputs of your model.
Step 4: Configure dependencies
cog.yaml file defines the CUDA and Python versions, and depenendencies for the model. This file tells Cog how to package the model.
Replace the contents of the
cog.yaml file with the following:
build: gpu: true python_version: "3.10" python_packages: - "torch==1.12.1" - "transformers==4.26.1" - "sentencepiece==0.1.97" - "accelerate==0.16.0" predict: "predict.py:Predictor"
Step 5: Customize your predictor
Replace the contents of the
predict.py file with the following:
from typing import List, Optional from cog import BasePredictor, Input from transformers import T5ForConditionalGeneration, AutoTokenizer import torch CACHE_DIR = 'weights' # Shorthand identifier for a transformers model. # See https://huggingface.co/models?library=transformers for a list of models. MODEL_NAME = 'google/flan-t5-xl' class Predictor(BasePredictor): def setup(self): self.device = 'cuda' if torch.cuda.is_available() else 'cpu' self.model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR, local_files_only=True) self.model.to(self.device) self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR, local_files_only=True) def predict( self, prompt: str = Input(description=f"Text prompt to send to the model."), n: int = Input(description="Number of output sequences to generate", default=1, ge=1, le=5), max_length: int = Input( description="Maximum number of tokens to generate. A word is generally 2-3 tokens", ge=1, default=50 ), temperature: float = Input( description="Adjusts randomness of outputs, greater than 1 is random and 0 is deterministic, 0.75 is a good starting value.", ge=0.01, le=5, default=0.75, ), top_p: float = Input( description="When decoding text, samples from the top p percentage of most likely tokens; lower to ignore less likely tokens", ge=0.01, le=1.0, default=1.0 ), repetition_penalty: float = Input( description="Penalty for repeated words in generated text; 1 is no penalty, values greater than 1 discourage repetition, less than 1 encourage it.", ge=0.01, le=5, default=1 ) ) -> List[str]: input = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device) outputs = self.model.generate( input, num_return_sequences=n, max_length=max_length, do_sample=True, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty ) out = self.tokenizer.batch_decode(outputs, skip_special_tokens=True) return out
The AutoTokenizer used above should work for all Transformers models.
If you want to use a Transformers model other than Flan-T5, you’ll need to specify the
model class to use. For example, if you’re using a GPT-J model, you’ll want to use
AutoModelForCausalLM instead of
T5ForConditionalGeneration. See the Transformers docs for more details.
Step 6: Download weights
Next you’ll create a script that uses the
transformers library to download pretrained weights.
Create a file for the script:
mkdir script touch script/download_weights chmod +x script/download_weights # makes the file executable
Paste the following code into the
#!/usr/bin/env python import os import shutil from transformers import T5ForConditionalGeneration, T5Tokenizer CACHE_DIR = 'weights' if os.path.exists(CACHE_DIR): shutil.rmtree(CACHE_DIR) os.makedirs(CACHE_DIR) model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-xl", cache_dir=CACHE_DIR) tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-xl", cache_dir=CACHE_DIR)`
Run the script to download the weights:
cog run script/download_weights
This process will take a while to run but you’ll only need to run it once, as it will cache the downloaded dependencies on disk. Get up and stretch, grab yourself a snack, or use this opportunity to add metadata to the model page you created on Replicate in Step 1 by adding a title, README, GitHub repository URL, etc.
Step 7: Run your model
Now that you’ve downloaded the weights, you can run the model locally with Cog:
cog predict -i prompt="Q: Answer the following yes/no question by reasoning step-by-step. Can a dog drive a car?"
This will run the model locally on your GPU machine, and return output text.
Step 8: Push your model
Now that you’ve created your model, it’s time to push it to Replicate.
First you’ll need to authenticate:
Then push your model using the name you specified in Step 1:
cog push r8.im/<your-username>/<your-model-name>
Step 9: Use your model
Your model is now live! 🚀
You can run the model from the website by clicking the “Demo” tab on the model page, or you can use the HTTP API to run the model from your own code.
Click the “API” tab on your model page to see example code for running the model:
Now that you have your own model, see what else you can do with it!
To see what models you can use, check out the Transformers docs on Hugging Face.
If you need inspiration or guidance, jump into our Discord.