adapter_train.py

import numpy as np
from transformers import AutoConfig, AutoModelWithHeads
from transformers import TrainingArguments, Trainer, EvalPrediction


config = AutoConfig.from_pretrained(
    "distilbert-base-uncased",
    num_labels=2,
)
model = AutoModelWithHeads.from_pretrained(
    "distilbert-base-uncased",
    config=config,
)

# Add a new adapter
model.add_adapter("imdb_review")
# Add a matching classification head
model.add_classification_head(
    "imdb_review",
    num_labels=2,
    id2label={ 0:"Negative" , 1: "Positive"}
  )
# Activate the adapter
model.train_adapter("imdb_review")


training_args = TrainingArguments(
    learning_rate=1e-4,
    num_train_epochs=6,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    logging_steps=200,
    output_dir="./training_output",
    overwrite_output_dir=True,
    # The next line is important to ensure the dataset labels are properly passed to the model
    remove_unused_columns=False,
)

def compute_accuracy(p: EvalPrediction):
  preds = np.argmax(p.predictions, axis=1)
  return {"acc": (preds == p.label_ids).mean()}

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    compute_metrics=compute_accuracy)

trainer.train()