Update main.py

main.py

@@ -8,8 +8,8 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from typing import Dict, List, Optional
 
 # --- Configuration ---
-# ** FIX: Hardcoding the correct model name from the URL to bypass any issues with secrets. **
-MODEL_NAME = "Sp2503/Bertmodel" 
+# ** FIX: Loading the model from a local directory within the Space **
+LOCAL_MODEL_PATH = "./" 
 
 # The MongoDB URI is loaded from Hugging Face Space secrets for security
 MONGO_URI = os.getenv("MONGO_URI") 
@@ -54,17 +54,12 @@ def startup_event():
         return
 
     try:
-        cache_dir = "/tmp"
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=cache_dir)
-        num_labels = len(intent_map)
-        model = AutoModelForSequenceClassification.from_pretrained(
-            MODEL_NAME, 
-            num_labels=num_labels, 
-            cache_dir=cache_dir
-        )
-        print(f"✅ Model '{MODEL_NAME}' and tokenizer loaded successfully.")
+        # Load the model and tokenizer from the LOCAL directory
+        tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_PATH)
+        model = AutoModelForSequenceClassification.from_pretrained(LOCAL_MODEL_PATH)
+        print(f"✅ Model and tokenizer loaded successfully from '{LOCAL_MODEL_PATH}'.")
     except Exception as e:
-        print(f"❌ Critical Error loading Hugging Face model: {e}")
+        print(f"❌ Critical Error loading model from local directory: {e}")
         model = None 
 
     try:
@@ -93,7 +88,7 @@ def get_legal_solution(request: QueryRequest):
         logits = model(**inputs).logits
     
     prediction_id = torch.argmax(logits, dim=1).item()
-    predicted_intent = intent_map.get(prediction_id, "Unknown Intent")
+    predicted_intent = model.config.id2label.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."