dfourneaux/stackoverflow

1from fastapi import FastAPI2from pydantic import BaseModel3import joblib4from typing import List5 6app = FastAPI()7 8# Chargement du pipeline complet (vectorizer + modèle)9model = joblib.load("model_sgd_hashing.pkl")10 11# ---------- Schémas ----------12class TitleRequest(BaseModel):13    title: str14 15 16class TagsResponse(BaseModel):17    title: str18    predicted_tags: List[str]19 20 21class BatchRequest(BaseModel):22    titles: List[str]23 24 25class BatchItemResponse(BaseModel):26    title: str27    predicted_tags: List[str]28 29 30class BatchResponse(BaseModel):31    results: List[BatchItemResponse]32 33# ---------- Fonction de prédiction ----------34def predict_tags(title: str):35    prediction = model.predict([title])36    tags = prediction[0]37 38    # Si ton modèle renvoie une seule string39    if isinstance(tags, str):40        tags = [tags]41 42    # Si numpy array43    if hasattr(tags, "tolist"):44        tags = tags.tolist()45 46    return tags47 48# ---------- Endpoint simple ----------49@app.post("/predict-tags", response_model=TagsResponse)50def predict(request: TitleRequest):51    tags = predict_tags(request.title)52 53    return {54        "title": request.title,55        "predicted_tags": tags56    }57 58# ---------- Endpoint batch ----------59@app.post("/predict-tags-batch", response_model=BatchResponse)60def predict_batch(request: BatchRequest):61 62    results = []63 64    # Prédiction en batch (plus propre et plus rapide)65    predictions = model.predict(request.titles)66 67    for title, pred in zip(request.titles, predictions):68 69        # Normalisation du format70        if isinstance(pred, str):71            tags = [pred]72        elif hasattr(pred, "tolist"):73            tags = pred.tolist()74        else:75            tags = pred76 77        results.append({78            "title": title,79            "predicted_tags": tags80        })81 82    return {"results": results}83