Custom frameworks

Build and deploy a model from any framework with the Baseten Python client.
While Baseten's deploy function supports some model frameworks out of the box with no need for customization, you can use the custom model deployment flow to deploy a model written in any framework. This gives you complete control over the Python environment and execution of your model.
Baseten uses Truss for custom models.
You can create an empty Truss using the Baseten cli:
!baseten truss init custom_model
or from your notebook/ipython directly:
import truss
truss_handle = truss.init('custom_model')
This will create a Truss, in other words your custom model, in the custom_model folder. You can update the file custom_model/model/model.py to configure your custom model.
Your model must be a Python class that implements two methods:
load, a method that will be called upon initialization of the model in the deployment environment.
predict, a method which consumes deserialized JSON input from a web request or from a Baseten worklet. This is the integration point for the underlying model object. It must return data in a JSON-serializable format.
For example:
model.py
import pickle
​
​
class Model:
    def __init__(self, data_dir):
        self.model = None
        self.encoder = None
        self.model_path = data_dir / 'my_model.pkl'
        self.encoder_path = data_dir / 'my_encoder.pkl'
​
    def load(self):
        self.model = pickle.load(self.model_path.open('rb'))
        self.encoder = pickle.load(self.model_path.open('rb'))
​
    def predict(self, request: dict) -> dict:
        inputs = request['inputs']
        # Encode the inputs
        model_inputs = self.encoder(inputs)
        # Run predict on the encoded inputs 
        predictions = self.model.predict(model_inputs)
        return {
            "predictions": predictions,
            "encoded_values": model_inputs,
        }
​
When a model is invoked in a Baseten application, either via a model block or model.predict in Python code, the input is wrapped as the value of the 'inputs' key of a dictionary. So, the predict function above should expect input in that form, as in:
{"inputs": ["JSON", "serializable", "values"]}
Also, Baseten applications expect the model to return the predictions as value of 'predictions' key in a dictionary, such as:
{"predictions": ["JSON", "serializable", "values"]}
Serialized objects such as model binaries, embeddings, and datasets can be placed under the Truss's data directory i.e. custom_model/data. The Path (pathlib.Path) to the data directory is made available to the Model constructor as the parameter data_dir.
Any supporting code can be put anywhere under the model directory custom_model/model and used in model.py.
Any Python requirements can be placed in the requirements field of the config file custom_model/config.yaml. They should be in the pip requirements format, with one line of requirements.txt per one list item in the requirements key of the config. Here's a quick example:
requirements:
  - torchvision==0.9.1
  - openai==0.6.3
In order to deploy a custom model, first get a handle to the truss using truss.load() and then deploy the truss using baseten.deploy().
import baseten, truss
​
baseten.login("*** INSERT API KEY ***")  # https://docs.baseten.co/settings/api-keys
truss_handle = truss.load('custom_model')
baseten.deploy(
    truss_handle,
    model_name='encoder model',
)
Baseten will package your files and deploy your custom model.
For guidance on how to interact with your model check out Interacting with models.
For a more in-depth look at the technology powering our model serving, please look at Truss.
XGBoost
Pre-trained models
Last modified 18d ago