|
import gradio as gr |
|
|
|
from transformers import pipeline |
|
get_completion = pipeline("ner", model="dslim/bert-base-NER") |
|
|
|
def merge_tokens(tokens): |
|
merged_tokens = [] |
|
for token in tokens: |
|
if (merged_tokens and token['word'].startswith('##')) or (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 |
|
merged_tokens[-1] = last_token |
|
|
|
else: |
|
|
|
merged_tokens.append(token) |
|
|
|
return merged_tokens |
|
|
|
def ner_merged(input): |
|
output = get_completion(input) |
|
merged_tokens = merge_tokens(output) |
|
return {"text": input, "entities": merged_tokens} |
|
|
|
demo = gr.Interface(fn=ner_merged, |
|
|
|
|
|
|
|
|
|
inputs=[gr.Textbox(label="Type or paste text to find Named Entities or even select and submit below examples", lines=2)], |
|
outputs=[gr.HighlightedText(label="Text with Named Entities identified")], |
|
title="Named Entity Recognition test and demo app by Srinivas.V ", |
|
description="Find entities", |
|
allow_flagging="never", |
|
examples=["My name is Srinivas and I live in Dubai, United Arab Emirates. I love DeepLearningAI", |
|
"I am a Data Scientist and I am a citizen of Bharat"]) |
|
demo.launch() |
