Update app.py

app.py

@@ -0,0 +1,65 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import spaces
+
+# Model name
+model_name = "medalpaca/medalpaca-7b"
+
+# Load tokenizer and model globally for efficiency
+print(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"GPU device count: {torch.cuda.device_count()}")
+    print(f"GPU device name: {torch.cuda.get_device_name(0)}")
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto",  # Use GPU if available
+    load_in_8bit=torch.cuda.is_available()  # 8-bit quantization for GPU
+)
+
+def format_prompt(message, chat_history):
+    prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n"
+    if chat_history:
+        prompt += "Previous conversation:\n"
+    for turn in chat_history:
+        user_message, assistant_message = turn
+        prompt += f"Human: {user_message}\nAssistant: {assistant_message}\n\n"
+    prompt += f"Human: {message}\nAssistant:"
+    return prompt
+
+@spaces.GPU  # <--- This is REQUIRED for ZeroGPU!
+def generate_response(message, chat_history):
+    prompt = format_prompt(message, chat_history)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        generation_output = model.generate(
+            input_ids=inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+        )
+    full_output = tokenizer.decode(generation_output[0], skip_special_tokens=True)
+    response = full_output.split("Assistant:")[-1].strip()
+    chat_history.append((message, response))
+    return "", chat_history
+
+with gr.Blocks(css="footer {visibility: hidden}") as demo:
+    gr.Markdown("# MedAlpaca Medical Chatbot")
+    gr.Markdown("A specialized medical chatbot powered by MedAlpaca-7B.")
+    gr.Markdown("Ask medical questions and get responses from a model trained on medical data.")
+    
+    chatbot = gr.Chatbot(type="messages")
+    msg = gr.Textbox(placeholder="Type your medical question here...")
+    clear = gr.Button("Clear")
+    
+    msg.submit(generate_response, [msg, chatbot], [msg, chatbot])  # Pass GPU-decorated function!
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+if __name__ == "__main__":
+    print("Starting Gradio app...")
+    demo.launch(server_name="0.0.0.0")