from fastapi import FastAPI
from pydantic import BaseModel
import spacy

app = FastAPI(title="spaCy NER API")

# Load transformer model
nlp = spacy.load(
    "en_core_web_trf",
    disable=["tagger", "parser", "lemmatizer", "attribute_ruler"]
)

class TextInput(BaseModel):
    text: str


@app.post("/ner")
def extract_ner(data: TextInput):
    doc = nlp(data.text)

    entities = []
    for ent in doc.ents:
        entities.append({
            "text": ent.text,
            "label": ent.label_
        })

    return {"entities": entities}