import gradio as gr
from transformers import pipeline
import os
from dotenv import load_dotenv
load_dotenv()

share = os.getenv("SHARE", False)

pipe = pipeline("token-classification", model="bnsapa/cybersecurity-ner")
def tag(input):
    k = pipe(input)
    for i in k:
        input = input.replace(i["word"], i["word"] + "[" + i["entity"] + " Score: " + str(i["score"]) +"]")
    return input

iface = gr.Interface(fn=tag, inputs="text", outputs="text", title="Cybersecurity NER", description="Named Entity Recognition for Cybersecurity")

if __name__ == "__main__":
    if share:
        server = "0.0.0.0"
    else:
        server = "127.0.0.1"
    iface.launch(server_name = server)