Update app.py

app.py

@@ -2,30 +2,39 @@ import gradio as gr
 from transformers import pipeline
 from datasets import load_dataset
 
-# Load model from Hugging Face Hub
-model_name = "Qwen/Qwen1.5-7B-Chat"  # or any other LLM
-pipe = pipeline("text-generation", model=model_name)
+# Choose your model here – change if needed
+model_name = "Qwen/Qwen1.5-7B-Chat"
 
-# Optionally load GTA queries
+# Load the text generation pipeline
+pipe = pipeline("text-generation", model=model_name, device=0)
+
+# Load GTA dataset (optional)
 gta_data = load_dataset("open-compass/GTA", split="test")
 
-def run_model(user_query, use_sample):
-    if use_sample:
-        question = gta_data[int(user_query)]["question"]
+# Inference function
+def run_model(input_text, use_gta_index):
+    if use_gta_index:
+        try:
+            index = int(input_text)
+            question = gta_data[index]["question"]
+        except Exception as e:
+            return f"❌ Error: Invalid index — {str(e)}"
     else:
-        question = user_query
+        question = input_text
 
     output = pipe(question, max_new_tokens=256, do_sample=True)
-    return f"**Input**: {question}\n\n**Output**:\n{output[0]['generated_text']}"
+    return f"**Question:** {question}\n\n**Response:**\n{output[0]['generated_text']}"
 
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 GTA Tool Reasoning Demo with Hugging Face Models")
+    gr.Markdown("# 🤖 GTA-style Reasoning with Hugging Face Models")
+    gr.Markdown("Enter a question or choose a sample from the GTA benchmark.")
     with gr.Row():
-        user_input = gr.Textbox(label="Enter your query or GTA index (if using sample)")
-        use_sample = gr.Checkbox(label="Use sample from GTA dataset (enter index)", value=False)
-    run_button = gr.Button("Run")
-    output = gr.Markdown()
+        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()
     
-    run_button.click(run_model, inputs=[user_input, use_sample], outputs=output)
+    run_button.click(run_model, inputs=[input_text, use_sample], outputs=output_text)
 
-demo.launch()
+demo.launch()