from transformers import BertTokenizer,AutoModel
from transformers.pipelines import pipeline
from register import register
import gradio as gr
from huggingface_hub import login
import os
register()
login(os.environ["HF_Token"])
tokenizer = BertTokenizer.from_pretrained("minskiter/resume_token_classification",use_auth_token=True)
model = AutoModel.from_pretrained("minskiter/resume_token_classification",use_auth_token=True)
ner_predictor = pipeline(
    "ner_predictor", 
    model=model, 
    tokenizer=tokenizer,
    device="cpu"
)

def ner_predictor_gradio(input):
    return ner_predictor(input)

demo = gr.Interface(fn=ner_predictor_gradio, inputs="text", outputs="text")
demo.launch()