Module tagger_train 

Step 2: tagger finetuning in Rust (candle) — the port of the reference spo_tagger.py::train.

Trains the multi-head tagger whose design is fixed in [crate::tagger_data]:

• Head A — BIO span typing over the spec’s semantic facets + structural kinds (dims 1, 3, 4)
• Head B — the 4-way epistemic reading (dim 5), which is also the infon polarity i the Dempster-Shafer layer consumes

The shared encoder is finetuned, not frozen: the pretrained tensors are named bert.*, so building the model under vb.pp("bert") lets the checkpoint load straight into the trainable VarMap, and the two new heads simply have no counterpart in the file (loaded per-tensor, missing names skipped).

Loss is token-level cross-entropy on both heads with an ignore mask: padding and sub-token continuations that carry no label are marked -100 and dropped by gathering valid positions before the CE, which candle’s cross_entropy does not do for us.

Head C (the biaffine relation scorer) is deliberately not here — it needs span pooling over the encoder output, so it lands as a second stage once Head A’s spans are reliable.

Structs

Encoded

One encoded example: token ids, attention mask, and per-token targets for both heads.

EpochReport

PredictedSpan

A span predicted by the tuned tagger: byte offsets, its facet (Head A) and epistemic reading (Head B, majority vote over the span’s tokens → the infon polarity i).

TrainConfig

TrainReport

TunedTagger

A tuned tagger ready for inference: weights + the label spaces they were trained with.

Functions

encode

Tokenize + project char-span labels onto tokens, padding to max_len. Reuses the offset projection already proven in [crate::tagger_data], so alignment semantics are identical between data generation and training.

hrm_config_from

Build the HRM config from a HuggingFace base config.json (only the embedding-table shapes are used; the encoder settings are irrelevant because the encoder is discarded).

save

Persist the finetuned tagger: weights plus the label spaces needed to decode it.

train

Finetune Head A + Head B. Returns the report; the trained weights are left in varmap for saving.

Type Aliases

MultiHeadTagger

The tagger is the reference HRM architecture: BERT embeddings only + the two-timescale reasoning core + independent heads (crate::hrm). Earlier revisions of this file used BERT’s full encoder with heads attached, which is a different model — capacity there comes from stacked layers, whereas the reference gets it from recurrent refinement at ~4M parameters.