View example on GitHub

This guide walks through deploying Mistral-7B, a powerful large language model (LLM), using Truss. You’ll configure the model, set up inference, allocate resources, and deploy it as an API endpoint.

Set up your model

Start by importing the necessary libraries:

model/model.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

Specify the Hugging Face model checkpoint:

model/model.py

CHECKPOINT = "mistralai/Mistral-7B-v0.1"

Define the model class

Create a Model class that loads Mistral-7B and its tokenizer when the server starts:

model/model.py

class Model:
    def __init__(self, **kwargs) -> None:
        self.tokenizer = None
        self.model = None

    def load(self):
        self.model = AutoModelForCausalLM.from_pretrained(
            CHECKPOINT, dtype=torch.float16, device_map="auto"
        )
        self.tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)

Implement inference

The predict function handles inference by tokenizing input, generating text, and decoding the output.

model/model.py

    def predict(self, request: dict):
        prompt = request.pop("prompt")
        generate_args = {
            "max_new_tokens": request.get("max_new_tokens", 128),
            "temperature": request.get("temperature", 1.0),
            "top_p": request.get("top_p", 0.95),
            "top_k": request.get("top_k", 50),
            "repetition_penalty": 1.0,
            "use_cache": True,
            "do_sample": True,
            "eos_token_id": self.tokenizer.eos_token_id,
            "pad_token_id": self.tokenizer.pad_token_id,
        }

        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.cuda()

        with torch.no_grad():
            output = self.model.generate(input_ids=input_ids, **generate_args)
            return self.tokenizer.decode(output[0])

Configure your deployment

Define dependencies

Specify the necessary Python packages in config.yaml:

config.yaml

model_name: Mistral 7B
python_version: py311
requirements:
  - transformers==5.17.0
  - sentencepiece==0.2.2
  - accelerate==1.15.0
  - torch==2.14.0
  - numpy==2.4.6

Allocate compute resources

Mistral-7B requires an NVIDIA A10G GPU for efficient inference:

config.yaml

resources:
  accelerator: A10G
  use_gpu: true

Deploy the model

Install the Baseten CLI and sign in:

Install the Baseten CLI

macOS or Linux
Windows

Terminal

brew tap basetenlabs/baseten
brew trust --formula basetenlabs/baseten/baseten
brew install baseten

Download and extract the Windows archive for the latest Baseten CLI release:

Terminal

$release = Invoke-RestMethod https://api.github.com/repos/basetenlabs/baseten-cli/releases/latest
$asset = $release.assets | Where-Object { $_.name -like "*_windows_amd64.zip" }
Invoke-WebRequest $asset.browser_download_url -OutFile baseten.zip
Expand-Archive -Force baseten.zip .

Move baseten.exe to a directory on your PATH.

For other platforms or a specific version, see the Baseten CLI install reference.

Sign in

baseten auth login

From the directory that contains config.yaml, push your model to Baseten:

baseten model push

You should see output like:

Pushing model "Mistral 7B"...
Uploading model...
Uploaded model in 0s
✨ Model Mistral 7B was successfully pushed ✨

  Model:       Mistral 7B (abc1d2ef)
  Deployment:  xyz123
  Environment: production

🪵 View logs:
   deployment:   baseten model deployment logs --model-id abc1d2ef --deployment-id xyz123
   environment:  baseten model environment logs --model-id abc1d2ef --environment production  (once deployed)
   app:          https://app.baseten.co/models/abc1d2ef/logs/xyz123

🚀 Invoke your model:
   URL:  https://model-abc1d2ef.api.baseten.co/deployment/xyz123/predict
   CLI:  baseten model predict --model-id abc1d2ef

Copy the model ID from the Model: line (abc1d2ef in the example). The first deploy takes several minutes while Baseten builds the image and the model loads its weights from Hugging Face. Once the deployment is active, call the model with the Baseten CLI:

baseten model predict --model-id abc1d2ef --data '{"prompt": "What is a large language model?"}'

To call the model’s endpoint from Python instead, create a personal key in API keys and set it in the terminal where you’ll run the request:

export BASETEN_API_KEY="YOUR_API_KEY"

$env:BASETEN_API_KEY = "YOUR_API_KEY"

Then send the request:

call_model.py

import os
import requests

model_id = "abc1d2ef"

response = requests.post(
    f"https://model-{model_id}.api.baseten.co/production/predict",
    headers={"Authorization": f"Bearer {os.environ['BASETEN_API_KEY']}"},
    json={"prompt": "What is a large language model?"},
)

print(response.json())

Replace abc1d2ef with your model ID. To use the Truss CLI instead, sign in with truss login, deploy with truss push, and call the model with truss predict --published -d '{"prompt": "What is a large language model?"}'.

Check for optimized engine support

For faster inference, check whether your model runs on a Baseten engine such as Engine-Builder-LLM, which builds TensorRT-LLM engines, or on an open-source server such as vLLM or SGLang. For an overview, see Engines.

Examples

Models

Engines

Custom Docker servers

Custom Python models

Chains

Post-training

Baseten for Model Labs

Build and deploy an LLM

View example on GitHub

Set up your model

Define the model class

Implement inference

Configure your deployment

Define dependencies

Allocate compute resources

Deploy the model

Check for optimized engine support

View example on GitHub

​Set up your model

​Define the model class

​Implement inference

​Configure your deployment

​Define dependencies

​Allocate compute resources

​Deploy the model

​Check for optimized engine support

Set up your model

Define the model class

Implement inference

Configure your deployment

Define dependencies

Allocate compute resources

Deploy the model

Check for optimized engine support