fixed.

app.py

@@ -1,90 +1,52 @@
 import gradio as gr
-import pandas as pd
-import faiss
-import time
-import numpy as np
-import torch
+import pandas as pd, faiss, torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from sentence_transformers import SentenceTransformer
 
-# Load retrieval corpus & FAISS index
+# —— Load data & embedding model ——
 df = pd.read_csv("retrieval_corpus.csv")
 index = faiss.read_index("faiss_index.bin")
-
-# Load embedding model
 embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
-model_id = "stanford-crfm/BioMedLM"
 
+# —— Quantized BioMedLM with CPU offload ——
+model_id = "stanford-crfm/BioMedLM"
 bnb_config = BitsAndBytesConfig(
     load_in_8bit=True,
     llm_int8_threshold=6.0,
+    llm_int8_enable_fp32_cpu_offload=True,
 )
-
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 tokenizer.pad_token = tokenizer.eos_token
-
 generation_model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="auto",
     quantization_config=bnb_config,
+    device_map={"": "cpu"},
 )
 
-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()])
-    return 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}
+def retrieve_top_k(q, k=5):
+    emb = embedding_model.encode([q]).astype("float32")
+    D,I = index.search(emb, k)
+    res = df.iloc[I[0]].copy(); res["score"]=D[0]; return res
 
-### Clinical Context:
-{context_text}
-
-### Diagnostic Explanation:
-[/INST]
-"""
+def build_prompt(q, docs):
+    ctx = "\n".join(f"- {d['text']}" for _,d in docs.iterrows())
+    return f"""[INST] <<SYS>>…[/INST]"""  # your existing template
 
 def generate_local_answer(prompt, max_new_tokens=512):
+    import time
     device = torch.device("cpu")
-    print(f"Using device: {device}")
     start = time.time()
-
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
-    input_ids = inputs["input_ids"].to(device)
-    attention_mask = inputs["attention_mask"].to(device)
-
-    output = generation_model.generate(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(device)
+    out = generation_model.generate(
+        input_ids=inputs.input_ids,
+        attention_mask=inputs.attention_mask,
         max_new_tokens=max_new_tokens,
-        do_sample=False,    # ← GREEDY
+        do_sample=False,
         num_beams=1,
     )
+    print(f"Gen time: {time.time()-start:.2f}s")
+    return tokenizer.decode(out[0], skip_special_tokens=True)
 
-    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-    print(f"Time taken: {time.time() - start:.2f}s")
-    return decoded.split("### Diagnostic Explanation:")[-1].strip()
-
-def rag_chat(query):
-    top_docs = retrieve_top_k(query, k=5)
-    prompt = build_prompt(query, top_docs)
-    return generate_local_answer(prompt)
-
-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 = gr.Interface(fn=lambda q: generate_local_answer(build_prompt(q, retrieve_top_k(q))),
+                     inputs="text", outputs="text")
 iface.launch()