LitServe

LitServe is a lightweight and fast inference server for machine learning models. It can be a good alternative to FastAPI if you are looking for dynamic batching support. It also has higher level abstractions for authentication, middleware, OpenAI compatible spec, streaming etc built on top of FastAPI.

In this example, we will deploy a simple Whisper model (speech to text) using faster-whisper and Litserve. You can find the code for this example here.

The key files are:

• whisper_server.py: Contains the WhisperLitAPI that implements the LitAPI interface.
• requirements.txt: Contains the dependencies.

​

How to write the inference function in LitServe

The whisper_server.py file contains the WhisperLitAPI class that implements the LitAPI interface.

...
def _get_model_dir():
    if "MODEL_DIR" not in os.environ:
        raise Exception(
            "MODEL_DIR environment variable is not set. Please set it to the directory containing the model."
        )
    return os.environ["MODEL_DIR"]

MODEL_DIR = _get_model_dir()

class WhisperLitAPI(ls.LitAPI):
    def setup(self, device):
        # Load the OpenAI Whisper model
        self.model = WhisperModel(MODEL_DIR, device="cpu")

    def decode_request(self, request: UploadFile):
        return request.file

    def predict(self, audio_stream):
        # Process the audio stream bytes and return the transcription result
        output = []
        segments, _ = self.model.transcribe(audio_stream)
        for segment in segments:
            output.append({"start": segment.start, "end": segment.end, "text": segment.text})
        return output

    def encode_response(self, output):
        # Output is already a list of dicts. So we can return it as is.
        return output

Mainly we inherit the LitAPI class and implement the setup, decode_request, predict and encode_response methods.

• setup: Load the model.
• decode_request: Decodes and transforms the request body to the input format expected by the model.
• predict: Processes the output of decode_request and runs model inference.
• encode_response: Formats the response. Can perform any postprocessing on the response.

See LitServe Docs for reference and more examples.

​

Running the server locally

1. Install the dependencies

pip install -r requirements.txt

2. Run the server

export MODEL_DIR="Systran/faster-whisper-tiny"
python whisper_server.py

3. Test the server

curl -X POST http://0.0.0.0:8000/predict -F "request=@./audio.mp3"

We’ll see something like this:

[
    {"start":0.0,"end":5.0,"text":" Oh, you think darkness is your ally."},
    {"start":5.0,"end":8.0,"text":" Are you merely adopted the dark?"},
    {"start":8.0,"end":11.0,"text":" I was born in it."},
    {"start":11.0,"end":14.0,"text":" More lit by it."},
    {"start":14.0,"end":17.0,"text":" I didn't see the light until I was already a man,"},
    {"start":17.0,"end":20.0,"text":" but then it was nothing to me but brightened."}
]

​

Deploying to TrueFoundry

Since the models are being pulled from HuggingFace Hub, we can directly deploy the code to TrueFoundry and use Artifacts Download feature to automatically download the model.