Update app.py

app.py

@@ -1,15 +1,22 @@
-import gradio as gr
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 import os
+import gradio as gr
 
-# Load Hugging Face token from secrets
-hf_token = os.environ.get("HF_TOKEN", "").strip()
+
+hf_token = os.environ.get("HF_TOKEN").strip()
 model_name = "jeevanrushi07/gemma-medical-assistant"
 
-# Load model and tokenizer
+# Directory to offload layers to CPU
+offload_dir = "/tmp/model_offload"
+
 tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    offload_folder=offload_dir,  # This is required for large models
+    token=hf_token
+)
 
 def generate_report(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
@@ -21,16 +28,14 @@ def generate_report(prompt):
         top_p=0.9,
         eos_token_id=tokenizer.eos_token_id
     )
-    text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-    return text
+    return tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
 
-# Gradio UI
 iface = gr.Interface(
     fn=generate_report,
-    inputs=gr.Textbox(lines=10, placeholder="Enter patient info and symptoms..."),
+    inputs=gr.Textbox(lines=10, placeholder="Enter patient info..."),
     outputs="textbox",
-    title="Medical Chatbot",
-    description="Enter patient data and receive AI-generated medical report."
+    title="Medical Chatbot"
 )
 
 iface.launch()
+