Update app.py

app.py

@@ -1,103 +1,61 @@
-# 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 to your real repo
-# TASK_PREFIX = os.getenv("TASK_PREFIX", "summarize: ")
-
-# HF_TOKEN = os.getenv("HF_TOKEN")
-
-# # -----------------------------
-# # LOAD MODELS (NEW API)
-# # -----------------------------
-# base_pipe = pipeline(
-#     "text2text-generation",
-#     model=BASE_MODEL_ID,
-#     token=HF_TOKEN,
-# )
-
-# finetuned_pipe = pipeline(
-#     "text2text-generation",
-#     model=FINETUNED_MODEL_ID,
-#     token=HF_TOKEN,
-# )
-
-# # -----------------------------
-# # FUNCTION
-# # -----------------------------
-# def summarize(text):
-#     prompt = TASK_PREFIX + text
-
-#     base_output = base_pipe(prompt, max_length=60, min_length=10, do_sample=False)
-#     finetuned_output = finetuned_pipe(prompt, max_length=60, min_length=10, do_sample=False)
-
-#     base_summary = base_output[0]["generated_text"]
-#     finetuned_summary = finetuned_output[0]["generated_text"]
-
-#     return base_summary, finetuned_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",
-# )
-
-# # -----------------------------
-# # RUN
-# # -----------------------------
-# if __name__ == "__main__":
-#     iface.launch()
-
-
-"""Gradio app comparing base FLAN-T5 vs the fine-tuned checkpoint."""
-
 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")
+FINETUNED_MODEL_ID = os.getenv("FINETUNED_MODEL_ID", "KB-Infinity-Tech/t5-samsum-mini")  # <-- CHANGE THIS
 TASK_PREFIX = os.getenv("TASK_PREFIX", "summarize: ")
 
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-base_pipe = pipeline("summarization", model=BASE_MODEL_ID)
-finetuned_pipe = pipeline("summarization", model=FINETUNED_MODEL_ID)
+# -----------------------------
+# LOAD MODELS (MODERN API)
+# -----------------------------
+base_pipe = pipeline(
+    "text2text-generation",
+    model=BASE_MODEL_ID,
+    token=HF_TOKEN,
+)
 
+finetuned_pipe = pipeline(
+    "text2text-generation",
+    model=FINETUNED_MODEL_ID,
+    token=HF_TOKEN,
+)
 
-def summarize(text: str) -> tuple[str, str]:
+# -----------------------------
+# FUNCTION
+# -----------------------------
+def summarize(text):
     prompt = TASK_PREFIX + text
-    base_summary = base_pipe(prompt, max_length=60, min_length=10, do_sample=False)[0]["summary_text"]
-    finetuned_summary = finetuned_pipe(prompt, max_length=60, min_length=10, do_sample=False)[0][
-        "summary_text"
-    ]
-    return base_summary, finetuned_summary
 
+    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_summary = base_output[0]["generated_text"]
+    fine_summary = fine_output[0]["generated_text"]
+
+    return base_summary, fine_summary
 
+# -----------------------------
+# UI
+# -----------------------------
 iface = gr.Interface(
     fn=summarize,
-    inputs=gr.Textbox(lines=8, label="Dialogue", placeholder="Paste a conversation here"),
+    inputs=gr.Textbox(lines=8, label="Dialogue"),
     outputs=[
-        gr.Textbox(label=f"Base model ({BASE_MODEL_ID})"),
-        gr.Textbox(label=f"Fine-tuned model ({FINETUNED_MODEL_ID})"),
+        gr.Textbox(label=f"Base ({BASE_MODEL_ID})"),
+        gr.Textbox(label=f"Fine-tuned ({FINETUNED_MODEL_ID})"),
     ],
     title="T5 SAMSum Demo",
-    description="Compare the original FLAN-T5-small summarization with the fine-tuned checkpoint.",
+    description="Compare base vs fine-tuned summarization",
 )
 
-
+# -----------------------------
+# RUN
+# -----------------------------
 if __name__ == "__main__":
-    iface.launch()
-
+    iface.launch()