Spaces:
Running
Running
File size: 2,591 Bytes
4f40928 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import gradio as gr
import re
from transformers import (
AutoTokenizer,
AutoModelForSeq2SeqLM,
)
def clean_text(text):
text = text.encode("ascii", errors="ignore").decode(
"ascii"
) # remove non-ascii, Chinese characters
text = re.sub(r"\n", " ", text)
text = re.sub(r"\n\n", " ", text)
text = re.sub(r"\t", " ", text)
text = re.sub(r"ADVERTISEMENT", " ", text)
text = re.sub(r"ADVERTISING", " ", text)
text = text.strip(" ")
text = re.sub(
" +", " ", text
).strip() # get rid of multiple spaces and replace with a single
return text
def newsroom_hd(hdchoice, text):
if hdchoice == "Singapore News":
modchoice = "chinhon/pegasus-newsroom-headline_writer_oct22"
elif hdchoice == "International News":
modchoice = "chinhon/pegasus-newsroom_wires_hdwriter42k"
elif hdchoice == "Commentary":
modchoice = "chinhon/bart-large-commentaries_hdwriter"
elif hdchoice == "News in Malay":
modchoice = "chinhon/pegasus-newsroom-malay_headlines"
else:
modchoice = "chinhon/pegasus-newsroom-headline_writer_oct22"
input_text = clean_text(text)
tokenizer = AutoTokenizer.from_pretrained(modchoice)
model = AutoModelForSeq2SeqLM.from_pretrained(modchoice)
with tokenizer.as_target_tokenizer():
batch = tokenizer(
input_text, truncation=True, padding="longest", return_tensors="pt"
)
raw = model.generate(**batch)
headline = tokenizer.batch_decode(raw, skip_special_tokens=True)
return headline[0]
gradio_ui = gr.Interface(
fn=newsroom_hd,
title="Generate Newsroom Headlines With AI",
description="**How to use**: Select the type of headline you wish to generate, paste in a relevant amount of text, and click submit.",
article="**Note**: Paste in as much text as you think necessary, though there's an automatic cut-off of about 500 words for some models and about 850 words for others. If you copy-and-paste directly from a website, take note to remove unrelated text such as those for advertisements and recommended links.",
inputs=[
gr.Dropdown(
label="Select the type of headlines you would like to generate",
choices=[
"Singapore News",
"International News",
"Commentary",
"News in Malay",
],
value="Singapore News",
),
gr.Textbox(label="Paste text here"),
],
outputs=gr.Textbox(label="Suggested Headline"),
)
gradio_ui.queue().launch()
|