Update app.py

app.py

@@ -7,98 +7,83 @@ import google.generativeai as genai
 import re
 import os
 
-# Load data and FAISS index
-def load_data_and_index():
-    docs_df = pd.read_pickle("data.pkl")  # Adjust path for HF Spaces
-    embeddings = np.array(docs_df['embeddings'].tolist(), dtype=np.float32)
-    dimension = embeddings.shape[1]
-    index = faiss.IndexFlatL2(dimension)
-    index.add(embeddings)
-    return docs_df, index
+# Load documents and FAISS index
+def load_index_and_data():
+    df = pd.read_pickle("data.pkl")
+    vecs = np.array(df['embeddings'].tolist(), dtype=np.float32)
+    idx = faiss.IndexFlatL2(vecs.shape[1])
+    idx.add(vecs)
+    return df, idx
 
-docs_df, index = load_data_and_index()
+docs_df, index = load_index_and_data()
 
-# Load SentenceTransformer
-minilm = SentenceTransformer('all-MiniLM-L6-v2')
+# Embedding model and Gemini setup
+encoder = SentenceTransformer("all-MiniLM-L6-v2")
 
-# Configure Gemini API using Hugging Face Secrets
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-if not GEMINI_API_KEY:
-    raise ValueError("Gemini API key not found. Please set it in Hugging Face Spaces secrets.")
-genai.configure(api_key=GEMINI_API_KEY)
-model = genai.GenerativeModel('gemini-2.0-flash')
+API_KEY = os.getenv("GEMINI_API_KEY")
+if not API_KEY:
+    raise EnvironmentError("Missing Gemini API key.")
+genai.configure(api_key=API_KEY)
+llm = genai.GenerativeModel("gemini-2.0-flash")
 
-# Preprocess text function
-def preprocess_text(text):
+# Clean text input
+def clean_text(text):
     text = text.lower()
-    text = text.replace('\n', ' ').replace('\t', ' ')
-    text = re.sub(r'[^\w\s.,;:>-]', ' ', text)
-    text = ' '.join(text.split()).strip()
-    return text
+    text = re.sub(r"[^\w\s.,]", " ", text)
+    return " ".join(text.split())
 
-# Retrieve documents
-def retrieve_docs(query, k=5):
-    query_embedding = minilm.encode([query], show_progress_bar=False)[0].astype(np.float32)
-    distances, indices = index.search(np.array([query_embedding]), k)
-    retrieved_docs = docs_df.iloc[indices[0]][['label', 'text', 'source']]
-    retrieved_docs['distance'] = distances[0]
-    return retrieved_docs
+# Retrieve relevant document context
+def get_context(query, k=5):
+    q_vec = encoder.encode([query])[0].astype(np.float32)
+    _, indices = index.search(np.array([q_vec]), k)
+    return "\n".join(docs_df.iloc[indices[0]]["text"].tolist())
 
-# RAG pipeline integrated into respond function
-def respond(message, system_message, max_tokens, temperature):
-    # Preprocess the user message
-    preprocessed_query = preprocess_text(message)
+# RAG-based Gemini response generation
+def generate_answer(user_input, system_note, max_tokens, temp):
+    query = clean_text(user_input)
+    context = get_context(query)
     
-    # Retrieve relevant documents
-    retrieved_docs = retrieve_docs(preprocessed_query, k=5)
-    context = "\n".join(retrieved_docs['text'].tolist())
-    
-    # Construct the prompt with system message and RAG context, asking for structured response
-    prompt = f"{system_message}\n\n"
-    prompt += (
-        f"Query: {message}\n"
-        f"Relevant Context: {context}\n"
-        f"Generate a short, concise response to the query based only on the provided context. "
-        f"Format the response as a structured list (e.g., bullet points or numbered items) instead of a paragraph."
+    prompt = (
+        f"Role Description:\n{system_note}\n\n"
+        f"User Question:\n{user_input}\n\n"
+        f"Knowledge Extracted From Records:\n{context}\n\n"
+        f"Instructions:\n"
+        f"- Analyze the user's query using ONLY the above context.\n"
+        f"- Do NOT add external or made-up information.\n"
+        f"- Begin with a brief summary of the identified condition or concern.\n"
+        f"- Provide detailed reasoning and explanation in bullet points:\n"
+        f"   • Include possible causes, symptoms, and diagnostic considerations.\n"
+        f"   • Mention relevant terms or observations from context.\n"
+        f"   • Explain how the context supports the conclusions.\n"
+        f"- End with a short, clear recommendation (if context permits).\n"
+        f"- Avoid medical advice unless the context contains it."
     )
-    
-    # Generate response with Gemini
-    response = model.generate_content(
+
+    result = llm.generate_content(
         prompt,
         generation_config=genai.types.GenerationConfig(
             max_output_tokens=max_tokens,
-            temperature=temperature
+            temperature=temp
         )
     )
-    answer = response.text.strip()
-    if not answer.endswith('.'):
-        last_period = answer.rfind('.')
-        if last_period != -1:
-            answer = answer[:last_period + 1]
-        else:
-            answer += "."
-    
-    return answer
-
-# Simple Gradio Interface
-def chatbot_interface(message, system_message, max_tokens, temperature):
-    return respond(message, system_message, max_tokens, temperature)
+    return result.text.strip()
 
+# Gradio interface
 demo = gr.Interface(
-    fn=chatbot_interface,
+    fn=generate_answer,
     inputs=[
-        gr.Textbox(label="Your Query", placeholder="Enter your medical question here..."),
+        gr.Textbox(label="Ask Something", placeholder="Describe your symptom or condition..."),
         gr.Textbox(
-            value="You are a medical AI assistant diagnosing patients based on their query, using relevant context from past records of other patients.",
-            label="System Message"
+            value="You are a virtual medical assistant using past medical records to respond intelligently.",
+            label="System Role"
         ),
-        gr.Slider(minimum=1, maximum=2048, value=150, step=1, label="Max Tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.75, step=0.1, label="Temperature"),
+        gr.Slider(50, 500, value=300, step=10, label="Max Tokens"),
+        gr.Slider(0.0, 1.0, value=0.4, step=0.1, label="Creativity (Temperature)")
     ],
-    outputs=gr.Textbox(label="Response"),
-    title="🏥 Medical Chat Assistant",
-    description="A simple medical assistant that diagnoses patient queries using AI and past records, providing structured responses."
+    outputs=gr.Textbox(label="AI Diagnosis"),
+    title="🩺 Smart Medical Query Assistant",
+    description="Submit a health-related question. The assistant analyzes similar past records to respond accurately and clearly."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()