Job-Tagger / app.py
usman-yello's picture
add slider
79539fd
Raw
History Blame Contribute Delete
859 Bytes
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()