JobTagger-V2 / app.py
usman-yello's picture
init
adf8279
Raw
History Blame Contribute Delete
1.69 kB
import re
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=device)
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSequenceClassification.from_pretrained("yello-co/ys-bert-base-uncased-summary-multi-label")
def remove_html(text):
html_pattern = re.compile(r'<.*?>|&([a-z0-9]+|#[0-9]{1,6}|#x[0-9a-f]{1,6});')
return re.sub(html_pattern, '', text)
def remove_markdown(text):
markdown_pattern = re.compile(r'(\*\*|__|`|~~|#|\*|_|-|\+|!|\[.*?\]\(.*?\))')
return re.sub(markdown_pattern, '', text)
def clean_text(input_text):
html_removed = remove_html(input_text)
html_markdown_removed = remove_markdown(html_removed)
return html_markdown_removed
def tag_job_description(position_title, job_description):
job_description = clean_text(job_description)
summary = summarizer(job_description, truncation=True)[0]["summary_text"]
text = position_title + ": " + summary
inputs = tokenizer(text, truncation=True, max_length=512, return_tensors="pt")
outputs = model(**inputs)
probabilities = torch.sigmoid(torch.Tensor(outputs.logits))
predictions = {model.config.id2label[idx]: p for idx, p in enumerate(probabilities[0])}
return predictions
job_tagger = gr.Interface(fn=tag_job_description, inputs=["text", "text"],
outputs=gr.Label(label="Predicted Tags"))
job_tagger.launch()