Update app.py

app.py

@@ -1,67 +1,45 @@
 import spaces
 import os
 import torch
-import random
 import numpy as np
 import pandas as pd
 import gradio as gr
-from pathlib import Path
-
 from transformers import pipeline, set_seed
 from sentence_transformers import SentenceTransformer
 from datasets import load_dataset
 
-# 1. Initialize LLM
-print("Loading BioMistral LLM...")
-llm = pipeline('text-generation', model='BioMistral/BioMistral-7B')
+# 1. Keep LLM instance uninitialized globally
+llm = None
+
+# 2. Keep lightweight embeddings on CPU during startup
+print("Indexing dataset embeddings on CPU...")
+embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
 
-# 2. Load Datasets
-print("Loading datasets...")
+# Load and process datasets at startup
 dataset = load_dataset("cbasu/Med-EASi", split="train")
 dataset2 = load_dataset("cbasu/Med-EASi", split="validation")
 dataset3 = load_dataset("cbasu/Med-EASi", split="test")
 dataset4 = load_dataset("csv", data_files="custom_terms.csv", split="train")
 
-# 3. Format dataset rows into strings
 texts = []
-for row in dataset:
-    content = f"Medical Context: {row['Expert']}\nSimplified Explanation: {row['Simple']}\nTerms: {row['expert_terms']}, {row['layman_terms']}"
-    texts.append(content)
-for row in dataset2:
-    content = f"Medical Context: {row['Expert']}\nSimplified Explanation: {row['Simple']}\nTerms: {row['expert_terms']}, {row['layman_terms']}"
-    texts.append(content)
-for row in dataset3:
-    content = f"Medical Context: {row['Expert']}\nSimplified Explanation: {row['Simple']}\nTerms: {row['expert_terms']}, {row['layman_terms']}"
-    texts.append(content)
-for row in dataset4:
-    content = f"Medical Context: {row['Expert']}\nSimplified Explanation: {row['Simple']}\nTerms: {row['expert_terms']}, {row['layman_terms']}"
-    texts.append(content)
-
-print(f"Total processed documents: {len(texts)}")
-
-# 4. Generate Embeddings & Index
-print("Indexing dataset embeddings...")
-embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
-embeddings = embedding_model.encode(texts)
+for ds in [dataset, dataset2, dataset3, dataset4]:
+    for row in ds:
+        content = f"Medical Context: {row['Expert']}\nSimplified Explanation: {row['Simple']}\nTerms: {row['expert_terms']}, {row['layman_terms']}"
+        texts.append(content)
 
-df = pd.DataFrame({
-    "Document": texts,
-    "Embedding": list(embeddings)
-})
+embeddings = embedding_model.encode(texts)
+df = pd.DataFrame({"Document": texts, "Embedding": list(embeddings)})
 
-# 5. Retrieval Helper Function
 def retrieve_with_pandas(query, top_k=3):
     query_embedding = embedding_model.encode([query])[0]
     df['Similarity'] = df['Embedding'].apply(
         lambda x: np.dot(query_embedding, x) / (np.linalg.norm(query_embedding) * np.linalg.norm(x))
     )
-    results = df.sort_values(by="Similarity", ascending=False).head(top_k)
-    return results["Document"].tolist()
+    return df.sort_values(by="Similarity", ascending=False).head(top_k)["Document"].tolist()
 
-# 6. Prompt Builder Helper Function
 def build_prompt(context_docs: list, user_query: str) -> str:
     context = "\n\n".join(context_docs)
-    prompt = f"""<s>[INST] You are a helpful medical assistant.
+    return f"""<s>[INST] You are a helpful medical assistant.
 Use the following retrieved medical context to answer the user's question clearly and concisely.
 Explain it in very simple terms as if the user has the vocabulary of an elementary school student.
 If the context does not contain enough information to answer, state that you don't know based on the provided text.
@@ -72,56 +50,45 @@ Context:
 
 Question:
 {user_query} [/INST]"""
-    return prompt
 
-# 7. Answer Generation Helper Function
-def generate_answer(prompt: str, max_new_tokens: int = 150) -> str:
-    set_seed(42)
+# 3. Decorate the master generation function
+@spaces.GPU(duration=60)
+def respond_to_query(user_query: str) -> str:
+    global llm
+    if not user_query.strip():
+        return "Please enter a valid medical question."
+
+    # Lazy-load BioMistral ONLY inside the GPU context
+    if llm is None:
+        print("Initializing BioMistral on ZeroGPU...")
+        llm = pipeline(
+            'text-generation',
+            model='BioMistral/BioMistral-7B',
+            torch_dtype=torch.bfloat16,
+            device_map="auto"
+        )
 
+    context_docs = retrieve_with_pandas(user_query, top_k=3)
+    prompt = build_prompt(context_docs, user_query)
+    
+    set_seed(42)
     if llm.model.config.pad_token_id is None:
         llm.model.config.pad_token_id = llm.model.config.eos_token_id
 
     output = llm(
         prompt,
-        max_new_tokens=max_new_tokens,
+        max_new_tokens=120,
         num_return_sequences=1,
         do_sample=False,
     )[0]["generated_text"]
 
-    # Strip out the prompt text to leave only the model's new response
-    answer = output[len(prompt):].strip()
-    return answer
-
-# 8. RAG Handler Function
-@spaces.GPU(duration=60)
-def respond_to_query(user_query: str) -> str:
-    if not user_query.strip():
-        return "Please enter a valid medical question or term."
-    
-    # Step A: Retrieve context documents based on user input
-    context_docs = retrieve_with_pandas(user_query, top_k=3)
-    
-    # Step B: Build prompt with retrieved context
-    prompt = build_prompt(context_docs, user_query)
-    
-    # Step C: Generate model answer
-    answer = generate_answer(prompt, max_new_tokens=80)
-    
-    return answer
+    return output[len(prompt):].strip()
 
-# 9. Gradio Web Interface
 demo = gr.Interface(
     fn=respond_to_query,
-    inputs=gr.Textbox(lines=2, placeholder="e.g., What is necrosis? or What is cholera?"),
-    outputs=gr.Textbox(label="BioMistral RAG Response", lines=4),
-    title="🩺 Medical Terminology RAG Assistant",
-    description="Ask a medical question to retrieve context from Med-EASi + custom terms and generate simplified explanations.",
-    examples=[
-        ["What causes necrosis?"],
-        ["What is cholera?"],
-        ["What is asbestosis?"],
-        ["How can syphilis be treated?"]
-    ]
+    inputs=gr.Textbox(lines=2, placeholder="e.g., What causes necrosis?"),
+    outputs=gr.Textbox(label="BioMistral Response", lines=4),
+    title="🩺 Medical Terminology RAG Assistant"
 )
 
 if __name__ == "__main__":