Update app.py

app.py

@@ -4,7 +4,7 @@ from sentence_transformers import SentenceTransformer
 import faiss
 import pandas as pd
 import random
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModelForCausalLM
 
 app = FastAPI()
 
@@ -14,6 +14,11 @@ embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
 summarization_model = AutoModelForSeq2SeqLM.from_pretrained("google/long-t5-tglobal-base")
 summarization_tokenizer = AutoTokenizer.from_pretrained("google/long-t5-tglobal-base")
 
+# New: Load Local LLM for Dynamic Emotional Responses (Mistral/Llama)
+response_model_name = "mistralai/Mistral-7B-Instruct"
+response_tokenizer = AutoTokenizer.from_pretrained(response_model_name)
+response_model = AutoModelForCausalLM.from_pretrained(response_model_name)
+
 # Load datasets
 recommendations_df = pd.read_csv("treatment_recommendations.csv")
 questions_df = pd.read_csv("symptom_questions.csv")
@@ -36,6 +41,36 @@ class SummaryRequest(BaseModel):
     chat_history: list  # List of messages
 
 
+@app.post("/get_questions")
+def get_recommended_questions(request: ChatRequest):
+    """Retrieve the most relevant diagnostic questions with a dynamically generated conversational response."""
+    
+    # Step 1: Encode the input message for FAISS search
+    input_embedding = embedding_model.encode([request.message], convert_to_numpy=True)
+    distances, indices = question_index.search(input_embedding, 3)
+
+    # Step 2: Retrieve the top 3 relevant questions
+    retrieved_questions = [questions_df["Questions"].iloc[i] for i in indices[0]]
+
+    # Step 3: Use a local LLM to generate context-aware empathetic responses
+    prompt = f"""
+    User: {request.message}
+    
+    You are a compassionate psychiatric assistant. Before asking a diagnostic question, respond empathetically.
+    
+    Questions:
+    1. {retrieved_questions[0]}
+    2. {retrieved_questions[1]}
+    3. {retrieved_questions[2]}
+    
+    Generate a conversational response that introduces each question naturally.
+    """
+
+    inputs = response_tokenizer(prompt, return_tensors="pt")
+    output = response_model.generate(**inputs, max_length=300)
+    enhanced_responses = response_tokenizer.decode(output[0], skip_special_tokens=True).split("\n")
+
+    return {"questions": enhanced_responses}
 
 
 @app.post("/summarize_chat")
@@ -47,6 +82,7 @@ def summarize_chat(request: SummaryRequest):
     summary = summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     return {"summary": summary}
 
+
 @app.post("/detect_disorders")
 def detect_disorders(request: SummaryRequest):
     """Detect psychiatric disorders from full chat history at the end."""
@@ -56,6 +92,7 @@ def detect_disorders(request: SummaryRequest):
     disorders = [recommendations_df["Disorder"].iloc[i] for i in indices[0]]
     return {"disorders": disorders}
 
+
 @app.post("/get_treatment")
 def get_treatment(request: SummaryRequest):
     """Retrieve treatment recommendations based on detected disorders."""