NLP API Systems

from fastapi import FastAPI
from pydantic import BaseModel, Field
from sentence_transformers import SentenceTransformer
import numpy as np

app = FastAPI()
model = SentenceTransformer("intfloat/multilingual-e5-base", device="cuda")

class EmbedReq(BaseModel):
    texts: list[str] = Field(min_length=1, max_length=128)

@app.post("/embed")
def embed(req: EmbedReq):
    vecs = model.encode(req.texts, batch_size=32, normalize_embeddings=True)
    return {"dim": vecs.shape[1], "vectors": vecs.tolist()}

from transformers import pipeline

clf = pipeline(
    "text-classification",
    model="cardiffnlp/twitter-xlm-roberta-base-sentiment",
    device=0,                    # GPU
    truncation=True,
    max_length=256,
)

@app.post("/sentiment")
def sentiment(req: BatchReq):
    out = clf(req.texts, batch_size=32)
    return [{"label": o["label"], "score": round(o["score"], 4)} for o in out]

# collect requests for 5ms, run as single forward pass
import asyncio
QUEUE: asyncio.Queue = asyncio.Queue()

async def batcher():
    while True:
        items = [await QUEUE.get()]
        try:
            while len(items) < 32:
                items.append(QUEUE.get_nowait())
        except asyncio.QueueEmpty:
            pass
        texts = [i["text"] for i in items]
        results = clf(texts, batch_size=len(texts))
        for it, r in zip(items, results):
            it["fut"].set_result(r)

@app.on_event("startup")
async def start():
    asyncio.create_task(batcher())

@app.post("/sentiment")
async def sentiment(req: TextReq):
    fut = asyncio.get_event_loop().create_future()
    await QUEUE.put({"text": req.text, "fut": fut})
    return await asyncio.wait_for(fut, timeout=5)