Update app.py

app.py

@@ -1,32 +1,22 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# Load model
 tokenizer = AutoTokenizer.from_pretrained("./mtpe-model")
 model = AutoModelForSeq2SeqLM.from_pretrained("./mtpe-model")
 pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
-def predict(task, prompt, context, auto_cot):
+def predict(task, prompt, context="", auto_cot=False):
     input_str = f"[TASK: {task.upper()}] {prompt}"
     if context:
         input_str += f" Context: {context}"
     if auto_cot:
         input_str += "\nLet's think step by step."
+    return pipe(input_str, max_new_tokens=128)[0]["generated_text"]
 
-    output = pipe(input_str, max_new_tokens=128)[0]["generated_text"]
-    return output
-
-iface = gr.Interface(
+app = gr.Interface(
     fn=predict,
-    inputs=[
-        gr.Textbox(label="Task", value="qa"),
-        gr.Textbox(label="Prompt"),
-        gr.Textbox(label="Context (optional)", lines=2),
-        gr.Checkbox(label="Enable Auto-CoT")
-    ],
-    outputs="text",
-    title="Prompt Playground Inference API",
-    description="Runs your trained mtpe-model from HF Spaces"
+    inputs=["text", "text", "text", "checkbox"],
+    outputs="text"
 )
 
-iface.launch()
+app.launch()