AI-Engineering Performance Cases

import time
from contextlib import contextmanager

timings: dict[str, float] = {}


@contextmanager
def stage(name: str):
    wall0, cpu0 = time.perf_counter(), time.process_time()
    try:
        yield
    finally:
        wall, cpu = time.perf_counter() - wall0, time.process_time() - cpu0
        timings[name] = wall
        print(f"{name:<10} wall {wall:6.3f}s  cpu {cpu:6.3f}s  cores busy {cpu / wall if wall else 0:4.2f}")


with stage("fetch"):
    time.sleep(0.2)                                    # network: cores busy ~0
with stage("clean"):
    docs = [" ".join(f"tok{i}".split()) for i in range(300_000)]   # Python CPU: cores busy ~1

import os

os.environ.setdefault("TOKENIZERS_PARALLELISM", "true")   # single process: let Rust use the cores

from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("bert-base-uncased", use_fast=True)


def token_counts(texts: list[str], batch: int = 1024) -> list[int]:
    counts: list[int] = []
    for i in range(0, len(texts), batch):
        enc = tok(texts[i:i + batch], add_special_tokens=False, truncation=False)
        counts.extend(len(ids) for ids in enc["input_ids"])
    return counts

import json

raw_outputs = ['{"answer": "Paris", "confidence": 0.93, "citations": [1, 4]}'] * 50_000

try:
    import orjson                                     # Rust; typically several times faster than json

    loads = orjson.loads
except ImportError:                                   # stdlib fallback keeps the code runnable
    loads = json.loads

parsed = [loads(s) for s in raw_outputs]
print(len(parsed), parsed[0]["answer"])

import numpy as np

rng = np.random.default_rng(0)
batch = rng.standard_normal((4096, 768)).astype(np.float32)

# Slow: Python loop per vector (holds the GIL, ~100x slower)
# normed = [v / np.linalg.norm(v) for v in batch]

normed = batch / np.linalg.norm(batch, axis=1, keepdims=True)          # one call
sims = normed @ normed[:8].T                                           # BLAS, multi-core
near_dupes = np.argwhere(np.triu(sims[:8], k=1) > 0.95)                 # pairs among the first 8
scale = np.abs(normed).max(axis=1, keepdims=True) / 127.0
int8 = np.round(normed / scale).astype(np.int8)                        # int8 quantisation
print(normed.shape, int8.dtype, len(near_dupes))

import asyncio
from concurrent.futures import ThreadPoolExecutor

model_pool = ThreadPoolExecutor(max_workers=1)        # model calls release the GIL inside torch/onnx


def model_forward(texts: list[str]) -> list[int]:      # stands in for tokenizer(texts) + model(**batch)
    return [len(t) for t in texts]


class MicroBatcher:
    def __init__(self, max_batch: int = 32, max_wait_s: float = 0.005):
        self.queue: asyncio.Queue = asyncio.Queue()
        self.max_batch, self.max_wait_s = max_batch, max_wait_s

    async def submit(self, text: str) -> int:
        fut = asyncio.get_running_loop().create_future()
        await self.queue.put((text, fut))
        return await fut

    async def run(self) -> None:
        loop = asyncio.get_running_loop()
        while True:
            text, fut = await self.queue.get()
            batch = [(text, fut)]
            deadline = loop.time() + self.max_wait_s
            while len(batch) < self.max_batch and (timeout := deadline - loop.time()) > 0:
                try:
                    batch.append(await asyncio.wait_for(self.queue.get(), timeout))
                except asyncio.TimeoutError:
                    break
            try:
                outs = await loop.run_in_executor(model_pool, model_forward, [t for t, _ in batch])
            except Exception as err:                    # fail every waiter, keep the batcher alive
                for _, f in batch:
                    if not f.done():
                        f.set_exception(err)
                continue
            for (_, f), out in zip(batch, outs):
                if not f.done():
                    f.set_result(out)


async def main() -> None:
    batcher = MicroBatcher()
    runner = asyncio.create_task(batcher.run())
    results = await asyncio.gather(*(batcher.submit(f"request {i}") for i in range(100)))
    runner.cancel()
    print(len(results), results[:3])


asyncio.run(main())

import numpy as np
import pandas as pd

df = pd.DataFrame({
    "prompt": ["Summarise this", "Translate: hola", "Summarise that report"],
    "tokens_in": [120, 15, 900],
    "tokens_out": [80, 5, 300],
})

# Before: df["cost"] = df.apply(lambda r: r.tokens_in * 3e-6 + r.tokens_out * 15e-6, axis=1)
df["cost"] = df["tokens_in"] * 3e-6 + df["tokens_out"] * 15e-6
df["task"] = np.select(
    [df["prompt"].str.startswith("Summarise"), df["prompt"].str.startswith("Translate")],
    ["summarise", "translate"],
    default="other",
)
df["long"] = df["tokens_in"] > 500
print(df[["task", "cost", "long"]])

import re

