fixed the urgencey related error

app.py

@@ -1,6 +1,8 @@
 import torch
+import os
 from fastapi import FastAPI
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig
+from transformers import AutoTokenizer
+from huggingface_hub import hf_hub_download
 
 from models import (
     QueryRequest,
@@ -8,49 +10,100 @@ from models import (
     CategoryPrediction,
     UrgencyPrediction
 )
+from multi_task_model_class import MultiTaskModel
 from rag import generate_answer
 from memory import get_conversation, add_message
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 CLASSIFIER_MODEL_ID = "Sandei/tech-support-classifier"
+ENCODER_NAME = "distilbert-base-uncased"
+
+tag_classes = ['Email & Communication', 'Classroom/Lab Support', 'Software & Applications', 'Classroom/Lab Support', 'Classroom/Lab Support', 'Network & Connectivity', 'General IT Support', 'Data Management', 'Classroom/Lab Support', 'Security & Compliance']
 
-tag_classes = [
-    "Billing",
-    "Network & Connectivity",
-    "Account Access",
-    "Hardware",
-    "Other"
-]
 
 urgency_encoder = {
     0: "low",
     1: "medium",
-    2: "high"
+    2: "high",
+    3: "critical"  # Added 4th level
 }
 
-tokenizer = AutoTokenizer.from_pretrained(CLASSIFIER_MODEL_ID)
-config = AutoConfig.from_pretrained(CLASSIFIER_MODEL_ID)
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(CLASSIFIER_MODEL_ID, trust_remote_code=True)
 
-model = AutoModelForSequenceClassification.from_pretrained(
-    CLASSIFIER_MODEL_ID,
-    config=config,
-    trust_remote_code=True
-).to(DEVICE)
+print("Initializing model structure...")
+model = MultiTaskModel(
+    encoder_name=ENCODER_NAME,
+    num_category_labels=len(tag_classes),
+    num_urgency_labels=4
+)
 
+# Load model weights
+print("Downloading model weights...")
+try:
+    model_path = hf_hub_download(
+        repo_id=CLASSIFIER_MODEL_ID,
+        filename="pytorch_model.bin",
+        token=None,  # Set to your HF token if repo is private
+    )
+    print(f"✓ Model downloaded to: {model_path}")
+    
+    print("Loading model weights...")
+    state_dict = torch.load(model_path, map_location=DEVICE, weights_only=False)
+    model.load_state_dict(state_dict)
+    print("✓ Model weights loaded successfully")
+    
+except Exception as e:
+    print(f"✗ Error downloading from Hugging Face: {e}")
+    print("\nTrying alternative methods...")
+    
+    # Method 2: Try loading from cache
+    from huggingface_hub import try_to_load_from_cache
+    cache_path = try_to_load_from_cache(
+        repo_id=CLASSIFIER_MODEL_ID,
+        filename="pytorch_model.bin"
+    )
+    
+    if cache_path and os.path.exists(cache_path):
+        print(f"✓ Found in cache: {cache_path}")
+        state_dict = torch.load(cache_path, map_location=DEVICE, weights_only=False)
+        model.load_state_dict(state_dict)
+        print("✓ Model loaded from cache")
+    else:
+        print("\n" + "="*60)
+        print("ERROR: Could not load model weights")
+        print("="*60)
+        print("\nPossible solutions:")
+        print("1. Login to Hugging Face:")
+        print("   huggingface-cli login")
+        print("\n2. Or download manually:")
+        print(f"   Visit: https://huggingface.co/{CLASSIFIER_MODEL_ID}/tree/main")
+        print(f"   Download 'pytorch_model.bin' to: ./Sandei/tech-support-classifier/")
+        print("\n3. Check your internet connection")
+        print("="*60)
+        raise
+
+model.to(DEVICE)
 model.eval()
 
+print(f"\n✓ Model ready on {DEVICE}\n")
+
 app = FastAPI(title="RAG + Conversation Memory API")
 
 # ---------------------
 # CLASSIFIER
 # ---------------------
 def classify_text(text: str, threshold: float = 0.5):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True).to(DEVICE)
+    """
+    Classify input text into categories and urgency level.
+    """
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(DEVICE)
 
     with torch.no_grad():
         outputs = model(**inputs)
 
+    # Category predictions (multi-label)
     category_probs = torch.sigmoid(outputs.category_logits)[0].cpu().numpy()
 
     categories = [
@@ -62,6 +115,7 @@ def classify_text(text: str, threshold: float = 0.5):
         if category_probs[i] >= threshold
     ]
 
+    # Urgency prediction (multi-class)
     urgency_probs = torch.softmax(outputs.urgency_logits, dim=-1)[0].cpu().numpy()
     urgency_idx = int(torch.argmax(outputs.urgency_logits, dim=-1)[0])
 
@@ -74,6 +128,9 @@ def classify_text(text: str, threshold: float = 0.5):
 
 
 def retrieve_documents(query: str):
+    """
+    Retrieve relevant documents for RAG.
+    """
     return [
         "Restarting the router fixes most connectivity issues.",
         "Check for planned ISP maintenance.",
@@ -81,19 +138,32 @@ def retrieve_documents(query: str):
     ]
 
 
+@app.get("/")
+def root():
+    """Health check endpoint"""
+    return {
+        "status": "running",
+        "device": DEVICE,
+        "model": CLASSIFIER_MODEL_ID
+    }
+
+
 @app.post("/query", response_model=QueryResponse)
 def query_endpoint(req: QueryRequest):
-    # ---- Load conversation
+    """
+    Main query endpoint.
+    """
+    # Load conversation history
     history = get_conversation(req.user_id)
 
-    # ---- Classification
+    # Classification
     categories, urgency = classify_text(req.query)
 
-    # ---- RAG
+    # RAG
     docs = retrieve_documents(req.query)
     answer = generate_answer(req.query, docs, history)
 
-    # ---- Update memory
+    # Update conversation memory
     add_message(req.user_id, "user", req.query)
     add_message(req.user_id, "assistant", answer)
 
@@ -105,3 +175,22 @@ def query_endpoint(req: QueryRequest):
         urgency=urgency,
         conversation=get_conversation(req.user_id)
     )
+
+
+@app.post("/classify")
+def classify_endpoint(req: QueryRequest):
+    """
+    Standalone classification endpoint.
+    """
+    categories, urgency = classify_text(req.query)
+    
+    return {
+        "query": req.query,
+        "categories": categories,
+        "urgency": urgency
+    }
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

multi_task_model_class.py

@@ -0,0 +1,32 @@
+import torch.nn as nn
+from transformers import AutoModel
+
+
+class MultiTaskModel(nn.Module):
+    def __init__(self, encoder_name, num_category_labels, num_urgency_labels):
+        super().__init__()
+
+        self.encoder = AutoModel.from_pretrained(encoder_name)
+        hidden_size = self.encoder.config.hidden_size
+
+        # Changed from category_head to category_classifier
+        self.category_classifier = nn.Linear(hidden_size, num_category_labels)
+        # Changed from urgency_head to urgency_classifier
+        self.urgency_classifier = nn.Linear(hidden_size, num_urgency_labels)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask
+        )
+
+        pooled = outputs.last_hidden_state[:, 0]
+
+        return type(
+            "Output",
+            (),
+            {
+                "category_logits": self.category_classifier(pooled),
+                "urgency_logits": self.urgency_classifier(pooled),
+            }
+        )()