daniB2112's picture
Create app.py
4e9a493 verified
Raw
History Blame Contribute Delete
1.12 kB
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
MODEL_ID = "daniB2112/bart-news-summarizer"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
def summarize(text, max_length=150, min_length=40):
inputs = tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
summary_ids = model.generate(
inputs["input_ids"],
max_length=max_length,
min_length=min_length,
length_penalty=2.0,
num_beams=4,
early_stopping=True
)
return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
demo = gr.Interface(
fn=summarize,
inputs=[
gr.Textbox(lines=10, placeholder="Paste news article here...", label="Article"),
gr.Slider(50, 300, value=150, label="Max Summary Length"),
gr.Slider(10, 100, value=40, label="Min Summary Length"),
],
outputs=gr.Textbox(label="Summary"),
title="BART News Summarizer",
description="Summarize news articles using daniB2112/bart-news-summarizer"
)
demo.launch()