Serving Data Science & AI

import io
from typing import Annotated

import pandas as pd
from fastapi import FastAPI, HTTPException, Query, UploadFile
from fastapi.testclient import TestClient
from pydantic import BaseModel

REVIEWS = pd.DataFrame(
    {
        "book_id": [1, 1, 1, 2, 2, 3],
        "reviewer": ["ana", "bo", "cy", "ana", "dee", "bo"],
        "stars": [5, 4, 4, 2, 3, None],
    }
)

app = FastAPI()


class BookRating(BaseModel):
    book_id: int
    reviews: int
    mean_stars: float | None


@app.get("/stats/ratings")
def ratings(min_reviews: Annotated[int, Query(ge=1)] = 1) -> list[BookRating]:
    grouped = (
        REVIEWS.groupby("book_id")
        .agg(reviews=("stars", "size"), mean_stars=("stars", "mean"))
        .reset_index()
    )
    grouped = grouped[grouped["reviews"] >= min_reviews]
    grouped = grouped.astype(object).where(grouped.notna(), None)
    return [BookRating(**row) for row in grouped.to_dict(orient="records")]


class UploadSummary(BaseModel):
    count: int
    mean: float | None


@app.post("/stats/upload")
def upload_summary(file: UploadFile) -> UploadSummary:
    try:
        frame = pd.read_csv(file.file, usecols=["stars"])
    except (ValueError, pd.errors.ParserError) as exc:
        raise HTTPException(status_code=422, detail="Expected a CSV with a stars column") from exc
    stars = frame["stars"]
    return UploadSummary(count=int(stars.count()), mean=None if stars.empty else float(stars.mean()))


@app.get("/stats/naive")
def naive():
    # NaN in the data and numpy scalars in the payload: the pitfall
    return REVIEWS.groupby("book_id")["stars"].mean().to_dict()

curl 'http://127.0.0.1:8000/stats/ratings?min_reviews=2'
curl -F 'file=@reviews.csv;type=text/csv' http://127.0.0.1:8000/stats/upload

[
  {"book_id": 1, "reviews": 3, "mean_stars": 4.333333333333333},
  {"book_id": 2, "reviews": 2, "mean_stars": 2.5},
  {"book_id": 3, "reviews": 1, "mean_stars": null}
]

from typing import Annotated, Literal

import pandas as pd
import plotly.express as px
from fastapi import FastAPI, Query, Response
from fastapi.testclient import TestClient

REVIEWS = pd.DataFrame(
    {
        "book": ["Dune", "Dune", "Dune", "Emma", "Emma", "Ulysses"],
        "stars": [5, 4, 4, 2, 3, 1],
    }
)

app = FastAPI()

MEDIA_TYPES = {"png": "image/png", "svg": "image/svg+xml"}


@app.get(
    "/charts/stars.{fmt}",
    response_class=Response,
    responses={200: {"content": {"image/png": {}, "image/svg+xml": {}}, "description": "The chart image"}},
)
def stars_chart(
    fmt: Literal["png", "svg"],
    width: Annotated[int, Query(ge=200, le=1600)] = 640,
    height: Annotated[int, Query(ge=150, le=1200)] = 400,
) -> Response:
    counts = REVIEWS.groupby("book", as_index=False)["stars"].mean()
    fig = px.bar(counts, x="book", y="stars", title="Mean stars per book", range_y=[0, 5])
    image = fig.to_image(format=fmt, width=width, height=height)
    return Response(
        content=image,
        media_type=MEDIA_TYPES[fmt],
        headers={"Cache-Control": "public, max-age=300"},
    )

{"description": "The chart image", "content": {"image/png": {}, "image/svg+xml": {}}}

import os
import threading
from collections.abc import Iterator

from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer


MODEL_ID = os.environ.get("MODEL_ID", "sshleifer/tiny-gpt2")


