Spaces:
Sleeping
Sleeping
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:]): | |
# If current token continues the entity of the last one, merge them | |
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: | |
# Otherwise, add the token to the list | |
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() | |