Skip to content

DistilBERT Workflow

Sudhanshu edited this page Nov 11, 2026 · 1 revision

DistilBERT Training Workflow

This is the executable-in-principle classification pipeline in DistilBERT_Colab_Fixed_2026.py. It starts model training immediately when the script is executed; there is no CLI argument parser or API endpoint.

sequenceDiagram
    autonumber
    participant User as Python/Jupyter user
    participant Data as Dataset + split
    participant Tok as AutoTokenizer
    participant Model as DistilBERT classifier
    participant Tr as HF Trainer
    participant Disk as Local artifacts
    User->>Data: Create 30 labeled examples; stratified 80/20 split
    Data->>Tok: map(tokenize_function)
    Tok-->>Data: token IDs and attention masks
    User->>Model: from_pretrained(distilbert-base-uncased, 4 labels)
    User->>Tr: create_trainer(model, args, datasets, collator, metrics)
    loop 4 epochs; batches of 8
      Tr->>Model: forward(batch)
      Model-->>Tr: logits / loss
      Tr->>Model: backward + optimizer step
      opt epoch evaluation
        Tr->>Model: validation forward
        Model-->>Tr: predictions
      end
    end
    Tr->>Disk: save_model(best_model)
    User->>Tok: save_pretrained(best_model)
    User->>Tr: evaluate(), predict()
    Tr-->>User: accuracy, weighted F1, report/confusion matrix
Loading

Stages

  1. Set seed 42 and detect CUDA/BF16/FP16 availability.
  2. Build a 30-example educational set across billing, technical, account and feature requests.
  3. Stratify to 24 train and 6 validation examples.
  4. Download the distilbert-base-uncased tokenizer, tokenize to maximum sequence length 128, and dynamically pad batches.
  5. Instantiate AutoModelForSequenceClassification with four labels.
  6. Run four epochs of full fine-tuning through Trainer.
  7. Save the selected model and tokenizer under support_classifier/best_model/.
  8. Evaluate with accuracy/F1, classification report and confusion matrix; run five sample queries.
  9. Build a second DistilBERT classifier and freeze only the DistilBERT encoder parameters; train the head for six epochs and compare metrics.

See Training-Configuration, Evaluation-and-Inference, and Frozen-Backbone-Comparison.

Clone this wiki locally