bgk commited on
Commit
e770f7d
·
1 Parent(s): e02d7b6

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +49 -0
app.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from simpletransformers.ner import NERModel
3
+ import string
4
+ labels = ["O",
5
+ "B-FOOD_QUANTITY",
6
+ "B-FOOD_SIZE",
7
+ "B-FOOD",
8
+ "I-FOOD",
9
+ "B-FOOD_INGREDIENTS",
10
+ "I-FOOD_INGREDIENTS",
11
+ "B-DRINK_SIZE",
12
+ "B-DRINK_QUANTITY",
13
+ "B-DRINK",
14
+ "B-PAYMENT",
15
+ "I-PAYMENT",
16
+ "B-DELIVERY_ADDRESS",
17
+ "I-DRINK_SIZE",
18
+ "I-DRINK",
19
+ "I-FOOD_SIZE",
20
+ "I-ING",
21
+ "I-DELIVERY_ADDRESS"]
22
+ model = NERModel(
23
+ "albert",
24
+ "bgk/lodosalberttr", labels=labels,
25
+ use_cuda=False,
26
+ ignore_mismatched_sizes=True
27
+ )
28
+ examples=['I want two hamburgers and one sprite and one milkshake send it to my workplace' , ' I want to order two large pizzas, two medium coke, send it to my home, I will pay with cash' ]
29
+ def ner(text):
30
+ trans_table = text.maketrans('', '', string.punctuation)
31
+ text = text.translate(trans_table)
32
+ text=text.lower()
33
+
34
+ prediction, model_output = model.predict([text])
35
+ entities = prediction
36
+ filtered_output = (({v: k} for d in sublist for k, v in d.items() if (v.startswith("B-") or v.startswith("I-"))) for sublist in prediction)
37
+ entities = []
38
+ for sublist in filtered_output:
39
+ for d in sublist:
40
+ for k, v in d.items():
41
+ label = k.split("-")[1]
42
+ entities.extend([(label, v)])
43
+
44
+ return entities # prediction
45
+ demo = gr.Interface(ner,
46
+ gr.Textbox(placeholder="Enter sentence here..."),
47
+ gr.HighlightedText(),
48
+ examples=examples)
49
+ demo.launch()