Spaces:
Sleeping
Sleeping
File size: 1,485 Bytes
e770f7d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import gradio as gr
from simpletransformers.ner import NERModel
import string
labels = ["O",
"B-FOOD_QUANTITY",
"B-FOOD_SIZE",
"B-FOOD",
"I-FOOD",
"B-FOOD_INGREDIENTS",
"I-FOOD_INGREDIENTS",
"B-DRINK_SIZE",
"B-DRINK_QUANTITY",
"B-DRINK",
"B-PAYMENT",
"I-PAYMENT",
"B-DELIVERY_ADDRESS",
"I-DRINK_SIZE",
"I-DRINK",
"I-FOOD_SIZE",
"I-ING",
"I-DELIVERY_ADDRESS"]
model = NERModel(
"albert",
"bgk/lodosalberttr", labels=labels,
use_cuda=False,
ignore_mismatched_sizes=True
)
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' ]
def ner(text):
trans_table = text.maketrans('', '', string.punctuation)
text = text.translate(trans_table)
text=text.lower()
prediction, model_output = model.predict([text])
entities = prediction
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)
entities = []
for sublist in filtered_output:
for d in sublist:
for k, v in d.items():
label = k.split("-")[1]
entities.extend([(label, v)])
return entities # prediction
demo = gr.Interface(ner,
gr.Textbox(placeholder="Enter sentence here..."),
gr.HighlightedText(),
examples=examples)
demo.launch()
|