The simplest way to serve AI/ML models in production
- Write once, run anywhere: Package and test model code, weights, and dependencies with a model server that behaves the same in development and production.
- Fast developer loop: Implement your model with fast feedback from a live reload server, and skip Docker and Kubernetes configuration with a batteries-included model serving environment.
- Support for all Python frameworks: From
transformersanddiffuserstoPyTorchandTensorFlowtoTensorRTandTriton, Truss supports models created and served with any framework.
See Trusses for popular models including:
- 🦙 Llama 2 7B (13B) (70B)
- 🎨 Stable Diffusion XL
- 🗣 Whisper
and dozens more examples.
Install Truss with:
pip install --upgrade truss
As a quick example, we'll package a text classification pipeline from the open-source transformers package.
To get started, create a Truss with the following terminal command:
truss init text-classificationWhen prompted, give your Truss a name like Text classification.
Then, navigate to the newly created directory:
cd text-classificationOne of the two essential files in a Truss is model/model.py. In this file, you write a Model class: an interface between the ML model that you're packaging and the model server that you're running it on.
There are two member functions that you must implement in the Model class:
load()loads the model onto the model server. It runs exactly once when the model server is spun up or patched.predict()handles model inference. It runs every time the model server is called.
Here's the complete model/model.py for the text classification model:
from transformers import pipeline
class Model:
def __init__(self, **kwargs):
self._model = None
def load(self):
self._model = pipeline("text-classification")
def predict(self, model_input):
return self._model(model_input)The other essential file in a Truss is config.yaml, which configures the model serving environment. For a complete list of the config options, see the config reference.
The pipeline model relies on Transformers and PyTorch. These dependencies must be specified in the Truss config.
In config.yaml, find the line requirements. Replace the empty list with:
requirements:
- torch==2.0.1
- transformers==4.30.0No other configuration is needed.
Truss is maintained by Baseten, which provides infrastructure for running ML models in production. We'll use Baseten as the remote host for your model.
Other remotes are coming soon, starting with AWS SageMaker.
To set up the Baseten remote, you'll need a Baseten API key. If you don't have a Baseten account, no worries, just sign up for an account and you'll be issued plenty of free credits to get you started.
With your Baseten API key ready to paste when prompted, you can deploy your model:
truss pushYou can monitor your model deployment from your model dashboard on Baseten.
After the model has finished deploying, you can invoke it from the terminal.
Invocation
truss predict -d '"Truss is awesome!"'Response
[
{
"label": "POSITIVE",
"score": 0.999873161315918
}
]Truss is backed by Baseten and built in collaboration with ML engineers worldwide. Special thanks to Stephan Auerhahn @ stability.ai and Daniel Sarfati @ Salad Technologies for their contributions.
We enthusiastically welcome contributions in accordance with our contributors' guide and code of conduct.
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent