feat(v4-2): stage_train tokenizes parquet 'text' column via tokenizer_path

Closes V4-2 / cppmega-mlx-862. Fixes G2 backend half: V3-2 only read
raw input_ids — Tokenizer Playground selection had no effect on
training because backend never tokenized anything.

New _tokenize_parquet_text(parquet, tokenizer, n_tokens) helper:
loads HuggingFace Tokenizer.from_file(tokenizer_path), encodes the
'text' column rows until n_tokens collected. Returns ([], None) on any
failure for clean degradation.

stage_train dispatch (priority order):
  1. tokenizer_path + parquet[text] → data_source='parquet_tokenized',
     tokenizer_used=basename
  2. parquet[input_ids|token_ids|tokens] → data_source='parquet' (V3-2)
  3. anything else → data_source='synthetic'

extras gains tokenizer_used: str | null.

5 new pytest (text column path / no-text-column fallback / missing
tokenizer fallback / corrupted parquet / vocab clip). Regression
90/90 across V3-1/V3-2/V3-5/V4-2/schedules/runner.

V4-1 UI already forwards tokenizer_path via stage_options when set;
V4-3 will add the Tokenizer Playground 'Use for training' button so
the user can actually pick one.

Author: apstenku123

cppmega_v4/runner/stages.py

@@ -423,24 +423,38 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
         rng_key = mx.random.key(0)
         data_source = "synthetic"
         token_count = 0
+        tokenizer_used: str | None = None
         parquet_path = opts.get("parquet_path")
+        tokenizer_path = opts.get("tokenizer_path")
+        targets = mx.random.randint(0, vocab_size, shape=(batch, seq))
         if parquet_path:
-            try:
-                tokens = _read_first_n_tokens(parquet_path, n=batch * seq)
+            # V4-2: prefer tokenize(text) path when both tokenizer and a
+            # 'text' column are present; fall back to raw-int input_ids
+            # column (V3-2) when not; fall through to synthetic otherwise.
+            if tokenizer_path:
+                tokens, used = _tokenize_parquet_text(
+                    parquet_path, tokenizer_path, n_tokens=batch * seq)
                 if len(tokens) >= batch * seq:
-                    tokens = [int(t) % vocab_size for t in tokens[:batch * seq]]
+                    tokens = [int(t) % vocab_size
+                              for t in tokens[:batch * seq]]
                     targets = mx.array(tokens, dtype=mx.int32).reshape(
                         batch, seq)
-                    data_source = "parquet"
+                    data_source = "parquet_tokenized"
                     token_count = len(tokens)
-                else:
-                    targets = mx.random.randint(
-                        0, vocab_size, shape=(batch, seq))
-            except Exception:
-                targets = mx.random.randint(
-                    0, vocab_size, shape=(batch, seq))
-        else:
-            targets = mx.random.randint(0, vocab_size, shape=(batch, seq))
+                    tokenizer_used = used
+            if data_source == "synthetic":
+                try:
+                    tokens = _read_first_n_tokens(
+                        parquet_path, n=batch * seq)
+                    if len(tokens) >= batch * seq:
+                        tokens = [int(t) % vocab_size
+                                  for t in tokens[:batch * seq]]
+                        targets = mx.array(tokens, dtype=mx.int32).reshape(
+                            batch, seq)
+                        data_source = "parquet"
+                        token_count = len(tokens)
+                except Exception:
+                    pass
         all_modules = nn.Sequential(*modules, lm_head)
         opt, optimizer_kind = _build_optimizer(spec_optim, lr)
         loss_and_grad = nn.value_and_grad(all_modules, loss_fn)
@@ -532,6 +546,7 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
                 "optimizer_kind": optimizer_kind,
                 "data_source": data_source,
                 "token_count": token_count,
+                "tokenizer_used": tokenizer_used,
                 "model_summary": _summarize_model(
                     ctx.spec, optimizer_kind, schedule_kind_label),
             },
@@ -545,6 +560,41 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
 # ---------------------------------------------------------------------------
 
 
