[inureyes]

Summary

Adds the shared per-token output surface from epic #2391 and lands it end to end with the BERT / XLM-RoBERTa token-classification head: mlxcel tag tags and redacts entities with Rampart and matches mlx-vlm span for span, and mlxcel extract carries the span-extraction contract that GLiNER2.5 (#2390) will implement.

What changed

• src/token_classification/: TokenClassifierModel, DecodeMode and the windowed TokenClassifier<M> pipeline (mod.rs); a one-to-one port of upstream's label, span, word-grouping, placeholder and redaction helpers with Unicode-scalar offsets (spans.rs); SpanExtractorModel, ExtractOptions, EntityResult and ClassifyTask with upstream's JSON shapes (extract.rs); config-only detection with TokenModelKind and a detector table, one variant plus one arm per family (loader.rs); the BERT pipeline model (bert.rs).
• BertTokenClassifier in src/models/bert_heads.rs: the encoder plus a per-position classifier projection, loaded like BertSequenceClassifier.
• mlxcel tag and mlxcel extract wired into the mlxcel binary (src/main.rs, src/commands/tag.rs, src/commands/extract.rs). A checkpoint of the other shape fails naming the right subcommand.
• Docs: new docs/token-classification.md, a token-classification section in docs/supported-models.md, README, docs index and changelog. The repo has no CLAUDE.md, so the reference-table row from the issue's scope was skipped.

Contract notes for #2387 / #2390 (also in the trait docs): log_probs receives one wrapped window and returns one row per id; the pipeline strips the special rows and applies an f32 log-softmax, so returning logits is fine. decode_modes()[0] is the default, --argmax forces Argmax, and a mode not listed is refused. The id2label vs classifier width check runs at load against the tensor, since detection sees only config.json.

Routing is unchanged: get_model_type still detects Rampart as the BERT embedder (its model_type is encoder-only), so mlxcel embed behaves as before; tag / extract use their own detection. XLM-RoBERTa uses the addressable position count (max_position_embeddings - pad - 1) as its context, where upstream's raw max_position_embeddings would index past the table.

