import gradio as gr
import time
import torch
from transformers import BertTokenizer, BertForSequenceClassification

label_dict = {"Urgency": 0, "Not Dark Pattern": 1, "Scarcity": 2, "Misdirection": 3, "Social Proof": 4, "Obstruction": 5, "Sneaking": 6, "Forced Action": 7}
model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=len(label_dict))
fine_tuned_model_path = "models/finetuned_BERT_5k_epoch_5.model"
model.load_state_dict(torch.load(fine_tuned_model_path, map_location=torch.device('cpu')))
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)

def get_dark_pattern_name(label):
    reverse_label_dict = {v: k for k, v in label_dict.items()}
    return reverse_label_dict[label]

def find_dark_pattern(text_predict):
    encoded_text = tokenizer.encode_plus(
        text_predict,
        add_special_tokens=True,
        return_attention_mask=True,
        pad_to_max_length=True,
        max_length=256,
        return_tensors='pt'
    )

    model.eval()

    with torch.no_grad():
        inputs = {
            'input_ids': encoded_text['input_ids'],
            'attention_mask': encoded_text['attention_mask']
        }
        outputs = model(**inputs)

    predictions = outputs.logits

    probabilities = torch.nn.functional.softmax(predictions, dim=1)
    predicted_label = torch.argmax(probabilities, dim=1).item()

    return get_dark_pattern_name(predicted_label)

def predict(text_to_predict):
    start_time = time.time()
    print("Predicting Dark Pattern...")
    for i in range(10):
        predicted_darkp = find_dark_pattern(text_to_predict)
        time.sleep(0.5)
    end_time = time.time()
    total_time = end_time - start_time
    return predicted_darkp

demo = gr.Interface(fn=predict, inputs="text", outputs="text")
demo.launch(share=True)