Add ADR-002: Dataset selection for Phase 3 demos — research findings, rationale, phased plan

docs/adr/ADR-002-dataset-selection.md

@@ -0,0 +1,322 @@
+# ADR-002: Dataset Selection for Phase 3 Domain Demos
+
+> **Status:** Accepted
+> **Date:** April 30, 2026
+> **Decision:** Start with `mindweave/bank-transactions-us` for pipeline validation, then scale to Sparkov (finance), REES46 (e-commerce), and Synthea (healthcare)
+
+---
+
+## 1. Context
+
+Phase 2 delivered a complete library (v0.4.0, 139 tests) with tokenizers, models, and training pipelines — all validated on synthetic data generated in test fixtures. Phase 3 requires running the full pipeline on **real public datasets** to produce trained models and benchmark against baselines.
+
+We need datasets for three domains matching our predefined schemas:
+
+| Schema | Required Fields | Minimum Scale for Demo |
+|--------|----------------|----------------------|
+| `FINANCE_SCHEMA` | timestamp, signed amount, text description | 100+ users × 10+ events |
+| `ECOMMERCE_SCHEMA` | timestamp, price, event type, category, product text | 1000+ users × 10+ events |
+| `HEALTHCARE_SCHEMA` | timestamp, event type, severity, cost, clinical text | 1000+ patients × 10+ events |
+
+The strategy: **start small to validate the pipeline end-to-end, then scale to production-sized datasets.**
+
+---
+
+## 2. Dataset Analysis
+
+### 2.1 Candidates Evaluated
+
+We evaluated 8 datasets across 3 domains. Each was checked for: HuggingFace Hub availability, schema compatibility with our field types, scale (users × events), licensing, and accessibility (instant download vs. gated/external).
+
+#### Finance Candidates
+
+| Dataset | Source | Users | Events | Schema Fit | Access |
+|---------|--------|-------|--------|-----------|--------|
+| **mindweave/bank-transactions-us** | HF Hub | ~20 accounts | ~400 | ✅ Perfect | Instant |
+| **Sparkov CC Fraud** (kartik2112) | Kaggle | ~1,000 | 1.3M | ✅ Excellent | Kaggle account |
+| IBM AML Transactions | GitHub | Thousands | 550K–55M | ✅ Good | Direct download |
+
+**mindweave/bank-transactions-us** — inspected in detail:
+- Config `bank_transactions`: 11 columns, 0.4MB Parquet
+- `transaction_date` (string, `"2024-01-04"`) → maps to `FINANCE_SCHEMA.timestamp` ✅
+- `amount` (float64, signed: `-17584.14` for payroll, `+1413.94` for deposits) → maps to `amount` and `amount_sign` ✅
+- `description` (string, `"Payroll - net wages"`, `"Customer payment received"`) → maps to `description` ✅
+- `source_module` (string, `"payroll"`, `"sales"`, `"purchases"`) → bonus categorical field ✅
+- `transaction_type` (string, `"withdrawal"`, `"deposit"`) → redundant with sign, but useful for validation
+- `bank_account_id` (UUID) → user grouping key ✅
+- Linked `bank_accounts` table has company/bank metadata for potential tabular features
+
+**Scale limitation:** ~20 accounts × ~20 transactions each = ~400 total. This is too small for meaningful pre-training, but **schema-perfect for pipeline validation**: every field maps directly to `FINANCE_SCHEMA` without transformation. The model will overfit immediately, but that's exactly what confirms the pipeline works.
+
+**Sparkov CC Fraud** — the scale-up target:
+- ~1,000 cardholders × ~1,300 transactions each = 1.3M total events
+- `trans_date_trans_time`, `amt`, `merchant`, `category`, `cc_num`, `is_fraud`
+- CC0 license (public domain)
+- `is_fraud` provides a natural fine-tuning label (binary classification)
+- Requires Kaggle account for download (free, instant)
+
+#### E-Commerce Candidates
+
+| Dataset | Source | Users | Events | Schema Fit | Access |
+|---------|--------|-------|--------|-----------|--------|
+| **REES46 Behavioral** | HF Hub | Millions | 42M | ✅ Perfect | Instant |
+| Amazon Reviews 2023 | HF Hub (gated) | 33M | 571M | ⚠️ No price in reviews | HF token |
+
+**REES46** (`kevykibbz/ecommerce-behavior-data-from-multi-category-store_oct-nov_2019`) — inspected:
+- 2GB Parquet (10 files), fully accessible on HF Hub
+- `event_time` (ISO 8601), `event_type` (`"view"`, `"cart"`, `"purchase"`), `product_id`, `category_code` (`"electronics.smartphone"`), `brand`, `price` (float64), `user_id`
+- Every ECOMMERCE_SCHEMA field maps directly
+- For pre-training: filter to `purchase` events for clean transaction sequences
+- Scale: can subsample to 10K–100K users for demo, millions available for production
+
+#### Healthcare Candidates
+
+| Dataset | Source | Patients | Events | Schema Fit | Access |
+|---------|--------|----------|--------|-----------|--------|
+| **Synthea 575K** | HF Hub | 575K | Millions | ✅ Excellent | Instant |
+| Synthea Direct | synthea.mitre.org | 100K–1M | Millions | ✅ Same | Direct download |
+| MIMIC-IV | PhysioNet | 40K+ ICU | Millions | ✅ Gold standard | 1-2 day DUA |
+
+**Synthea 575K** (`richardyoung/synthea-575k-patients`) — inspected:
+- 136GB total across 18 Parquet files (allergies, conditions, encounters, medications, observations, procedures, etc.)
+- Default config shows allergies table: `START` (date), `PATIENT` (UUID), `DESCRIPTION`, `TYPE`, `CATEGORY`, `SEVERITY1`
+- For richer sequences: load `encounters.parquet` (5.1GB) with `Start`, `DESCRIPTION`, `Base_Cost`, `REASONDESCRIPTION`
+- Fully synthetic — no IRB, no access restrictions, MIT/Apache 2.0 license
+
+### 2.2 Schema Mapping Verification
+
+Direct field mapping from `mindweave/bank-transactions-us` to `FINANCE_SCHEMA`:
+
+```
+Dataset Column          →  FINANCE_SCHEMA Field     →  Tokenizer
+─────────────────────────────────────────────────────────────────
+amount (sign)           →  amount_sign              →  SignTokenizer (2 tokens)
+amount (magnitude)      →  amount                   →  MagnitudeBucketTokenizer (21 bins)
+transaction_date        →  timestamp                →  CalendarTokenizer (month/dow/dom/hour)
+description             →  description              →  BPE subword tokenizer
+─────────────────────────────────────────────────────────────────
+bank_account_id         →  (user grouping key)      →  group-by for user sequences
+source_module           →  (bonus: not in schema)   →  could extend schema
+transaction_type        →  (redundant with sign)    →  validation check
+```
+
+**Zero transformation needed.** The `amount` field is already signed (negative = withdrawal, positive = deposit). The `description` field contains natural text suitable for BPE. The `transaction_date` is a standard date string. This is the cleanest possible mapping to our schema.
+
+---
+
+## 3. Decision
+
+### Phased approach: validate small → scale up
+
+| Phase | Dataset | Purpose | Scale |
+|-------|---------|---------|-------|
+| **3.0: Pipeline Validation** | `mindweave/bank-transactions-us` | Verify end-to-end: load → tokenize → pack → train → loss decreases | ~400 events, ~20 accounts |
+| **3.1: Finance Demo** | Sparkov CC Fraud (Kaggle) | Train 24M model, fine-tune fraud detection, benchmark vs LightGBM | 1.3M events, 1K users |
+| **3.2: E-Commerce Demo** | REES46 (HF Hub) | Train 24M model, next-purchase prediction | 42M events, subsample to 100K users |
+| **3.3: Healthcare Demo** | Synthea 575K (HF Hub) | Train 24M model, condition prediction | 575K patients, subsample encounters |
+
+### Rationale
+
+1. **Start with mindweave because it's schema-perfect and instant.** No data cleaning, no field renaming, no Kaggle credentials needed. The pipeline either works or it doesn't — this dataset tells us in minutes.
+
+2. **The model will overfit on 400 events — that's the point.** If loss doesn't decrease on 400 events, the pipeline is broken. If it does, the pipeline works and we can scale with confidence.
+
+3. **Sparkov is the real finance demo.** 1,000 users × 1,300 events is the exact scale where a 24M-parameter model should learn meaningful patterns. The `is_fraud` label enables a direct comparison with LightGBM on the same data.
+
+4. **REES46 is the flagship demo.** Millions of events, real behavioral data, perfect schema fit, instant HF download. This is the dataset that demonstrates domainTokenizer's value proposition most compellingly.
+
+5. **Synthea is the healthcare proof point.** Fully synthetic (no access barriers), massive scale, multiple event types. Validates that the domain tokenizer approach generalizes beyond finance and e-commerce.
+
+---
+
+## 4. Implementation
+
+### 4.1 Phase 3.0: Pipeline Validation with mindweave
+
+**Goal:** Run the complete pipeline end-to-end on real data, verify loss decreases, confirm no bugs.
+
+**Step 1: Load and explore the data**
+
+```python
+from datasets import load_dataset
+import pandas as pd
+
+# Load bank transactions
+ds = load_dataset("mindweave/bank-transactions-us", "bank_transactions", split="train")
+df = ds.to_pandas()
+
+# Basic stats
+print(f"Total transactions: {len(df)}")
+print(f"Unique accounts: {df['bank_account_id'].nunique()}")
+print(f"Date range: {df['transaction_date'].min()} to {df['transaction_date'].max()}")
+print(f"Amount range: {df['amount'].min():.2f} to {df['amount'].max():.2f}")
+print(f"Descriptions: {df['description'].nunique()} unique")
+print(f"Source modules: {df['source_module'].value_counts().to_dict()}")
+```
+
+**Step 2: Convert to domainTokenizer event format**
+
+```python
+from datetime import datetime
+
+def row_to_event(row):
+    """Convert a DataFrame row to a FINANCE_SCHEMA event dict."""
+    return {
+        "amount_sign": row["amount"],         # SignTokenizer reads the sign
+        "amount": row["amount"],               # MagnitudeBucketTokenizer reads abs value
+        "timestamp": datetime.strptime(row["transaction_date"], "%Y-%m-%d"),
+        "description": row["description"],     # BPE tokenizer
+    }
+
+# Group by account → list of event sequences
+user_sequences = []
+for account_id, group in df.sort_values("transaction_date").groupby("bank_account_id"):
+    events = [row_to_event(row) for _, row in group.iterrows()]
+    user_sequences.append(events)
+
+print(f"Users: {len(user_sequences)}")
+print(f"Events per user: {[len(s) for s in user_sequences]}")
+```
+
+**Step 3: Build tokenizer, prepare data, train**
+
+```python
+from domain_tokenizer import (
+    DomainTokenizerBuilder, DomainTransformerConfig,
+    DomainTransformerForCausalLM, prepare_clm_dataset, pretrain_domain_model,
+)
+from domain_tokenizer.schemas import FINANCE_SCHEMA
+
+# Build tokenizer
+all_events = [e for seq in user_sequences for e in seq]
+builder = DomainTokenizerBuilder(FINANCE_SCHEMA)
+builder.fit(all_events)
+hf_tokenizer = builder.build(
+    text_corpus=[e["description"] for e in all_events],
+    bpe_vocab_size=500,  # small vocab for small dataset
+)
+
+# Prepare packed dataset
+dataset = prepare_clm_dataset(user_sequences, builder, hf_tokenizer, block_size=128)
+print(f"Packed blocks: {len(dataset)} × 128 tokens")
+
+# Create tiny model (for validation, not real training)
+config = DomainTransformerConfig(
+    vocab_size=hf_tokenizer.vocab_size,
+    hidden_size=128, num_hidden_layers=4, num_attention_heads=4,
+    intermediate_size=512,
+)
+model = DomainTransformerForCausalLM(config)
+print(f"Model params: {sum(p.numel() for p in model.parameters()):,}")
+
+# Train — expect loss to decrease rapidly (overfitting on small data = pipeline works)
+pretrain_domain_model(
+    model, hf_tokenizer, dataset,
+    num_epochs=20,
+    per_device_batch_size=4,
+    gradient_accumulation_steps=1,
+    learning_rate=3e-4,
+    warmup_steps=10,
+    logging_steps=5,
+    save_strategy="no",
+    report_to="none",
+)
+```
+
+**Expected outcome:** Loss should drop from ~6.0 to <2.0 within 20 epochs on 400 events. If it does, the pipeline is validated. If it doesn't, there's a bug in tokenization, packing, or model architecture.
+
+**Validation checks after training:**
+- [ ] Loss decreased monotonically (overfitting expected and desired)
+- [ ] No NaN/inf in loss or gradients
+- [ ] Token distribution is reasonable (no >50% UNK tokens)
+- [ ] `builder.tokenize_event()` produces expected token strings for sample events
+- [ ] `hf_tokenizer.decode()` on model output produces recognizable token strings
+
+### 4.2 Phase 3.1: Finance Demo with Sparkov (After Validation)
+
+```bash
+# Download from Kaggle
+kaggle datasets download kartik2112/fraud-detection -p data/
+unzip data/fraud-detection.zip -d data/sparkov/
+```
+
+```python
+import pandas as pd
+
+df = pd.read_csv("data/sparkov/fraudTrain.csv")
+
+def sparkov_to_event(row):
+    return {
+        "amount_sign": row["amt"],           # always positive in Sparkov; sign from context
+        "amount": row["amt"],
+        "timestamp": datetime.strptime(row["trans_date_trans_time"], "%Y-%m-%d %H:%M:%S"),
+        "description": f"{row['merchant']} {row['category']}",
+    }
+
+# Group by cardholder
+user_sequences = []
+labels = []  # for fine-tuning: any fraud in user's history?
+for cc_num, group in df.sort_values("trans_date_trans_time").groupby("cc_num"):
+    events = [sparkov_to_event(row) for _, row in group.iterrows()]
+    user_sequences.append(events)
+    labels.append(int(group["is_fraud"].any()))
+
+# Pre-train 24M model on 1K users × 1.3K events
+config = DomainTransformerConfig.from_preset("24m", vocab_size=hf_tokenizer.vocab_size)
+model = DomainTransformerForCausalLM(config)
+# ... pretrain_domain_model(model, ..., bf16=True)  # requires GPU
+
+# Fine-tune for fraud detection
+# ... finetune_domain_model(fusion_model, ft_dataset, ...)
+```
+
+**Hardware:** a10g-large (24GB VRAM), ~2-3 hours for 24M model on 1.3M events.
+
+### 4.3 Phase 3.2: E-Commerce Demo with REES46
+
+```python
+from datasets import load_dataset
+
+ds = load_dataset(
+    "kevykibbz/ecommerce-behavior-data-from-multi-category-store_oct-nov_2019",
+    split="train",
+)
+
+# Filter to purchases and subsample users
+purchases = ds.filter(lambda x: x["event_type"] == "purchase")
+# Group by user_id, take top 100K users by event count
+# ... build ECOMMERCE_SCHEMA tokenizer, train 24M model
+```
+
+### 4.4 Phase 3.3: Healthcare Demo with Synthea
+
+```python
+from huggingface_hub import hf_hub_download
+import pandas as pd
+
+encounters = pd.read_parquet(hf_hub_download(
+    "richardyoung/synthea-575k-patients",
+    "data/encounters.parquet",
+    repo_type="dataset",
+))
+
+# Group by PATIENT, sort by Start date
+# Map: Start→timestamp, Base_Cost→amount, DESCRIPTION→description
+# ... build HEALTHCARE_SCHEMA tokenizer, train 24M model
+```
+
+---
+
+## 5. Risks and Mitigations
+
+| Risk | Impact | Mitigation |
+|------|--------|-----------|
+| mindweave too small to catch scale bugs | Bugs only surface at 1M+ events | Run Sparkov immediately after validation passes |
+| Sparkov has no negative amounts | `SignTokenizer` always produces `[AMT_SIGN_POS]` | Concatenate merchant+category as description; test sign tokenizer separately on mindweave (which has signed amounts) |
+| REES46 2GB download slow | Delays e-commerce demo | Stream via HF datasets `streaming=True` or subsample first |
+| Synthea encounters lack numerical values | `MagnitudeBucketTokenizer` underutilized | Use `Base_Cost` for cost binning; join with `observations.parquet` for lab values |
+| Model overfits on 400 events | Expected — not a bug | Overfitting on validation set = pipeline works. Move to Sparkov for real training. |
+
+---
+
+*This ADR will be updated with results from each phase as demos are completed.*