Step 5: Implement model load

Load an ML model into your Truss

The other essential file 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.

The code to load and invoke a model in a Jupyter notebook or Python script maps directly to the code used in model/model.py.

We’ll go line-by-line through the code. Open model/model.py in your text editor.

Import transformers

Import transformers.pipeline at the top of model/model.py:

Code

model/model.py

from transformers import pipeline

Load the model

The Model.load() function runs exactly once when the model server is spun up or patched and loads the model onto the model server.

Update load() to bring in the text-classification model from transformers.pipeline:

Code

model/model.py

def load(self):
    self._model = pipeline("text-classification")

You should see this change patched onto the model server in your truss watch terminal tab.

model/model.py ●

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 model_input