4darsh-Dev/dark_pattern_detector_app_v2

1import gradio as gr2import time3import torch4from transformers import BertTokenizer, BertForSequenceClassification5 6label_dict = {"Urgency": 0, "Not Dark Pattern": 1, "Scarcity": 2, "Misdirection": 3, "Social Proof": 4, "Obstruction": 5, "Sneaking": 6, "Forced Action": 7}7model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=len(label_dict))8fine_tuned_model_path = "models/finetuned_BERT_5k_epoch_5.model"9model.load_state_dict(torch.load(fine_tuned_model_path, map_location=torch.device('cpu')))10tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)11 12def get_dark_pattern_name(label):13    reverse_label_dict = {v: k for k, v in label_dict.items()}14    return reverse_label_dict[label]15 16def find_dark_pattern(text_predict):17    encoded_text = tokenizer.encode_plus(18        text_predict,19        add_special_tokens=True,20        return_attention_mask=True,21        pad_to_max_length=True,22        max_length=256,23        return_tensors='pt'24    )25 26    model.eval()27 28    with torch.no_grad():29        inputs = {30            'input_ids': encoded_text['input_ids'],31            'attention_mask': encoded_text['attention_mask']32        }33        outputs = model(**inputs)34 35    predictions = outputs.logits36 37    probabilities = torch.nn.functional.softmax(predictions, dim=1)38    predicted_label = torch.argmax(probabilities, dim=1).item()39 40    return get_dark_pattern_name(predicted_label)41 42def predict(text_to_predict):43    start_time = time.time()44    print("Predicting Dark Pattern...")45    for i in range(10):46        predicted_darkp = find_dark_pattern(text_to_predict)47        time.sleep(0.5)48    end_time = time.time()49    total_time = end_time - start_time50    return predicted_darkp51 52demo = gr.Interface(fn=predict, inputs="text", outputs="text")53demo.launch(share=True)54 55 56