Updated.

app.py

@@ -1,106 +1,103 @@
 import gradio as gr
 import pandas as pd
 import faiss
-import torch
 import numpy as np
-
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ===============================
-# Load Retrieval Components
-# ===============================
-print("Loading corpus and FAISS index...")
+# ----------------------
+# Load Retrieval Corpus & FAISS Index
+# ----------------------
 df = pd.read_csv("retrieval_corpus.csv")
 index = faiss.read_index("faiss_index.bin")
 
-print("Loading embedding model...")
+# ----------------------
+# Load Embedding Model
+# ----------------------
 embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
 
-# ===============================
-# Load LLM on CPU
-# ===============================
+# ----------------------
+# Load HuggingFace LLM (Nous-Hermes)
+# ----------------------
 model_id = "BioMistral/BioMistral-7B"
 
-print(f"Loading tokenizer and model: {model_id}")
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+)
+
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
+generation_model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype=torch.float32,
-    low_cpu_mem_usage=True
-).to("cpu")
-
-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:"
-    )
+    torch_dtype=torch.float16,
+    device_map="auto",
+    quantization_config=bnb_config
+)
+
+# ----------------------
+# 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
 
-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()
+def generate_local_answer(prompt, max_new_tokens=512):
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
+    output = generation_model.generate(
+        input_ids=input_ids,
+        max_new_tokens=max_new_tokens,
+        temperature=0.5,
+        do_sample=True,
+        top_k=50,
+        top_p=0.95,
+    )
+    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
+
+iface = gr.Interface(
+    fn=rag_chat,
+    inputs=gr.Textbox(lines=3, placeholder="Enter a clinical query..."),
+    outputs="text",
+    title="🩺 Clinical Reasoning RAG Assistant",
+    description="Ask a medical question based on MIMIC-IV-Ext-DiReCT's diagnostic knowledge.",
+    allow_flagging="never"
+)
+
+iface.launch()

requirements.txt

@@ -4,4 +4,6 @@ faiss-cpu
 torch
 gradio
 accelerate
-sentencepiece
+sentencepiece
+bitsandbytes
+blobfile