App.py updated.

app.py

@@ -1,134 +1,106 @@
 import gradio as gr
 import pandas as pd
 import faiss
-import numpy as np
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import numpy as np
+
 from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ----------------------
-# Load Retrieval Corpus & FAISS Index
-# ----------------------
+# ===============================
+# Load Retrieval Components
+# ===============================
+print("Loading corpus and FAISS index...")
 df = pd.read_csv("retrieval_corpus.csv")
 index = faiss.read_index("faiss_index.bin")
 
-# ----------------------
-# Load Embedding Model (very lightweight)
-# ----------------------
+print("Loading embedding model...")
 embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
 
-# ----------------------
-# Load HuggingFace LLM (BioMistral-7B, 8bit CPU)
-# ----------------------
+# ===============================
+# Load LLM on CPU
+# ===============================
 model_id = "BioMistral/BioMistral-7B"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-tokenizer.pad_token = tokenizer.eos_token
 
-generation_model = AutoModelForCausalLM.from_pretrained(
+print(f"Loading tokenizer and model: {model_id}")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="auto",
-    offload_folder="offload",
-    offload_state_dict=True,
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True,
-)
-
-# ----------------------
-# RAG Functions
-# ----------------------
-
-def retrieve_top_k(query, k=5):
-    query_embedding = embedding_model.encode([query]).astype("float32")
-    D, I = index.search(query_embedding, k)
-    results = df.iloc[I[0]].copy()
-    results["score"] = D[0]
-    return results
-
-def build_prompt(query, retrieved_docs):
-    context_text = "\n".join([
-        f"- {doc['text']}" for _, doc in retrieved_docs.iterrows()
-    ])
-    prompt = f"""[INST] <<SYS>>
-You are a medical assistant trained on clinical reasoning data. Given the following patient query and related clinical observations, generate a diagnostic explanation or suggestion based on the context.
-<</SYS>>
-
-### Patient Query:
-{query}
-
-### Clinical Context:
-{context_text}
-
-### Diagnostic Explanation:
-[/INST]
-"""
-    return prompt
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True
+).to("cpu")
 
-def generate_local_answer(prompt, max_new_tokens=256):  # ✅ Reduced token budget
-    tokens = tokenizer(
-        prompt,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=1024
-    )
-    input_ids = tokens["input_ids"]
-    attention_mask = tokens["attention_mask"]
-
-    output = generation_model.generate(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        max_new_tokens=max_new_tokens,
-        temperature=0.5,
-        do_sample=True,
-        top_k=50,
-        top_p=0.95,
-        pad_token_id=tokenizer.pad_token_id
+tokenizer.pad_token = tokenizer.eos_token
+
+# ===============================
+# RAG Pipeline
+# ===============================
+def get_top_k_chunks(query, k=5):
+    query_embedding = embedding_model.encode([query])
+    scores, indices = index.search(np.array(query_embedding).astype("float32"), k)
+    return df.iloc[indices[0]]["text"].tolist()
+
+def build_prompt(query, chunks):
+    context = "\n".join(f"{i+1}. {chunk}" for i, chunk in enumerate(chunks))
+    prompt = (
+        "You are a clinical reasoning assistant. Based on the following medical information, "
+        "answer the query with a detailed explanation.\n\n"
+        f"Context:\n{context}\n\n"
+        f"Query: {query}\n"
+        "Answer:"
     )
+    return prompt
 
-    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-    return decoded.split("### Diagnostic Explanation:")[-1].strip()
-
-# ----------------------
-# Gradio Interface
-# ----------------------
-
-def rag_chat(query):
-    top_docs = retrieve_top_k(query, k=5)
-    prompt = build_prompt(query, top_docs)
-    answer = generate_local_answer(prompt)
-    return answer
-
-custom_css = """
-textarea, .input_textbox {
-    font-size: 1.05rem !important;
-}
-.output-markdown {
-    font-size: 1.08rem !important;
-}
-"""
-
-with gr.Blocks(css=custom_css, theme=gr.themes.Default(primary_hue="blue")) as demo:
-    gr.Markdown("""
-# 🩺 RAGnosis — Clinical Reasoning Assistant
-
-Enter a natural-language query describing your patient's condition to receive an AI-generated diagnostic reasoning response.
-
-**Example:**  
-*Patient has shortness of breath, fatigue, and leg swelling.*
-""")
-
-    with gr.Row():
-        with gr.Column():
-            query_input = gr.Textbox(
-                lines=4,
-                label="📝 Patient Query",
-                placeholder="Enter patient symptoms or findings..."
-            )
-            submit_btn = gr.Button("🔍 Generate Diagnosis")
-
-        with gr.Column():
-            output = gr.Markdown(label="���� Diagnostic Reasoning")
-
-    submit_btn.click(fn=rag_chat, inputs=query_input, outputs=output)
-
-demo.launch()
+def generate_diagnosis(query):
+    chunks = get_top_k_chunks(query)
+    prompt = build_prompt(query, chunks)
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
+    input_ids = inputs.input_ids.to("cpu")
+
+    with torch.no_grad():
+        output = model.generate(
+            input_ids=input_ids,
+            max_new_tokens=256,
+            do_sample=True,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.7,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    answer = generated_text.split("Answer:")[-1].strip()
+    return answer, "\n\n".join(chunks)
+
+# ===============================
+# Gradio UI
+# ===============================
+def run_interface():
+    with gr.Blocks(theme=gr.themes.Soft()) as demo:
+        gr.Markdown("## 🧠 Clinical Diagnosis Assistant (RAG)")
+        gr.Markdown("Enter a clinical query. The assistant retrieves relevant medical facts and generates a diagnostic explanation.")
+
+        with gr.Row():
+            query_input = gr.Textbox(label="Clinical Query", placeholder="e.g. 65-year-old male with shortness of breath...")
+            generate_btn = gr.Button("Generate Diagnosis")
+
+        with gr.Accordion("📄 Retrieved Context", open=False):
+            context_output = gr.Textbox(label="Top-5 Retrieved Chunks", lines=10, interactive=False)
+
+        answer_output = gr.Textbox(label="Generated Diagnosis", lines=8)
+
+        generate_btn.click(
+            fn=generate_diagnosis,
+            inputs=query_input,
+            outputs=[answer_output, context_output]
+        )
+
+    return demo
+
+# ===============================
+# Launch App
+# ===============================
+if __name__ == "__main__":
+    demo = run_interface()
+    demo.launch()