Spaces:
Runtime error
Runtime error
| 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"http\S+", "", text) | |
| text = re.sub(r"\n", " ", text) | |
| text = re.sub(r"\n\n", " ", text) | |
| text = re.sub(r"\t", " ", text) | |
| text = text.strip(" ") | |
| text = re.sub( | |
| " +", " ", text | |
| ).strip() # get rid of multiple spaces and replace with a single | |
| return text | |
| model_name = "chinhon/headline_writer" | |
| def headline_writer(text): | |
| input_text = clean_text(text) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
| with tokenizer.as_target_tokenizer(): | |
| batch = tokenizer( | |
| input_text, | |
| truncation=True, | |
| max_length=1024, | |
| padding="longest", | |
| return_tensors="pt", | |
| ) | |
| raw_write = model.generate(**batch) | |
| headline = tokenizer.batch_decode( | |
| raw_write, skip_special_tokens=True, min_length=200, length_penalty=50.5 | |
| ) | |
| return headline[0] | |
| gradio_ui = gr.Interface( | |
| fn=headline_writer, | |
| title="Generate News Headlines with AI", | |
| description="Too busy or tired to write a headline? Try this instead.", | |
| inputs=gr.inputs.Textbox( | |
| lines=20, label="Paste the first few paras of your news story here" | |
| ), | |
| outputs=gr.outputs.Textbox(label="Suggested Headline"), | |
| theme="darkdefault" | |
| ) | |
| gradio_ui.launch() | |