Update app.py

app.py

@@ -7,76 +7,92 @@ import google.generativeai as genai
 import re
 import os
 
-# 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
+# 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
 
-docs_df, index = load_index_and_data()
+docs_df, index = load_data_and_index()
 
-# Embedding model and Gemini setup
-encoder = SentenceTransformer("all-MiniLM-L6-v2")
+# Load SentenceTransformer
+minilm = SentenceTransformer('all-MiniLM-L6-v2')
 
-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")
+# 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')
 
-# Clean text input
-def clean_text(text):
+# Preprocess text function
+def preprocess_text(text):
     text = text.lower()
-    text = re.sub(r"[^\w\s.,]", " ", text)
-    return " ".join(text.split())
+    text = text.replace('\n', ' ').replace('\t', ' ')
+    text = re.sub(r'[^\w\s.,;:>-]', ' ', text)
+    text = ' '.join(text.split()).strip()
+    return text
 
-# 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())
+# 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
 
-# RAG-based Gemini response generation
-def generate_answer(user_input, system_note, max_tokens, temp):
-    query = clean_text(user_input)
-    context = get_context(query)
+# RAG pipeline integrated into respond function
+def respond(message, system_message, max_tokens, temperature):
+    # Preprocess the user message
+    preprocessed_query = preprocess_text(message)
     
-    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"   • 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."
+    # 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 with headings and information write in the form of points not paragraph"
     )
-
-    result = llm.generate_content(
+    
+    # Generate response with Gemini
+    response = model.generate_content(
         prompt,
         generation_config=genai.types.GenerationConfig(
             max_output_tokens=max_tokens,
-            temperature=temp
+            temperature=temperature
         )
     )
-    return result.text.strip()
+    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)
 
-# Gradio interface
 demo = gr.Interface(
-    fn=generate_answer,
+    fn=chatbot_interface,
     inputs=[
-        gr.Textbox(label="Ask Something", placeholder="Describe your symptom or condition..."),
+        gr.Textbox(label="Your Query", placeholder="Enter your medical question here..."),
     ],
-    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."
+    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."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()