Update model.py

model.py

@@ -1,6 +1,5 @@
 import chromadb
 import traceback
-import os
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from retriever import retrieve
 from utils import build_prompt, refine_response
@@ -12,24 +11,18 @@ from utils import build_prompt, refine_response
 
 _vector_store = None
 _finetuned_llm = None
+_base_model = None
 
 def get_vector_store():
     """Load vector store (lazy-loaded on first use)"""
     global _vector_store
-    
     if _vector_store is None:
-        # Create directory if it doesn't exist
-        db_path = "./MedQuAD_db"
-        os.makedirs(db_path, exist_ok=True)
-        
-        db_client = chromadb.PersistentClient(path=db_path)
-        
+        db_client = chromadb.PersistentClient(path="./MedQuAD_db")
         try:
             _vector_store = db_client.get_collection("medical_rag")
         except:
-            # Collection doesn't exist yet - create it
+            # If collection doesn't exist, create it
             _vector_store = db_client.create_collection(name="medical_rag")
-    
     return _vector_store
 
 def get_finetuned_llm():
@@ -49,8 +42,13 @@ def get_finetuned_llm():
     return _finetuned_llm
 
 
+# ============================================================================
+# MAIN RAG FUNCTION
+# ============================================================================
+
 def rag(user_query):
     """Main RAG function: retrieve context and generate answer.
+    Takes a question string and returns an answer string with confidence.
     Returns: str(generated_answer)"""
     
     try:
@@ -58,22 +56,58 @@ def rag(user_query):
         vector_store = get_vector_store()
         finetuned_llm = get_finetuned_llm()
         
-        # Check for emergency keywords
-        emergency_keywords = ["emergency", "severe pain", "bleeding", "blind", "lose consciousness", "pass out"]
+        # 1. Check for emergency keywords
+        emergency_keywords = ["emergency", "severe pain", "bleeding",
+                            "blind", "lose consciousness", "pass out"]
         
         if any(keyword in user_query.lower() for keyword in emergency_keywords):
             emergency_msg = """I am an AI and cannot provide medical advice for emergencies.
 PLEASE contact emergency services or a medical professional immediately."""
-            return emergency_msg
+            
+            try:
+                # Still generate answer for context
+                contexts = retrieve(vector_store, user_query, top_k=3, use_reranking=True)
+                
+                if not contexts:
+                    return f"{emergency_msg}\n\nNo relevant information found for your query."
+                
+                prompt = build_prompt(user_query, contexts)
+                result = finetuned_llm(
+                    prompt,
+                    max_new_tokens=70,
+                    num_beams=3,
+                    early_stopping=True,
+                    do_sample=False,
+                    repetition_penalty=1.4,
+                    eos_token_id=finetuned_llm.tokenizer.eos_token_id
+                )
+                
+                answer = result[0]['generated_text'].strip()
+                answer = refine_response(answer)
+                
+                # Calculate confidence
+                if contexts:
+                    avg_distance = sum(c.get('chroma_distance', 1.0) for c in contexts) / len(contexts)
+                    confidence_score = (1 - avg_distance) * 100
+                    confidence_score = max(0, min(100, confidence_score))
+                else:
+                    confidence_score = 0
+                
+                return f"{emergency_msg}\n\n[Confidence: {confidence_score:.1f}%]\n\n{answer}"
+            
+            except Exception as e:
+                return f"{emergency_msg}\n\nError generating answer: {str(e)}"
         
-        # Retrieve relevant contexts
+        # 2. Retrieve relevant contexts
         contexts = retrieve(vector_store, user_query, top_k=3, use_reranking=True)
         
         if not contexts:
             return "I'm not confident about my answer (0%).\n\nCouldn't find relevant information to answer your question."
         
-        # Build prompt and generate answer
+        # 3. Build prompt with context
         prompt = build_prompt(user_query, contexts)
+        
+        # 4. Generate answer
         result = finetuned_llm(
             prompt,
             max_new_tokens=70,
@@ -87,18 +121,23 @@ PLEASE contact emergency services or a medical professional immediately."""
         answer = result[0]['generated_text'].strip()
         answer = refine_response(answer)
         
-        # Calculate confidence score
+        # 5. Calculate confidence score based on retrieval quality
         if contexts and len(contexts) > 0:
             avg_distance = sum(c.get('chroma_distance', 1.0) for c in contexts) / len(contexts)
             confidence_score = (1 - avg_distance) * 100
             confidence_score = max(0, min(100, confidence_score))
             
+            # Build final response with confidence
             if confidence_score < 40:
-                return f"I'm not confident ({confidence_score:.1f}%).\n\n{answer}"
+                final_response = f"I'm not confident about my answer ({confidence_score:.1f}%).\n\n{answer}"
             else:
-                return f"{answer}\n\n[Confidence: {confidence_score:.1f}%]"
+                final_response = f"{answer}\n\n[Confidence: {confidence_score:.1f}%]"
+        else:
+            final_response = "I'm not confident about my answer (0%).\n\n" + answer
         
-        return answer
+        return final_response
     
     except Exception as e:
-        return f"ERROR: {str(e)}"
+        error_msg = f"ERROR in RAG pipeline: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+        print(error_msg)
+        return error_msg