Pulastya0's picture
Update app.py
025579b verified
import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel , Extra
from transformers import pipeline
# -------------------------------
# Set Hugging Face cache to writable directory
# -------------------------------
os.environ["HF_HOME"] = "/tmp/huggingface"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
os.environ["HF_DATASETS_CACHE"] = "/tmp/huggingface"
# -------------------------------
# FastAPI app
# -------------------------------
app = FastAPI(title="Routing Service - Space 2")
# -------------------------------
# Request Model
# -------------------------------
class RoutingRequest(BaseModel):
text: str
class Config:
extra = Extra.ignore # allow extra fields like user_email
# -------------------------------
# Load Zero-Shot Classifier
# -------------------------------
DEPARTMENTS = ['Account', 'Software', 'Network', 'Security', 'Hardware',
'Infrastructure', 'Licensing', 'Communication', 'RemoteWork',
'Training', 'Performance']
classifier = pipeline(
"zero-shot-classification",
model="MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli"
)
# -------------------------------
# Routing Endpoint
# -------------------------------
@app.post("/route")
async def route_ticket(req: RoutingRequest):
text = req.text
if not text:
raise HTTPException(status_code=400, detail="Text cannot be empty")
result = classifier(text, DEPARTMENTS)
department = result["labels"][0] # highest scoring department
return {"department": department}
# -------------------------------
# Health Check
# -------------------------------
@app.get("/health")
async def health():
return {"status": "ok"}