Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import torch
 import gradio as gr
+import pandas as pd
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
 REPO_ID = os.environ.get("HF_MODEL_ID", "HuyTran1301/constrative_cont_so_phase2_SI")
@@ -11,27 +12,31 @@ torch.set_num_threads(int(os.environ.get("TORCH_NUM_THREADS", "1")))
 tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
 model = AutoModelForSeq2SeqLM.from_pretrained(REPO_ID)
 
-summarizer = pipeline(
-    task="summarization",
-    model=model,
-    tokenizer=tokenizer,
-    device=-1
-)
-
 def summarize_one(lang: str, desc: str, code: str):
     if not any([lang.strip(), desc.strip(), code.strip()]):
-        return ""
+        return pd.DataFrame([["", ""]], columns=["#","Summary"])
     merged_text = f"{lang.strip()}: {desc.strip()} <code> {code.strip()}"
-    summary = summarizer(
+    input_ids = tokenizer(
         merged_text,
+        return_tensors="pt",
         truncation=True,
-        max_length=GEN_MAX_LENGTH,
-        min_length=5,
-        do_sample=False
-    )[0]["summary_text"]
-    return summary
-
-with gr.Blocks(title="Code Summarizer") as demo:
+        max_length=MAX_LENGTH
+    ).input_ids
+
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids,
+            max_length=GEN_MAX_LENGTH,
+            num_beams=5,
+            num_return_sequences=5,
+            early_stopping=True,
+        )
+
+    summaries = [tokenizer.decode(o, skip_special_tokens=True).strip() for o in outputs]
+    df = pd.DataFrame(list(enumerate(summaries, start=1)), columns=["#", "Summary"])
+    return df
+
+with gr.Blocks(title="Code Summarization") as demo:
     gr.Markdown("# Code Summarization")
 
     with gr.Row():
@@ -39,16 +44,16 @@ with gr.Blocks(title="Code Summarizer") as demo:
         desc = gr.Textbox(label="Description", placeholder="What does the code do?")
     code = gr.Textbox(lines=8, label="Code", placeholder="Paste your code here...")
 
-    btn = gr.Button("Generate Summary")
-    out_summary = gr.Textbox(label="Generated Summary", lines=4)
+    btn = gr.Button("Generate Summaries (5 beams)")
+    out_table = gr.Dataframe(headers=["#", "Summary"], label="Generated Summaries", interactive=False)
 
     btn.click(
         summarize_one,
         inputs=[lang, desc, code],
-        outputs=[out_summary]
+        outputs=[out_table]
     )
 
-    gr.Markdown(f"**Model:** `{REPO_ID}`  •  **Input max length:** {MAX_LENGTH}  •  **Output max length:** {GEN_MAX_LENGTH}")
+    gr.Markdown(f"**Model:** `{REPO_ID}`  •  **Input max length:** {MAX_LENGTH}  •  **Output max length:** {GEN_MAX_LENGTH}  •  **num_beams:** 5")
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))