Update app.py

app.py

@@ -2,39 +2,52 @@ import gradio as gr
 import fitz  # PyMuPDF
 import re
 import numpy as np
-from chromadb import Client
-from sklearn.metrics.pairwise import cosine_similarity
+import faiss
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.preprocessing import normalize
 import google.generativeai as gemini
 
 # Step 1: Configure Gemini API
 gemini.configure(api_key="AIzaSyCOxpeeq4qUMjZje8sNtwnQZiQ9xVShLd0")
 
-# Step 2: Set up ChromaDB for Knowledge Retrieval
-client = Client()
-collection = client.get_or_create_collection("medical_knowledge")
+# Step 2: Define Medical Knowledge Base
+medical_knowledge = [
+    "Normal hemoglobin levels are 13.8 to 17.2 g/dL for men and 12.1 to 15.1 g/dL for women.",
+    "Low hemoglobin levels indicate anemia, causing fatigue and weakness.",
+    "High hemoglobin levels may suggest polycythemia or dehydration.",
+    "Normal fasting blood glucose levels are 70 to 99 mg/dL.",
+    "Elevated glucose levels indicate diabetes or prediabetes and require further testing.",
+]
 
-# Load medical knowledge into ChromaDB (run once)
-def load_medical_knowledge():
-    knowledge = [
-        {"name": "hemoglobin_normal", "text": "Normal hemoglobin levels are 13.8-17.2 g/dL for men and 12.1-15.1 g/dL for women."},
-        {"name": "hemoglobin_low", "text": "Low hemoglobin levels indicate anemia, causing fatigue and weakness."},
-        {"name": "glucose_normal", "text": "Normal fasting blood glucose levels are 70-99 mg/dL."},
-        {"name": "glucose_high", "text": "High glucose levels suggest diabetes or prediabetes and need further testing."},
-    ]
-    for item in knowledge:
-        collection.add(documents=[item["text"]], metadatas={"name": item["name"]}, ids=[item["name"]])
-    print("Medical knowledge loaded.")
+# Step 3: Build FAISS Index
+vectorizer = TfidfVectorizer()
+knowledge_vectors = vectorizer.fit_transform(medical_knowledge).toarray()
+knowledge_vectors = normalize(knowledge_vectors)  # Normalize vectors for cosine similarity
 
-# Step 3: Extract Text from PDF
+# Initialize FAISS Index
+dimension = knowledge_vectors.shape[1]
+faiss_index = faiss.IndexFlatL2(dimension)
+faiss_index.add(knowledge_vectors)
+
+def retrieve_medical_knowledge(parameter):
+    """Retrieve relevant knowledge using FAISS."""
+    query_vector = vectorizer.transform([parameter]).toarray()
+    query_vector = normalize(query_vector)  # Normalize the query vector
+    _, indices = faiss_index.search(query_vector, 1)  # Retrieve top 1 result
+    return medical_knowledge[indices[0][0]]
+
+# Step 4: Extract Text from PDF
 def extract_text_from_pdf(pdf_file):
+    """Extract text from the uploaded PDF file."""
     text = ""
     with fitz.open(stream=pdf_file.read(), filetype="pdf") as pdf:
         for page in pdf:
             text += page.get_text()
     return text
 
-# Step 4: Parse Blood Test Results
+# Step 5: Parse Blood Test Results
 def parse_blood_test_results(text):
+    """Parse blood test results from the extracted text."""
     results = {}
     hemoglobin_match = re.search(r'Hemoglobin:\s*(\d+\.\d+)\s*g/dL', text, re.IGNORECASE)
     glucose_match = re.search(r'Glucose:\s*(\d+)\s*mg/dL', text, re.IGNORECASE)
@@ -44,19 +57,15 @@ def parse_blood_test_results(text):
         results['Glucose'] = int(glucose_match.group(1))
     return results
 
-# Step 5: Retrieve Knowledge Dynamically from ChromaDB
-def retrieve_medical_knowledge(parameter):
-    results = collection.query(query_texts=[parameter], n_results=1)
-    return results['documents'][0] if results['documents'] else "No relevant knowledge found."
-
 # Step 6: Generate Personalized Advice
 def generate_advice(test_results):
+    """Generate personalized health advice using Gemini API."""
     advice = {}
     for parameter, value in test_results.items():
-        medical_knowledge = retrieve_medical_knowledge(parameter)
+        medical_info = retrieve_medical_knowledge(parameter)
         prompt = (
-            f"The patient's {parameter} level is {value}. {medical_knowledge} "
-            "Provide a clear, concise health recommendation."
+            f"The patient's {parameter} level is {value}. {medical_info} "
+            "Provide clear, concise health advice."
         )
         response = gemini.generate_text(prompt)
         advice[parameter] = response.result
@@ -64,6 +73,7 @@ def generate_advice(test_results):
 
 # Step 7: Main Function for Gradio Interface
 def analyze_blood_test(pdf_file):
+    """Main function to analyze the uploaded blood test PDF."""
     text = extract_text_from_pdf(pdf_file)
     test_results = parse_blood_test_results(text)
     if not test_results:
@@ -71,15 +81,14 @@ def analyze_blood_test(pdf_file):
     advice = generate_advice(test_results)
     return advice
 
-# Gradio interface
+# Gradio Interface
 iface = gr.Interface(
     fn=analyze_blood_test,
     inputs=gr.inputs.File(label="Upload Blood Test PDF"),
     outputs="json",
-    title="Blood Test Analysis with Full RAG Implementation",
+    title="Blood Test Analysis with RAG and Gemini (FAISS)",
     description="Upload a PDF with blood test results to receive personalized health advice."
 )
 
 if __name__ == "__main__":
-    load_medical_knowledge()  # Run once to load knowledge into ChromaDB
     iface.launch()