Update main.py

main.py

@@ -1,116 +1,56 @@
-# main.py
 
 import os
-import torch
-import pandas as pd
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
-from langdetect import detect
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import pandas as pd
 
-# ========== CONFIG ==========
-FINAL_MODEL_PATH = './final_bert_model_pdf'  # Local fine-tuned English model
-MURIL_MODEL_ID = 'Sp2503/Muril-Model'        # Hugging Face multilingual model
+# --- Configuration ---
+FINAL_MODEL_PATH = './final_bert_model_pdf' 
 SOLUTIONS_DATASET_PATH = 'qa_dataset_detailed_answers.csv'
-# ✅ Fix cache permissions for Hugging Face Spaces
-HF_CACHE_DIR = "/tmp/hf_cache"
-os.environ["HF_HOME"] = HF_CACHE_DIR
-os.environ["TRANSFORMERS_CACHE"] = HF_CACHE_DIR
-os.makedirs(HF_CACHE_DIR, exist_ok=True)
 
-
-
-# ========== LOAD MODELS ==========
+# --- Load Models and Data ---
 def load_resources():
     try:
-        # Load English model
         tokenizer = AutoTokenizer.from_pretrained(FINAL_MODEL_PATH)
         model = AutoModelForSequenceClassification.from_pretrained(FINAL_MODEL_PATH)
-        
-        # Load multilingual MuRIL model for non-English
-        muril_pipeline = pipeline("text-classification", model=MURIL_MODEL_ID)
-
-        # Load answers dataset
-        df = pd.read_csv(SOLUTIONS_DATASET_PATH)
-        solution_db = df.set_index('Intent')['Answer'].to_dict()
-
-        print("✅ All models & data loaded successfully!")
-        return model, tokenizer, muril_pipeline, solution_db
+        solutions_df = pd.read_csv(SOLUTIONS_DATASET_PATH)
+        solution_database = solutions_df.set_index('Intent')['Answer'].to_dict()
+        print("✅ Resources loaded successfully!")
+        return model, tokenizer, solution_database
     except Exception as e:
-        print(f"❌ Error loading models or data: {e}")
-        return None, None, None, None
+        print(f"❌ Critical Error loading resources: {e}")
+        return None, None, None
 
-model, tokenizer, muril_pipeline, solution_db = load_resources()
+model, tokenizer, solution_database = load_resources()
 
-# ========== FASTAPI APP ==========
-app = FastAPI(title="AI LegalAid Chatbot Server")
+# --- Initialize FastAPI ---
+app = FastAPI(title="Legal Aid API")
 
-# Request / Response Schemas
+# --- API Data Models ---
 class QueryRequest(BaseModel):
     question: str
 
 class SolutionResponse(BaseModel):
     predicted_intent: str
     solution: str
-    model_used: str
 
-# ========== LOGIC ==========
+# --- API Endpoints ---
 @app.post("/get-solution", response_model=SolutionResponse)
 def get_legal_solution(request: QueryRequest):
-    if not model or not tokenizer:
-        return {
-            "predicted_intent": "Error",
-            "solution": "Model not loaded properly.",
-            "model_used": "None"
-        }
-
-    question = request.question.strip()
-
-    # Detect language
-    try:
-        lang = detect(question)
-    except:
-        lang = "en"
-
-    # If not English, use MuRIL model
-    if lang != "en":
-        try:
-            muril_result = muril_pipeline(question)
-            predicted_intent = muril_result[0]['label']
-            solution = solution_db.get(predicted_intent, "No solution found for this intent.")
-            return {
-                "predicted_intent": predicted_intent,
-                "solution": solution,
-                "model_used": "MuRIL"
-            }
-        except Exception as e:
-            return {
-                "predicted_intent": "Error",
-                "solution": f"MuRIL model failed: {e}",
-                "model_used": "MuRIL"
-            }
-
-    # For English questions → use fine-tuned BERT model
-    try:
-        inputs = tokenizer(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 = model.config.id2label[prediction_id]
-        solution = solution_db.get(predicted_intent, "No solution found for this intent.")
-        return {
-            "predicted_intent": predicted_intent,
-            "solution": solution,
-            "model_used": "English BERT"
-        }
-    except Exception as e:
-        return {
-            "predicted_intent": "Error",
-            "solution": f"English model failed: {e}",
-            "model_used": "English BERT"
-        }
+    if not model:
+        return {"predicted_intent": "Error", "solution": "Model not loaded."}
+
+    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 = model.config.id2label[prediction_id]
+    solution = solution_database.get(predicted_intent, "No solution found.")
+    
+    return {"predicted_intent": predicted_intent, "solution": solution}
 
 @app.get("/")
-def root():
-    ready = all([model, tokenizer, muril_pipeline])
-    return {"status": "✅ AI LegalAid Chatbot Running", "models_ready": ready}
+def read_root():
+    return {"status": "Legal Aid API is running."}