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

app = FastAPI()

# Load model once
classifier = pipeline(
    "zero-shot-classification",
    model="MoritzLaurer/mDeBERTa-v3-base-mnli-xnli"
)

# Your classes
CANDIDATE_LABELS = [
    "Garbage issue",
    "Streetlight not working",
    "Road damage",
    "Water supply issue",
    "Noise pollution",
    "Flooding",
    "Corruption",
    "Other"
]

class Query(BaseModel):
    text: str

@app.post("/predict")
def predict(query: Query):
    result = classifier(
        query.text,
        candidate_labels=CANDIDATE_LABELS,
        multi_label=False   # <-- single best class
    )

    # Top-1 predicted label
    predicted_label = result["labels"][0]

    return {"label": predicted_label}