|
import gradio as gr |
|
from transformers import pipeline |
|
|
|
ner = pipeline('ner') |
|
|
|
def merge_tokens(tokens): |
|
merged_tokens = [] |
|
for token in tokens: |
|
if merged_tokens and token['entity'].startswith('I-') and merged_tokens[-1]['entity'].endswith(token['entity'][2:]): |
|
|
|
last_token = merged_tokens[-1] |
|
last_token['word'] += token['word'].replace('##', '') |
|
last_token['end'] = token['end'] |
|
last_token['score'] = (last_token['score'] + token['score']) / 2 |
|
else: |
|
|
|
merged_tokens.append(token) |
|
|
|
return merged_tokens |
|
|
|
def named(input): |
|
output = ner(input) |
|
merged_word = merge_tokens(output) |
|
return {'text': input, 'entities': merged_word} |
|
|
|
a = gr.Interface(fn=named, |
|
inputs=[gr.Textbox(label="Text input", lines= 2)], |
|
outputs=[gr.HighlightedText(label='Text with entities')], |
|
title='Named Entity Recognition', examples=["My name is Andrew, I'm building DeeplearningAI and I live in California", "My name is Poli, I live in Vienna and work at HuggingFace"]) |
|
a.launch() |
|
|