Update app.py

app.py

@@ -1,63 +1,45 @@
+import os
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-
-import os
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 from huggingface_hub import login
 
+# --- Authenticate with HF token (from Spaces Secrets) ---
 login(os.environ["HF_TOKEN"])
 
+# --- Model setup ---
+MODEL_ID = "google/medgemma-4b-it"
+
+# 4-bit quantization config
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+# Load model + tokenizer with quantization
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    quantization_config=bnb_config,
+    device_map="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+# --- Gradio app ---
+def medgemma_chat(prompt):
+    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)
+    return outputs[0]["generated_text"]
 
-# Load MedGemma
-MODEL_NAME = "google/medgemma-4b-it"  # choose the model size you want. here
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")  # GPU if available
-
-# Function to generate response
-def medgemma_chat(prompt, max_length=200):
-    """
-    Generates medical responses from MedGemma.
-    
-    Args:
-        prompt (str): Medical question, lab results, or patient info.
-        max_length (int): Max number of tokens to generate.
-    
-    Returns:
-        str: AI-generated response.
-    """
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_length=max_length,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
-        )
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response
-
-# Gradio UI
-iface = gr.Interface(
+demo = gr.Interface(
     fn=medgemma_chat,
-    inputs=[
-        gr.Textbox(lines=5, placeholder="Enter patient's info, lab results, or medical question here...", label="Input")
-    ],
-    outputs=[
-        gr.Textbox(label="MedGemma Response")
-    ],
-    title="MedGemma Medical Assistant",
-    description=(
-        "Ask questions or provide patient information. "
-        "MedGemma generates medical insights, summaries, and guidance. "
-        "⚠️ For educational and research purposes only — not a substitute for professional medical advice."
-    ),
-    examples=[
-        ["Patient: 45-year-old male, BMI 28, blood pressure 140/90, glucose 7.5 mmol/L. Suggest possible conditions."],
-        ["Summarize the following lab report: Hemoglobin 11 g/dL, WBC 9 x10^9/L, Platelets 200 x10^9/L."]
-    ],
-    allow_flagging="never"
+    inputs=gr.Textbox(label="Enter medical question", lines=4, placeholder="e.g. What are symptoms of malaria?"),
+    outputs=gr.Textbox(label="MedGemma Response"),
+    title="🧠 MedGemma (4-bit Quantized)",
+    description="Ask medical questions (research/demo use only). Running in 4-bit quantized mode for efficiency."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)