HugoHE commited on
Commit
c82548e
1 Parent(s): 5a0f4cc

Create new file

Browse files
Files changed (1) hide show
  1. app.py +35 -0
app.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import transformers
3
+ from transformers import pipeline, BertForSequenceClassification, BertTokenizer
4
+
5
+ def classify(input_text):
6
+ tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
7
+ model = BertForSequenceClassification.from_pretrained('./bert-cls')
8
+ classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, top_k=3)
9
+ class_dict = {0:'story',
10
+ 1:'culture',
11
+ 2:'entertainment',
12
+ 3:'sports',
13
+ 4:'finance',
14
+ 6:'house',
15
+ 7:'car',
16
+ 8:'edu',
17
+ 9:'tech',
18
+ 10:'military',
19
+ 12:'travel',
20
+ 13:'world',
21
+ 14:'stock',
22
+ 15:'argriculture',
23
+ 16:'game'}
24
+ output = classifier([input_text])
25
+ idx_list = [output[0][i]['label'].split('_')[1] for i in range(len(output[0]))]
26
+ label_list = [class_dict[int(idx)] for idx in idx_list]
27
+ score_list = [output[0][i]['score'] for i in range(len(output[0]))]
28
+ return dict(zip(label_list, score_list))
29
+ label = gr.Label()
30
+ iface = gr.Interface(fn = classify,
31
+ inputs = "text",
32
+ outputs = label,
33
+ title = 'chinese news classification')
34
+
35
+ iface.launch(inline = False)