# heytamjid/babylm25 architecture

> Architecture diagram of the GitHub repository heytamjid/babylm25, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-27). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/heytamjid/babylm25
- Repository: https://github.com/heytamjid/babylm25

## Overview

1. **Repository purpose**
   - This is a BabyLM reproduction repo centered on **pretraining small language models** and **evaluating them on BLiMP-style minimal pairs**.
   - The code supports both standard **masked language modeling (MLM)** and **ELECTRA-style replaced-token detection (RTD)**, plus an optional hybrid path that adds an MLM head on top of an ELECTRA discriminator.

2. **Main execution path**
   - The primary entrypoint is **`train_mask.py`**, which orchestrates training, validation, checkpointing, and optional logging.
   - **`train_mask_basic.py`** appears to be a simpler or older training variant, likely for baseline MLM runs without the more advanced masking/ELECTRA options.

3. **Model architecture layer**
   - The default backbone is a **DeBERTa-v3-style masked LM setup** unless overridden.
   - **`electra_model.py`** contains the custom architecture used for the ELECTRA-derived MLM head (`ElectraHeadForMaskedLM` per the README), combining a pretrained discriminator backbone with a newly trained prediction head.
   - The repository also appears to rely on Hugging Face model loading conventions, including `AutoModel` / `AutoModelForMaskedLM` and `trust_remote_code` for custom classes.

4. **Training subsystem**
   - Training consumes:
     - a training text corpus,
     - a validation corpus,
     - a tokenizer,
     - and a configurable model/backbone size.
   - The training loop supports common LM knobs: batch size, gradient accumulation, learning rate, weight decay, dropout, seed, CPU data workers, checkpoint cadence, and evaluation cadence.
   - It also supports **dynamic masking**, with scheduled mask updates and optional mask decay over training.

5. **ELECTRA branch**
   - Passing `--custom electra` switches the objective from MLM to **RTD**.
   - In this mode the model is split into:
     - a **generator** that proposes replacement tokens,
     - and a **discriminator** that learns to detect which tokens were replaced.
   - The README makes clear that the **discriminator checkpoint is the downstream artifact** saved for later fine-tuning or evaluation, and that it is a plain `DebertaV2Model` at save time.

6. **ELECTRA-to-MLM recovery path**
   - Passing `--custom electra_head` builds a **second-stage MLM model** on top of a previously trained ELECTRA discriminator.
   - This stage is additive: it reads the discriminator checkpoint, optionally unfreezes the top N transformer layers, and trains a new MLM head with configurable depth.
   - This is the repo’s bridge from an RTD-only checkpoint back to a model that can be scored with standard pseudo-log-likelihood methods.

7. **Masking and sequence-length control**
   - The training pipeline supports several masking behaviors:
     - standard BERT-style MLM masking probabilities,
     - forced full masking for ELECTRA generator inputs,
     - soft masking,
     - and a mask decay schedule.
   - Sequence length can vary by epoch via the `--max_seq_len` schedule, which is important for BabyLM-style curriculum or compute budgeting.

8. **Data and preprocessing boundary**
   - The repository separates raw/processed corpus files from the training code:
     - `data/bb24.train`
     - `data/bb25_small.dev`
   - **`preprocessing.py`** is the boundary for turning raw text into model-ready inputs or preparing corpus formats, while **`tokenizers/bb24.model`** is the trained tokenizer artifact consumed by the training and scoring scripts.
   - This makes the data pipeline largely external to the core training loop: the scripts assume prebuilt text corpora and a SentencePiece-style tokenizer model already exist.

9. **Evaluation subsystem**
   - **`score_blimp.py`** is the evaluation entrypoint for BLiMP-style minimal-pair scoring.
   - It supports multiple scoring modes:
     - **MLM pseudo-log-likelihood** for standard masked LMs,
     - **RTD-realness scoring** for ELECTRA discriminators,
     - and the MLM-head-on-discriminator variant.
   - The script consumes a model checkpoint plus a minimal-pair dataset in JSONL-like form, and compares “good” vs “bad” sentences using configurable field names.

10. **Model registration and reloadability**
   - **`register_model.py`** exists to make custom architectures reloadable through Hugging Face’s model registry / remote-code mechanism.
   - This is important because the repo includes nonstandard model variants beyond vanilla Transformers, and the architecture must survive checkpoint save/load cycles for downstream use.

11. **Artifacts and checkpoints**
   - Training outputs are saved under user-specified output directories, typically with **`checkpoint-N/`** subdirectories.
   - The repo distinguishes between:
     - a model checkpoint intended for downstream fine-tuning,
     - and a separate checkpoint or head-training stage used only for zero-shot scoring conversion.
   - The trained model collection is also published externally on Hugging Face, indicating that checkpoints are a first-class output of the pipeline.

