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()