import gradio as gr
from transformers import pipeline

token_skill_classifier = pipeline(model="jjzha/jobbert_skill_extraction", aggregation_strategy="simple")
token_knowledge_classifier = pipeline(model="jjzha/jobbert_knowledge_extraction", aggregation_strategy="simple")


examples = [
        "Knowing Python is a plus.",
        ]


def ner(text):
    output_skills = token_skill_classifier(text)
    for result in output_skills:
        if result.get("entity_group"):
            tag = result["entity_group"]
            result["entity"] = tag + "-Skill"
            del result["entity_group"]

    output_knowledge = token_knowledge_classifier(text)
    for result in output_knowledge:
        if result.get("entity_group"):
            tag = result["entity_group"]
            result["entity"] = tag + "-Knowledge"
            del result["entity_group"]

    output = output_skills + output_knowledge
    return {"text": text, "entities": output}


demo = gr.Interface(fn=ner,
                    inputs=gr.Textbox(placeholder="Enter sentence here..."),
                    outputs=gr.HighlightedText(),
                    examples=examples)

demo.launch()