+def _tokenize_parquet_text(
+    parquet_path: str, tokenizer_path: str, n_tokens: int,
+) -> tuple[list[int], str | None]:
+    """V4-2: encode the parquet ``text`` column through a real tokenizer.
+
+    Returns (token_ids, tokenizer_basename) on success, ([], None) on any
+    failure so stage_train can fall through cleanly to the V3-2 raw-int
+    path or the synthetic fallback. We deliberately swallow all exceptions
+    here — the goal is non-fatal degradation, not surfacing tokenizer
+    bugs through the train pipeline.
+    """
+    try:
+        import pyarrow.parquet as pq
+        from tokenizers import Tokenizer
+        from pathlib import Path
+        table = pq.read_table(parquet_path)
+        if "text" not in table.column_names:
+            return [], None
+        tok = Tokenizer.from_file(str(tokenizer_path))
+        out: list[int] = []
+        col = table.column("text")
+        for chunk in col.chunks:
+            for cell in chunk.to_pylist():
+                if cell is None:
+                    continue
+                enc = tok.encode(str(cell))
+                for tid in enc.ids:
+                    out.append(int(tid))
+                    if len(out) >= n_tokens:
+                        return out, Path(tokenizer_path).name
+        return out, Path(tokenizer_path).name
+    except Exception:
+        return [], None
+
+
 def _read_first_n_tokens(parquet_path: str, n: int) -> list[int]:
     """V3-2: read the first ``n`` token ids from a parquet shard so
     stage_train can train against real corpus tokens instead of a

tests/v4/test_stage_train_tokenize.py

@@ -0,0 +1,166 @@
+"""V4-2: stage_train tokenizes parquet text column when tokenizer_path supplied.
+
+The previous V3-2 path only read raw input_ids columns. Now if the
+parquet has a 'text' column and a tokenizer_path is supplied,
+stage_train encodes via the tokenizer and reports
+extras.data_source='parquet_tokenized' + extras.tokenizer_used.
+Falls back cleanly to V3-2 raw-int path or synthetic.
+"""
+
+from __future__ import annotations
+
+import pathlib
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+from tokenizers import Tokenizer, models, pre_tokenizers, trainers
+
+from cppmega_v4.jsonrpc.schema import VerifyParams
+from cppmega_v4.runner import Pipeline, run_pipeline
+
+
+def _verify_params() -> VerifyParams:
+    return VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "attn", "kind": "attention", "params": {}},
+                {"id": "mlp", "kind": "mlp",
+                 "params": {"intermediate_size": 64, "activation": "swiglu"}},
+            ],
+            "edges": [{"src": "attn", "dst": "mlp"}],
+        },
+        "dim_env": {"B": 1, "S": 8, "H": 32, "nh": 2, "nkv": 1,
+                    "head_dim": 16},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["mlp"]},
+        "optim": {"kind": "adamw",
+                  "groups": [{"matcher": "all", "lr": 1e-3,
+                              "weight_decay": 0.01,
+                              "betas": [0.9, 0.95]}]},
+    })
+
+
+@pytest.fixture(scope="module")
+def trained_tokenizer(tmp_path_factory) -> str:
+    """Tiny BPE tokenizer trained on a 4-line synthetic corpus.
+    Returns path to tokenizer.json. Vocab ≈ 64 tokens."""
+    tmp = tmp_path_factory.mktemp("tok")
+    corpus = tmp / "corpus.txt"
+    corpus.write_text(
+        "hello world foo bar baz qux\n"
+        "the quick brown fox jumps over the lazy dog\n"
+        "lorem ipsum dolor sit amet consectetur\n"
+        "abcdefghij klmnopqrst uvwxyz\n"
+    )
+    tok = Tokenizer(models.BPE(unk_token="<unk>"))
+    tok.pre_tokenizer = pre_tokenizers.Whitespace()
+    trainer = trainers.BpeTrainer(
+        vocab_size=64, min_frequency=1,
+        special_tokens=["<unk>"],
+    )
+    tok.train([str(corpus)], trainer)
+    out = tmp / "tokenizer.json"
+    tok.save(str(out))
+    return str(out)
+
+
+def _write_text_parquet(tmp_path: pathlib.Path,
+                        rows: list[str] | None = None) -> str:
+    rows = rows or [
+        "hello world hello world hello world hello world hello world",
+        "the quick brown fox jumps over the lazy dog the quick brown",
+        "lorem ipsum dolor sit amet consectetur adipiscing elit sed do",
+        "foo bar baz qux foo bar baz qux foo bar baz qux foo bar baz",
+    ]
+    table = pa.table({"text": rows})
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    path = tmp_path / "text.parquet"
+    pq.write_table(table, path)
+    return str(path)
+
+
+def _run(stage_options: dict) -> dict:
+    spec = _verify_params()
+    report = run_pipeline(spec, Pipeline.from_dict({
+        "stages": ["parse", "verify_build_spec", "build_model", "train"],
+        "stage_options": stage_options,
+    }))
+    train = next(s for s in report.stages if s.name == "train")
+    assert train.status == "ok", f"stage_train failed: {train.error}"
+    return train.extras
+
+
+def test_stage_train_tokenizes_text_column(tmp_path, trained_tokenizer):
+    """tokenizer_path + parquet[text] → data_source='parquet_tokenized'
+    + tokenizer_used reports the tokenizer basename."""
+    parquet = _write_text_parquet(tmp_path)
+    extras = _run(stage_options={"train": {
+        "num_steps": 2,
+        "parquet_path": parquet,
+        "tokenizer_path": trained_tokenizer,
+    }})
+    assert extras["data_source"] == "parquet_tokenized"
+    assert extras["tokenizer_used"] == "tokenizer.json"
+    assert extras["token_count"] >= 8  # batch*seq = 1*8
+
+
+def test_stage_train_falls_back_to_raw_when_no_text_column(
+    tmp_path, trained_tokenizer,
+):
+    """No 'text' column → falls through to V3-2 input_ids path."""
+    table = pa.table({"input_ids": [list(range(64))]})
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    p = tmp_path / "ids.parquet"
+    pq.write_table(table, p)
+    extras = _run(stage_options={"train": {
+        "num_steps": 2,
+        "parquet_path": str(p),
+        "tokenizer_path": trained_tokenizer,
+    }})
+    # Tokenizer was supplied but text column absent → V3-2 path wins.
+    assert extras["data_source"] == "parquet"
+    assert extras["tokenizer_used"] is None
+    assert extras["token_count"] == 8
+
+
+def test_stage_train_falls_back_to_synthetic_when_tokenizer_missing(
+    tmp_path,
+):
+    """Bad tokenizer_path AND no input_ids column → synthetic."""
+    parquet = _write_text_parquet(tmp_path)
+    extras = _run(stage_options={"train": {
+        "num_steps": 2,
+        "parquet_path": parquet,
+        "tokenizer_path": "/nonexistent/tok.json",
+    }})
+    assert extras["data_source"] == "synthetic"
+    assert extras["tokenizer_used"] is None
+
+
+def test_stage_train_falls_back_when_corrupted_parquet(
+    tmp_path, trained_tokenizer,
+):
+    """Corrupt parquet bytes → synthetic; no crash."""
+    tmp_path.mkdir(parents=True, exist_ok=True)
+    bad = tmp_path / "bad.parquet"
+    bad.write_bytes(b"not a parquet file")
+    extras = _run(stage_options={"train": {
+        "num_steps": 2,
+        "parquet_path": str(bad),
+        "tokenizer_path": trained_tokenizer,
+    }})
+    assert extras["data_source"] == "synthetic"
+    assert extras["tokenizer_used"] is None
+
+
+def test_stage_train_clips_token_ids_to_vocab(tmp_path, trained_tokenizer):
+    """Vocab clip applies to tokenized ids — large id values wrap mod vocab."""
+    parquet = _write_text_parquet(tmp_path)
+    extras = _run(stage_options={"train": {
+        "num_steps": 2, "vocab_size": 32,  # smaller than tokenizer vocab
+        "parquet_path": parquet,
+        "tokenizer_path": trained_tokenizer,
+    }})
+    assert extras["data_source"] == "parquet_tokenized"
+    # Loss must stay finite even though tokens were clipped.
+    assert all(l == l and -1e10 < l < 1e10 for l in extras["losses"])