Advanced_BART / app.py
vangru's picture
Update app.py
ad7b5db verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
MODEL_NAME = "facebook/bart-large-cnn"
# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
def summarize(text, max_length, min_length):
if not text.strip():
return "Please enter some text to summarize."
inputs = tokenizer(
text,
return_tensors="pt",
max_length=1024,
truncation=True
).to(device)
summary_ids = model.generate(
inputs["input_ids"],
num_beams=4,
max_length=max_length,
min_length=min_length,
early_stopping=True
)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return summary
with gr.Blocks(title="Advanced BART Summarizer") as demo:
gr.Markdown("# Advanced BART Text Summarizer")
gr.Markdown("Summarization using facebook/bart-large-cnn")
input_text = gr.Textbox(
lines=12,
placeholder="Enter long article or paragraph here..."
)
with gr.Row():
max_len = gr.Slider(50, 300, value=150, label="Max Summary Length")
min_len = gr.Slider(20, 100, value=40, label="Min Summary Length")
output_text = gr.Textbox(
lines=8,
label="Generated Summary"
)
summarize_btn = gr.Button("Summarize")
summarize_btn.click(
summarize,
inputs=[input_text, max_len, min_len],
outputs=output_text
)
demo.launch()