Hugging Face is a supported framework on Truss. To package a Hugging Face model, follow the steps below or run this Google Colab notebook.
If you're using a Jupyter notebook, add a line to install the
trusspackages. Otherwise, ensure the packages are installed in your Python environment.
!pip install --upgrade transformers truss
Truss officially supports
transformersversion 4.21.0 or higher. Especially if you're using an online notebook environment like Google Colab or a bundle of packages like Anaconda, ensure that the version you are using is supported. If it's not, use the
--upgradeflag and pip will install the most recent version.
This is the part you want to replace with your own code. Using a Hugging Face transformer, build a machine learning model and keep it in-memory. In this example we're using bert-base-uncased, which will fill in the missing word in a sentence.
All Hugging Face models must be wrapped as a pipeline.
from transformers import pipeline
model = pipeline('fill-mask', model='bert-base-uncased')
createcommand to package your model into a Truss.
from truss import create
tr = create(model, target_directory="huggingface_truss")
Check the target directory to see your new Truss!
To get a prediction from the Truss, try running:
tr.predict("Donatello is a teenage mutant [MASK] turtle")