Spaces:
Running
Running
| 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() | |