Step 5: Implement model load

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

environment_variables: {}
model_name: My First Truss
requirements:
  - torch==2.0.1
  - transformers==4.30.0
resources:
  accelerator: null
  cpu: "1"
  memory: 2Gi
  use_gpu: false
secrets: {}
system_packages: []

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
Diff

model/model.py

from transformers import pipeline

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
Diff

model/model.py

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

model/model.py

def load(self):
-    # Load model here and assign to self._model.
-    pass
+    self._model = pipeline("text-classification")

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

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

environment_variables: {}
model_name: My First Truss
requirements:
  - torch==2.0.1
  - transformers==4.30.0
resources:
  accelerator: null
  cpu: "1"
  memory: 2Gi
  use_gpu: false
secrets: {}
system_packages: []

Learn

Truss 101: Model serving

Step 5: Implement model load

Import transformers

Load the model

​Import transformers

​Load the model

Import transformers

Load the model