Update main.py

main.py

@@ -1,102 +1,57 @@
+# main.py
+
 import os
 from fastapi import FastAPI
 from pydantic import BaseModel
-from pymongo import MongoClient
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from sentence_transformers import SentenceTransformer
-from typing import List, Optional
+import pandas as pd
 
 # --- Configuration ---
-MONGO_URI = os.getenv("MONGO_URI", "mongodb+srv://saisunil22ecs:9m2ajd0GxVn43Fbu@majorproject.g0g1as0.mongodb.net/?retryWrites=true&w=majority&appName=MajorProject")
-DB_NAME = os.getenv("MONGO_DB", "legal_chatbot_db")
-COLLECTION_NAME = os.getenv("MONGO_COLLECTION", "datasets")
-MODEL_PATH = os.getenv("MODEL_PATH", "./final_bert_model_pdf")
-EMBED_MODEL = os.getenv("EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+FINAL_MODEL_PATH = './final_bert_model_pdf' 
+SOLUTIONS_DATASET_PATH = 'qa_dataset_detailed_answers.csv'
 
-# --- Resource Loading ---
+# --- Load Models and Data ---
 def load_resources():
     try:
-        print("🔄 Loading intent classification model...")
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-        intent_model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
-        print("✅ Intent model loaded.")
-
-        print("🔄 Loading embedding model for vector search...")
-        
-        # --- THIS IS THE FIX ---
-        # We specify a local cache directory where the app has write permissions.
-        cache_dir = "./model_cache"
-        embedding_model = SentenceTransformer(EMBED_MODEL, cache_folder=cache_dir)
-        # ---
-        
-        print("✅ Embedding model loaded.")
-
-        print("🔄 Connecting to MongoDB Atlas...")
-        client = MongoClient(MONGO_URI)
-        db = client[DB_NAME]
-        collection = db[COLLECTION_NAME]
-        db.command('ping') # Verify connection
-        print("✅ MongoDB connection successful!")
-        
-        return tokenizer, intent_model, embedding_model, collection
+        tokenizer = AutoTokenizer.from_pretrained(FINAL_MODEL_PATH)
+        model = AutoModelForSequenceClassification.from_pretrained(FINAL_MODEL_PATH)
+        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"❌ Critical Error during startup: {e}")
-        return None, None, None, None
+        print(f"❌ Critical Error loading resources: {e}")
+        return None, None, None
 
-tokenizer, intent_model, embedding_model, collection = load_resources()
-# --- FastAPI App ---
-app = FastAPI(title="Legal Aid Chatbot API")
+model, tokenizer, solution_database = load_resources()
+
+# --- Initialize FastAPI ---
+app = FastAPI(title="Legal Aid API")
 
 # --- API Data Models ---
-class ChatRequest(BaseModel):
-    query: str
-    top_k: Optional[int] = 3
+class QueryRequest(BaseModel):
+    question: str
 
-class ChatResponse(BaseModel):
-    intent: str
-    answers: List[dict]
+class SolutionResponse(BaseModel):
+    predicted_intent: str
+    solution: str
 
 # --- API Endpoints ---
-@app.get("/")
-def health_check():
-    """A simple endpoint to check if the API is running."""
-    return {"status": "ok", "resources_loaded": all([tokenizer, intent_model, embedding_model, collection])}
-
-@app.post("/chat", response_model=ChatResponse)
-def chat(req: ChatRequest):
-    """
-    Main chat endpoint that performs intent classification and vector search.
-    """
-    if not all([tokenizer, intent_model, embedding_model, collection]):
-        return {"intent": "Error", "answers": [{"answer": "Server is not ready. Resources could not be loaded."}]}
+@app.post("/get-solution", response_model=SolutionResponse)
+def get_legal_solution(request: QueryRequest):
+    if not model:
+        return {"predicted_intent": "Error", "solution": "Model not loaded."}
 
-    # Step 1: Intent Classification
-    inputs = tokenizer(req.query, return_tensors="pt", truncation=True)
+    inputs = tokenizer(request.question, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
-        logits = intent_model(**inputs).logits
-    pred_id = torch.argmax(logits, dim=1).item()
-    intent = intent_model.config.id2label[pred_id]
-
-    # Step 2: Vector Search in MongoDB to find the most relevant documents
-    query_embedding = embedding_model.encode(req.query, normalize_embeddings=True).tolist()
-    pipeline = [
-        {
-            '$vectorSearch': {
-                'index': 'kb_vector_index', # Ensure this index name matches your MongoDB Atlas index
-                'path': 'embedding',
-                'queryVector': query_embedding,
-                'numCandidates': 100,
-                'limit': req.top_k
-            }
-        },
-        {'$project': {'_id': 0, 'answer': '$Answer', 'question': '$Question', 'intent': '$Intent', 'score': {'$meta': 'vectorSearchScore'}}}
-    ]
-    
-    try:
-        results = list(collection.aggregate(pipeline))
-    except Exception as e:
-        print(f"Error during vector search: {e}")
-        return {"intent": intent, "answers": [{"answer": "Could not retrieve documents from the knowledge base."}]}
+        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 {"intent": intent, "answers": results}
+    return {"predicted_intent": predicted_intent, "solution": solution}
+
+@app.get("/")
+def read_root():
+    return {"status": "Legal Aid API is running."}