Review hardening: redact takes one forward pass over the sorted, disjoint spans decode_spans produces (upstream's right-to-left splice is kept for any other span list), the token head refuses a position table shorter than max_position_embeddings at load, the tagger refuses a token id outside the word-embedding table, and the docs note that text / spans[].text are not redacted.

Test plan

All runs on GB10 (CUDA, --profile test-fast --features cuda), test binaries under gpu-lock.

• mlxcel lib tests, filter token_classification: 43 passed (span ports, windowing, detection, extract contract, Rampart real-checkpoint and parity), with MLX_ENABLE_TF32=0 and again with TF32 on.
• lib tests, filter bert: 56 passed, including the new BertTokenClassifier head tests; embedding and rerank BERT tests unchanged.
• mlxcel bin tests, filters extract and tag: 6 passed (flag parsing, stub SpanExtractorModel dispatch, help listing).
• Real binary: mlxcel tag -m models/mlx/rampart-mlx-fp16 --keep CITY,STATE,ZIP_CODE "I'm Sarah Connor, ..." prints I'm <GIVEN_NAME> <SURNAME>, <BUILDING_NUMBER> <STREET_NAME>, Los Angeles CA 90012, <EMAIL>, the model card's output; rampart-mlx-4bit (8-bit embeddings, 4-bit linears) prints the same 8 spans with the same offsets.
• Parity fixture tests/fixtures/token_classification/rampart_mlx_vlm_spans.json, from mlx-vlm 0.7.8 on mlx CPU (git 73d366ec; the three sentences are identical under b8bf04c6): three sentences (README example, 617-555-0142 merge, multi-byte German address) plus a 706-token two-window case with the boundary inside a word. Labels, offsets, text and redaction match exactly. The generator rejects any case with a top-1/top-2 gap under 0.05; an earlier repeated-template long case had an exact tie on one token that CUDA broke the other way.
• Not applicable: logit traces (new feature, no existing output moves).

Closes #2383

[inureyes]

Implementation Review Summary

Intent

Add the per-token output surface (TokenClassifierModel pipeline and SpanExtractorModel contract) and land it end to end with the BERT / XLM-RoBERTa token-classification head behind mlxcel tag and mlxcel extract (#2383).

Findings Addressed

• None needed: no CRITICAL or HIGH findings. The span, word-grouping, placeholder and redaction helpers read line for line against upstream mlx_vlm/token_classification.py, and the contract types match the issue body verbatim.

Remaining Items (LOW, not fixed)

• --keep trims entries and drops whitespace-only ones, where upstream run_cli only drops empty strings (LOW), harmless divergence.
• TEXT of - reads stdin, a CLI addition the issue does not specify (LOW), worth a docs line or keep as is.
• group_words falls back to the first tag when an entity has no I state, and a duplicate non-background O continues the open run; upstream raises KeyError on that corner (LOW), safer than upstream.
• BertTokenTagger::log_probs applies a log-softmax and the pipeline applies it again (LOW), idempotent, small redundant compute.
• redact rebuilds the char vector per span, O(text x spans) (LOW), fine at CLI sizes.
• The CLAUDE.md reference-table row from the issue scope was skipped because the repo has no CLAUDE.md (LOW), justified.

Verification

• All stated requirements implemented
• No placeholder/mock code remaining (the extract loader refusing every current kind is the intended contract until GLiNER2.5 lands)
• Integrated into project code flow (mlxcel tag / mlxcel extract subcommands, cli_model_dir, docs)
• Project conventions followed
• Existing modules reused where applicable (BertEncoder, load_encoder / sanitize / quant path, resolve_model_source_with_override, read_embedding_config)
• No unintended structural changes (BertSequenceClassifier::load only extracted read_bert_config; embed / rerank routing untouched)
• Tests pass

Re-run at head 3574a0a9 on GB10 (CUDA, test-fast, MLX_ENABLE_TF32=0): lib filter token_classification 42 passed, bert 55 passed; the Rampart real-checkpoint gates ran (not skipped): README redaction exact on fp16 and 4-bit, fixture parity on 23 / 12 / 42 / 706-token cases. Extra spot check (with the mlxcel binary built at 586393d6; later commits only annotate types, rename a test helper, adjust a test and rustfmt): Call Maria Garcia on 617-555-0142, Müller. with --replacement "[X]" --keep NOPE gives the same 5 spans and offsets as mlx-vlm on both fp16 and 4-bit; Rampart passed to extract fails naming mlxcel tag.

[inureyes]

Security and performance review (head 6d308b1)

No CRITICAL or HIGH findings, so nothing was changed. I audited id2label parsing, the span and offset indexing in decode_spans / token_spans / group_words, the score row-count and row-width checks, the context arithmetic, and --labels-json / --classify parsing. None of them panics on malformed input. There is no HTTP route, so all input comes from the local user or a downloaded checkpoint.

MEDIUM

• src/token_classification/spans.rs:401: redact rebuilds the whole char vector for every span, so it costs O(text_chars x spans). 1 MB of entity-dense stdin (about 1e5 spans) means about 1e11 copies, which is an effective hang. Upstream has the same cost. The fix is one forward pass over the sorted, non-overlapping spans, and it is worth doing before an HTTP route lands.
• src/token_classification/bert.rs:157: the context comes from config.json max_position_embeddings, and nothing checks it (or the vocab size) against the position_embeddings / word_embeddings row counts. A crafted checkpoint (repo-ids auto-download) can declare a huge context. Long text then becomes one window with O(L^2) attention memory, plus out-of-bounds gathers past the tables. The BertEncoder trunk already behaves this way, but tag fills the window by design. Clamp the context to the table's row count at load.

LOW

• src/token_classification/extract.rs:160-174: ClassifyTask::Spec does not range-check its threshold and accepts unknown fields, so a typo such as "multilabel": true is silently ignored.
• src/token_classification/mod.rs:262: literal [SEP] / [CLS] / </s> in the input still tokenize as specials, which lets adversarial text perturb the tagging and slip past redaction. Upstream behaves the same.
• docs/token-classification.md: only redacted_text is redacted. text and spans[].text keep the original PII, and the docs should say so.
• src/commands/tag.rs:104: model_type from config.json goes to stderr unescaped, which allows terminal escape injection. The rest of the CLI does the same.