Update app.py

app.py

@@ -1,63 +1,103 @@
+# 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")  # <-- change to your real repo
+FINETUNED_MODEL_ID = os.getenv("FINETUNED_MODEL_ID", "KB-Infinity-Tech/t5-samsum-mini")
 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,
-)
+base_pipe = pipeline("summarization", model=BASE_MODEL_ID)
+finetuned_pipe = pipeline("summarization", model=FINETUNED_MODEL_ID)
 
-finetuned_pipe = pipeline(
-    "text2text-generation",
-    model=FINETUNED_MODEL_ID,
-    token=HF_TOKEN,
-)
 
-# -----------------------------
-# FUNCTION
-# -----------------------------
-def summarize(text):
+def summarize(text: str) -> tuple[str, str]:
     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"]
-
+    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
 
-# -----------------------------
-# UI
-# -----------------------------
+
 iface = gr.Interface(
     fn=summarize,
-    inputs=gr.Textbox(lines=8, label="Dialogue"),
+    inputs=gr.Textbox(lines=8, label="Dialogue", placeholder="Paste a conversation here"),
     outputs=[
-        gr.Textbox(label=f"Base ({BASE_MODEL_ID})"),
-        gr.Textbox(label=f"Fine-tuned ({FINETUNED_MODEL_ID})"),
+        gr.Textbox(label=f"Base model ({BASE_MODEL_ID})"),
+        gr.Textbox(label=f"Fine-tuned model ({FINETUNED_MODEL_ID})"),
     ],
     title="T5 SAMSum Demo",
-    description="Compare base vs fine-tuned summarization",
+    description="Compare the original FLAN-T5-small summarization with the fine-tuned checkpoint.",
 )
 
-# -----------------------------
-# RUN
-# -----------------------------
+
 if __name__ == "__main__":
     iface.launch()
 
-