app.py

"""GLiNER NER Server for HuggingFace Spaces."""
import os
import json
from fastapi import FastAPI
from pydantic import BaseModel
from gliner import GLiNER

MODEL_NAME = os.environ.get("MODEL_NAME", "urchade/gliner_multi-v2.1")
print(f"[GLiNER] Loading model: {MODEL_NAME}...", flush=True)
model = GLiNER.from_pretrained(MODEL_NAME)
print("[GLiNER] Model loaded.", flush=True)

app = FastAPI()

DEFAULT_LABELS = [
    "person", "organization", "location", "date", "event",
    "PERSONAL_FACT", "PREFERENCE", "RELATIONSHIP", "PET",
    "HOBBY", "HEALTH", "OCCUPATION", "FAMILY_MEMBER", "EVENT",
]


class ExtractRequest(BaseModel):
    text: str
    labels: list[str] | None = None
    entity_types: list[str] | None = None  # alias used by KikoBot adapter
    threshold: float = 0.3


@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_NAME}


@app.post("/extract")
def extract(req: ExtractRequest):
    labels = req.labels or req.entity_types or DEFAULT_LABELS
    results = model.predict_entities(req.text, labels, threshold=req.threshold)
    entities = []
    for ent in results:
        entities.append({
            "text": ent["text"],
            "type": ent["label"],
            "score": round(ent["score"], 4),
            "start": ent["start"],
            "end": ent["end"],
        })
    return {"entities": entities, "model": MODEL_NAME}