We’ll begin by importing a number of useful libraries.
from datasets import DatasetDict, Dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification,
TrainingArguments, Coach
import consider
import numpy as np
from transformers import DataCollatorWithPadding
Subsequent, we’ll load the coaching dataset. It consists of three,000 text-label pairs with a 70–15–15 train-test-validation cut up. The information are initially from right here (open database license).
dataset_dict = load_dataset("shawhin/phishing-site-classification")
The Transformer library makes it tremendous straightforward to load and adapt pre-trained fashions. Right here’s what that appears like for the BERT mannequin.
# outline pre-trained mannequin path
model_path = "google-bert/bert-base-uncased"# load mannequin tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path)
# load mannequin with binary classification head
id2label = {0: "Secure", 1: "Not Secure"}
label2id = {"Secure": 0, "Not Secure": 1}
mannequin = AutoModelForSequenceClassification.from_pretrained(model_path,
num_labels=2,
id2label=id2label,
label2id=label2id,)
Once we load a mannequin like this, all of the parameters will probably be set as trainable by default. Nevertheless, coaching all 110M parameters will probably be computationally pricey and doubtlessly pointless.
As a substitute, we are able to freeze many of the mannequin parameters and solely prepare the mannequin’s closing layer and classification head.
# freeze all base mannequin parameters
for title, param in mannequin.base_model.named_parameters():
param.requires_grad = False# unfreeze base mannequin pooling layers
for title, param in mannequin.base_model.named_parameters():
if "pooler" in title:
param.requires_grad = True
Subsequent, we might want to preprocess our information. It will include two key operations: tokenizing the URLs (i.e., changing them into integers) and truncating them.
# outline textual content preprocessing
def preprocess_function(examples):
# return tokenized textual content with truncation
return tokenizer(examples["text"], truncation=True)# preprocess all datasets
tokenized_data = dataset_dict.map(preprocess_function, batched=True)
One other vital step is making a information collator that may dynamically pad token sequences in a batch throughout coaching so that they have the identical size. We will do that in a single line of code.
# create information collator
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
As a closing step earlier than coaching, we are able to outline a operate to compute a set of metrics to assist us monitor coaching progress. Right here, we’ll think about mannequin accuracy and AUC.
# load metrics
accuracy = consider.load("accuracy")
auc_score = consider.load("roc_auc")def compute_metrics(eval_pred):
# get predictions
predictions, labels = eval_pred
# apply softmax to get possibilities
possibilities = np.exp(predictions) / np.exp(predictions).sum(-1,
keepdims=True)
# use possibilities of the constructive class for ROC AUC
positive_class_probs = possibilities[:, 1]
# compute auc
auc = np.spherical(auc_score.compute(prediction_scores=positive_class_probs,
references=labels)['roc_auc'],3)
# predict most possible class
predicted_classes = np.argmax(predictions, axis=1)
# compute accuracy
acc = np.spherical(accuracy.compute(predictions=predicted_classes,
references=labels)['accuracy'],3)
return {"Accuracy": acc, "AUC": auc}
Now, we’re able to fine-tune our mannequin. We begin by defining hyperparameters and different coaching arguments.
# hyperparameters
lr = 2e-4
batch_size = 8
num_epochs = 10training_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,
)
Then, we go our coaching arguments right into a coach class and prepare the mannequin.
coach = Coach(
mannequin=mannequin,
args=training_args,
train_dataset=tokenized_data["train"],
eval_dataset=tokenized_data["test"],
tokenizer=tokenizer,
data_collator=data_collator,
compute_metrics=compute_metrics,
)coach.prepare()
The coaching outcomes are proven beneath. We will see that the coaching and validation loss are monotonically lowering whereas the accuracy and AUC enhance with every epoch.
As a closing check, we are able to consider the efficiency of the mannequin on the impartial validation information, i.e., information not used for coaching or setting hyperparameters.
# apply mannequin to validation dataset
predictions = coach.predict(tokenized_data["validation"])# Extract the logits and labels from the predictions object
logits = predictions.predictions
labels = predictions.label_ids
# Use your compute_metrics operate
metrics = compute_metrics((logits, labels))
print(metrics)
# >> {'Accuracy': 0.889, 'AUC': 0.946}
Bonus: Though a 110M parameter mannequin is tiny in comparison with fashionable language fashions, we are able to cut back its computational necessities utilizing mannequin compression strategies. I cowl methods to cut back the reminiscence footprint mannequin by 7X within the article beneath.
Positive-tuning pre-trained fashions is a robust paradigm for growing higher fashions at a decrease price than coaching them from scratch. Right here, we noticed how to do that with BERT utilizing the Hugging Face Transformers library.
Whereas the instance code was for URL classification, it may be readily tailored to different textual content classification duties.
Extra on LLMs 👇