12. **External dependencies and tooling**
   - The project is Python-based and depends on the Hugging Face ecosystem for model loading and checkpoint interoperability.
   - It also optionally integrates with **Weights & Biases** for experiment tracking.
   - The README mentions **LAMB vs AdamW**, FLOPs logging, and checkpoint evaluation at BabyLM-required token-count milestones, which suggests the training system is designed for reproducible benchmark submission rather than general-purpose serving.

13. **Overall architecture in one line**
   - The repo forms a pipeline of **text corpus + tokenizer → masked/ELECTRA pretraining → checkpoint export → BLiMP-style scoring**, with a custom model-registration layer enabling nonstandard architectures to be saved and reloaded cleanly.

Main entry points: [`train_mask.py`](https://github.com/heytamjid/babylm25/blob/electra_v1/train_mask.py) (Train mask), [`score_blimp.py`](https://github.com/heytamjid/babylm25/blob/electra_v1/score_blimp.py) (Score BLiMP), [`electra_model.py`](https://github.com/heytamjid/babylm25/blob/electra_v1/electra_model.py) (ELECTRA head), [`tokenizers/bb24.model`](https://github.com/heytamjid/babylm25/blob/electra_v1/tokenizers/bb24.model) (Tokenizer)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_group_data["Data & preprocessing"]
  node_node_corpus_train["Train corpus<br/>text corpus<br/>[bb24.train]"]
  node_node_corpus_dev["Dev corpus<br/>validation corpus<br/>[bb25_small.dev]"]
  node_node_tokenizer["Tokenizer<br/>sentencepiece model<br/>[bb24.model]"]
  node_node_preprocessing["Preprocess<br/>data prep<br/>[preprocessing.py]"]
end

subgraph group_group_train["Pretraining"]
  node_node_train_mask["Train mask<br/>main trainer<br/>[train_mask.py]"]
  node_node_train_basic["Basic train<br/>baseline trainer"]
  node_node_electra_model["ELECTRA model<br/>custom model<br/>[electra_model.py]"]
  node_node_mlm_path{{"MLM path<br/>masked LM"}}
  node_node_rtd_path{{"RTD path<br/>electra rtd"}}
  node_node_electra_head["ELECTRA head<br/>hybrid stage<br/>[electra_model.py]"]
  node_node_checkpoints[("Checkpoints<br/>artifacts")]
  node_node_wandb(("W&B<br/>tracking"))
end

subgraph group_group_eval["Evaluation"]
  node_node_register_model["Register model<br/>hf registry<br/>[register_model.py]"]
  node_node_score_blimp["Score BLiMP<br/>evaluator<br/>[score_blimp.py]"]
  node_node_hf_models(("HF loading<br/>transformers api"))
end

node_node_preprocessing -->|"prepares"| node_node_corpus_train
node_node_preprocessing -->|"prepares"| node_node_corpus_dev
node_node_tokenizer -->|"consumed by"| node_node_train_mask
node_node_tokenizer -->|"consumed by"| node_node_score_blimp
node_node_corpus_train -->|"input"| node_node_train_mask
node_node_corpus_dev -->|"validate"| node_node_train_mask
node_node_train_mask -->|"orchestrates"| node_node_mlm_path
node_node_train_mask -->|"orchestrates"| node_node_rtd_path
node_node_train_basic -->|"baseline"| node_node_mlm_path
node_node_rtd_path -->|"saves"| node_node_checkpoints
node_node_mlm_path -->|"saves"| node_node_checkpoints
node_node_checkpoints -->|"evaluated by"| node_node_score_blimp
node_node_register_model -->|"enables"| node_node_hf_models
node_node_electra_model -->|"loads via"| node_node_hf_models
node_node_train_mask -.->|"optional"| node_node_wandb
node_node_electra_model -->|"extends"| node_node_electra_head
node_node_electra_head -->|"exports"| node_node_checkpoints
node_node_score_blimp -->|"scores"| node_node_mlm_path
node_node_score_blimp -->|"scores"| node_node_rtd_path
node_node_score_blimp -->|"scores"| node_node_electra_head

click node_node_corpus_train "https://github.com/heytamjid/babylm25/blob/electra_v1/data/bb24.train"
click node_node_corpus_dev "https://github.com/heytamjid/babylm25/blob/electra_v1/data/bb25_small.dev"
click node_node_tokenizer "https://github.com/heytamjid/babylm25/blob/electra_v1/tokenizers/bb24.model"
click node_node_preprocessing "https://github.com/heytamjid/babylm25/blob/electra_v1/preprocessing.py"
click node_node_train_mask "https://github.com/heytamjid/babylm25/blob/electra_v1/train_mask.py"
click node_node_train_basic "https://github.com/heytamjid/babylm25/blob/electra_v1/train_mask_basic.py"
click node_node_electra_model "https://github.com/heytamjid/babylm25/blob/electra_v1/electra_model.py"
click node_node_register_model "https://github.com/heytamjid/babylm25/blob/electra_v1/register_model.py"
click node_node_score_blimp "https://github.com/heytamjid/babylm25/blob/electra_v1/score_blimp.py"
click node_node_electra_head "https://github.com/heytamjid/babylm25/blob/electra_v1/electra_model.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_node_corpus_train,node_node_corpus_dev,node_node_tokenizer,node_node_preprocessing toneBlue
class node_node_train_mask,node_node_train_basic,node_node_electra_model,node_node_mlm_path,node_node_rtd_path,node_node_electra_head,node_node_checkpoints,node_node_wandb toneAmber
class node_node_register_model,node_node_score_blimp,node_node_hf_models toneMint
```

## Components

### Data & preprocessing

Corpus and tokenizer artifacts plus preprocessing boundary.

- [Train corpus](https://github.com/heytamjid/babylm25/blob/electra_v1/data/bb24.train): text corpus · Primary pretraining text input. · `data/bb24.train`
- [Dev corpus](https://github.com/heytamjid/babylm25/blob/electra_v1/data/bb25_small.dev): validation corpus · Development text used for validation during training. · `data/bb25_small.dev`
- [Tokenizer](https://github.com/heytamjid/babylm25/blob/electra_v1/tokenizers/bb24.model): sentencepiece model · Shared tokenizer artifact consumed by training and scoring. · `tokenizers/bb24.model`
- [Preprocess](https://github.com/heytamjid/babylm25/blob/electra_v1/preprocessing.py): data prep · Boundary for corpus formatting and input preparation. · `preprocessing.py`

### Pretraining

Masked LM and ELECTRA-style training paths.

- [Train mask](https://github.com/heytamjid/babylm25/blob/electra_v1/train_mask.py): main trainer · Primary training entrypoint with masking, checkpoints, and logging. · `train_mask.py`
- [Basic train](https://github.com/heytamjid/babylm25/blob/electra_v1/train_mask_basic.py): baseline trainer · Simpler MLM training variant. · `train_mask_basic.py`
- [ELECTRA model](https://github.com/heytamjid/babylm25/blob/electra_v1/electra_model.py): custom model · Custom ELECTRA-derived MLM head and discriminator-backed architecture. · `electra_model.py`
- **MLM path**: masked LM · Standard masked language modeling training and scoring route.
- **RTD path**: electra rtd · Replaced-token detection branch with generator and discriminator.
- [ELECTRA head](https://github.com/heytamjid/babylm25/blob/electra_v1/electra_model.py): hybrid stage · Adds a new MLM head on top of a trained ELECTRA discriminator. · `electra_model.py`
- **Checkpoints**: artifacts · Saved model outputs for downstream use and scoring.
- **W&B**: tracking · Optional experiment tracking integration.

### Evaluation

BLiMP-style scoring and model reload support.

- [Register model](https://github.com/heytamjid/babylm25/blob/electra_v1/register_model.py): hf registry · Makes custom architectures reloadable through Hugging Face conventions. · `register_model.py`
- [Score BLiMP](https://github.com/heytamjid/babylm25/blob/electra_v1/score_blimp.py): evaluator · Evaluates checkpoints on BLiMP-style minimal pairs with multiple scoring modes. · `score_blimp.py`
- **HF loading**: transformers api · AutoModel and AutoModelForMaskedLM loading plus remote-code support.

## Connections

- Preprocess → Train corpus: prepares · Turns raw or external text into training-ready corpus files.
- Preprocess → Dev corpus: prepares · Turns raw or external text into validation corpus files.
- Tokenizer → Train mask: consumed by · Tokenizer is required for training inputs.
- Tokenizer → Score BLiMP: consumed by · Tokenizer is used during scoring and tokenization.
- Train corpus → Train mask: input · Primary training corpus.
- Dev corpus → Train mask: validate · Validation corpus for periodic evaluation.
- Train mask → MLM path: orchestrates · Main trainer drives standard MLM runs.
- Train mask → RTD path: orchestrates · Main trainer also drives ELECTRA-style RTD runs.
- Basic train → MLM path: baseline · Simpler MLM-only training route.
- RTD path → Checkpoints: saves · Exports discriminator checkpoints for downstream use.
- MLM path → Checkpoints: saves · Exports masked LM checkpoints.
- Checkpoints → Score BLiMP: evaluated by · Saved checkpoints are scored on minimal-pair tasks.
- Register model → HF loading: enables · Registers custom classes for reloadable HF checkpoints.
- ELECTRA model → HF loading: loads via · Custom ELECTRA-derived models rely on HF remote-code loading.
- Train mask → W&B: optional · Training can emit experiment logs to W&B.
- ELECTRA model → ELECTRA head: extends · Hybrid stage adds an MLM head on top of discriminator weights.
- ELECTRA head → Checkpoints: exports · Produces a scored model artifact for pseudo-log-likelihood evaluation.
- Score BLiMP → MLM path: scores · Uses MLM pseudo-log-likelihood for masked language models.
- Score BLiMP → RTD path: scores · Uses realness scoring for ELECTRA discriminators.
- Score BLiMP → ELECTRA head: scores · Supports MLM-head-on-discriminator evaluation mode.