Skip to content

Data Pipeline

Sudhanshu edited this page Nov 11, 2026 · 1 revision

Data Pipeline and Batch 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

Three-way split strategy

The smoke path generates 80 synthetic samples and partitions 64/8/8. In real-data mode, load_real_data selects shuffled, length-filtered candidates from the original Yelp training split, reserving the last val_count selected examples for validation. The official Yelp test split provides held-out test examples.

The script limits the candidate review pool to roughly 3× the requested size before filtering, so sufficiently long reviews may reduce actual returned sample counts; it also checks whether validation can be formed.

Collation and device movement

  • The collate_fn tokenizes text, pads/truncates to max_len, confirms token IDs fit the embedding vocabulary, and uses frozen F.embedding on CPU.
  • It emits both attention_mask (1 valid, 0 padding) and pad_mask (True where padding).
  • The training/evaluation loop then moves embeddings, masks and labels onto CPU or CUDA.

Why the frozen matrix is not trainable

emb_weight is a standalone CPU tensor captured by the collation closure rather than an nn.Embedding parameter registered in the custom model. Gradient-based optimizer updates therefore affect only the Transformer encoder and classification head.

Clone this wiki locally