AWS Lambda Python Hugging Face

Tip

This page is best viewed through the site search or through the AI.

Uses a Python Lambda container image to deploy a lightweight Hugging Face model.

Uses the transformers library to generate text using the TinyStories-33M model. The backend is the pytorch cpu runtime.

Note

This is not a production ready example.

This example also shows how it is possible to use custom index resolution to get dependencies from a private pypi server such as the pytorch cpu link. This example also shows how to use a custom Dockerfile to handle complex builds such as installing pytorch and pruning the build size.

sst.config.ts

new sst.aws.Function("PythonFunction", {
  python: {
    container: true,
  },
  handler: "functions/src/functions/api.handler",
  runtime: "python3.12",
  memory: "2048 MB",
  timeout: "120 seconds",
  url: true,
});

functions/src/functions/api.py

from transformers import pipeline
import json

# Initialize the pipeline once (outside the handler) to avoid reloading the model on each request
pipe = pipeline("text-generation", model="roneneldan/TinyStories-1M")


def handler(event, context):
    # Define the prompt for text generation
    prompt = "Write a short story about the magical framework SST that makes the cloud so easy to use!"

    # Generate text using the pipeline
    response = pipe(
        prompt,
        max_length=150,
        num_return_sequences=1,
        temperature=0.6,
        top_k=50,
        do_sample=True,
    )

    # Extract the generated text
    generated_text = response[0]["generated_text"]

    # Return the response with a status code and JSON-encoded body
    return {
        "statusCode": 200,
        "body": json.dumps({"story": generated_text}),
    }

View the full example.