from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

app = FastAPI()

# Load NER pipeline
ner_pipeline = pipeline(
    "ner",
    model="dslim/bert-large-NER",
    aggregation_strategy="simple"
)

class RequestData(BaseModel):
    sentence: str

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

@app.post("/predict")
def predict(data: RequestData):

    predictions = ner_pipeline(data.sentence)

    allowed = {"PER", "ORG", "LOC", "MISC"}

    entities = []
    seen = set()

    for pred in predictions:

        label = pred["entity_group"]

        if label not in allowed:
            continue

        start = pred["start"]
        end = pred["end"]

        key = (start, end)

        if key in seen:
            continue

        seen.add(key)

        entities.append({
            "text": pred["word"],
            "start": start,
            "end": end,
            "label": label,
            "score": float(pred["score"])
        })

    return {
        "entities": entities
    }