class TextGenerator:
    """Owns a tokenizer and a model; both are loaded once."""

    def __init__(self, model_id: str) -> None:
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model = AutoModelForCausalLM.from_pretrained(model_id)
        self.model.eval()
        self._lock = threading.Lock()  # one generation at a time on this model

    def _kwargs(self, prompt: str, max_new_tokens: int) -> dict:
        inputs = self.tokenizer(prompt, return_tensors="pt")
        return dict(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=self.tokenizer.eos_token_id,
        )

    def generate(self, prompt: str, max_new_tokens: int) -> str:
        kwargs = self._kwargs(prompt, max_new_tokens)
        with self._lock:
            output = self.model.generate(**kwargs)
        new_tokens = output[0][kwargs["input_ids"].shape[1]:]
        return self.tokenizer.decode(new_tokens, skip_special_tokens=True)

    def stream(self, prompt: str, max_new_tokens: int) -> Iterator[str]:
        kwargs = self._kwargs(prompt, max_new_tokens)
        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, timeout=30.0)

        def run() -> None:
            with self._lock:
                self.model.generate(**kwargs, streamer=streamer)

        worker = threading.Thread(target=run, daemon=True)
        worker.start()
        try:
            yield from streamer
        finally:
            worker.join()

import hashlib
import re
from contextlib import asynccontextmanager
from typing import Annotated, Protocol

import numpy as np
from fastapi import Depends, FastAPI, Request
from fastapi.testclient import TestClient
from pydantic import BaseModel, Field

DIM = 256


class Embedder(Protocol):
    def embed(self, texts: list[str]) -> np.ndarray: ...


class HashingEmbedder:
    """A stand-in with the interface of a real embedding model.

    It hashes words into a fixed-size vector, so it matches shared words, not meaning.
    Replace it with a real embedding model in production; nothing else changes.
    """

    def embed(self, texts: list[str]) -> np.ndarray:
        out = np.zeros((len(texts), DIM), dtype=np.float32)
        for row, text in enumerate(texts):
            for word in re.findall(r"[a-z']+", text.lower()):
                bucket = int.from_bytes(hashlib.sha256(word.encode()).digest()[:4], "big") % DIM
                out[row, bucket] += 1.0
        norms = np.linalg.norm(out, axis=1, keepdims=True)
        return out / np.maximum(norms, 1e-9)  # unit length: dot product == cosine similarity


REVIEWS = [
    {"id": 1, "book_id": 1, "text": "A slow and political start that rewards patience"},
    {"id": 2, "book_id": 1, "text": "Gripping desert adventure with a political edge"},
    {"id": 3, "book_id": 2, "text": "A witty comedy of manners"},
    {"id": 4, "book_id": 3, "text": "Dense and difficult, but a political and literary landmark"},
]


class Index:
    def __init__(self, embedder: Embedder, reviews: list[dict]) -> None:
        self.embedder = embedder
        self.reviews = reviews
        self.matrix = embedder.embed([r["text"] for r in reviews])  # shape (n, DIM)

    def search(self, query: str, limit: int) -> list[tuple[dict, float]]:
        vector = self.embedder.embed([query])[0]
        scores = self.matrix @ vector
        best = np.argsort(-scores)[:limit]
        return [(self.reviews[i], float(scores[i])) for i in best]


@asynccontextmanager
async def lifespan(app: FastAPI):
    yield {"index": Index(HashingEmbedder(), REVIEWS)}


def get_index(request: Request) -> Index:
    return request.state.index


app = FastAPI(lifespan=lifespan)


class SearchRequest(BaseModel):
    query: str = Field(min_length=1, max_length=500)
    limit: int = Field(default=3, ge=1, le=20)


class Hit(BaseModel):
    review_id: int
    book_id: int
    text: str
    score: float


@app.post("/reviews/search")
def search_reviews(body: SearchRequest, index: Annotated[Index, Depends(get_index)]) -> list[Hit]:
    return [
        Hit(review_id=r["id"], book_id=r["book_id"], text=r["text"], score=round(score, 3))
        for r, score in index.search(body.query, body.limit)
    ]

curl -X POST http://127.0.0.1:8000/reviews/search \
  -H 'Content-Type: application/json' \
  -d '{"query": "political", "limit": 2}'

[
  {"review_id": 2, "book_id": 1, "text": "Gripping desert adventure with a political edge", "score": 0.378},
  {"review_id": 1, "book_id": 1, "text": "A slow and political start that rewards patience", "score": 0.354}
]