chinhon commited on
Commit
e37d0f5
1 Parent(s): 2c642a9

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +64 -0
app.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import re
3
+
4
+ from transformers import (
5
+ BertTokenizer,
6
+ BartForConditionalGeneration,
7
+ pipeline
8
+ )
9
+
10
+ #DEFINE TEXT CLEANING FUNCTION
11
+ def clean_text(text):
12
+ text = re.sub(r"http\S+", "", text)
13
+ text = re.sub(r"ADVERTISEMENT", " ", text)
14
+ text = re.sub(r"ADVERTISING", " ", text)
15
+ text = re.sub(r"\n", " ", text)
16
+ text = re.sub(r"\n\n", " ", text)
17
+ text = re.sub(r"\t", " ", text)
18
+ text = text.strip(" ")
19
+ text = re.sub(
20
+ " +", " ", text
21
+ ).strip() # get rid of multiple spaces and replace with a single
22
+ return text
23
+
24
+ #Define headlne writer functon
25
+ model_name = "chinhon/bart-large-chinese-cnhdwriter"
26
+
27
+ tokenizer = BertTokenizer.from_pretrained(model_name, model_max_length=512)
28
+
29
+ model = BartForConditionalGeneration.from_pretrained(model_name)
30
+
31
+ text2text_generator = pipeline(
32
+ "text2text-generation",
33
+ model=model,
34
+ tokenizer=tokenizer,
35
+ truncation=True
36
+ )
37
+
38
+ def cn_text(text):
39
+ input_text = clean_text(text)
40
+
41
+ prediction = text2text_generator(
42
+ input_text,
43
+ max_length=128,
44
+ length_penalty=50.,
45
+ )
46
+
47
+ pred_text = [x.get("generated_text") for x in prediction]
48
+
49
+ return pred_text[0]
50
+
51
+
52
+ #3 Define Gradio UI
53
+ gradio_ui = gr.Interface(
54
+ fn=cn_text,
55
+ title="Generate Chinese News Headlines with AI",
56
+ description="Too busy or tired to write a headline for your Chinese news story? Try this instead.",
57
+ inputs=gr.inputs.Textbox(
58
+ lines=20, label="Paste Chinese text here"
59
+ ),
60
+ outputs=gr.outputs.Textbox(label="Suggested Headline"),
61
+ theme="huggingface",
62
+ )
63
+
64
+ gradio_ui.launch(enable_queue=True)