System Architecture

System Architecture — High-Level Design (HLD)

Purpose and system boundary

The system is a single-process command-line training/evaluation application, not a hosted API or distributed training service. It runs synthetic CPU-only smoke tests or Hugging Face-backed Yelp training, produces sentiment logits, stores a best-validation checkpoint, and evaluates a held-out test split.

System context

flowchart LR
    User["Developer / Researcher"] --> CLI["Python CLI: transformer_from_scratch.py"]
    CLI --> Mode{"Execution mode"}
    Mode -->|"--smoke-test"| Synthetic["Synthetic reviews + SmokeTokenizer"]
    Mode -->|"--train / --full"| Yelp["Hugging Face Yelp Polarity dataset"]
    Mode -->|"--train / --full"| MiniLM["MiniLM tokenizer + frozen embedding weights"]
    Synthetic --> Collator["Tokenizer + frozen embedding lookup on CPU"]
    Yelp --> Collator
    MiniLM --> Collator
    Collator --> Loaders["PyTorch DataLoader: train / validation / test"]
    Loaders --> Trainer["Training / evaluation loop"]
    Trainer <--> Model["Custom Transformer Encoder + Classification Head"]
    Trainer --> Checkpoint["Best validation checkpoint: .pt"]
    Checkpoint --> Test["Reload checkpoint + held-out test evaluation"]
    CLI -.->|"--visualize optional"| Graphviz["torchview + Graphviz PNG"]

Loading

Component pipeline

flowchart TD
    Text["Input review text"] --> Tokenizer["Tokenizer"]
    Tokenizer --> IDs["Token IDs and attention mask: B x S"]
    IDs --> Frozen["Frozen embedding lookup on CPU: B x S x D"]
    Frozen --> Scale["Multiply by sqrt(D)"]
    Scale --> Pos["Sinusoidal positional encoding + dropout"]
    Pos --> Stack["Pre-LN Transformer encoder layer x N"]
    Stack --> Norm["Final LayerNorm"]
    Norm --> Pool["Mask-aware mean pool across S: B x D"]
    IDs -.->|"attention mask"| Pool
    Pool --> Head["Linear classification head: B x 2"]
    Head --> Logits["Logits: negative / positive"]
    subgraph EncoderDetail["Each encoder layer"]
        LN1["LayerNorm"] --> Attn["Multi-head self-attention"]
        Attn --> Add1["Residual connection + dropout"]
        Add1 --> LN2["LayerNorm"]
        LN2 --> FFN["Linear -> GELU -> Dropout -> Linear"]
        FFN --> Add2["Residual connection + dropout"]
    end
    Stack -.-> EncoderDetail

Loading

External systems and artifacts

Component	Responsible for	Location / lifecycle
Hugging Face Datasets	fancyzhx/yelp_polarity train/test reviews	Downloaded/cached by datasets; real-data mode only
Hugging Face Transformers	Tokenizer and frozen input embedding weights from MiniLM	Downloaded/cached by transformers; real-data mode only
YelpTextDataset + DataLoader	Batch source, shuffling, collating	Local process, in memory
TransformerClassifier	Custom Transformer layers and binary logits	Local PyTorch process, CPU/CUDA
AdamW and validation	Weight updates and best-model selection	Local process
.pt checkpoint	Trainable model weights and metadata, not embedding matrix	Local output directory
torchview + Graphviz	Optional PNG architecture rendering	Optional packages / installed executable

Important constraints

• The core model is an encoder-only classifier. No generation, autoregressive decoder, retrieval system, API endpoint, or vector database is implemented.
• Frozen MiniLM embeddings are computed by F.embedding in the DataLoader collate_fn on CPU; embedding tensors are later transferred to the selected training device.
• Training uses a single CUDA GPU when available, otherwise CPU; no multi-GPU or DDP implementation is present.
• --sdpa selects PyTorch scaled dot-product attention rather than the explicit score/softmax implementation.
• Dataset, pretrained MiniLM weights, and Graphviz executable are not included in this project ZIP.