spb/trouve-ka

1# Trouve-KA — serveur d'embeddings2# Author: Simon-Pierre Boucher3# Contact: contact@spboucher.ai45"""Serveur HTTP d'embeddings multilingues.67Modèle : intfloat/multilingual-e5-small (384 dims) — espace vectoriel aligné8FR/EN : « thermopompe » et « heat pump » sont voisins sans traduction (§9).9e5 exige les préfixes "query: " / "passage: " — le paramètre `kind` les gère.1011Volontairement autonome (fastapi + sentence-transformers seulement) : il tourne12dans son propre venv sur un node satellite, PAS dans l'image Docker principale13(torch pèse ~1 Go et n'a rien à faire dans le chemin de crawl).1415Lancement :  uvicorn services.embedding.server:app --host 0.0.0.0 --port 809116"""1718import os19import threading20from typing import Literal2122from fastapi import FastAPI23from pydantic import BaseModel, Field2425MODEL_NAME = os.environ.get("EMBEDDING_MODEL", "intfloat/multilingual-e5-small")26RERANK_MODEL_NAME = os.environ.get(27    "RERANK_MODEL", "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"28)29MAX_TEXTS = int(os.environ.get("EMBEDDING_MAX_TEXTS", "256"))30MAX_CHARS = int(os.environ.get("EMBEDDING_MAX_CHARS", "1200"))3132app = FastAPI(title="Trouve-KA Embeddings", docs_url=None, redoc_url=None)3334_model = None35_reranker = None36_model_lock = threading.Lock()373839def _get_model():40    global _model41    if _model is None:42        with _model_lock:43            if _model is None:44                from sentence_transformers import SentenceTransformer4546                _model = SentenceTransformer(MODEL_NAME, device="cpu")47    return _model484950def _get_reranker():51    global _reranker52    if _reranker is None:53        with _model_lock:54            if _reranker is None:55                from sentence_transformers import CrossEncoder5657                _reranker = CrossEncoder(RERANK_MODEL_NAME, device="cpu")58    return _reranker596061class EmbedRequest(BaseModel):62    texts: list[str] = Field(..., min_length=1, max_length=MAX_TEXTS)63    kind: Literal["query", "passage"] = "passage"646566class EmbedResponse(BaseModel):67    vectors: list[list[float]]68    dim: int69    model: str707172@app.get("/health")73def health() -> dict:74    return {"ok": True, "model": MODEL_NAME, "loaded": _model is not None}757677@app.post("/embed", response_model=EmbedResponse)78def embed(req: EmbedRequest) -> EmbedResponse:79    model = _get_model()80    prefix = "query: " if req.kind == "query" else "passage: "81    texts = [prefix + t[:MAX_CHARS] for t in req.texts]82    vectors = model.encode(texts, normalize_embeddings=True, batch_size=64)83    return EmbedResponse(vectors=vectors.tolist(), dim=len(vectors[0]), model=MODEL_NAME)848586class RerankRequest(BaseModel):87    query: str = Field(..., min_length=1, max_length=300)88    texts: list[str] = Field(..., min_length=1, max_length=64)899091class RerankResponse(BaseModel):92    scores: list[float]93    model: str949596@app.post("/rerank", response_model=RerankResponse)97def rerank(req: RerankRequest) -> RerankResponse:98    """Cross-encoder (requête, passage) → score de pertinence fin.99100    Réservé au top-N (coût quadratique en attention) : l'API n'envoie que ~20101    candidats déjà filtrés par BM25+kNN.102    """103    model = _get_reranker()104    pairs = [(req.query, t[:MAX_CHARS]) for t in req.texts]105    scores = model.predict(pairs, batch_size=32)106    return RerankResponse(scores=[float(s) for s in scores], model=RERANK_MODEL_NAME)107