Update main.py

main.py

@@ -1,61 +1,49 @@
+# main.py
+
 import os
 import torch
 import pandas as pd
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 from langdetect import detect
 
-# ============================================================
-# ✅ Environment setup — avoids permission errors on Hugging Face
-# ============================================================
+# ========== CONFIG ==========
+FINAL_MODEL_PATH = './final_bert_model_pdf'  # Local fine-tuned English model
+MURIL_MODEL_ID = 'Sp2503/Muril-Model'        # Hugging Face multilingual model
+SOLUTIONS_DATASET_PATH = 'qa_dataset_detailed_answers.csv'
+
+# Fix cache permissions for Spaces
 os.environ["HF_HOME"] = "/app/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
-os.environ["HF_DATASETS_CACHE"] = "/app/hf_cache"
-os.environ["TORCH_HOME"] = "/app/hf_cache"
-
-# ============================================================
-# ✅ Configuration
-# ============================================================
-FINAL_MODEL_PATH = './final_bert_model_pdf'
-SOLUTIONS_DATASET_PATH = 'qa_dataset_detailed_answers.csv'
-MURIL_MODEL_NAME = 'Sp2503/Muril-Model'   # Your public HF model
+os.makedirs("/app/hf_cache", exist_ok=True)
 
-# ============================================================
-# ✅ Load models and data
-# ============================================================
+# ========== LOAD MODELS ==========
 def load_resources():
     try:
         # Load English model
-        tokenizer_en = AutoTokenizer.from_pretrained(FINAL_MODEL_PATH)
-        model_en = AutoModelForSequenceClassification.from_pretrained(FINAL_MODEL_PATH)
-
-        # Load MuRIL multilingual model (for non-English)
-        tokenizer_muril = AutoTokenizer.from_pretrained(MURIL_MODEL_NAME)
-        model_muril = AutoModelForSequenceClassification.from_pretrained(MURIL_MODEL_NAME)
-
-        # Load Q&A dataset
-        solutions_df = pd.read_csv(SOLUTIONS_DATASET_PATH)
-        solution_database = solutions_df.set_index('Intent')['Answer'].to_dict()
-
-        print("✅ All models and data loaded successfully!")
-        return model_en, tokenizer_en, model_muril, tokenizer_muril, solution_database
-
+        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
     except Exception as e:
         print(f"❌ Error loading models or data: {e}")
-        return None, None, None, None, None
-
+        return None, None, None, None
 
-model_en, tokenizer_en, model_muril, tokenizer_muril, solution_database = load_resources()
+model, tokenizer, muril_pipeline, solution_db = load_resources()
 
-# ============================================================
-# ✅ FastAPI app setup
-# ============================================================
-app = FastAPI(title="Legal Aid API")
+# ========== FASTAPI APP ==========
+app = FastAPI(title="AI LegalAid Chatbot Server")
 
-# ============================================================
-# ✅ Request and Response Models
-# ============================================================
+# Request / Response Schemas
 class QueryRequest(BaseModel):
     question: str
 
@@ -64,55 +52,63 @@ class SolutionResponse(BaseModel):
     solution: str
     model_used: str
 
-    model_config = {
-        "protected_namespaces": ()  # suppress Pydantic warning
-    }
-
-# ============================================================
-# ✅ Helper: Detect if question is English
-# ============================================================
-def is_english(text: str) -> bool:
-    try:
-        lang = detect(text)
-        return lang == "en"
-    except:
-        return True  # default fallback to English
-
-# ============================================================
-# ✅ Main API Endpoint
-# ============================================================
+# ========== LOGIC ==========
 @app.post("/get-solution", response_model=SolutionResponse)
 def get_legal_solution(request: QueryRequest):
-    if not model_en or not model_muril:
-        return {"predicted_intent": "Error", "solution": "Models not loaded.", "model_used": "None"}
+    if not model or not tokenizer:
+        return {
+            "predicted_intent": "Error",
+            "solution": "Model not loaded properly.",
+            "model_used": "None"
+        }
 
     question = request.question.strip()
-    use_english = is_english(question)
-
-    # Select model based on language
-    model = model_en if use_english else model_muril
-    tokenizer = tokenizer_en if use_english else tokenizer_muril
-    model_name = "BERT-English" if use_english else "MuRIL-Multilingual"
-
-    # Tokenize and predict
-    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.get(prediction_id, "Unknown")
-
-    # Fetch solution
-    solution = solution_database.get(predicted_intent, "No solution found in database.")
-
-    return {
-        "predicted_intent": predicted_intent,
-        "solution": solution,
-        "model_used": model_name
-    }
-
-# ============================================================
-# ✅ Root Endpoint
-# ============================================================
+
+    # 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"
+        }
+
 @app.get("/")
-def read_root():
-    return {"status": "✅ Legal Aid API is running with English + MuRIL multilingual support."}
+def root():
+    ready = all([model, tokenizer, muril_pipeline])
+    return {"status": "✅ AI LegalAid Chatbot Running", "models_ready": ready}