!pip install datasets evaluate peft torch --upgrade

from datasets import DatasetDict, Dataset, load_dataset
from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification, DataCollatorWithPadding, TrainingArguments, Trainer
from peft import PeftModel, PeftConfig, get_peft_model, LoraConfig

import evaluate
import torch
import numpy as np

model_checkpoint = "distilbert-base-uncased"

id2label = {0: "Negative", 1: "Positive"}
label2id = {"Negative": 0, "Positive": 1}

model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, 
                                                           num_labels=2, 
                                                           id2label=id2label, 
                                                           label2id=label2id,).to("mps")

dataset_dict = load_dataset("shawhin/imdb-truncated")

tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True)

def tokenize_function(examples):
    text = examples["text"]
    tokenizer.truncation_side = "left"
    tokenized_inputs = tokenizer(text, truncation=True, return_tensors="np", max_length=512)
    return tokenized_inputs

if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({"pad_token": "[PAD]"})
    model.resize_token_embeddings(len(tokenizer))

tokenized_dataset = dataset_dict.map(tokenize_function, batched=True)

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return {"accuracy": accuracy.compute(predictions=predictions, references=labels)}

text_list = ["I love this movie", "I hate this movie", "Not a fan, don't recommend", "This one is a pass"]

print("untrained model predictions:")
print("-"*20)

for text in text_list:
    inputs = tokenizer.encode(text, return_tensors="pt").to("mps")
    logits = model(inputs).logits
    predictions = torch.argmax(logits)
    print(f"{text} -> {id2label[predictions.tolist()]}")

# Now let's fine-tune the model with LoRA
peft_config = LoraConfig(task_type="SEQ_CLS", r=4, lora_alpha=32, target_modules=["q_lin"], lora_dropout=0.01)

model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

lr = 1e-3
batch_size = 4
num_epochs = 1

