Spaces:
Running
Running
| 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 | |
| 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} | |