Update app.py

app.py

@@ -10,7 +10,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 os.environ["BITSANDBYTES_NOWELCOME"] = "1"
 os.environ["DISABLE_BITSANDBYTES"] = "1"
 
-# Hugging Face token login via env variable
+# Hugging Face token login via environment variable
 from huggingface_hub import login
 login(os.environ.get("HF_TOKEN", ""))
 
@@ -19,7 +19,7 @@ login(os.environ.get("HF_TOKEN", ""))
 # -------------------------------
 MODEL_NAME = "google/medgemma-4b-it"  # or lighter if CPU only
 
-# auto-detect device
+# Auto-detect device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load tokenizer
@@ -28,8 +28,8 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 # Load model safely
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.float32,  # safer than 4-bit on CPU
-    device_map="auto" if device=="cuda" else None
+    torch_dtype=torch.float32,  # safer on CPU
+    device_map="auto" if device == "cuda" else None
 )
 
 # -------------------------------
@@ -59,11 +59,21 @@ def medgemma_generate(prompt):
 # -------------------------------
 demo = gr.Interface(
     fn=medgemma_generate,
-    inputs=gr.Textbox(lines=4, placeholder="Enter your medical prompt..."),
-    outputs="text",
+    inputs=gr.Textbox(
+        lines=4,
+        placeholder="Enter your medical prompt...",
+        label="Prompt"
+    ),
+    outputs=gr.Textbox(
+        lines=15,        # start with 15 lines
+        max_lines=100,   # auto-expand up to 100 lines
+        interactive=False,
+        label="Generated Answer"
+    ),
     title="MedGemma Q&A",
     description="Ask medical questions (English). Safe generation config prevents NaNs on CPU."
 )
 
+# Launch the app
 if __name__ == "__main__":
     demo.launch()