Summarizer / app.py
IotaCluster's picture
Update app.py
069ec03 verified
import re
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
# Choose the distilled BART checkpoint
model_name = "sshleifer/distilbart-cnn-12-6"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# Pre‑build the HF summarization pipeline (faster in repeated calls)
summarizer = pipeline(
"summarization",
model=model,
tokenizer=tokenizer,
device=-1, # set to -1 for CPU, or the GPU ID if available
truncation = True
)
def remove_confidentiality(text: str) -> str:
pattern = r"\*\*CONFIDENTIALITY NOTICE:.*"
return re.split(pattern, text, flags=re.DOTALL)[0].strip()
def summarize_text(text):
cleaned = remove_confidentiality(text)
# pipeline will chunk long inputs automatically if you pass `max_length` and `min_length`
summary_list = summarizer(
cleaned,
max_length=250, # shorter target length for punchier summaries
min_length=50,
length_penalty=1.0,
num_beams=3,
temperature=0.7,
early_stopping=True
)
return summary_list[0]["summary_text"]
iface = gr.Interface(
fn=summarize_text,
inputs=gr.Textbox(lines=15, placeholder="Paste your text here..."),
outputs=gr.Textbox(label="Summary"),
title="Fast & Accurate Summarizer",
description="Using the distilled BART model for quicker, high-quality summaries."
)
if __name__ == "__main__":
iface.launch()