PATTERNS = [
    (re.compile(r"<script.*?</script>", re.S | re.I), " "),
    (re.compile(r"<[^>]+>"), " "),
    (re.compile(r"https?://\S+"), " <url> "),
    (re.compile(r"[ \t ]+"), " "),
    (re.compile(r"\n{3,}"), "\n\n"),
]
BOILER = re.compile(r"(?:cookie policy|subscribe to our newsletter|all rights reserved)", re.I)


def clean(doc: str) -> str:
    if "<" not in doc and "http" not in doc:          # cheap prefilter skips most regex work
        return " ".join(doc.split())
    for rx, repl in PATTERNS:
        doc = rx.sub(repl, doc)
    return BOILER.sub("", doc).strip()


print(clean("<p>Hello   <b>world</b> see https://x.y/z</p>\n\n\n\nAll rights reserved"))

import numpy as np

docs = ["vector search with hnsw", "bm25 keyword search", "reranking with cross encoders", "hnsw graph index"]
vocab = {w: i for i, w in enumerate(sorted({w for d in docs for w in d.split()}))}
tf = np.zeros((len(docs), len(vocab)), dtype=np.float32)
for r, d in enumerate(docs):                      # build once at index time
    for w in d.split():
        tf[r, vocab[w]] += 1
doc_len = tf.sum(axis=1)
idf = np.log(1 + (len(docs) - (tf > 0).sum(0) + 0.5) / ((tf > 0).sum(0) + 0.5))


def bm25(query: str, k1: float = 1.5, b: float = 0.75) -> np.ndarray:
    cols = [vocab[w] for w in query.split() if w in vocab]
    f = tf[:, cols]
    denom = f + k1 * (1 - b + b * doc_len[:, None] / doc_len.mean())
    return (idf[cols] * f * (k1 + 1) / denom).sum(axis=1)


print(np.argsort(-bm25("hnsw search"))[:2])

import os

import torch
from torch.utils.data import DataLoader, Dataset


class TextDataset(Dataset):
    def __init__(self, paths: list[str]):
        self.paths = paths                       # store paths, not loaded data: cheap to send to workers

    def __len__(self) -> int:
        return len(self.paths)

    def __getitem__(self, i: int) -> torch.Tensor:
        with open(self.paths[i], "rb") as fh:
            data = fh.read()
        return torch.frombuffer(bytearray(data[:512].ljust(512, b"\0")), dtype=torch.uint8)


def worker_init(_worker_id: int) -> None:
    torch.set_num_threads(1)                     # one intra-op thread per worker
    os.environ["TOKENIZERS_PARALLELISM"] = "false"


def make_loader(paths: list[str]) -> DataLoader:
    return DataLoader(TextDataset(paths), batch_size=64, num_workers=4, worker_init_fn=worker_init,
                      persistent_workers=True, prefetch_factor=4, pin_memory=torch.cuda.is_available())

import asyncio
from concurrent.futures import ProcessPoolExecutor
from contextlib import asynccontextmanager

from fastapi import FastAPI
from pydantic import BaseModel


def chunk_document(text: str, size: int) -> list[str]:        # module-level: picklable
    words = text.split()
    return [" ".join(words[i:i + size]) for i in range(0, len(words), size)]


@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.cpu_pool = ProcessPoolExecutor(max_workers=4)   # created once, reused
    try:
        yield
    finally:
        app.state.cpu_pool.shutdown(wait=True, cancel_futures=True)


app = FastAPI(lifespan=lifespan)


class ChunkRequest(BaseModel):
    text: str
    size: int = 200


@app.post("/chunk")
async def chunk(req: ChunkRequest) -> dict:
    loop = asyncio.get_running_loop()
    chunks = await loop.run_in_executor(app.state.cpu_pool, chunk_document, req.text, req.size)
    return {"n_chunks": len(chunks)}

import asyncio
from concurrent.futures import ProcessPoolExecutor


def score(pred: str, gold: str) -> float:            # CPU: normalise + token-overlap F1
    p, g = pred.lower().split(), gold.lower().split()
    common = len(set(p) & set(g))
    if not common:
        return 0.0
    prec, rec = common / len(p), common / len(g)
    return 2 * prec * rec / (prec + rec)


async def generate(prompt: str) -> str:              # I/O: stands in for the API call
    await asyncio.sleep(0.01)
    return prompt.replace("Q:", "A:")


async def run_eval(items: list[tuple[str, str]], pool: ProcessPoolExecutor) -> float:
    sem = asyncio.Semaphore(32)

    async def one(prompt: str, gold: str) -> float:
        async with sem:
            pred = await generate(prompt)
        return await asyncio.get_running_loop().run_in_executor(pool, score, pred, gold)

    scores = await asyncio.gather(*(one(p, g) for p, g in items))
    return sum(scores) / len(scores)


if __name__ == "__main__":
    data = [(f"Q: capital {i}", f"A: capital {i}") for i in range(200)]
    with ProcessPoolExecutor(max_workers=2) as pool:
        print(f"mean F1 {asyncio.run(run_eval(data, pool)):.2f}")