embedding-server.py

#!/usr/bin/env python3
"""OpenAI-compatible embedding server wrapping Qwen3-Embedding-4B on GPU."""

import time
import torch
from contextlib import asynccontextmanager
from typing import Union

from fastapi import FastAPI
from pydantic import BaseModel, Field
from sentence_transformers import SentenceTransformer

MODEL_NAME = "Qwen/Qwen3-Embedding-4B"
EMBEDDING_DIM = 1024  # MRL truncation

model: SentenceTransformer | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    print(f"Loading {MODEL_NAME} onto GPU...")
    t0 = time.time()
    model = SentenceTransformer(
        MODEL_NAME,
        trust_remote_code=True,
        model_kwargs={"torch_dtype": torch.float16},
    )
    model.max_seq_length = 8192
    model.truncate_dim = EMBEDDING_DIM
    print(f"Model loaded in {time.time() - t0:.1f}s")
    yield
    del model
    torch.cuda.empty_cache()


app = FastAPI(title="Kit Embeddings", lifespan=lifespan)


class EmbeddingRequest(BaseModel):
    input: Union[str, list[str]]
    model: str = "qwen3-embedding-4b"
    encoding_format: str = "float"


class EmbeddingData(BaseModel):
    object: str = "embedding"
    embedding: list[float]
    index: int


class Usage(BaseModel):
    prompt_tokens: int
    total_tokens: int


class EmbeddingResponse(BaseModel):
    object: str = "list"
    data: list[EmbeddingData]
    model: str = "qwen3-embedding-4b"
    usage: Usage


@app.post("/v1/embeddings")
async def create_embeddings(req: EmbeddingRequest) -> EmbeddingResponse:
    texts = [req.input] if isinstance(req.input, str) else req.input

    embeddings = model.encode(
        texts,
        batch_size=32,
        normalize_embeddings=True,
        show_progress_bar=False,
    )

    # Estimate token count (~4 chars per token)
    total_chars = sum(len(t) for t in texts)
    est_tokens = max(1, total_chars // 4)

    data = [
        EmbeddingData(embedding=emb.tolist(), index=i)
        for i, emb in enumerate(embeddings)
    ]

    return EmbeddingResponse(
        data=data,
        usage=Usage(prompt_tokens=est_tokens, total_tokens=est_tokens),
    )


@app.get("/health")
async def health():
    return {"status": "ok", "model": MODEL_NAME, "dim": EMBEDDING_DIM, "gpu": torch.cuda.is_available()}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8678)