liamvbetts's picture
random button
e673826
raw
history blame
1.32 kB
import gradio as gr
import random
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from datasets import load_dataset
tokenizer = AutoTokenizer.from_pretrained("liamvbetts/bart-large-cnn-v4")
model = AutoModelForSeq2SeqLM.from_pretrained("liamvbetts/bart-large-cnn-v4")
dataset = load_dataset("cnn_dailymail", "3.0.0")
def summarize(article):
inputs = tokenizer(article, return_tensors="pt").input_ids
outputs = model.generate(inputs, max_new_tokens=128, do_sample=False)
summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
return summary
def get_random_article():
random.seed()
val_example = dataset["validation"].shuffle().select(range(1))
val_article = val_example['article'][0][:512]
return val_article
# Create Gradio interface
input_text = gr.Textbox(lines=10, label="Input Text")
output_text = gr.Textbox(label="Summary")
random_article_button = gr.Button("Load Random Article")
def update_input_text():
return get_random_article()
gr.Interface(
fn=summarize,
inputs=[input_text, gr.components.Button("Load Random Article").click(update_input_text, [], input_text)],
outputs=output_text,
title="News Summary App",
description="Enter a news text and get its summary, or load a random article from the validation set."
).launch()