Spaces:
Runtime error
Runtime error
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() |