import re
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("usman-yello/job-tagging")


def tag_job_description(job_description, num_tags):
    job_description = re.sub(r'[^a-zA-Z]', ' ', job_description)
    inputs = tokenizer(job_description, truncation=True, max_length=512, return_tensors="pt")
    outputs = model(**inputs)
    topk_values, topk_indices = torch.topk(outputs.logits, num_tags)
    predicted_tags = [model.config.id2label[i.item()] for i in topk_indices[0]]
    return predicted_tags


job_tagger = gr.Interface(fn=tag_job_description, inputs=["text", gr.Slider(value=1, minimum=1, maximum=10, step=1)], outputs="text")
job_tagger.launch()