from transformers import pipeline
import os
import gradio as gr

os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

ner_pipeline = pipeline(task="ner", model="Tirendaz/roberta-base-NER", framework="pt")


def ner(text):
    output = ner_pipeline(text, aggregation_strategy="simple")
    return {"text": text, "entities": output}


examples = \
    [
        "My name is Tim and I live in California",
        "Ich arbeite bei Google in Berlin",
        "Ali, Ankara'lı mı?"
    ]

demo = (gr.Interface(ner, gr.Textbox(placeholder="Enter sentence here..."), gr.Highlightedtext(), examples))

demo.launch(share=True)