Update app.py

app.py

@@ -1,15 +1,14 @@
 import gradio as gr
+import json
 from transformers import pipeline
-from datasets import load_dataset
 
-# ✅ Use smaller, memory-friendly model
+# Load Qwen2.5-3B (should fit in 16GB with CPU)
 model_name = "Qwen/Qwen2.5-3B"
-
-# Load model pipeline
 pipe = pipeline("text-generation", model=model_name, device=0)
 
-# Load GTA benchmark (optional)
-gta_data = load_dataset("open-compass/GTA", split="test")
+# ✅ Load local GTA dataset (test set)
+with open("gta_test.json", "r", encoding="utf-8") as f:
+    gta_data = json.load(f)
 
 def run_model(input_text, use_gta_index):
     if use_gta_index:
@@ -17,23 +16,22 @@ def run_model(input_text, use_gta_index):
             index = int(input_text)
             question = gta_data[index]["question"]
         except Exception as e:
-            return f"❌ Error: Invalid index — {str(e)}"
+            return f"❌ Invalid index: {e}"
     else:
         question = input_text
 
     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-style Reasoning with Qwen 2.5-3B")
-    gr.Markdown("Test queries from the GTA benchmark or enter your own.")
+    gr.Markdown("# 🤖 GTA Reasoning with Qwen2.5-3B")
+    gr.Markdown("Use a GTA query by index or enter your own.")
     with gr.Row():
-        input_text = gr.Textbox(label="Enter a question or GTA index (e.g., 0–228)")
-        use_sample = gr.Checkbox(label="Use as GTA index", value=False)
-    run_button = gr.Button("Generate")
-    output_text = gr.Markdown()
+        input_text = gr.Textbox(label="Enter a question or GTA index (0–228)")
+        use_index = gr.Checkbox(label="Treat input as GTA index", value=False)
+    run_btn = gr.Button("Generate")
+    output_md = gr.Markdown()
     
-    run_button.click(run_model, inputs=[input_text, use_sample], outputs=output_text)
+    run_btn.click(run_model, inputs=[input_text, use_index], outputs=output_md)
 
 demo.launch()