app.py

"""
ML Backend Label Studio - Detection de concepts medicaux avec GLiNER-BioMed.
Surligne automatiquement les entites medicales (maladies, medicaments, procedures...)
dans les notes cliniques pour faciliter le codage CIM-10.
"""

import os
import uuid
import logging
from contextlib import asynccontextmanager
from typing import Any, Optional, List, Dict

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from model import MedicalNERModel

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

ner_model = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global ner_model
    logger.info("Chargement du modele GLiNER-BioMed...")
    ner_model = MedicalNERModel()
    logger.info("Modele pret.")
    yield
    logger.info("Arret du serveur.")


app = FastAPI(title="CIM-10 Medical NER Backend", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ------------------------------------------------------------------ #
# Configuration : adapter a votre interface Label Studio              #
#   <Labels name="label" toName="text">  -> FROM_NAME / TO_NAME      #
#   <Text name="text" value="$text"/>    -> DATA_FIELD                #
# ------------------------------------------------------------------ #
FROM_NAME = os.environ.get("FROM_NAME", "label")
TO_NAME = os.environ.get("TO_NAME", "text")
DATA_FIELD = os.environ.get("DATA_FIELD", "text")
MODEL_VERSION = "gliner-biomed-small-v1.0"


# -- Pydantic schemas --
class PredictRequest(BaseModel):
    tasks: List[Dict[str, Any]]
    label_config: Optional[str] = None


class WebhookRequest(BaseModel):
    action: Optional[str] = None
    project: Optional[Dict[str, Any]] = None


# -- Endpoints --

@app.get("/health")
def health():
    return {
        "status": "UP",
        "model": MODEL_VERSION,
        "model_loaded": ner_model is not None,
    }


@app.post("/setup")
def setup(payload: Dict[str, Any] = {}):
    return {"model_version": MODEL_VERSION}


@app.post("/predict")
def predict(req: PredictRequest):
    """
    Endpoint principal. Recoit des taches, renvoie des predictions NER.
    Chaque concept medical detecte devient un span surligne dans Label Studio.

    Exemple de sortie pour une tache :
    {
      "result": [
        {"id":"a1","from_name":"label","to_name":"text","type":"labels",
         "value":{"start":45,"end":72,"text":"diabete de type 2",
                  "score":0.92,"labels":["Maladie"]}},
        {"id":"b2","from_name":"label","to_name":"text","type":"labels",
         "value":{"start":120,"end":132,"text":"metformine",
                  "score":0.88,"labels":["Medicament"]}}
      ],
      "score": 0.90,
      "model_version": "gliner-biomed-small-v1.0"
    }
    """
    predictions = []

    for task in req.tasks:
        text = task.get("data", {}).get(DATA_FIELD, "")
        if not text:
            predictions.append({"result": [], "model_version": MODEL_VERSION})
            continue

        entities = ner_model.predict(text)

        results = []
        for ent in entities:
            results.append({
                "id": str(uuid.uuid4())[:8],
                "from_name": FROM_NAME,
                "to_name": TO_NAME,
                "type": "labels",
                "value": {
                    "start": ent["start"],
                    "end": ent["end"],
                    "text": ent["text"],
                    "score": round(ent["score"], 4),
                    "labels": [ent["label"]],
                },
            })

        avg_score = (
            round(sum(e["score"] for e in entities) / len(entities), 4)
            if entities
            else 0.0
        )
        predictions.append({
            "result": results,
            "score": avg_score,
            "model_version": MODEL_VERSION,
        })

    return {"results": predictions}


@app.post("/webhook")
def webhook(req: WebhookRequest):
    logger.info("Webhook recu : action=%s", req.action)
    return {"status": "ok"}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)