Model Components

Model and framework component design

This diagram distinguishes framework classes from source-defined helper functions; the script itself does not declare custom Python model classes.

classDiagram
    class AutoTokenizer {
      +from_pretrained(model_name)
      +__call__(text, truncation, max_length)
      +save_pretrained(path)
    }
    class DatasetDict {
      +train
      +validation
    }
    class DataCollatorWithPadding {
      +__call__(features)
    }
    class DistilBertForSequenceClassification {
      +distilbert
      +pre_classifier
      +classifier
      +forward(input_ids, attention_mask)
    }
    class Trainer {
      +train()
      +evaluate()
      +predict(dataset)
      +save_model(path)
    }
    class TrainingArguments {
      +output_dir
      +learning_rate
      +num_train_epochs
    }
    DatasetDict --> AutoTokenizer : tokenized by
    AutoTokenizer --> DataCollatorWithPadding : vocabulary and padding
    Trainer --> DatasetDict : iterates
    Trainer --> DataCollatorWithPadding : collates
    Trainer --> TrainingArguments : config
    Trainer --> DistilBertForSequenceClassification : optimizes

Loading

DistilBERT classification head

The sequence-classification wrapper includes the distilbert encoder, a pre_classifier transformation, and a final classifier layer. All components are trainable during full fine-tuning. In the comparison, the code sets requires_grad=False on the encoder parameters and trains the head.

The architecture uses a pretrained encoder and a randomly initialized four-class head, not a model trained from random initialization.