docs/ARCHITECTURE.md

# Architectural Documentation - Small Language Model (SLM)

## System Architecture

`lawslm` is an industrial-grade, decoder-only Small Language Model (SLM) ecosystem built completely from scratch using Python and PyTorch tensor primitives.

```
+-----------------------------------------------------------------------+
|                             SLM ENGINE                                |
+-----------------------------------------------------------------------+
|  [REST API / FastAPI]                  [CLI Interface / main.py]      |
+-----------------------------------++----------------------------------+
                                    ||
+-----------------------------------vv----------------------------------+
|                           Text Generator                              |
|       (Greedy, Temp, Top-K, Top-P, Penalties, Streaming Callbacks)    |
+-----------------------------------++----------------------------------+
                                    ||
+-----------------------------------vv----------------------------------+
|                         SLMForCausalLM                                |
|  Token Embeddings -> RoPE -> N x TransformerBlock -> RMSNorm -> Head  |
+-----------------------------------++----------------------------------+
                                    ||
+-----------------------------------vv----------------------------------+
|                        Transformer Block                              |
|  Pre-RMSNorm -> Multi-Head Causal Attention -> Pre-RMSNorm -> SwiGLU  |
+-----------------------------------------------------------------------+
```

## Modular Breakdown

1. **Tokenizer (`slm/tokenizer/`)**:
   - Pure Python Byte-Pair Encoding (BPE) subword algorithm (`bpe.py`).
   - Character tokenizer fallback (`char_tokenizer.py`).
   - Vocabulary manager, frequency dictionary, and JSON serialization (`vocab.py`).

2. **Dataset & Cleaning (`slm/dataset/`)**:
   - Text cleaning, NFC Unicode normalization, HTML stripping, deduplication (`cleaner.py`).
   - Ingestion parsers for TXT, CSV, JSON, JSONL, MD, HTML, XML (`readers.py`).
   - Sliding causal context window PyTorch Dataset (`dataset.py`).

3. **Embeddings & Normalization (`slm/embeddings/`, `slm/normalization/`)**:
   - Rotary Position Embeddings (RoPE) applied to Queries and Keys (`positional.py`).
   - Learnable and Sinusoidal positional embeddings option (`positional.py`).
   - Root Mean Square Layer Normalization (`rmsnorm.py`).
   - Layer Normalization (`layernorm.py`).

4. **Attention & FeedForward (`slm/attention/`, `slm/feedforward/`)**:
   - Scaled Dot-Product Attention with triangular causal mask (`causal_attention.py`).
   - Multi-Head Causal Attention with RoPE (`causal_attention.py`).
   - SwiGLU (Swish Gated Linear Unit) Feed-Forward Network (`mlp.py`).

5. **Optimizers & Schedulers (`slm/optimizer/`, `slm/scheduler/`)**:
   - Custom `AdamW` with decoupled weight decay (`adamw.py`).
   - Custom `Lion` (EvoLved Sign Momentum) optimizer (`lion.py`).
   - Cosine Annealing with Warmup learning rate scheduler (`schedulers.py`).

6. **Training Engine (`slm/training/`)**:
   - Teacher-forcing training loop with AMP FP16/BF16, gradient accumulation, gradient clipping, evaluation, and checkpoint manager integration.