Update main.py

main.py

@@ -76,7 +76,6 @@ class SolutionResponse(BaseModel):
 
 @app.post("/get-solution", response_model=SolutionResponse)
 def get_legal_solution(request: QueryRequest):
-    # ** FIX: Check for None explicitly as required by pymongo **
     if model is None or tokenizer is None or collection is None:
         raise HTTPException(status_code=503, detail="Server resources are not ready. Check startup logs for errors.")
 
@@ -85,7 +84,9 @@ def get_legal_solution(request: QueryRequest):
         logits = model(**inputs).logits
     
     prediction_id = torch.argmax(logits, dim=1).item()
-    predicted_intent = model.config.id2label.get(prediction_id, "Unknown Intent")
+    
+    # ** FIX: Use the intent_map created from CSVs to get the human-readable label **
+    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."