# ADR-001: Implementation Framework for domainTokenizer

> **Status:** Accepted
> **Date:** April 29, 2026
> **Decision:** PyTorch + HuggingFace Transformers as primary framework, with JAX/Flax NNX as future scaling path
> **Deciders:** domainTokenizer core team

---

## Table of Contents

1. [Context](#1-context)
2. [Goal](#2-goal)
3. [Options Evaluated](#3-options-evaluated)
4. [Decision](#4-decision)
5. [Trade-offs and Justification](#5-trade-offs-and-justification)
6. [Consequences](#6-consequences)
7. [Implementation Roadmap](#7-implementation-roadmap)
8. [Appendix A: Framework Usage Across Reference Papers](#appendix-a-framework-usage-across-reference-papers)
9. [Appendix B: Head-to-Head Comparison Matrix](#appendix-b-head-to-head-comparison-matrix)
10. [Appendix C: Key Code Patterns](#appendix-c-key-code-patterns)

---

## 1. Context

### What We're Building

domainTokenizer is a library for building **small Transformer models (24M–330M parameters)** that process **domain-specific tokens** — financial transactions, e-commerce events, healthcare records — instead of natural language text. The architecture follows the validated pattern from Nubank's nuFormer ([arXiv: 2507.23267](https://arxiv.org/abs/2507.23267)):

```
Domain Events → Custom Tokenizer → GPT-style Transformer → Foundation Model → Downstream Tasks
```

### The Implementation Question

A video from Google for Developers presents the **Keras 3 + JAX/Flax NNX integration** as a potential framework, citing:
- **Explicit state management** (Flax NNX) — useful for tracking sequential transaction state
- **Custom training loops** — Keras structure + JAX/Optax for domain-specific gradient control
- **JIT compilation** (`@nnx.jit`) — high-performance processing of millions of transactions
- **Paradigm mixing** — Keras layers for standard components + NNX for custom sequential encoders

The question: **Is Keras + JAX/Flax NNX the right framework for domainTokenizer, or is there a better choice?**

### Constraints

1. **Custom tokenizer required:** We need a tokenizer that maps structured fields (amounts, dates, categories) to special tokens — not a standard text tokenizer
2. **Small models:** 24M–330M parameters, not 70B+ — framework overhead matters less than developer velocity
3. **Production deployment:** Models must be servable with low latency for real-time applications (fraud detection, recommendations)
4. **GPU hardware:** Development on A100/A10G GPUs, not TPUs (standard cloud environment)
5. **Team context:** ML engineers familiar with Python, PyTorch, and the HuggingFace ecosystem
6. **Iteration speed:** Need to prototype quickly across multiple domains (finance, e-commerce, healthcare)

---

## 2. Goal

Choose an implementation framework that:

1. **Minimizes time from research to working prototype** — weeks, not months
2. **Supports custom domain tokenizers** as first-class citizens
3. **Integrates with the HuggingFace Hub** for model sharing, versioning, and community
4. **Enables production deployment** via standard serving infrastructure (ONNX, TGI, vLLM, etc.)
5. **Scales to 330M parameters** on 4–8 GPUs without heroic engineering
6. **Does not preclude future migration** to JAX/TPU if we need to scale beyond 1B parameters

---

## 3. Options Evaluated

### Option A: PyTorch + HuggingFace Transformers

The dominant ecosystem for custom NLP/sequential models. Provides `PreTrainedModel`, `PreTrainedTokenizerFast`, `Trainer`, `push_to_hub`, ONNX export, and integration with TRL, PEFT, Accelerate, DeepSpeed.

### Option B: Keras 3 + JAX Backend + Flax NNX

Google's multi-backend framework. Keras provides high-level APIs; JAX provides XLA compilation and functional transforms; Flax NNX provides PyTorch-like stateful modules on top of JAX.

### Option C: Pure JAX + Flax NNX + Optax

Skip Keras entirely. Use Flax NNX for model definition, Optax for optimization, Orbax for checkpointing, and Grain/tf.data for data loading. Google's MaxText framework follows this pattern.

### Option D: PyTorch + Custom (no HuggingFace)

Use PyTorch directly without the HuggingFace abstraction layer. Full control but no ecosystem integration.

---

## 4. Decision

### Primary: PyTorch + HuggingFace Transformers (Option A)

### Future scaling path: JAX/Flax NNX (Option C) — if and when we need TPU training at >1B parameters

---

## 5. Trade-offs and Justification

### 5.1 What the Reference Papers Actually Use

We audited the frameworks used by every paper in the domainTokenizer research corpus. The result is overwhelming:

| Paper | Framework | Confidence |
|-------|-----------|------------|
| **nuFormer** (Nubank) | PyTorch + HF Transformers (inferred) | ~90% |
| **TIGER** (Google) | JAX + T5X (official); PyTorch (community reimpl) | 100% |
| **ActionPiece** (Google DeepMind) | **PyTorch + HF Transformers** (stated verbatim in paper) | 100% |
| **RecFormer** (UCSD/Amazon) | **PyTorch + HF Transformers (Longformer)** (stated verbatim) | 100% |
| **Banking Transaction Flow** | **PyTorch** (stated verbatim in appendix) | 100% |
| **PLR Embeddings** (Yandex) | **PyTorch** + scikit-learn + Optuna | 100% |

**5 of 6 papers use PyTorch.** The sole JAX user (TIGER) was a Google-internal project using T5X, and even its most popular community reimplementation (781⭐) is in PyTorch.

Even **Google DeepMind's own ActionPiece** — the paper most relevant to our domain tokenization approach — uses PyTorch + HuggingFace. This is the strongest signal possible.

### 5.2 Custom Tokenizer Story

This is the **decisive factor**. domainTokenizer's core innovation is the tokenizer itself. The framework must provide first-class support for custom token vocabularies.

**PyTorch + HuggingFace:**
- Train custom BPE tokenizer via `tokenizers` library (Rust-backed, fast)
- Wrap in `PreTrainedTokenizerFast` → full Trainer compatibility
- Add domain special tokens via `add_special_tokens()` → auto-resize embeddings
- Push tokenizer to Hub: `tokenizer.push_to_hub("org/my-tokenizer")`
- Load anywhere: `AutoTokenizer.from_pretrained("org/my-tokenizer")`
- **KL3M** ([arXiv: 2503.17247](https://arxiv.org/abs/2503.17247)) — the gold standard for financial domain tokenizers — is built entirely on this stack

**Keras + JAX/Flax NNX:**
- No equivalent to `PreTrainedTokenizerFast`
- No Hub-integrated tokenizer format
- Must build custom tokenizer from scratch with no ecosystem support
- No standard serialization/deserialization for domain vocabularies

**Verdict:** PyTorch/HF has a **complete, production-tested** custom tokenizer pipeline. JAX/Keras has **nothing** — you'd build everything from scratch.

### 5.3 Production Deployment

| Path | PyTorch | JAX/Keras |
|------|---------|-----------|
| ONNX export | `torch.onnx.export()` — one line | Requires TF backend intermediate or experimental `jax.export` |
| TensorRT | ONNX → TRT (standard) | Multi-hop, fragile |
| TGI (HuggingFace inference) | First-class | Not supported |
| vLLM | First-class | Not supported |
| Triton Inference Server | Direct ONNX/TorchScript | Via ONNX (workaround) |
| BentoML | Supported | Supported |
| Model Hub sharing | `push_to_hub()` → `from_pretrained()` | Works but fragmented (`.msgpack` weights, no Trainer compat) |

**Verdict:** PyTorch has **direct, tested paths** to every major serving framework. JAX requires **multiple intermediate conversions**, each introducing failure points.

### 5.4 Training Speed

At our scale (24M–330M parameters on 4–8 A100s):

| Scenario | PyTorch | JAX |
|----------|---------|-----|
| Steady-state training throughput | **Comparable** (`torch.compile`) | **Comparable** (XLA JIT) |
| Variable-length sequences | **Native** — dynamic shapes | **Problematic** — recompiles on new shapes; must pad to buckets |
| Multi-GPU (FSDP) | `accelerate` + FSDP2 — mature | `pmap`/`shard_map` — works but harder to configure |
| First-run compilation | Instant (eager mode) | 5–20s JIT compilation overhead |
| Debugging | Standard Python debugger | `print` debugging; cryptic XLA errors |

**Verdict:** At 330M parameters, training speed is a **wash**. JAX's advantages (XLA kernel fusion, TPU native) only matter at 10B+ parameters on 256+ accelerators. At our scale, **developer velocity dominates throughput**.

### 5.5 The JAX Advantage: When It Would Win

JAX/Flax NNX would be the right choice **if**:

1. **Training exclusively on Google TPUs** — JAX is the native TPU compiler; PyTorch/XLA is a port with overhead
2. **Models >1B parameters** — XLA's whole-program optimization shines at scale
3. **Fixed-shape workloads** — images, fixed-length token sequences (no variable-length padding issues)
4. **Need functional transforms** — `vmap` (per-sample gradients), `pmap` (data parallelism), `grad` (higher-order derivatives)
5. **Google Cloud infrastructure** — Vertex AI, TPU VMs, GCS integration

For domainTokenizer's current scope (24M–330M, GPU, variable-length sequences, fast iteration), **none of these conditions apply**.

### 5.6 The Keras + JAX Mixing Argument

The Google for Developers video argues for mixing Keras layers (high-level) with NNX modules (custom, high-performance). In theory, this lets you:
- Use Keras for standard Transformer layers
- Use NNX for custom sequential transaction encoders
- Get JIT compilation on the NNX parts

**In practice, this creates problems:**
1. **Two mental models:** Keras (layer-oriented, `fit/compile`) vs. NNX (functional, explicit state) — context switching slows development
2. **Limited interop documentation:** Keras ↔ NNX examples are thin; edge cases are poorly documented
3. **No HF ecosystem integration:** You lose Trainer, push_to_hub, PEFT, TRL, Accelerate — the entire ecosystem Nubank and ActionPiece rely on
4. **Debugging complexity:** Errors in the Keras↔NNX boundary are hard to diagnose

**Better approach with PyTorch:** Use `torch.compile()` on performance-critical modules to get JIT compilation benefits without leaving the PyTorch ecosystem. Write custom `nn.Module` subclasses for domain-specific components. This gives you the same "standard parts + custom parts" architecture without framework mixing.

---

## 6. Consequences

### What We Gain

1. **Immediate access to the entire HuggingFace ecosystem:** Trainer, Accelerate, PEFT (LoRA), TRL, Evaluate, push_to_hub, from_pretrained, ONNX export, TGI serving
2. **Copy-paste from reference implementations:** ActionPiece, RecFormer, Banking TF, and PLR embeddings are all PyTorch — we can directly reuse their code
3. **KL3M tokenizer as starting point:** The best financial domain tokenizer already exists in PyTorch/HF format at `alea-institute/kl3m-004-128k-cased`
4. **Standard production deployment:** ONNX → TensorRT → Triton, or direct TGI/vLLM serving
5. **Community and hiring:** PyTorch is the dominant ML framework; finding contributors and documentation is easy
6. **`torch.compile()` for performance:** When we need JIT compilation on hot paths, `torch.compile()` provides 10–30% speedups without leaving the ecosystem

### What We Accept

1. **No native TPU support:** If we later need to train on Google TPUs, we'll need PyTorch/XLA (slower than native JAX) or migrate the model code
2. **No functional transforms:** `vmap` (per-sample gradients) isn't available without `functorch` (experimental). If we need advanced gradient manipulation for meta-learning or Nested Learning (HOPE-style), JAX would be better
3. **Potential future migration cost:** If we scale beyond 1B parameters and move to TPUs, we'll need to rewrite model code in Flax NNX. This is mitigated by keeping model definitions clean and modular

### Migration Strategy (If Needed Later)

If domainTokenizer grows to >1B parameters and we need TPU training:

1. **Tokenizer layer stays in Python/HF:** Tokenizer is framework-agnostic — it produces integer sequences regardless of whether the model is PyTorch or JAX
2. **Model architecture translates 1:1:** PyTorch `nn.Module` → Flax NNX `nnx.Module` mapping is straightforward for standard Transformer components
3. **Training loop changes:** PyTorch Trainer → custom Flax NNX training loop with Optax
4. **Reference:** Google's MaxText (`github.com/google/maxtext`) provides production-grade JAX Transformer patterns we can follow

**Estimated migration effort:** 2–4 weeks for a clean, well-separated codebase.

---

## 7. Implementation Roadmap

### Phase 2A: Core Tokenizer Library (Weeks 1–3)

#### Step 1: Domain Schema Definition

Create a declarative schema format that describes the fields in a domain's event data:

```python
# src/tokenizers/schema.py

from dataclasses import dataclass, field
from enum import Enum
from typing import List, Optional

class FieldType(Enum):
    NUMERICAL_CONTINUOUS = "numerical_continuous"    # prices, amounts → magnitude bins
    NUMERICAL_DISCRETE = "numerical_discrete"        # quantities → small fixed vocab
    CATEGORICAL_FIXED = "categorical_fixed"          # categories, days of week → direct mapping
    CATEGORICAL_ENTITY = "categorical_entity"        # products, merchants → Semantic IDs (RQ-VAE)
    TEMPORAL = "temporal"                            # timestamps → calendar decomposition
    TEXT = "text"                                     # descriptions → BPE subwords
    SIGN = "sign"                                    # credit/debit → 2 tokens

@dataclass
class FieldSpec:
    name: str
    field_type: FieldType
    vocab_size: Optional[int] = None          # for fixed categorical
    n_bins: int = 21                          # for numerical (Nubank uses 21)
    calendar_fields: List[str] = field(       # for temporal
        default_factory=lambda: ["month", "dow", "dom", "hour"]
    )

@dataclass
class DomainSchema:
    name: str                                 # e.g., "ecommerce", "finance"
    fields: List[FieldSpec]                   # ordered list of fields per event
    
    @property
    def special_token_count(self) -> int:
        """Total domain-specific special tokens needed."""
        count = 0
        for f in self.fields:
            if f.field_type == FieldType.SIGN:
                count += 2
            elif f.field_type == FieldType.NUMERICAL_CONTINUOUS:
                count += f.n_bins
            elif f.field_type == FieldType.CATEGORICAL_FIXED:
                count += f.vocab_size
            elif f.field_type == FieldType.TEMPORAL:
                count += sum({
                    "month": 12, "dow": 7, "dom": 31, "hour": 24,
                    "quarter": 4, "year": 10
                }.get(cf, 0) for cf in f.calendar_fields)
        return count

# Example: Nubank-style financial schema
FINANCE_SCHEMA = DomainSchema(
    name="finance",
    fields=[
        FieldSpec("amount_sign", FieldType.SIGN),
        FieldSpec("amount", FieldType.NUMERICAL_CONTINUOUS, n_bins=21),
        FieldSpec("timestamp", FieldType.TEMPORAL,
                  calendar_fields=["month", "dow", "dom", "hour"]),
        FieldSpec("description", FieldType.TEXT),
    ]
)

# Example: E-commerce schema  
ECOMMERCE_SCHEMA = DomainSchema(
    name="ecommerce",
    fields=[
        FieldSpec("event_type", FieldType.CATEGORICAL_FIXED, vocab_size=5),
        FieldSpec("price", FieldType.NUMERICAL_CONTINUOUS, n_bins=21),
        FieldSpec("quantity", FieldType.NUMERICAL_DISCRETE, vocab_size=11),
        FieldSpec("category_l1", FieldType.CATEGORICAL_FIXED, vocab_size=30),
        FieldSpec("category_l2", FieldType.CATEGORICAL_FIXED, vocab_size=200),
        FieldSpec("timestamp", FieldType.TEMPORAL,
                  calendar_fields=["month", "dow", "dom", "hour"]),
        FieldSpec("product_title", FieldType.TEXT),
    ]
)
```

#### Step 2: Per-Field Tokenizers

Implement each field type tokenizer as a standalone module:

```python
# src/tokenizers/field_tokenizers.py

import numpy as np
from typing import List

class SignTokenizer:
    """Tokenizes sign of a numerical value (credit/debit, inflow/outflow)."""
    
    def __init__(self, prefix: str = "SIGN"):
        self.tokens = [f"[{prefix}_POS]", f"[{prefix}_NEG]"]
    
    def __call__(self, value: float) -> str:
        return self.tokens[0] if value >= 0 else self.tokens[1]
    
    @property
    def vocab(self) -> List[str]:
        return self.tokens


class MagnitudeBucketTokenizer:
    """Quantizes continuous values into bins (Nubank-style).
    
    Uses quantile-based binning on the training distribution.
    Follows the Relative Magnitude Tokenization principle from TP-BERTa.
    """
    
    def __init__(self, n_bins: int = 21, prefix: str = "AMT"):
        self.n_bins = n_bins
        self.prefix = prefix
        self.bin_edges = None  # fitted from data
    
    def fit(self, values: np.ndarray):
        """Compute bin edges from training data using quantiles."""
        # Use absolute values for magnitude binning
        abs_vals = np.abs(values[~np.isnan(values)])
        quantiles = np.linspace(0, 100, self.n_bins + 1)
        self.bin_edges = np.percentile(abs_vals, quantiles)
        return self
    
    def __call__(self, value: float) -> str:
        if self.bin_edges is None:
            raise ValueError("Tokenizer not fitted. Call .fit() first.")
        bin_idx = np.searchsorted(self.bin_edges[1:-1], abs(value))
        return f"[{self.prefix}_{bin_idx:02d}]"
    
    @property
    def vocab(self) -> List[str]:
        return [f"[{self.prefix}_{i:02d}]" for i in range(self.n_bins)]


class CalendarTokenizer:
    """Decomposes timestamps into calendar components (Nubank-style)."""
    
    FIELD_VOCABS = {
        "month": ([f"[MON_{i:02d}]" for i in range(1, 13)], lambda dt: dt.month - 1),
        "dow":   ([f"[DOW_{i}]" for i in range(7)], lambda dt: dt.weekday()),
        "dom":   ([f"[DOM_{i:02d}]" for i in range(1, 32)], lambda dt: dt.day - 1),
        "hour":  ([f"[HOUR_{i:02d}]" for i in range(24)], lambda dt: dt.hour),
    }
    
    def __init__(self, fields: List[str] = None):
        self.fields = fields or ["month", "dow", "dom", "hour"]
    
    def __call__(self, timestamp) -> List[str]:
        tokens = []
        for field_name in self.fields:
            vocab, extractor = self.FIELD_VOCABS[field_name]
            idx = extractor(timestamp)
            tokens.append(vocab[idx])
        return tokens
    
    @property
    def vocab(self) -> List[str]:
        all_tokens = []
        for field_name in self.fields:
            all_tokens.extend(self.FIELD_VOCABS[field_name][0])
        return all_tokens


class CategoricalTokenizer:
    """Maps categorical values to fixed vocabulary tokens."""
    
    def __init__(self, categories: List[str], prefix: str = "CAT"):
        self.prefix = prefix
        self.token_map = {cat: f"[{prefix}_{i:03d}]" for i, cat in enumerate(categories)}
        self.unk_token = f"[{prefix}_UNK]"
    
    def __call__(self, value: str) -> str:
        return self.token_map.get(value, self.unk_token)
    
    @property
    def vocab(self) -> List[str]:
        return list(self.token_map.values()) + [self.unk_token]
```

#### Step 3: Composite Domain Tokenizer

Assemble per-field tokenizers into a complete domain tokenizer, wrapped as `PreTrainedTokenizerFast`:

```python
# src/tokenizers/domain_tokenizer.py

from tokenizers import Tokenizer, models, trainers, pre_tokenizers
from transformers import PreTrainedTokenizerFast

class DomainTokenizerBuilder:
    """Builds a HuggingFace-compatible tokenizer from a DomainSchema."""
    
    def __init__(self, schema: DomainSchema):
        self.schema = schema
        self.field_tokenizers = {}  # name → field tokenizer instance
        self._build_field_tokenizers()
    
    def _build_field_tokenizers(self):
        for field_spec in self.schema.fields:
            if field_spec.field_type == FieldType.SIGN:
                self.field_tokenizers[field_spec.name] = SignTokenizer(field_spec.name.upper())
            elif field_spec.field_type == FieldType.NUMERICAL_CONTINUOUS:
                self.field_tokenizers[field_spec.name] = MagnitudeBucketTokenizer(
                    n_bins=field_spec.n_bins, prefix=field_spec.name.upper()
                )
            elif field_spec.field_type == FieldType.TEMPORAL:
                self.field_tokenizers[field_spec.name] = CalendarTokenizer(field_spec.calendar_fields)
            # ... other types
    
    def fit(self, data):
        """Fit data-dependent tokenizers (magnitude bins, etc.)."""
        for field_spec in self.schema.fields:
            if field_spec.field_type == FieldType.NUMERICAL_CONTINUOUS:
                values = [getattr(event, field_spec.name) for event in data]
                self.field_tokenizers[field_spec.name].fit(np.array(values))
        return self
    
    def build_hf_tokenizer(self, text_corpus=None, bpe_vocab_size=8000) -> PreTrainedTokenizerFast:
        """Build a complete HuggingFace tokenizer.
        
        1. Collect all domain special tokens
        2. Train BPE on text fields (if any)
        3. Merge into a single PreTrainedTokenizerFast
        """
        # Collect all special tokens from field tokenizers
        all_special_tokens = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "[BOS]", "[EOS]"]
        for name, tok in self.field_tokenizers.items():
            if hasattr(tok, 'vocab'):
                all_special_tokens.extend(tok.vocab)
        
        # Train BPE on text fields
        bpe_tokenizer = Tokenizer(models.BPE(unk_token="[UNK]"))
        bpe_tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
        trainer = trainers.BpeTrainer(
            vocab_size=bpe_vocab_size,
            special_tokens=all_special_tokens,
            min_frequency=2,
        )
        if text_corpus:
            bpe_tokenizer.train_from_iterator(text_corpus, trainer=trainer)
        
        # Wrap as HuggingFace tokenizer
        hf_tokenizer = PreTrainedTokenizerFast(
            tokenizer_object=bpe_tokenizer,
            bos_token="[BOS]",
            eos_token="[EOS]",
            pad_token="[PAD]",
            unk_token="[UNK]",
            mask_token="[MASK]",
        )
        
        return hf_tokenizer
    
    def tokenize_event(self, event) -> List[str]:
        """Convert a single domain event into a list of token strings."""
        tokens = []
        for field_spec in self.schema.fields:
            value = getattr(event, field_spec.name, None)
            if value is None:
                tokens.append("[UNK]")
                continue
            tok = self.field_tokenizers[field_spec.name]
            result = tok(value)
            if isinstance(result, list):
                tokens.extend(result)
            else:
                tokens.append(result)
        return tokens
```

### Phase 2B: Model Architecture (Weeks 3–5)

#### Step 4: GPT-style Causal Transformer (NoPE)

Implement as a HuggingFace-compatible `PreTrainedModel`:

```python
# src/models/configuration_domain_transformer.py

from transformers import PretrainedConfig

class DomainTransformerConfig(PretrainedConfig):
    model_type = "domain_transformer"
    
    def __init__(
        self,
        vocab_size: int = 32000,
        hidden_size: int = 256,        # 256 = 24M params, 1024 = 330M (Nubank sizes)
        num_hidden_layers: int = 24,    # Nubank uses 24 for both sizes
        num_attention_heads: int = 16,  # Nubank uses 16 for both sizes
        intermediate_size: int = None,  # defaults to 4 * hidden_size
        max_position_embeddings: int = 2048,
        dropout: float = 0.1,
        use_positional_encoding: bool = False,  # NoPE by default!
        **kwargs
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size or 4 * hidden_size
        self.max_position_embeddings = max_position_embeddings
        self.dropout = dropout
        self.use_positional_encoding = use_positional_encoding
        super().__init__(**kwargs)
```

```python
# src/models/modeling_domain_transformer.py

import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

class DomainTransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config.hidden_size)
        self.attn = nn.MultiheadAttention(
            config.hidden_size, config.num_attention_heads,
            dropout=config.dropout, batch_first=True
        )
        self.ln2 = nn.LayerNorm(config.hidden_size)
        self.mlp = nn.Sequential(
            nn.Linear(config.hidden_size, config.intermediate_size),
            nn.GELU(),
            nn.Linear(config.intermediate_size, config.hidden_size),
            nn.Dropout(config.dropout),
        )
    
    def forward(self, x, attn_mask=None):
        # Pre-norm architecture
        h = self.ln1(x)
        h, _ = self.attn(h, h, h, attn_mask=attn_mask, is_causal=True)
        x = x + h
        x = x + self.mlp(self.ln2(x))
        return x


class DomainTransformerForCausalLM(PreTrainedModel):
    config_class = DomainTransformerConfig
    
    def __init__(self, config):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        
        # NoPE: no positional encoding by default (Kazemnejad et al. 2023)
        if config.use_positional_encoding:
            self.embed_positions = nn.Embedding(
                config.max_position_embeddings, config.hidden_size
            )
        else:
            self.embed_positions = None
        
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([
            DomainTransformerBlock(config)
            for _ in range(config.num_hidden_layers)
        ])
        self.ln_f = nn.LayerNorm(config.hidden_size)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        # Weight tying (standard for small models)
        self.lm_head.weight = self.embed_tokens.weight
        
        self.post_init()
    
    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        x = self.embed_tokens(input_ids)
        
        if self.embed_positions is not None:
            positions = torch.arange(input_ids.size(1), device=input_ids.device)
            x = x + self.embed_positions(positions)
        
        x = self.drop(x)
        
        for block in self.blocks:
            x = block(x, attn_mask=attention_mask)
        
        x = self.ln_f(x)
        logits = self.lm_head(x)
        
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = nn.functional.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        
        return CausalLMOutputWithPast(loss=loss, logits=logits)

# Register with AutoClass for Hub compatibility
DomainTransformerConfig.register_for_auto_class()
DomainTransformerForCausalLM.register_for_auto_class("AutoModelForCausalLM")
```

#### Step 5: PLR Numerical Embeddings (for Joint Fusion)

Port from Yandex's implementation:

```python
# src/models/plr_embeddings.py

import torch
import torch.nn as nn
import math

class PeriodicLinearReLU(nn.Module):
    """PLR numerical embeddings (Gorishniy et al. 2022).
    
    Maps scalar x → [sin(2π·w·x + b), cos(2π·w·x + b)] → Linear → ReLU
    Frequencies w and phases b are LEARNED parameters.
    """
    
    def __init__(self, n_features: int, n_frequencies: int = 64, embedding_dim: int = 64):
        super().__init__()
        self.n_features = n_features
        self.n_frequencies = n_frequencies
        
        # Learnable frequencies and phases (per feature)
        self.frequencies = nn.Parameter(
            torch.randn(n_features, n_frequencies) * 0.01
        )
        self.phases = nn.Parameter(
            torch.zeros(n_features, n_frequencies)
        )
        
        # Linear projection: 2*n_frequencies → embedding_dim
        self.linear = nn.Linear(2 * n_frequencies, embedding_dim)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (batch, n_features) — raw scalar feature values
        Returns:
            (batch, n_features, embedding_dim)
        """
        # x: (B, F) → (B, F, 1)
        x = x.unsqueeze(-1)
        
        # Periodic encoding: (B, F, n_freq)
        angles = 2 * math.pi * self.frequencies.unsqueeze(0) * x + self.phases.unsqueeze(0)
        periodic = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)  # (B, F, 2*n_freq)
        
        # Linear + ReLU: (B, F, embedding_dim)
        return torch.relu(self.linear(periodic))
```

### Phase 2C: Pre-training (Weeks 5–7)

#### Step 6: Data Pipeline

```python
# src/training/data_pipeline.py

from torch.utils.data import Dataset
from typing import List

class DomainSequenceDataset(Dataset):
    """Converts user event sequences into token sequences for CLM training."""
    
    def __init__(self, user_sequences, tokenizer_builder, hf_tokenizer, max_length=2048):
        self.user_sequences = user_sequences
        self.tokenizer_builder = tokenizer_builder
        self.hf_tokenizer = hf_tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.user_sequences)
    
    def __getitem__(self, idx):
        events = self.user_sequences[idx]
        
        # Tokenize each event into token strings
        token_strings = []
        for event in events:
            event_tokens = self.tokenizer_builder.tokenize_event(event)
            token_strings.extend(event_tokens)
        
        # Convert token strings to IDs via HF tokenizer
        encoding = self.hf_tokenizer(
            " ".join(token_strings),
            max_length=self.max_length,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )
        
        input_ids = encoding["input_ids"].squeeze(0)
        
        return {
            "input_ids": input_ids,
            "labels": input_ids.clone(),  # CLM: labels = input shifted by 1
            "attention_mask": encoding["attention_mask"].squeeze(0),
        }
```

#### Step 7: Pre-training with HuggingFace Trainer

```python
# src/training/pretrain.py

from transformers import Trainer, TrainingArguments

def pretrain_domain_model(
    model,
    train_dataset,
    eval_dataset=None,
    output_dir="./checkpoints",
    hub_model_id="org/domain-model-24m",
    num_epochs=3,
    batch_size=64,
    learning_rate=3e-4,
    context_length=2048,
):
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=4,
        learning_rate=learning_rate,
        lr_scheduler_type="cosine",
        warmup_ratio=0.05,
        weight_decay=0.01,
        logging_strategy="steps",
        logging_steps=100,
        logging_first_step=True,
        disable_tqdm=True,           # plain text logging for cloud
        eval_strategy="steps" if eval_dataset else "no",
        eval_steps=500,
        save_strategy="steps",
        save_steps=1000,
        save_total_limit=3,
        push_to_hub=True,
        hub_model_id=hub_model_id,
        bf16=True,
        dataloader_num_workers=4,
        report_to="trackio",
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
    )
    
    trainer.train()
    trainer.push_to_hub()
```

### Phase 2D: Joint Fusion Fine-tuning (Weeks 7–9)

#### Step 8: nuFormer-style Joint Fusion

```python
# src/models/joint_fusion.py

import torch
import torch.nn as nn

class DCNv2CrossLayer(nn.Module):
    """Single cross layer from DCN V2 (Wang et al. 2021)."""
    
    def __init__(self, dim):
        super().__init__()
        self.weight = nn.Linear(dim, dim, bias=True)
    
    def forward(self, x0, x):
        return x0 * self.weight(x) + x  # element-wise product with anchor


class JointFusionModel(nn.Module):
    """nuFormer-style: Transaction Transformer + DCNv2(PLR) → Joint Prediction.
    
    Architecture:
        Transaction Sequence → Pre-trained DomainTransformer → user_embedding
        Tabular Features → PLR → DCNv2 → tab_embedding
        Concatenate → MLP Head → prediction
    """
    
    def __init__(self, transformer_model, n_tabular_features, n_classes=1,
                 plr_frequencies=64, dcn_layers=3, hidden_dim=256):
        super().__init__()
        
        self.transformer = transformer_model  # pre-trained, unfrozen for fine-tuning
        transformer_dim = transformer_model.config.hidden_size
        
        # Tabular branch: PLR → DCNv2
        self.plr = PeriodicLinearReLU(n_tabular_features, plr_frequencies, hidden_dim)
        tab_input_dim = n_tabular_features * hidden_dim
        
        self.dcn_layers = nn.ModuleList([
            DCNv2CrossLayer(tab_input_dim) for _ in range(dcn_layers)
        ])
        self.dcn_deep = nn.Sequential(
            nn.Linear(tab_input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        
        # Joint head
        self.head = nn.Sequential(
            nn.Linear(transformer_dim + hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, n_classes),
        )
    
    def forward(self, input_ids, attention_mask, tabular_features, labels=None):
        # Transaction branch: get last-token embedding
        transformer_output = self.transformer(input_ids, attention_mask=attention_mask)
        user_embedding = transformer_output.logits[:, -1, :]  # last token representation
        
        # Tabular branch: PLR → flatten → DCNv2
        tab_embedded = self.plr(tabular_features)  # (B, F, D)
        tab_flat = tab_embedded.view(tab_embedded.size(0), -1)  # (B, F*D)
        
        x0 = tab_flat
        x = tab_flat
        for cross_layer in self.dcn_layers:
            x = cross_layer(x0, x)
        tab_output = self.dcn_deep(x)  # (B, hidden_dim)
        
        # Joint fusion
        combined = torch.cat([user_embedding, tab_output], dim=-1)
        logits = self.head(combined)
        
        loss = None
        if labels is not None:
            loss = nn.functional.binary_cross_entropy_with_logits(logits.squeeze(-1), labels.float())
        
        return {"loss": loss, "logits": logits}
```

### Phase 3: Domain Demos (Weeks 9–12)

| Week | Deliverable | Hardware |
|------|------------|----------|
| 9–10 | **Finance demo:** Transaction tokenizer + 24M model pre-trained on synthetic/public financial data + fraud detection fine-tuning | a10g-large |
| 10–11 | **E-commerce demo:** Event tokenizer + 24M model pre-trained on Amazon review sequences + next-purchase prediction | a10g-large |
| 11–12 | **Evaluation & benchmarking:** Compare domain tokenizer vs. text serialization vs. LightGBM baselines on each domain | a10g-large |

### Phase 4: Scale & Optimize (Weeks 12+)

| Task | Details |
|------|---------|
| Scale to 330M params | Increase `hidden_size` to 1024, train on a100-large |
| `torch.compile()` | Apply to attention and MLP blocks for 10–30% speedup |
| ONNX export | `torch.onnx.export()` for production serving |
| Context window experiments | Ablate 512/1024/2048/4096 context lengths |
| Data source ablation | Test impact of different event types (Nubank found adding low-signal sources hurts) |
| ActionPiece vocabulary | Implement BPE-like cross-field merging on top of per-field tokens |

---

## Appendix A: Framework Usage Across Reference Papers

| Paper | ArXiv | Framework | Verbatim Evidence |
|-------|-------|-----------|-------------------|
| nuFormer (Nubank) | 2507.23267 | PyTorch + HF (inferred) | All dependencies are PyTorch-based |
| TIGER (Google) | 2305.05065 | JAX + T5X | "We use the open-sourced T5X framework" |
| ActionPiece (DeepMind) | 2502.13581 | PyTorch + HF | "HuggingFace Transformers and PyTorch" (Appendix H) |
| RecFormer | 2305.13731 | PyTorch + HF Longformer | "Longformer implemented by Huggingface" (§3.1.4) |
| Banking TF | 2410.08243 | PyTorch | "Pytorch backend is used" (Appendix B) |
| PLR Embeddings (Yandex) | 2203.05556 | PyTorch | Repository: pure PyTorch + scikit-learn |
| KL3M Tokenizers | 2503.17247 | HF `tokenizers` + PyTorch | "tokenizers" BPE for HF compatibility |

---

## Appendix B: Head-to-Head Comparison Matrix

| Criterion | PyTorch + HF | JAX/Flax NNX | Keras 3 + JAX |
|-----------|-------------|-------------|---------------|
| Custom domain tokenizer | ✅ `PreTrainedTokenizerFast` | ❌ Build from scratch | ❌ Build from scratch |
| HF Trainer integration | ✅ Native | ❌ Not compatible | ❌ Not compatible |
| Hub push/pull | ✅ `push_to_hub()` | ⚠️ Works, fragmented | ⚠️ Limited |
| PEFT/LoRA | ✅ Drop-in | ❌ Manual | ❌ Manual |
| ONNX export | ✅ One-line | ❌ Multi-hop | ⚠️ TF backend required |
| TGI/vLLM serving | ✅ First-class | ❌ Not supported | ❌ Not supported |
| TPU training | ⚠️ PyTorch/XLA (overhead) | ✅ Native | ✅ Native |
| JIT compilation | ✅ `torch.compile()` | ✅ `@nnx.jit` | ✅ XLA via JAX |
| Dynamic shapes (NLP) | ✅ Native | ❌ Recompiles | ❌ Recompiles |
| Debugging | ✅ Eager mode, std debugger | ⚠️ Challenging | ⚠️ Challenging |
| Reference implementations | 5/6 papers | 1/6 papers | 0/6 papers |
| Community/hiring pool | 🟢 Large | 🟡 Small | 🟡 Small |

---

## Appendix C: Key Code Patterns

### Adding Domain Special Tokens to an Existing Tokenizer

```python
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")  # or any base tokenizer

# Add all domain special tokens
special_tokens = {
    "additional_special_tokens": [
        # Amount tokens (Nubank-style)
        "[AMT_POS]", "[AMT_NEG]",
        *[f"[AMT_{i:02d}]" for i in range(21)],
        # Calendar tokens  
        *[f"[MON_{i:02d}]" for i in range(1, 13)],
        *[f"[DOW_{i}]" for i in range(7)],
        *[f"[DOM_{i:02d}]" for i in range(1, 32)],
        *[f"[HOUR_{i:02d}]" for i in range(24)],
    ]
}
num_added = tokenizer.add_special_tokens(special_tokens)
print(f"Added {num_added} domain tokens. Vocab size: {len(tokenizer)}")

# CRITICAL: resize model embeddings
model.resize_token_embeddings(len(tokenizer))
```

### Registering a Custom Model for Hub Deployment

```python
# In your package's __init__.py or a registration script:
from transformers import AutoConfig, AutoModelForCausalLM

from .configuration_domain_transformer import DomainTransformerConfig
from .modeling_domain_transformer import DomainTransformerForCausalLM

# Register so AutoClass can find your model
AutoConfig.register("domain_transformer", DomainTransformerConfig)
AutoModelForCausalLM.register(DomainTransformerConfig, DomainTransformerForCausalLM)

# Enable push_to_hub with custom code
DomainTransformerConfig.register_for_auto_class()
DomainTransformerForCausalLM.register_for_auto_class("AutoModelForCausalLM")

# Push: uploads configuration.py, modeling.py, config.json, model.safetensors
model.push_to_hub("org/domain-transformer-24m")

# Load anywhere:
model = AutoModelForCausalLM.from_pretrained("org/domain-transformer-24m", trust_remote_code=True)
```

---

*This ADR is a living document and will be updated as implementation progresses.*