Update app.py

app.py

@@ -1,8 +1,7 @@
 import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
+from transformers import pipeline
 
 # -------------------------------
 # Set Hugging Face cache to writable directory
@@ -11,7 +10,6 @@ os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
 os.environ["HF_DATASETS_CACHE"] = "/tmp/huggingface"
 
-
 # -------------------------------
 # FastAPI app
 # -------------------------------
@@ -24,16 +22,16 @@ class RoutingRequest(BaseModel):
     text: str
 
 # -------------------------------
-# Load DeBERTa MNLI Model
+# Load Zero-Shot Classifier
 # -------------------------------
-MODEL_NAME = "MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-
-# Departments mapping (example, can adjust for hackathon)
 DEPARTMENTS = ['Account', 'Software', 'Network', 'Security', 'Hardware',
-       'Infrastructure', 'Licensing', 'Communication', 'RemoteWork',
-       'Training', 'Performance']
+               'Infrastructure', 'Licensing', 'Communication', 'RemoteWork',
+               'Training', 'Performance']
+
+classifier = pipeline(
+    "zero-shot-classification",
+    model="MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli"
+)
 
 # -------------------------------
 # Routing Endpoint
@@ -43,16 +41,11 @@ async def route_ticket(req: RoutingRequest):
     text = req.text
     if not text:
         raise HTTPException(status_code=400, detail="Text cannot be empty")
-    
-    inputs = tokenizer(text, return_tensors="pt", truncation=True)
-    outputs = model(**inputs)
-    logits = outputs.logits[0]
 
-    # Simple mapping: max logit → department
-    department_idx = torch.argmax(logits).item() % len(DEPARTMENTS)
-    department = DEPARTMENTS[department_idx]
+    result = classifier(text, DEPARTMENTS)
+    department = result["labels"][0]  # highest scoring department
 
-    return {"department": department}
+    return {"department": department, "scores": result["scores"]}
 
 # -------------------------------
 # Health Check