Update app.py

app.py

@@ -1,46 +1,38 @@
-import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Model name
+# Load the model and tokenizer
 model_name = "Amir230703/phi3-medmcqa-finetuned"
-
-# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Use float16 for speed if GPU available
-    device_map="auto"  # Moves model to GPU if available
-)
-
-# Function for generating responses
-def generate_response(input_text):
-    try:
-        input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(model.device)
-        
-        output = model.generate(
-            input_ids, 
-            max_length=200,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True
-        )
-
-        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-        return generated_text
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-    except Exception as e:
-        return f"Error: {str(e)}"
+def generate_answer(question):
+    # Tokenize the input question
+    input_ids = tokenizer(question, return_tensors="pt").input_ids.to(model.device)
+    
+    # Generate the answer
+    output = model.generate(
+        input_ids,
+        max_length=100,  # Reduced max_length for faster response
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True,
+        num_return_sequences=1  # Only return one answer
+    )
+    
+    # Decode the output
+    answer = tokenizer.decode(output[0], skip_special_tokens=True)
+    return answer
 
-# Create Gradio interface
+# Gradio Interface
 demo = gr.Interface(
-    fn=generate_response,
-    inputs=gr.Textbox(placeholder="Enter a medical question..."),
+    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.",
+    description="Enter a medical question, and the AI will provide an answer."
 )
 
-# Run the app
-if __name__ == "__main__":
-    demo.launch()
+# Launch the Gradio app
+demo.launch()