Update app.py

app.py

@@ -1,37 +1,36 @@
 import gradio as gr
-import json
 from transformers import pipeline
+from datasets import load_dataset
 
-# Load Qwen2.5-3B (should fit in 16GB with CPU)
+# Load smaller model to fit memory
 model_name = "Qwen/Qwen2.5-3B"
 pipe = pipeline("text-generation", model=model_name, device=0)
 
-# ✅ Load local GTA dataset (test set)
-with open("gta_test.json", "r", encoding="utf-8") as f:
-    gta_data = json.load(f)
+# Load the GTA dataset from Hugging Face
+gta = load_dataset("Jize1/GTA", split="test")
 
-def run_model(input_text, use_gta_index):
-    if use_gta_index:
+def run_model(input_text, use_gta_idx):
+    if use_gta_idx:
         try:
-            index = int(input_text)
-            question = gta_data[index]["question"]
+            idx = int(input_text)
+            question = gta[idx]["dialogs"][0]["content"].strip()
         except Exception as e:
-            return f"❌ Invalid index: {e}"
+            return f"❌ Invalid index (0–{len(gta)-1}): {e}"
     else:
-        question = input_text
+        question = input_text.strip()
 
     output = pipe(question, max_new_tokens=256, do_sample=True)
     return f"**Question:** {question}\n\n**Response:**\n{output[0]['generated_text']}"
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 GTA Reasoning with Qwen2.5-3B")
-    gr.Markdown("Use a GTA query by index or enter your own.")
+    gr.Markdown("# 🤖 GTA Reasoning Demo (Qwen2.5‑3B + GTA Dataset)")
+    gr.Markdown("Enter your own question or a GTA index (0–228).")
     with gr.Row():
-        input_text = gr.Textbox(label="Enter a question or GTA index (0–228)")
+        input_text = gr.Textbox(label="Your input or GTA index")
         use_index = gr.Checkbox(label="Treat input as GTA index", value=False)
     run_btn = gr.Button("Generate")
     output_md = gr.Markdown()
     
-    run_btn.click(run_model, inputs=[input_text, use_index], outputs=output_md)
+    run_btn.click(fn=run_model, inputs=[input_text, use_index], outputs=[output_md])
 
 demo.launch()