Update app.py

app.py

@@ -2,18 +2,19 @@ import gradio as gr
 from transformers import pipeline
 from datasets import load_dataset
 
-# Load smaller model to fit memory
+# ✅ Load the smaller model (fits in 16GB)
 model_name = "Qwen/Qwen2.5-3B"
 pipe = pipeline("text-generation", model=model_name, device=0)
 
-# Load the GTA dataset from Hugging Face
-gta = load_dataset("Jize1/GTA", split="test")
+# ✅ Load the GTA dataset (correct split is 'train')
+gta = load_dataset("Jize1/GTA", split="train")
 
+# ✅ Inference function
 def run_model(input_text, use_gta_idx):
     if use_gta_idx:
         try:
             idx = int(input_text)
-            question = gta[idx]["dialogs"][0]["content"].strip()
+            question = gta[idx]["dialogs"][0]["content"]
         except Exception as e:
             return f"❌ Invalid index (0–{len(gta)-1}): {e}"
     else:
@@ -22,9 +23,10 @@ def run_model(input_text, use_gta_idx):
     output = pipe(question, max_new_tokens=256, do_sample=True)
     return f"**Question:** {question}\n\n**Response:**\n{output[0]['generated_text']}"
 
+# ✅ Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# 🤖 GTA Reasoning Demo (Qwen2.5‑3B + GTA Dataset)")
-    gr.Markdown("Enter your own question or a GTA index (0–228).")
+    gr.Markdown("Enter a custom question or choose a sample from the GTA dataset (index 0–228).")
     with gr.Row():
         input_text = gr.Textbox(label="Your input or GTA index")
         use_index = gr.Checkbox(label="Treat input as GTA index", value=False)