Update main.py

main.py

@@ -8,8 +8,8 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from typing import Dict, List, Optional
 
 # --- Configuration using Environment Variables from Hugging Face Secrets ---
-MODEL_NAME = os.getenv("MODEL_NAME") 
-MONGO_URI = os.getenv("MONGO_URI") 
+MODEL_NAME = os.getenv("MODEL_NAME")
+MONGO_URI = os.getenv("MONGO_URI")
 DB_NAME = "legal_aid-chatbot"
 COLLECTION_NAME = "categories"
 
@@ -41,7 +41,7 @@ app = FastAPI(title="Legal Aid Chatbot API", version="1.0.0")
 def startup_event():
     """Loads all necessary resources when the FastAPI application starts."""
     global model, tokenizer, intent_map, collection
-    
+
     print("--- Loading resources on application startup ---")
 
     if not MONGO_URI:
@@ -50,7 +50,7 @@ def startup_event():
     if not MODEL_NAME:
         print("❌ Critical Error: MODEL_NAME secret is not set in Hugging Face Space settings.")
         return
-    
+
     intent_map = create_intent_map(['womens_legal_questions_20k.csv', 'legal_aid_chatbot_dataset_20k.csv'])
     if not intent_map:
         print("❌ Could not create intent map. API will not function correctly.")
@@ -60,18 +60,18 @@ def startup_event():
         # Use a writable cache directory for Hugging Face Spaces
         cache_dir = "/tmp"
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=cache_dir)
-        
+
         # Explicitly provide the number of labels for the classifier
         num_labels = len(intent_map)
         model = AutoModelForSequenceClassification.from_pretrained(
-            MODEL_NAME, 
-            num_labels=num_labels, 
+            MODEL_NAME,
+            num_labels=num_labels,
             cache_dir=cache_dir
         )
         print(f"✅ Model '{MODEL_NAME}' and tokenizer loaded successfully.")
     except Exception as e:
         print(f"❌ Critical Error loading Hugging Face model: {e}")
-        model = None 
+        model = None
 
     try:
         client = MongoClient(MONGO_URI)
@@ -100,12 +100,12 @@ def get_legal_solution(request: QueryRequest):
     inputs = tokenizer(request.question, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         logits = model(**inputs).logits
-    
+
     prediction_id = torch.argmax(logits, dim=1).item()
     predicted_intent = intent_map.get(prediction_id, "Unknown Intent")
 
     document = collection.find_one({"intent": predicted_intent})
-    
+
     solution = document["answer"] if document and "answer" in document else "No specific solution was found for this topic."
 
     return SolutionResponse(predicted_intent=predicted_intent, solution=solution)