Update app.py

app.py

@@ -1,19 +1,30 @@
+import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
+# Load the model and tokenizer
 model_name = "Amir230703/phi3-medmcqa-finetuned"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-input_text = "What are the symptoms of the flu?"
-input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+def generate_answer(question):
+    input_ids = tokenizer(question, return_tensors="pt").input_ids.to(model.device)
+    output = model.generate(
+        input_ids,
+        max_length=200,
+        temperature=0.7,  # Controls randomness
+        top_p=0.9,  # Nucleus sampling
+        do_sample=True  # Enables variability
+    )
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
-output = model.generate(
-    input_ids, 
-    max_length=200, 
-    temperature=0.7,  # Controls randomness
-    top_p=0.9,  # Nucleus sampling
-    do_sample=True  # Enables variability
+# Gradio Interface
+demo = gr.Interface(
+    fn=generate_answer,
+    inputs=gr.Textbox(placeholder="Enter a medical question here..."),
+    outputs=gr.Textbox(),
+    title="Medical QA Model",
+    description="Enter a medical question, and the AI will provide an answer."
 )
 
-generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-print(generated_text)
+demo.launch()