ner_comparation / app.py
Roland Szabo
Rename file
3f4e60f
raw
history blame
1.58 kB
import gradio as gr
import spacy
from transformers import pipeline
import boto3
nlp = spacy.load("en_core_web_sm")
ner_pipeline = pipeline("ner", model="Jean-Baptiste/roberta-large-ner-english", aggregation_strategy="simple", grouped_entities=True)
def greet(model_type, text):
if model_type == "Spacy":
doc = nlp(text)
pos_tokens = []
for token in doc:
if token.ent_type_ != "":
pos_tokens.append((token.text, token.ent_type_))
else:
pos_tokens.append((token.text, None))
return pos_tokens
elif model_type == "Roberta":
output = ner_pipeline(text)
print(output)
return {"text": text, "entities": [
{"word": entity["word"], "entity": entity["entity_group"], "start": entity['start'],
'end': entity['end']}
for entity in output]}
elif model_type == "AWS Comprehend":
client = boto3.client('comprehend')
response = client.detect_entities(
Text=text, LanguageCode='en')
print(response)
return {"text": text, "entities": [{"word": entity["Text"], "entity": entity["Type"], "start": entity['BeginOffset'], 'end': entity['EndOffset']}
for entity in response["Entities"]]}
demo = gr.Interface(fn=greet, inputs=[gr.Radio(["Spacy", "Roberta", "AWS Comprehend"]), "text"],
outputs="highlight", title="Key Matcher",
description=f"Find the best match for a key from the list of ",)
demo.launch()