Update app.py

app.py

@@ -2,60 +2,37 @@ import os
 import gradio as gr
 from transformers import pipeline
 
-# -----------------------------
-# CONFIG
-# -----------------------------
-BASE_MODEL_ID = os.getenv("BASE_MODEL_ID", "google/flan-t5-small")
-FINETUNED_MODEL_ID = os.getenv("FINETUNED_MODEL_ID", "KB-Infinity-Tech/t5-samsum-mini")  # <-- CHANGE THIS
-TASK_PREFIX = os.getenv("TASK_PREFIX", "summarize: ")
+BASE_MODEL_ID = "google/flan-t5-small"
+FINETUNED_MODEL_ID = "kbsha/t5-samsum-mini"  # change to yours
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# -----------------------------
-# LOAD MODELS (MODERN API)
-# -----------------------------
 base_pipe = pipeline(
-    "text2text-generation",
+    "text-generation",
     model=BASE_MODEL_ID,
     token=HF_TOKEN,
 )
 
 finetuned_pipe = pipeline(
-    "text2text-generation",
+    "text-generation",
     model=FINETUNED_MODEL_ID,
     token=HF_TOKEN,
 )
 
-# -----------------------------
-# FUNCTION
-# -----------------------------
 def summarize(text):
-    prompt = TASK_PREFIX + text
+    prompt = "summarize: " + text
 
-    base_output = base_pipe(prompt, max_length=60, min_length=10, do_sample=False)
-    fine_output = finetuned_pipe(prompt, max_length=60, min_length=10, do_sample=False)
+    base = base_pipe(prompt, max_new_tokens=60)[0]["generated_text"]
+    fine = finetuned_pipe(prompt, max_new_tokens=60)[0]["generated_text"]
 
-    base_summary = base_output[0]["generated_text"]
-    fine_summary = fine_output[0]["generated_text"]
+    return base, fine
 
-    return base_summary, fine_summary
-
-# -----------------------------
-# UI
-# -----------------------------
 iface = gr.Interface(
     fn=summarize,
-    inputs=gr.Textbox(lines=8, label="Dialogue"),
-    outputs=[
-        gr.Textbox(label=f"Base ({BASE_MODEL_ID})"),
-        gr.Textbox(label=f"Fine-tuned ({FINETUNED_MODEL_ID})"),
-    ],
-    title="T5 SAMSum Demo",
-    description="Compare base vs fine-tuned summarization",
+    inputs=gr.Textbox(lines=8),
+    outputs=[gr.Textbox(), gr.Textbox()],
+    title="T5 Summarizer",
 )
 
-# -----------------------------
-# RUN
-# -----------------------------
 if __name__ == "__main__":
     iface.launch()