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()