training_args = TrainingArguments(
    output_dir=model_checkpoint + "-lora-text-classification",
    learning_rate=lr,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=num_epochs,
    weight_decay=0.01,
    save_strategy="epoch",
    evaluation_strategy="epoch",
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["validation"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

print("trained model predictions:")
print("-"*20)

for text in text_list:
    inputs = tokenizer.encode(text, return_tensors="pt").to("mps")
    logits = model(inputs).logits
    predictions = torch.argmax(logits)
    print(f"{text} -> {id2label[predictions.tolist()]}")

Some weights of DistilBertForSequenceClassification were not initialized from the model checkpoint at distilbert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight', 'pre_classifier.bias', 'pre_classifier.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

untrained model predictions:
--------------------
I love this movie -> Negative
I hate this movie -> Negative
Not a fan, don't recommend -> Negative
This one is a pass -> Negative
trainable params: 628,994 || all params: 67,584,004 || trainable%: 0.9307

/Users/LSoica/work/AI/blog/.venv/lib/python3.12/site-packages/transformers/training_args.py:1568: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead
  warnings.warn(
/var/folders/59/c32_bthx48jd9m2ym5m3tnpw0000j7/T/ipykernel_19744/1863821626.py:78: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `Trainer.__init__`. Use `processing_class` instead.
  trainer = Trainer(
                                                 
100%|██████████| 250/250 [01:47<00:00,  3.20it/s]

{'eval_loss': 0.30286046862602234, 'eval_accuracy': {'accuracy': 0.891}, 'eval_runtime': 34.387, 'eval_samples_per_second': 29.081, 'eval_steps_per_second': 7.27, 'epoch': 1.0}

100%|██████████| 250/250 [01:48<00:00,  2.31it/s]

{'train_runtime': 108.427, 'train_samples_per_second': 9.223, 'train_steps_per_second': 2.306, 'train_loss': 0.4968539123535156, 'epoch': 1.0}
trained model predictions:
--------------------
I love this movie -> Positive
I hate this movie -> Negative
Not a fan, don't recommend -> Negative
This one is a pass -> Negative

import pandas as pd
from datasets import DatasetDict, Dataset, load_dataset

from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer

import evaluate
import numpy as np
from transformers import DataCollatorWithPadding
from torch.utils.data import DataLoader
import torch

/Users/LSoica/work/AI/blog/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

dataset_dict = load_dataset("shawhin/phishing-site-classification")

# Load model directly
model_path = "google-bert/bert-base-uncased"

tokenizer = AutoTokenizer.from_pretrained(model_path)

id2label = {0: "Safe", 1: "Not Safe"}
label2id = {"Safe": 0, "Not Safe": 1}
model = AutoModelForSequenceClassification.from_pretrained(model_path, 
                                                           num_labels=2, 
                                                           id2label=id2label, 
                                                           label2id=label2id,).to("mps")

Some weights of BertForSequenceClassification were not initialized from the model checkpoint at google-bert/bert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

# freeze base model parameters
for name, param in model.base_model.named_parameters():
    param.requires_grad = False

# unfreeze base model pooling layers
for name, param in model.base_model.named_parameters():
    if "pooler" in name:
        param.requires_grad = True

# define text preprocessing
def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True)

tokenized_data = dataset_dict.map(preprocess_function, batched=True)

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

# load metrics
accuracy = evaluate.load("accuracy")
auc_score = evaluate.load("roc_auc")

def compute_metrics(eval_pred):
    # get predictions
    predictions, labels = eval_pred
    
    # apply softmax to get probabilities
    probabilities = np.exp(predictions) / np.exp(predictions).sum(-1, keepdims=True)
    # use probabilities of the positive class for ROC AUC
    positive_class_probs = probabilities[:, 1]
    # compute auc
    auc = np.round(auc_score.compute(prediction_scores=positive_class_probs, references=labels)['roc_auc'],3)
    
    # predict most probable class
    predicted_classes = np.argmax(predictions, axis=1)
    # compute accuracy
    acc = np.round(accuracy.compute(predictions=predicted_classes, references=labels)['accuracy'],3)
    
    return {"Accuracy": acc, "AUC": auc}

# hyperparameters
lr = 2e-4
batch_size = 8
num_epochs = 10

training_args = TrainingArguments(
    output_dir="bert-phishing-classifier_teacher",
    learning_rate=lr,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=num_epochs,
    logging_strategy="epoch",
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_data["train"],
    eval_dataset=tokenized_data["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

/var/folders/59/c32_bthx48jd9m2ym5m3tnpw0000j7/T/ipykernel_69699/2732273287.py:1: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `Trainer.__init__`. Use `processing_class` instead.
  trainer = Trainer(
 10%|█         | 263/2630 [00:21<03:28, 11.33it/s]

{'loss': 0.4922, 'grad_norm': 1.1272960901260376, 'learning_rate': 0.00018, 'epoch': 1.0}

 10%|█         | 263/2630 [00:26<03:28, 11.33it/s]

{'eval_loss': 0.4166680574417114, 'eval_Accuracy': 0.787, 'eval_AUC': 0.912, 'eval_runtime': 4.9842, 'eval_samples_per_second': 90.285, 'eval_steps_per_second': 11.436, 'epoch': 1.0}

 20%|██        | 526/2630 [00:47<02:45, 12.70it/s]

{'loss': 0.3905, 'grad_norm': 2.191509485244751, 'learning_rate': 0.00016, 'epoch': 2.0}

 20%|██        | 526/2630 [00:52<02:45, 12.70it/s]

{'eval_loss': 0.3610435724258423, 'eval_Accuracy': 0.813, 'eval_AUC': 0.931, 'eval_runtime': 5.0556, 'eval_samples_per_second': 89.01, 'eval_steps_per_second': 11.275, 'epoch': 2.0}

 30%|███       | 789/2630 [01:14<02:28, 12.42it/s]

{'loss': 0.3854, 'grad_norm': 0.820168137550354, 'learning_rate': 0.00014, 'epoch': 3.0}

# apply model to validation dataset
predictions = trainer.predict(tokenized_data["test"])

# Extract the logits and labels from the predictions object
logits = predictions.predictions
labels = predictions.label_ids

# Use your compute_metrics function
metrics = compute_metrics((logits, labels))
print(metrics)

100%|██████████| 57/57 [00:04<00:00, 12.49it/s]

{'Accuracy': np.float64(0.864), 'AUC': np.float64(0.951)}

## Infer on new data

urls = [
  "google.com",
  "yahoo.com",
  "www.yahoo.com",
  "https://www.yahoo.com",
  "https://microsoft.user-account.online/14e84edd29dc7302?l=861",
  "users11.jabry.com/reaseo/Aolupdate.htm",
  "www.allandmedia.com/opencart/system/Cielo/index.html",
  "mrterabit.com/remax/index.php",
  "phishing.org"
]
for url in urls:
  inputs = tokenizer(url, return_tensors="pt", truncation=True, padding=True)

  with torch.no_grad():
    outputs = trainer.model.forward(**inputs)
    probabilities = np.exp(outputs.logits) / np.exp(outputs.logits).sum(-1, keepdims=True)
    print("Safe" if probabilities[0][0].item() > 0.9 else "Not Safe", url)

/var/folders/59/c32_bthx48jd9m2ym5m3tnpw0000j7/T/ipykernel_69699/3933651874.py:17: DeprecationWarning: __array_wrap__ must accept context and return_scalar arguments (positionally) in the future. (Deprecated NumPy 2.0)
  probabilities = np.exp(outputs.logits) / np.exp(outputs.logits).sum(-1, keepdims=True)

Safe google.com
Safe yahoo.com
Safe www.yahoo.com
Safe https://www.yahoo.com
Not Safe https://microsoft.user-account.online/14e84edd29dc7302?l=861
Not Safe users11.jabry.com/reaseo/Aolupdate.htm
Not Safe www.allandmedia.com/opencart/system/Cielo/index.html
Not Safe mrterabit.com/remax/index.php
Not Safe phishing.org

BERT¶

Training a BERT model¶

Fine tuning¶

Fine tuning BERT with LoRA¶

Fine tuning for phishing link detection¶

Import¶

Train¶

Evaluate¶

References¶