Update app.py

app.py

@@ -2,20 +2,21 @@ import gradio as gr
 import torch
 import torch.nn.functional as F
 from transformers import AutoTokenizer, AutoModelForMultipleChoice
+import spaces  # <--- Added spaces import
 
-# 1. Load fine-tuned model and tokenizer from Hugging Face Hub
+# Load fine-tuned model and tokenizer from Hugging Face Hub
 MODEL_ID = "udaypratap/smart-mcq-solver"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForMultipleChoice.from_pretrained(MODEL_ID)
 model.eval()
 
-# 2. Define prediction function
+# Add the @spaces.GPU decorator right here
+@spaces.GPU
 def predict_mcq(prompt, option_a, option_b, option_c, option_d, option_e):
     options = [option_a, option_b, option_c, option_d, option_e]
     labels = ["A", "B", "C", "D", "E"]
     
-    # Format inputs for AutoModelForMultipleChoice
     first_sentences = [prompt] * 5
     second_sentences = options
     
@@ -28,7 +29,11 @@ def predict_mcq(prompt, option_a, option_b, option_c, option_d, option_e):
         return_tensors="pt"
     )
     
-    # Reshape input tensors for multiple choice model: (batch_size=1, num_choices=5, seq_len)
+    # Move inputs and model to the GPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    
     input_ids = inputs["input_ids"].unsqueeze(0)
     attention_mask = inputs["attention_mask"].unsqueeze(0)
     
@@ -37,13 +42,12 @@ def predict_mcq(prompt, option_a, option_b, option_c, option_d, option_e):
         logits = outputs.logits.squeeze(0)
         probs = F.softmax(logits, dim=-1)
         
-    # Get top 3 predicted choices
     top3_indices = torch.topk(probs, k=3).indices.tolist()
     top3_choices = [f"{labels[idx]} ({probs[idx].item():.2%})" for idx in top3_indices]
     
     return f"Top 3 Predicted Answers: {', '.join(top3_choices)}"
 
-# 3. Create Gradio Interface
+# Create Gradio Interface
 demo = gr.Interface(
     fn=predict_mcq,
     inputs=[