from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from gliner import GLiNER
import gradio as gr
import spaces
import time

app = FastAPI(title="SIGAP AI ML Service (HF Spaces)")

MODEL_HUB = "urchade/gliner_multi-v2.1"

print(f"Downloading model {MODEL_HUB} from HuggingFace...", flush=True)
model = GLiNER.from_pretrained(MODEL_HUB)
print("Model GLiNER loaded successfully!", flush=True)

LABELS = [
    "nama",
    "usia",
    "kondisi medis",
    "keterbatasan mobilitas",
    "asal lokasi",
    "anggota keluarga",
    "ketiadaan obat",
]

class ExtractRequest(BaseModel):
    teks: str

class Entity(BaseModel):
    label: str
    text: str
    confidence: float

@spaces.GPU
def predict_entities(teks: str):
    predictions = model.predict_entities(teks, LABELS, threshold=0.4)
    return [
        {
            "label": p["label"],
            "text": p["text"],
            "confidence": float(p["score"]),
        }
        for p in predictions
    ]

@app.post("/extract")
async def extract_entities(req: ExtractRequest):
    start_time = time.time()
    try:
        entities = predict_entities(req.teks)
        latency = (time.time() - start_time) * 1000
        return {"entities": entities, "latency_ms": latency}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {"status": "Model GLiNER Ready!"}

def ui_predict(teks):
    return predict_entities(teks)

demo = gr.Interface(
    fn=ui_predict,
    inputs="text",
    outputs="json",
    title="SIGAP AI ML Service"
)

app = gr.mount_gradio_app(app, demo, path="/ui")