Low Level Design

Low-Level Design (LLD)

Class and composition diagram

classDiagram
    direction TB
    class YelpTextDataset {
      +texts: list
      +labels: list
      +__len__() int
      +__getitem__(idx) tuple
    }
    class SmokeTokenizer {
      +vocab: dict
      +__call__(texts, padding, truncation, max_length, return_tensors) dict
    }
    class SinusoidalPositionalEncoding {
      +pe: Tensor
      +forward(x) Tensor
    }
    class MultiHeadSelfAttention {
      +d_model: int
      +n_heads: int
      +head_dim: int
      +use_sdpa: bool
      +forward(x, pad_mask) Tensor
    }
    class TransformerEncoderLayer {
      +ln1: LayerNorm
      +attn: MultiHeadSelfAttention
      +ln2: LayerNorm
      +ffn: Sequential
      +forward(x, pad_mask) Tensor
    }
    class TransformerClassifier {
      +pos_enc: SinusoidalPositionalEncoding
      +layers: ModuleList
      +ln_f: LayerNorm
      +head: Linear
      +forward(x_emb, attention_mask, pad_mask) Tensor
    }
    TransformerClassifier *-- SinusoidalPositionalEncoding : pos_enc
    TransformerClassifier *-- "N" TransformerEncoderLayer : layers
    TransformerEncoderLayer *-- MultiHeadSelfAttention : attn
    YelpTextDataset ..> TransformerClassifier : training batches
    SmokeTokenizer ..> YelpTextDataset : synthetic mode

Loading

Class responsibilities

Class / function	What the code actually does
YelpTextDataset	Stores review text/labels; supports indexed access for DataLoader
SmokeTokenizer	Offline whitespace vocabulary/token-ID/padding logic for the synthetic path
make_collate_with_static_embeddings	Closes over tokenizer and CPU frozen embedding matrix; creates embeddings, attention_mask, pad_mask, labels
SinusoidalPositionalEncoding	Adds fixed sinusoidal signal; supports odd/even model dimensions; buffer not stored in state dict
MultiHeadSelfAttention	Separate Q/K/V projections, split heads, optional SDPA, key padding masking, output projection
TransformerEncoderLayer	Pre-LN attention residual, then pre-LN GELU feed-forward residual
TransformerClassifier	Embedding scale, positional encoding, N encoder layers, final LN, masked mean pool, 2-logit head
select_reviews / load_real_data	Length-filtered Yelp selection and separate validation from original training split
evaluate	Inference-mode cross-entropy and accuracy
run	Parse mode, make data loaders/model, optimize, checkpoint best validation accuracy, reload, test

Tensor contracts

Let B = batch size, S = padded sequence length, D = embedding/model width, H = attention heads, Dh = D/H.

Stage	Shape
Token IDs, attention mask	B × S
Frozen token embeddings	B × S × D
Q, K, V after head split	B × H × S × Dh
Attention score matrix (manual path)	B × H × S × S
Concatenated head outputs	B × S × D
Masked mean pooled vector	B × D
Classification logits	B × 2

Default model shape (real-data --train)

D=384, H=6, Dh=64, N=3 custom encoder blocks, and num_classes=2. The smoke test instead sets D=48, H=4, N=2, S=24, batch_size=8.

Exact forward-pass order

1. x_emb * sqrt(D) → positional encoding → dropout.
2. For each encoder block: x + dropout(attn(LayerNorm(x), pad_mask)).
3. Within that same block: x + dropout(FFN(LayerNorm(x))).
4. Final LayerNorm; apply attention_mask to eliminate padded token contributions.
5. Mean across valid tokens (denominator clamped to at least 1); apply linear classification head.

Runtime sequence

sequenceDiagram
    autonumber
    participant Runner as run(args)
    participant Data as Data loader / synthetic sampler
    participant Tokenizer as Tokenizer
    participant Lookup as Frozen embedding matrix (CPU)
    participant Collate as collate(batch)
    participant Model as TransformerClassifier
    Runner->>Data: Obtain train / validation / test splits
    Runner->>Tokenizer: Create synthetic tokenizer or load MiniLM
    Runner->>Lookup: Random frozen matrix or MiniLM input embeddings
    Runner->>Collate: Build collate function with tokenizer and frozen weights
    Runner->>Data: Construct DataLoaders
    Data->>Collate: Batch of (review, label)
    Collate->>Tokenizer: Tokenize + pad + truncate
    Tokenizer-->>Collate: input_ids, attention_mask
    Collate->>Lookup: F.embedding(input_ids, emb_weight)
    Lookup-->>Collate: embeddings B x S x D
    Collate-->>Runner: embeddings, attention_mask, pad_mask, labels
    Runner->>Model: forward(embeddings, attention_mask, pad_mask)
    Model-->>Runner: logits B x 2

Loading

For attention internals and the branching SDPA path see Attention Mechanism.