import gradio as gr
from transformers import pipeline

get_completion = pipeline("ner", model="tororoin/longformer-8bitadam-2048-main")

def ner(input):
    output = get_completion(input)
    return {"text": input, "entities": output}

demo = gr.Interface(fn=ner,
                    inputs=[gr.Textbox(label="Text to find entities", lines=4)],
                    outputs=[gr.HighlightedText(label="Text with entities")],
                    title="Detecting PII",
                    description="Find PII in text data! A small demo",
                    allow_flagging="never",
                    examples=["My name is Andrew and I live in at Baker Str"])
demo.launch()