Update app.py

app.py

@@ -1,50 +1,49 @@
-import os
-import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
-from huggingface_hub import login
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
+import gradio as gr
 
 import os
-os.environ["BITSANDBYTES_NOWELCOME"] = "1"
-os.environ["DISABLE_BITSANDBYTES"] = "1"
-
+from huggingface_hub import login
 
-# --- Authenticate with HF token (from Spaces Secrets) ---
-login(os.environ["HF_TOKEN"])
+login(os.environ["HF_TOKEN"])  # use the token with gated repo access
 
-# --- Model setup ---
-MODEL_ID = "google/medgemma-4b-it"
+MODEL_ID = "google/med-gemma-2b"
 
-# 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 tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# Load model + tokenizer with quantization
+# Load model with 4-bit quantization (works on CPU)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    quantization_config=bnb_config,
-    device_map="auto"
+    device_map="cpu",
+    torch_dtype=torch.float32,   # stay safe from NaN in CPU mode
+    load_in_4bit=True            # quantize
 )
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+# Wrap in a pipeline
+pipe = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=-1)
 
-# --- Gradio app ---
+# Safe generation function
 def medgemma_chat(prompt):
-    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)
-    return outputs[0]["generated_text"]
-
-demo = gr.Interface(
-    fn=medgemma_chat,
-    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."
-)
+    try:
+        output = pipe(
+            prompt,
+            max_new_tokens=200,
+            temperature=1.0,   # stable
+            top_p=0.9,
+            do_sample=True
+        )
+        return output[0]["generated_text"]
+    except Exception as e:
+        return f"⚠️ Error: {str(e)}"
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# 🩺 MedGemma (Quantized, CPU-safe)")
+    inp = gr.Textbox(label="Enter patient info", placeholder="Example: Patient has fever and cough...")
+    out = gr.Textbox(label="Model Output")
+    btn = gr.Button("Generate")
+    btn.click(medgemma_chat, inp, out)
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()