From b453cf3d417c3a5cf59cce1105628cb61b73934c Mon Sep 17 00:00:00 2001 From: Stephen Solka Date: Tue, 18 Aug 2026 12:59:54 -0400 Subject: [PATCH 1/3] Add pdoc API reference --- .gitignore | 1 + .zed/tasks.json | 10 + README.md | 12 ++ mise.toml | 5 + pyproject.toml | 3 +- src/langset/__init__.py | 16 ++ src/langset/bridge_emission.py | 118 +++++----- src/langset/data.py | 20 +- src/langset/heads.py | 202 +++++++++-------- src/langset/loss.py | 29 ++- src/langset/masking.py | 161 +++++++++----- src/langset/modeling.py | 276 ++++++++++++++---------- src/langset/selection.py | 12 +- src/langset/st_module.py | 13 +- src/langset/strategies.py | 383 +++++++++++++++------------------ src/langset/trainer.py | 265 +++++++++++------------ src/langset/training_args.py | 14 +- uv.lock | 32 ++- 18 files changed, 899 insertions(+), 673 deletions(-) create mode 100644 .zed/tasks.json diff --git a/.gitignore b/.gitignore index f39acdd..cb1dc37 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ __pycache__/ .eggs/ build/ dist/ +docs/ # environments .venv/ diff --git a/.zed/tasks.json b/.zed/tasks.json new file mode 100644 index 0000000..70ab052 --- /dev/null +++ b/.zed/tasks.json @@ -0,0 +1,10 @@ +[ + { + "label": "Docs: serve live API reference", + "command": "uv run --extra docs pdoc src/langset/*.py", + "use_new_terminal": true, + "allow_concurrent_runs": false, + "reveal": "always", + "hide": "never" + } +] diff --git a/README.md b/README.md index 1c7a167..9454aad 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,18 @@ trajectory and the uncertainty inside each step — no hand-built simulator. Wha pip install langset ``` +## API documentation + +Generate the API reference locally with [pdoc](https://pdoc.dev/): + +```bash +uv sync --extra docs +mise run docs +open docs/langset.html +``` + +The generated `docs/` directory is local build output and is intentionally not committed. The reference includes the public package plus its implementation modules, including `langset.loss`. + ## Quickstart — a world model Rows are `input_text` (the current state) → `target_texts` (the **ordered ticks** to emit). A matching `concepts` diff --git a/mise.toml b/mise.toml index 4e1d01d..84e59c0 100644 --- a/mise.toml +++ b/mise.toml @@ -23,3 +23,8 @@ run = "uv run ruff format --check src tests" description = "Run the test suite" sources = ["src/**/*.py", "tests/**/*.py", "pyproject.toml"] run = "uv run pytest tests/" + +[tasks.docs] +description = "Generate API documentation with pdoc" +sources = ["src/**/*.py", "README.md", "pyproject.toml"] +run = "uv run --extra docs pdoc -o docs src/langset/*.py" diff --git a/pyproject.toml b/pyproject.toml index ea839e9..4ec5dcb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,8 @@ setfit = ["setfit>=1.1.0", "transformers>=4.41,<4.47", "torch>=2.2,<2.5", "sciki wandb = ["wandb>=0.18"] # world-model property probes (langset.probes): calibration_corr + linear_decodability probes = ["scikit-learn", "scipy>=1.10"] -dev = ["pytest", "ruff", "ty>=0.0.61", "scikit-learn", "scipy>=1.10"] +docs = ["pdoc>=15"] +dev = ["pytest", "ruff", "ty>=0.0.61", "scikit-learn", "scipy>=1.10", "pdoc>=15"] [tool.hatch.build.targets.wheel] packages = ["src/langset"] diff --git a/src/langset/__init__.py b/src/langset/__init__.py index 44a0c73..d04ee7e 100644 --- a/src/langset/__init__.py +++ b/src/langset/__init__.py @@ -3,6 +3,22 @@ Few-shot fine-tune a pretrained LLM to predict in latent space: emit a sequence of latents (a JEPA world model) that holds a calibrated superposition of next states — or a single latent as a bespoke embedding model, which is Sentence-Transformer-shaped and drops into SetFit as a `model_body`. + +## Module reference + +- [Modeling](langset/modeling.html) — model and emission heads +- [Training](langset/trainer.html) — training loop and orchestration +- [Training arguments](langset/training_args.html) — run configuration +- [Losses](langset/loss.html) — composable training objectives +- [Strategies](langset/strategies.html) — emission and target strategies +- [Heads](langset/heads.html) — auxiliary supervised heads +- [Masking](langset/masking.html) — input masking helpers +- [Probes](langset/probes.html) — representation-quality probes +- [Bridge emission](langset/bridge_emission.html) — continuous vector-set emission +- [Data](langset/data.html) — dataset helpers +- [Selection](langset/selection.html) — model selection utilities +- [SIGReg](langset/sigreg.html) — isotropic regularization +- [Sentence-transformer integration](langset/st_module.html) — embedding-model adapter """ from langset.heads import Head diff --git a/src/langset/bridge_emission.py b/src/langset/bridge_emission.py index 541bebf..01c8c4c 100644 --- a/src/langset/bridge_emission.py +++ b/src/langset/bridge_emission.py @@ -1,23 +1,20 @@ -"""QueryBridgeEmission — a NON-autoregressive, continuous, parallel-query emission family, ported from the -validated proof (ii3 `bridge_lightning.py`) into langset's `_EmissionObjective` seam. - -One forward: N learned query tokens cross-attend the (frozen) backbone's per-token hidden states → N vectors + -a validity/count logit each. Trained by DETR Hungarian matching to the row's target latents, then per-fact -InfoNCE (cross-row negatives) + validity BCE. It reuses the widened seams: `build_targets` (default positional -teacher-forcing supplies the target latents), the objective-before-optimizer order (its module registers on the -model, so `m.parameters()` trains it), and the canonical `EmissionOut` log keys. - -Inject with `TrainingArguments(emission=QueryBridgeEmission, freeze_backbone=True)`. Pair with a query/frozen -`_TargetSource` and set `lam_multi_nce=0` (its base loss already does the contrastive term). `n_queries` (default -16) is read off args via getattr — no TrainingArguments change required. - -HARD NEGATIVES: set `hard_neg_field=""` (a per-row list of confusable texts, e.g. adjacent-quarter queries) and -the family folds them into its OWN InfoNCE denominator — one softmax over [in-batch targets ++ hard-negs], the exact -mechanism that lets the emitted vector separate same-entity/wrong-time near-misses. Absent `hard_neg_field`, the bank is plain in-batch (byte-identical). - -Retrieval is preserved BY CONSTRUCTION: the backbone is frozen, so the base embedder's geometry is untouched; -the bridge is a pure add-on. Eval routes through `emit_infer` (one pass, validity-gated) instead of the AR -rollout. +"""Parallel-query continuous emission for multi-latent models. + +`QueryBridgeEmission` uses learned query vectors to cross-attend the backbone's token-level hidden states. Each +query produces an L2-normalized latent vector and a validity logit. During training, Hungarian matching assigns +query slots to target latents, and the objective combines an InfoNCE loss over matched vectors with binary +cross-entropy over validity logits. + +A typical configuration uses `QueryBridgeEmission` with `FrozenEncoderTarget`, `freeze_backbone=True`, and +`lam_multi_nce=0`. The separate multi-latent NCE term should be disabled because this objective already includes +a contrastive loss. `n_queries` controls the maximum number of vectors emitted per input. + +When `hard_neg_field` is configured, the corresponding texts are encoded and added to the InfoNCE target bank as +additional negatives. At inference time, the bridge emits query slots whose validity probability exceeds 0.5, +subject to `max_steps`, with a one-slot fallback when none pass the threshold. + +Freezing the backbone keeps its existing embedding parameters unchanged; this emission strategy does not freeze +the backbone automatically. """ from __future__ import annotations @@ -40,10 +37,14 @@ class FrozenEncoderTarget(_TargetSource): - """Targets = the FROZEN encoder's OWN embedding of the target texts, i.e. E(text). Pair with - `target_texts = queries` to reproduce the 'query-target' objective (emitted vectors trained to match E(query), - which is exactly what the bench scores). No EMA twin and no self-distillation — the frozen model IS the target, - so `twin = model` and `update` is a no-op. Retrieval geometry preserved by construction.""" + """Use the model's normalized text embeddings as target latents. + + `encode()` evaluates target texts with the same `LangSetModel` used for training and does not track gradients. + The target source has no EMA copy: `twin` refers to the model itself and `update()` is a no-op. + + Use `freeze_backbone=True` when the target embedding space must remain fixed. This target source does not freeze + model parameters itself. + """ suppresses_nce = False @@ -55,14 +56,14 @@ def __init__( dev: torch.device, ) -> None: self.m, self.a, self.tok, self.dev = model, args, tok, dev - self.twin = model # eval encodes its retrieval bank with the frozen model itself + self.twin = model # Evaluation encodes the retrieval bank with the same model. def encode(self, texts: list[str]) -> torch.Tensor: with torch.no_grad(): z = self.m.encode(texts, convert_to_numpy=False, normalize_embeddings=True) - return z.to(self.dev).float() # ty: ignore[unresolved-attribute] # encode(convert_to_numpy=False) -> Tensor + return z.to(self.dev).float() # ty: ignore[unresolved-attribute] # Tensor when convert_to_numpy=False - def update(self) -> None: # nothing to track — the target model is frozen + def update(self) -> None: # This target source has no EMA state to update. pass @@ -74,7 +75,10 @@ def _heads_for(d: int) -> int: class QueryBridge(nn.Module): - """N learned queries cross-attend the frozen substrate -> N L2-normalized vectors + per-query validity logit.""" + """Decode learned queries against token-level hidden states. + + The module returns one L2-normalized vector and one validity logit for each query slot. + """ def __init__(self, d: int, n_queries: int, n_layers: int = 2) -> None: super().__init__() @@ -93,7 +97,10 @@ def forward( class QueryBridgeEmission(_EmissionObjective): - """Parallel-query continuous emission (see module docstring). `codebook=False`: no codebook and no AR rollout.""" + """Train and infer a set of continuous latent vectors in one decoder pass. + + This objective uses no discrete codebook and implements non-autoregressive inference through `emit_infer()`. + """ codebook = False @@ -105,7 +112,7 @@ def __init__( trainer: "Trainer", ) -> None: super().__init__(model, args, dev, trainer) - try: # SciPy is an optional dep — only this emission family needs it (Hungarian matching) + try: # SciPy is required only for Hungarian matching in this emission strategy. from scipy.optimize import linear_sum_assignment except ModuleNotFoundError as e: # pragma: no cover - trivial guard raise ModuleNotFoundError( @@ -120,13 +127,12 @@ def __init__( self.lam_valid = float(getattr(args, "bridge_lam_valid", 1.0)) self.pos_weight = float(getattr(args, "bridge_pos_weight", 2.0)) bridge = QueryBridge(d, self.n_queries).to(dev) - # Register on the model so `m.parameters()` (built AFTER the objective now) trains it. Persistence via the - # aux-head plug is a follow-up; for now this makes the module optimizer-visible. + # Register the bridge on the model so optimization and checkpoint persistence include it. model.add_module("emission_bridge", bridge) pending = getattr(model, "_emission_bridge_state", None) if ( pending is not None - ): # restore weights persisted by save_pretrained (serve-time / resumed run) + ): # `LangSetModel.load()` stages persisted weights until this strategy is attached. bridge.load_state_dict(pending) self.bridge = bridge @@ -134,8 +140,10 @@ def parameters(self) -> Iterable[nn.Parameter]: return self.bridge.parameters() def _hard_negs(self, bidx: list[int]) -> list[str]: - """Flatten this batch's per-row hard-negative texts (from the trainer's `hard_neg_field` column) into one - pooled list. Empty when no `hard_neg_field` is configured (or a serve-time attach without a trainer).""" + """Collect nonblank hard-negative texts for the selected dataset rows. + + Returns an empty list when no trainer hard-negative column is available. + """ hn = getattr(self.trainer, "hard_neg_texts", None) if not hn: return [] @@ -156,28 +164,27 @@ def emit( m, dev = self.m, self.dev d = target_lat.size(-1) ids, mask = se["input_ids"], se["attention_mask"] - # per-token substrate from the (frozen when freeze_backbone) backbone — same read seed_hidden does, un-pooled + # Read unpooled token-level hidden states from the backbone. substrate = m._last_hidden(m._run_backbone(m.embed(ids), mask, ids, 0)).float() # [B, T, d] vecs, vlog = self.bridge(substrate, mask.bool()) # [B, nq, d], [B, nq] tgt = F.normalize(target_lat.float(), dim=-1) # [B, lmax, d] bank = tgt[ valid - ] # [Ntot, d] — cross-row (in-batch) InfoNCE negatives; positives index into THIS prefix - # HARD-NEG bank: adjacent-quarter (same-entity / wrong-time) confusables from `hard_neg_field`, encoded in the - # SAME (query) space as the targets and folded into the SAME InfoNCE denominator — the mechanism the standalone - # used to beat entity×time. Absent hard_neg_field -> no-op, byte-identical to the plain in-batch bank. + ] # Valid targets lead the InfoNCE bank; positive indices use this ordering. + # Encode configured hard negatives in the target embedding space and append them to the InfoNCE denominator. + # They are negatives only, so positive indices remain valid. hn_texts = self._hard_negs(bidx) if hn_texts: with torch.no_grad(): hn = self.m.encode(hn_texts, convert_to_numpy=False, normalize_embeddings=True) - bank = torch.cat([bank, hn.to(dev).float()], dim=0) # ty: ignore[unresolved-attribute] # -> Tensor + bank = torch.cat([bank, hn.to(dev).float()], dim=0) # ty: ignore[unresolved-attribute] # encode returns a Tensor recon = torch.zeros(b, lmax, d, device=dev) vlab = torch.zeros_like(vlog) matched_pred: list[torch.Tensor] = [] pos: list[int] = [] - nq = vecs.size(1) # number of query slots = MAX matchable emissions per row - # ONE device sync for the whole batch: query×target cosine, then Hungarian per row on CPU (SciPy is CPU-only). + nq = vecs.size(1) # Maximum number of targets that can be matched in each row. + # Transfer all query-target similarities to CPU once, then run SciPy matching per row. sims = torch.bmm(vecs, tgt.transpose(1, 2)).detach().float().cpu().numpy() # [B, nq, lmax] off = 0 for r in range(b): @@ -186,12 +193,10 @@ def emit( continue mi_eff = min( mi, nq - ) # guard: SciPy returns min(nq, mi) pairs — extra targets stay in the bank as negs - pr, tc = self._match(-sims[r, :, :mi_eff]) # cost = -cos over [nq, mi_eff] + ) # At most `nq` targets can be matched; additional targets remain negatives. + pr, tc = self._match(-sims[r, :, :mi_eff]) # Maximize cosine similarity. for p, c in zip(pr, tc): - recon[r, c] = vecs[ - r, p - ] # matched emission -> its target's slot (aligns recon with `valid`) + recon[r, c] = vecs[r, p] # Place the matched prediction in its target-aligned slot. matched_pred.append(vecs[r, p]) pos.append(off + int(c)) vlab[r, p] = 1.0 @@ -218,21 +223,28 @@ def emit( return EmissionOut( recon=recon, base_loss=base, - # reuse the canonical log keys (semantic map: recon_loss<-InfoNCE, loss_stop<-validity BCE, loss_dims<-0) + # Preserve canonical fields: `recon_loss` is InfoNCE, `loss_stop` is validity BCE, and `loss_dims` is unused. logs={"loss_stop": vloss.detach(), "loss_dims": zero, "recon_loss": nce.detach()}, code_logits=None, ) @torch.no_grad() def emit_infer(self, texts: list[str], max_steps: int) -> tuple[torch.Tensor, torch.Tensor]: - """One-pass inference: text -> a validity-gated fact-vector SET per row (NO autoregression). Returns - (lat [B, Lmax, d], lens [B]) padded — the same shape the AR rollout yields, so eval/retrieval is unchanged.""" + """Emit validity-gated latent vectors for a batch of texts. + + The bridge runs once per batch and keeps query slots whose validity probability exceeds 0.5. If no slot + passes the threshold, the highest-logit slot is retained. At most `max_steps` vectors are returned per input. + + Returns: + A pair `(latents, lengths)`. `latents` has shape `[batch, max_emitted, latent_dim]` and is zero-padded; + `lengths` contains the number of retained vectors for each input. + """ m, dev = self.m, self.dev e = m.tokenizer( texts, padding=True, truncation=True, - max_length=self.a.max_len, # honor the run's TrainingArguments.max_len (== training truncation) + max_length=self.a.max_len, # Match the truncation length used during training. padding_side="left", return_tensors="pt", ).to(dev) @@ -244,7 +256,7 @@ def emit_infer(self, texts: list[str], max_steps: int) -> tuple[torch.Tensor, to lens: list[int] = [] for r in range(vecs.size(0)): idx = keep[r].nonzero(as_tuple=True)[0] - if idx.numel() == 0: # fallback: never emit zero — keep the single most-confident query + if idx.numel() == 0: # Always retain the most confident slot. idx = vlog[r].argmax().unsqueeze(0) idx = idx[:max_steps] rows.append(vecs[r, idx]) diff --git a/src/langset/data.py b/src/langset/data.py index 3ac354d..cfbe66f 100644 --- a/src/langset/data.py +++ b/src/langset/data.py @@ -1,13 +1,17 @@ -"""Dataset contract for langset. +"""Dataset contracts and small helpers for preparing langset training data. -A training example is a row with: - - `input_text` : the text you'll have at inference (e.g. a name, a query, a review) - - `target_text` : a description of the SAME item that DEFINES where it should land. The geometry is whatever - these descriptions describe — point them at the axis you care about. Point them at something - `input_text` can't trivially regenerate, or you're just distilling a text encoder. +The core training contract is a row with two views of the same item: -Pass a `datasets.Dataset` or a `list[dict]` to `Trainer`; use `column_mapping` to rename your columns onto -`input_text` / `target_text`. +- ``input_text`` is the text available at inference time, such as a name, + query, or review. +- ``target_text`` is a description of that same item that defines the geometry + the model should learn. Choose a target that exposes the property of + interest and cannot be trivially regenerated from ``input_text``; otherwise + the objective mostly distills a text encoder. + +Pass a ``datasets.Dataset`` or ``list[dict]`` to :class:`langset.Trainer`. Use +``column_mapping`` when an existing dataset uses different column names, or +use :func:`from_records` to project arbitrary records onto the basic contract. """ from __future__ import annotations diff --git a/src/langset/heads.py b/src/langset/heads.py index 29f706e..206fea4 100644 --- a/src/langset/heads.py +++ b/src/langset/heads.py @@ -1,29 +1,18 @@ -"""Pluggable AUXILIARY-SUPERVISED HEADS — the generalization of the hard-coded phase head. - -The multi-latent trainer historically carried ONE domain concept baked into the library: a transient -`hidden->phase` CE classifier (`lam_phase`) that shaped the emitted geometry to be phase-separable. Baking -"phase" (and a would-be "value") into a domain-AGNOSTIC latent-prediction library is wrong: phase, value, time, -ordinal-rank are all the SAME mechanism — a small supervised head hung off the model that (a) injects a shaping -gradient into the backbone and/or (b) is itself a queryable readout. This module makes that mechanism a plug. - -A `Head` declares the four axes the phase head implicitly hard-coded: - 1. READ SITE (`reads`): "recon" = one prediction per EMITTED latent, in the token stream (what the phase head - reads — grad flows into the emitter and backbone, shaping the geometry). - "hidden" = the pooled/final backbone hidden, ONE prediction per sequence (a value/time - readout of the current state; grad flows into the backbone/LoRA). - 2. TARGET (`target`): a per-item dataset column supplying the supervision (like the existing `sup_field`). - 3. LOSS (`loss`): "ce" (classification) | "mse" (regression) | a CUSTOM CALLABLE (pred, target) -> scalar - (the seam a censored survival/hazard loss plugs into). - 4. LIFECYCLE (`transient`): transient (NOT saved; its only job is the shaping gradient — eval re-fits its own - probe, like today's phase head) OR PERSISTED (saved with the checkpoint AND reloadable - at inference via `LangSetModel.head_output(name, ...)` — the value/time head you query). -Plus `weight` (the `lam`) and optional `warmup` (ramp the weight 0->w over the first N epochs — a hidden->scalar -head on garbage early emissions destabilizes, the same reason the phase head can stick at chance; mirrors -`ss_warmup`). - -`lam_phase` is now a SHIM over this plug: `Head.phase_shim(sup_field, weight)` reconstructs the exact transient -recon+CE head, so the historical path is byte-identical (see the multi-latent golden, which trains with -lam_phase>0). +"""Configurable auxiliary supervision for multi-latent training. + +A `Head` attaches a linear prediction layer to one of two representations: + +- `reads="recon"` applies the head to each emitted latent. +- `reads="hidden"` applies the head to the pooled backbone hidden state. + +Each head selects a dataset column through `target` and uses classification (`loss="ce"`), regression +(`loss="mse"`), or a custom loss callable. `weight` scales the loss, and `warmup` can ramp that weight during early +epochs. + +Transient heads participate in training but are not saved. Persisted heads (`transient=False`) are stored with +the model and can be queried through `LangSetModel.head_output()` after loading the checkpoint. + +`Head.phase_shim()` provides compatibility with the `lam_phase` training option. """ from __future__ import annotations @@ -34,39 +23,43 @@ import torch import torch.nn.functional as F -# A custom loss reads (pred [N, dim], target [N] or [N, k]) and returns a scalar. The target packs whatever the -# loss needs — e.g. a censored survival loss reads target[:, 0]=Δt, target[:, 1]=event-observed (dim can be 1). +# A custom loss receives predictions shaped `[N, dim]` and targets shaped `[N]` or `[N, k]`, and returns a scalar. HeadLoss = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] _READS = ("recon", "hidden") _BUILTIN_LOSSES = ("ce", "mse") -# label tokens treated as "no supervision" for this item. +# Labels treated as missing supervision. _MISSING = ("", "unknown", "none", "nan") @dataclass class Head: - """One auxiliary supervised head. See the module docstring for the four axes. Instantiate and pass a list via - `TrainingArguments(heads=[Head(...), ...])`; each is built once and trained in every step alongside the - emission objective. Domain-agnostic: the library never names "phase" or "value" — you do, per Head.""" - - name: str # log key (`loss_{name}`) and, when persisted, the inference lookup key for `head_output(name, ...)` - reads: str # "recon" (per emitted latent) | "hidden" (pooled per-sequence backbone hidden) - target: str # per-item dataset column supplying the labels/values (like `sup_field`) - loss: Union[str, HeadLoss] # "ce" | "mse" | a custom callable(pred, target) -> scalar - dim: Optional[int] = ( - None # output width; None => infer for "ce" (#classes), else REQUIRED (mse/custom) - ) - transient: bool = ( - True # True: not saved (shaping-gradient only). False: PERSISTED + queryable at inference. - ) - weight: float = 1.0 # the `lam` — this head's loss weight in the total - warmup: int = ( - 0 # ramp `weight` 0->weight linearly over the first N epochs (0 = full weight from ep0) - ) - - # populated by resolution when `loss` is a custom callable, so the trainer can special-case target parsing - # (a custom loss may want a per-item VECTOR target, e.g. (Δt, censored)). Not user-set. + """Configuration for an auxiliary supervised head. + + Args: + name: Unique name used for logging and persisted-head lookup. + reads: Representation to read. `"recon"` predicts once per emitted latent; `"hidden"` predicts once per + input sequence. + target: Dataset column containing labels or numeric target values. + loss: `"ce"`, `"mse"`, or a callable accepting `(predictions, targets)` and returning a scalar tensor. + dim: Output width. Classification heads infer this from their labels; MSE and custom-loss heads must + specify it. + transient: If true, train the head without saving it in the model. If false, persist it for use with + `LangSetModel.head_output()`. + weight: Multiplier applied to this head's loss. + warmup: Number of epochs over which to ramp the effective weight from zero to `weight`. + """ + + name: str + reads: str + target: str + loss: Union[str, HeadLoss] + dim: Optional[int] = None # Infer the output width for CE; required for MSE and custom losses. + transient: bool = True # Persist the head only when false. + weight: float = 1.0 + warmup: int = 0 # Zero applies the full weight from epoch zero. + + # Set during validation when `loss` is a custom callable. _custom: bool = field(default=False, init=False, repr=False) def __post_init__(self) -> None: @@ -105,49 +98,58 @@ def loss_key(self) -> str: return f"loss_{self.name}" def eff_weight(self, ep: int) -> float: - """The warmup-ramped weight for epoch `ep` (mirrors ss_warmup: linear 0->weight over the first N epochs).""" + """Return this head's effective loss weight for epoch `ep`. + + With a positive `warmup`, the weight increases linearly from zero at epoch zero to `weight` at epoch + `warmup`. + """ if self.warmup <= 0: return self.weight return self.weight * min(1.0, ep / self.warmup) @classmethod def phase_shim(cls, sup_field: str, weight: float) -> "Head": - """Reconstruct the historical `lam_phase` head as a Head: a TRANSIENT recon+CE classifier over the - `sup_field` per-item stage labels. The resolver builds its module and class map identically to the old - inline phase head, so the training path stays byte-identical (guarded by the multi-latent golden).""" + """Create the transient classification head used by `lam_phase`. + + The returned head reads emitted latents, predicts labels from `sup_field`, and scales its cross-entropy + loss by `weight`. + """ return cls( name="phase", reads="recon", target=sup_field, loss="ce", - dim=None, # inferred #classes — same sorted/filtered class set as the old phase head + dim=None, # Infer the class count from labels. transient=True, weight=weight, ) def build_ce_classes(rows: list[list[str]]) -> tuple[list[str], dict[str, int]]: - """Class list + label->id map for a CE head, built EXACTLY like the old phase head: the sorted set of non-empty, - non-missing labels across every item of every row. Items whose label is missing map to -100 (ignored) at loss - time. Deterministic (sorted) so the head width and class indices are reproducible.""" + """Build a deterministic class vocabulary for categorical targets. + + Empty labels and labels matching the module's missing-value tokens are excluded case-insensitively. + + Returns: + A sorted class list and a mapping from each class label to its integer ID. + """ classes = sorted({lb for row in rows for lb in row if lb and lb.lower() not in _MISSING}) return classes, {c: i for i, c in enumerate(classes)} @dataclass class RtHead: - """A RESOLVED head: the user `Head` spec + its built `nn.Linear` module + (for CE) its class map + a reference - to the per-row target values. Trainer-internal — one per `Head`, built once in `_train_multi`.""" + """Resolved runtime state for a `Head`. + + This trainer-internal object combines the user specification, its linear module, optional classification + metadata, and the source target values. + """ spec: Head module: torch.nn.Linear - values: list[ - object - ] # per-row target values: a LIST per row for reads="recon", a scalar for reads="hidden" - classes: Optional[list[str]] = ( - None # CE only: index->label (also persisted so `head_output` can name argmaxes) - ) - class_map: Optional[dict[str, int]] = None # CE only: label->index + values: list[object] # One list per row for `recon`; one scalar or vector per row for `hidden`. + classes: Optional[list[str]] = None # CE index-to-label mapping persisted with the head. + class_map: Optional[dict[str, int]] = None # CE label-to-index mapping used during training. @property def out_dim(self) -> int: @@ -158,8 +160,7 @@ def in_dim(self) -> int: return int(self.module.in_features) def spec_dict(self) -> dict[str, object]: - """Serializable metadata for a PERSISTED head — written to config.json so `LangSetModel.load` can rebuild - the module and `head_output` knows how to read it back.""" + """Return the metadata needed to save and reconstruct a persisted head.""" loss = self.spec.loss if isinstance(self.spec.loss, str) else "custom" return { "name": self.spec.name, @@ -170,11 +171,12 @@ def spec_dict(self) -> dict[str, object]: "classes": self.classes, } - # --- per-step loss ------------------------------------------------------------------------------------------ + # Per-step loss helpers. def _flat_recon_targets(self, bidx: list[int], lens_l: list[int]) -> list[object]: - """Per-emitted-latent target values in `recon[valid]` row-major order (row0's items, then row1's, ...), - aligned exactly like the old phase head's `pf`: item j of row k is `values[k][j]`, or None past the row's - list length (-> ignored).""" + """Flatten selected row targets in the order used by `recon[valid]`. + + Targets missing beyond the end of a row's target list are represented by `None`. + """ out: list[object] = [] for r, k in enumerate(bidx): row = self.values[k] @@ -192,17 +194,23 @@ def _ce_ids(self, flat: list[object], dev: torch.device) -> torch.Tensor: def loss_on( self, pred: torch.Tensor, flat: list[object], dev: torch.device ) -> Optional[torch.Tensor]: - """This head's loss for predictions `pred` [N, out_dim] against the flat per-prediction targets. Returns - None when nothing in the batch supervises this head (all items missing) so the caller can skip it. - CE ignores missing items (-100); MSE masks to finite targets; a CUSTOM loss gets (pred, target) with - missing packed as NaN (its job to handle censoring). The CE branch is byte-identical to the old phase term.""" + """Compute this head's loss for flattened batch targets. + + For classification, missing targets are encoded as `-100` and ignored by cross-entropy. For MSE, rows + containing non-finite targets are excluded and `None` is returned if no rows remain. Custom losses receive + the complete prediction and target tensors, with missing or unparseable targets represented as `NaN`, and + are responsible for their own masking. + + Args: + pred: Predictions shaped `[N, out_dim]`. + flat: One target value per prediction. + dev: Device on which to construct the target tensor. + """ if self.spec.is_ce: return F.cross_entropy(pred, self._ce_ids(flat, dev), ignore_index=-100) tgt = _as_float_target(flat, dev) # [N] or [N, k]; NaN = missing if self.spec._custom: - return cast_loss(self.spec.loss)( - pred, tgt - ) # custom loss owns masking/censoring over the full batch + return cast_loss(self.spec.loss)(pred, tgt) # Custom losses handle their own masking. tgt2 = tgt.unsqueeze(1) if tgt.dim() == 1 else tgt keep = torch.isfinite(tgt2).all(dim=-1) if not bool(keep.any()): @@ -211,15 +219,17 @@ def loss_on( def cast_loss(loss: Union[str, HeadLoss]) -> HeadLoss: - """Narrow a Head.loss to the callable branch (only called when _custom is True, i.e. loss is a callable).""" + """Return `loss` narrowed to the custom callable type.""" assert callable(loss) return cast(HeadLoss, loss) def _as_float_target(flat: list[object], dev: torch.device) -> torch.Tensor: - """Parse per-item target values to a float tensor [N] (scalars) or [N, k] (list/tuple items, e.g. a survival - (Δt, event) pair). Missing / unparseable values become NaN so downstream masking (MSE) or the custom loss can - drop them.""" + """Convert scalar or vector targets to a float tensor. + + Missing and unparseable values become `NaN`. Scalar inputs produce shape `[N]`; consistently sized list or + tuple inputs produce shape `[N, k]`. + """ def _f(x: object) -> float: if x is None: @@ -241,9 +251,21 @@ def _f(x: object) -> float: def resolve_head( spec: Head, values: list[object], latent_dim: int, hidden_dim: int, dev: torch.device ) -> RtHead: - """Build the head's `nn.Linear` (its RNG draw happens HERE — for the phase shim this is at the exact point the - old inline phase head was built, keeping init byte-identical) + its CE class map. `values` is the per-row target - column (a list per row for reads="recon", a scalar per row for reads="hidden").""" + """Build the runtime module and metadata for a `Head`. + + For classification heads, classes are inferred from `values` and `dim`, when provided, must match the inferred + class count. Other loss types use the explicitly configured `dim`. + + Args: + spec: Head configuration to resolve. + values: Per-row target values. + latent_dim: Input width for `reads="recon"`. + hidden_dim: Input width for `reads="hidden"`. + dev: Device on which to create the linear module. + + Returns: + The resolved runtime head. + """ in_dim = latent_dim if spec.reads == "recon" else hidden_dim classes: Optional[list[str]] = None class_map: Optional[dict[str, int]] = None @@ -252,19 +274,19 @@ def resolve_head( rows = [[str(e) for e in v] if isinstance(v, (list, tuple)) else [str(v)] for v in values] classes, class_map = build_ce_classes(rows) n_cls = len(classes) - if n_cls == 0: # else F.cross_entropy fails opaquely on a [N, 0] logit + if n_cls == 0: # Avoid constructing a classifier with zero output classes. raise ValueError( f"CE head {spec.name!r} has no classes — all `target` labels are missing " f"(empty/unknown/none/nan). Supply real class labels, or use a non-CE loss." ) if ( spec.dim is not None and spec.dim != n_cls - ): # a CE head's width IS the class count — don't let it drift + ): # CE width must equal the inferred class count. raise ValueError( f"CE head {spec.name!r}: dim={spec.dim} conflicts with the {n_cls} inferred classes " f"({classes}). Leave dim=None for a CE head — its width is inferred from the labels." ) - out_dim = n_cls # CE width is always the inferred class count (phase shim keeps dim=None -> byte-identical) + out_dim = n_cls assert out_dim is not None # __post_init__ guarantees dim is set for non-CE losses module = torch.nn.Linear(in_dim, out_dim).to(dev) return RtHead(spec, module, values, classes, class_map) diff --git a/src/langset/loss.py b/src/langset/loss.py index 6b84d72..76eabff 100644 --- a/src/langset/loss.py +++ b/src/langset/loss.py @@ -1,4 +1,31 @@ -"""Loss functions.""" +"""Composable loss objectives used by langset's training strategies. + +This module separates the mathematical objectives from the training loop. +Each loss function accepts a typed context dataclass containing the tensors +and configuration it needs, and returns a :class:`Loss`. The trainer can then +combine the result with other objectives and call ``to_tensor()`` for the +scalar used in backpropagation. + +The objectives fall into four groups: + +- **Latent grounding:** :func:`recon_loss` teaches an emitted latent to + reconstruct its target text. +- **Representation alignment:** :func:`sl_loss` and :func:`info_nce_loss` + align paired embeddings and contrastive candidates. The supervised + contrastive objective groups emitted embeddings by labels. +- **Multi-latent supervision:** :func:`soft_target_cross_entropy_loss` + trains distributions over concepts or states, while :func:`stop_loss` + teaches autoregressive emissions when to terminate. +- **Text-generation support:** :func:`learn_loss` rehearses document-to-target + associations, and :func:`cot_loss` trains optional chain-of-thought + generation. + +The context classes are intentionally explicit: they document tensor shapes, +make each objective independently testable, and keep loss implementations +decoupled from batch construction and optimization orchestration. Most users +will select these objectives through a strategy or ``TrainingArguments`` +configuration rather than call them directly. +""" from __future__ import annotations diff --git a/src/langset/masking.py b/src/langset/masking.py index 394af12..822f02d 100644 --- a/src/langset/masking.py +++ b/src/langset/masking.py @@ -1,30 +1,21 @@ -"""Masked self-prediction (JEPA-style) data construction — langset's generic JEPA mode. +"""Utilities for masked self-prediction datasets and training views. -langset already trains the JEPA way on the *target* side: the default `EMATwinTarget` is a stop-grad EMA copy -of the model that supplies the target latents (BYOL/JEPA teacher, no collapse). What made a langset encoder -JEPA-vs-not was never a flag — it was WHAT you put in `target_text`: +A masker is a callable with the following interface: - * target_text = an external LABEL (e.g. "the legal moves here") -> contrastive-to-label; the encoder keeps - only label-relevant structure and DISCARDS the rest (this is how our chess affordance encoder collapsed - near-identical positions — one move barely changes the label). - * target_text = a HELD-OUT PART of the SAME text, predicted from the VISIBLE part -> real JEPA: no label, - the encoder is FORCED to represent what it can't see, so fine detail survives. + (text: str, rng: random.Random) -> (visible_text, hidden_text) -This module gives you the second one GENERICALLY. You hand it whole texts and a `Masker` (the algorithm that -decides what to hide); it returns langset rows `{input_text: visible, target_text: full-or-hidden}`. Train -those with the normal (single-latent, default EMATwinTarget) `Trainer` and you have a JEPA encoder for ANY -text domain — prose, structured records, board renders, code — by swapping the masker. +The built-in maskers hide contiguous spans, scattered tokens, or fields in a delimited record. `Trainer` can +apply a masker to a raw `text` column at the start of each epoch while using the original text as the target: - # PREFERRED: hand the Trainer RAW text; it masks fresh every epoch. No pre-built masked data. - Trainer(model, TrainingArguments(), [{"text": t} for t in texts]).train() # auto-masks (word @0.15) - # pick the algorithm / ratio: TrainingArguments(masker="span", mask_ratio=0.2) - # custom masker (e.g. protect chess move-numbers): TrainingArguments(masker=TokenMasker(0.15, protect=...)) + rows = [{"text": text} for text in texts] + args = TrainingArguments(masker="span", mask_ratio=0.2) + Trainer(model, args, rows).train() -The Trainer path re-masks EVERY epoch (unlimited diversity from a small corpus) — you never enumerate masked -copies. `build_masked` / `build_masked_pairs` below still exist for materializing a fixed masked dataset when -you want one, but for training just give `text`. +When no masker is specified for raw-text training, scattered-token masking is used with the configured mask +ratio. -A Masker is any callable `(text: str, rng: random.Random) -> (visible_text, hidden_text)`. +`build_masked()` and `build_masked_pairs()` materialize deterministic masked datasets for callers that need fixed +views instead of per-epoch masking. """ from __future__ import annotations @@ -40,8 +31,11 @@ def __call__(self, text: str, rng: random.Random) -> tuple[str, str]: ... class SpanMasker: - """Hide a CONTIGUOUS run of whitespace tokens — the general-text default (a phrase, a clause). The encoder - must infer the missing span from both sides, so it has to represent surrounding meaning, not just keywords.""" + """Replace one contiguous span of whitespace-delimited tokens with a sentinel. + + At least one token is hidden and at least one original token remains visible. Texts containing fewer than two + tokens are returned unchanged with an empty hidden string. + """ def __init__(self, ratio: float = 0.15, sentinel: str = "[MASK]") -> None: self.ratio, self.sentinel = ratio, sentinel @@ -58,12 +52,14 @@ def __call__(self, text: str, rng: random.Random) -> tuple[str, str]: class TokenMasker: - """Hide RANDOM scattered tokens (BERT/data2vec flavor). More diffuse pressure than a single span. + """Replace randomly selected whitespace-delimited tokens with a sentinel. - `protect` — an optional predicate `(token) -> bool`; a protected token is NEVER masked and never - counts toward the maskable pool. Use it to keep scaffolding visible while hiding only content, e.g. - for chess movetext `protect=lambda t: t.endswith(".")` masks moves but keeps the `1.`/`2.` numbering - so positions stay anchored (mask a random % of the actual MOVES, both players).""" + `ratio` is applied to the tokens eligible for masking. When `protect` is provided, tokens for which + `protect(token)` returns true remain visible and are excluded from the maskable population. + + The hidden text contains the selected original tokens in source order. Texts with no maskable tokens, or fewer + than two total tokens, are returned unchanged with an empty hidden string. + """ def __init__( self, @@ -86,10 +82,15 @@ def __call__(self, text: str, rng: random.Random) -> tuple[str, str]: class FieldMasker: - """Hide whole FIELDS of a delimited record (structured text / tabular / board renders). Splits on `sep`; - hides a fraction of fields. With `kv_sep` set, keeps each field's KEY and hides only its VALUE - (e.g. cells `e1:K` -> `e1:[MASK]`), so the model must infer WHAT is at a KNOWN slot — the sharp form of - "represent the fine detail." This is the masker for a chess board render (cells) or any key:value record.""" + """Mask randomly selected fields in a delimited record. + + Fields are split with `sep`, and `ratio` determines how many are selected. Without `kv_sep`, each selected field + is replaced by `sentinel`. With `kv_sep`, a selected key/value field keeps its key and replaces its value; for + example, `e1:K` becomes `e1:[MASK]`. + + The hidden text contains each selected field in its original form. Inputs with fewer than two fields are returned + unchanged with an empty hidden string. + """ def __init__( self, @@ -122,9 +123,19 @@ def __call__(self, text: str, rng: random.Random) -> tuple[str, str]: def resolve_masker(spec: Masker | str | None, ratio: float = 0.15) -> Masker: - """Turn a masker SPEC into a Masker. A callable passes through; a string picks a default algorithm; None - is the reasonable default ('word' = scattered-word masking). Used by the Trainer so callers can write - `masker="span"` / `masker=None` and never build a masker by hand.""" + """Resolve a masker specification. + + Args: + spec: A masker callable, `"word"`, `"token"`, `"span"`, `"field"`, or `None`. `None`, `"word"`, and + `"token"` select `TokenMasker`. + ratio: Masking ratio passed to a built-in masker. + + Returns: + The supplied callable or a newly constructed built-in masker. + + Raises: + ValueError: If `spec` is an unrecognized string. + """ if callable(spec): return cast("Masker", spec) if spec in (None, "word", "token"): @@ -139,10 +150,18 @@ def resolve_masker(spec: Masker | str | None, ratio: float = 0.15) -> Masker: def mask_view(texts: list[str], masker: Masker, rng: random.Random) -> list[str]: - """One FRESH masked view per text — the VISIBLE side; the target is the full text itself. This is what the - Trainer calls at the start of EVERY epoch so masks are never reused (unlimited diversity from raw text). - A text where the masker hides nothing (too short) passes through unmasked (that row just carries no signal - this epoch).""" + """Generate one masked visible view for each text. + + If a masker returns an empty or whitespace-only hidden string, the original text is preserved in the result. + + Args: + texts: Source texts to mask. + masker: Masking callable. + rng: Random generator controlling mask selection. + + Returns: + One visible text for each source text, in the same order. + """ out: list[str] = [] for t in texts: visible, hidden = masker(t, rng) @@ -153,9 +172,25 @@ def mask_view(texts: list[str], masker: Masker, rng: random.Random) -> list[str] def build_masked( texts: list[str], masker: Masker, views: int = 1, seed: int = 0, target_mode: str = "full" ) -> list[dict[str, str]]: - """Whole texts -> langset JEPA rows. Each text yields `views` masked pairs with different random masks - (mask diversity, like JEPA re-masking). target_mode: 'full' = predict the WHOLE text's latent from the - masked view (masked-view -> full-view alignment); 'hidden' = predict only the withheld content's latent.""" + """Build fixed masked input/target rows from whole texts. + + Each source text is masked `views` times using a deterministic random generator initialized from `seed`. + Degenerate views that hide no content are omitted. + + Args: + texts: Source texts. + masker: Masking callable. + views: Number of masking attempts per source text. + seed: Seed for deterministic mask selection. + target_mode: `"full"` uses the original text as `target_text`; `"hidden"` uses only the text returned as + hidden by the masker. + + Returns: + Dictionaries containing `input_text` and `target_text`. + + Raises: + ValueError: If `target_mode` is not `"full"` or `"hidden"`. + """ if target_mode not in ("full", "hidden"): raise ValueError("target_mode must be 'full' or 'hidden'") rng = random.Random(seed) @@ -164,7 +199,7 @@ def build_masked( for _ in range(views): visible, hidden = masker(t, rng) if not visible.strip() or not hidden.strip() or visible.strip() == t.strip(): - continue # skip degenerate (nothing hidden) + continue # Skip views that hide no usable content. out.append( {"input_text": visible, "target_text": t if target_mode == "full" else hidden} ) @@ -180,20 +215,28 @@ def build_masked_pairs( mask_region: str = "target", target_mode: str = "full", ) -> list[dict[str, str]]: - """FUSE (input, target) then mask — the "state + facets you want to predict" JEPA. - - Each pair is `(input, target)`: `input` is the state of the world you ALWAYS have; `target` is the - ADDITIONAL facets you want the latent to be able to recover. We MERGE them (`input + sep + target`), - run the masker on the merged text, and predict the masked content — so one encoder is forced to hold - BOTH the state and the extra facets, with no external label. - - mask_region: - * 'target' (default) — mask ONLY inside the target span; the whole input stays visible. This is the - pointed form: "given the full state, reconstruct the withheld facets." (The encoder learns to - predict the facets FROM the state.) - * 'all' — mask anywhere across the merged text (symmetric masked self-prediction over state+facets). - target_mode: 'full' = predict the whole merged (input+sep+target) latent; 'hidden' = predict only the - withheld content's latent. + """Build fixed masked rows from `(input, target)` text pairs. + + The two parts are joined as `input + sep + target`. With `mask_region="target"`, only the target part is passed + to the masker and the input part remains visible. With `mask_region="all"`, the masker is applied to the joined + text. + + Args: + pairs: Source `(input, target)` pairs. + masker: Masking callable. + views: Number of masking attempts per pair. + seed: Seed for deterministic mask selection. + sep: Separator inserted between the two parts. + mask_region: `"target"` to mask only the target part, or `"all"` to mask the joined text. + target_mode: `"full"` uses the unmasked joined text as `target_text`; `"hidden"` uses only the content + returned as hidden by the masker. + + Returns: + Dictionaries containing `input_text` and `target_text`. Degenerate views that hide no usable content are + omitted. + + Raises: + ValueError: If `mask_region` or `target_mode` is invalid. """ if mask_region not in ("target", "all"): raise ValueError("mask_region must be 'target' or 'all'") @@ -206,11 +249,11 @@ def build_masked_pairs( for _ in range(views): if mask_region == "all": visible, hidden = masker(merged, rng) - else: # mask only the target span; keep input whole + else: # Mask only the target part. vis_t, hidden = masker(tgt, rng) visible = f"{inp}{sep}{vis_t}" if not visible.strip() or not hidden.strip() or visible.strip() == merged.strip(): - continue # skip degenerate (nothing hidden) + continue # Skip views that hide no usable content. out.append( {"input_text": visible, "target_text": merged if target_mode == "full" else hidden} ) diff --git a/src/langset/modeling.py b/src/langset/modeling.py index be1e052..3b3d34c 100644 --- a/src/langset/modeling.py +++ b/src/langset/modeling.py @@ -1,8 +1,12 @@ -"""LangSetModel: an LLM (LoRA) + a learned emit head that maps input text -> a latent in a bespoke geometry. +"""Model and emission-head implementations for LangSet. -The latent lives in the model's OWN hidden space, and the output is Sentence-Transformer-shaped (`encode`, -`get_sentence_embedding_dimension`, `as_sentence_transformer`) so the trained model drops straight into SetFit -as a `model_body`. +:class:`LangSetModel` combines a language-model backbone with a learned emission interface. The emission head can +project backbone hidden states into a single embedding, multiple continuous latents, or autoregressive codebook +states. Optional LoRA adaptation supports parameter-efficient fine-tuning. + +The model provides Sentence-Transformer-compatible ``encode``, ``get_sentence_embedding_dimension``, and +``as_sentence_transformer`` methods. Codebook-configured multi-latent models additionally support autoregressive +state emission through ``rollout``. """ from __future__ import annotations @@ -33,10 +37,9 @@ class _HiddenOutput(Protocol): class _Backbone(Protocol): - """The structural surface `LangSetModel` uses on its (peft-wrapped, maybe text-tower-unwrapped) backbone — - one Protocol spanning plain Llama/Qwen `ForCausalLM`, text towers, and Gemma-E PLE models. The PLE-only - `get_per_layer_inputs` is deliberately NOT declared: it exists on Gemma-E alone and is reached through the - `_ple_dim` runtime guard, so pinning it here would wrongly exclude every non-PLE backbone.""" + """ + Wraps the backbone model, providing a common interface for accessing hidden states. + """ config: PretrainedConfig @@ -219,10 +222,11 @@ def _cfg_set(config: PretrainedConfig, name: str, val: object) -> None: def _text_tower(model: _Backbone) -> _Backbone: - """Descend a (peft-wrapped) causal/conditional-generation model to its TEXT transformer — the module that - returns hidden states directly, with NO lm_head. Skips the huge-vocab logits projection (Gemma's 262k-vocab - lm_head over a full sequence OOMs — we only ever read hidden states) and any vision tower. LoRA is injected - in-place on the language Linears, so calling the text tower directly still applies it.""" + """Return the text transformer inside a PEFT-wrapped generation model. + + Calling the text tower directly avoids the vocabulary projection and any vision tower while retaining LoRA + modules injected into the language layers. The returned module exposes hidden states rather than LM logits. + """ node = getattr( getattr(model, "base_model", model), "model", model ) # peft LoraModel -> underlying HF model @@ -257,8 +261,8 @@ def build_backbone( vocab_size: Optional[int] = None, ) -> _Backbone: def _top_k_layers(n_layers: int) -> Optional[list[int]]: - # LoRA ONLY the top-K transformer layers -> fewer adapters, smaller activation graph -> bigger batch. Emission - # reads the FINAL hidden state, so the top layers carry the task-shaping. 0 = all layers (default, unchanged). + # Restrict LoRA to the top K transformer layers when requested. A value of zero leaves all eligible layers + # adapted. Emission reads the final hidden state, so upper-layer adapters can directly shape its representation. return ( list(range(max(0, n_layers - lora_top_k), n_layers)) if lora_top_k and n_layers @@ -270,10 +274,8 @@ def _top_k_layers(n_layers: int) -> Optional[list[int]]: dt = torch.bfloat16 if bf16 else torch.float32 if not pretrained: - # RANDOM-INIT control arm: copy `llm_model`'s ARCHITECTURE (config) but NOT its weights, then train the whole - # net (no LoRA — a low-rank adapter over random weights is meaningless). `arch_overrides` shrinks the net; a - # decoupled tokenizer sets `vocab_size` so the fresh embedding table matches it. This is the "does pretraining - # matter" baseline: same emission/anti-collapse machinery, zero inherited knowledge. + # Build the requested architecture from configuration without loading pretrained weights. The complete + # backbone is trainable, and `vocab_size` keeps its new embedding table aligned with a decoupled tokenizer. from transformers import AutoConfig # type: ignore[import-untyped] cfg = AutoConfig.from_pretrained(llm_model) @@ -303,8 +305,8 @@ def _top_k_layers(n_layers: int) -> Optional[list[int]]: from peft import LoraConfig, get_peft_model # type: ignore[import-untyped] def _try_load(dtype_key: str, attn: Optional[str]) -> _Backbone: - # sdpa (default) avoids materializing the O(S^2) eager-attention score matrix — a long seed (3072 tokens) - # OOM'd a 0.6B model at 72GB on eager. attention_dropout is dropped for multimodal wrappers that reject it. + # SDPA avoids materializing the eager attention score matrix. Some multimodal wrappers reject + # `attention_dropout`, so retry without that keyword when necessary. kw: dict[str, Any] = {dtype_key: dt} if attn: kw["attn_implementation"] = attn @@ -321,9 +323,9 @@ def _load(attn: Optional[str]) -> _Backbone: except TypeError: return _try_load("torch_dtype", attn) - # A non-default impl (flash_attention_2, flex_attention, ...) is an EXPLICIT performance choice: refuse to silently - # downgrade it. sdpa/eager may still fall back (a model that can't do sdpa -> eager) since those are just defaults; - # with the default attn_implementation="sdpa", _strict is False so the load path is byte-identical to before. + # Treat an explicitly selected nonstandard attention implementation as a strict requirement. Default SDPA and + # eager modes may fall back to one another for model compatibility, but an explicit optimized kernel must not be + # silently replaced with a slower implementation. _strict = bool(attn_implementation) and attn_implementation not in ("sdpa", "eager") try: base = _load(attn_implementation or None) @@ -372,19 +374,17 @@ def _load(attn: Optional[str]) -> _Backbone: ], layers_to_transform=ltt, ) - # strip the lm_head: we only read hidden states, and computing the full-vocab logits - # ([B,S,vocab]) every forward OOMs — a 0.6B at bs48/384 hit 78GB purely on Qwen3's 152k-vocab projection. + # The model consumes hidden states only. Returning the text tower avoids allocating full-vocabulary logits with + # shape `[batch, sequence, vocabulary]` on every forward. peft = get_peft_model(base, lora) if train_base: - # KNOWLEDGE INJECTION: rank-16 LoRA on a FROZEN base can't STORE new facts (only re-style existing ones) — - # facts live in the base MLP weights. Unfreeze the whole base so next-token [LEARN] can actually rewrite - # "GrEStG=Grundgesetz" -> the real statute. Full-FT capacity; default off (frozen-LoRA, unchanged). + # `train_base` requests full fine-tuning rather than adapter-only training, so next-token replay can update + # all backbone parameters. for p in peft.parameters(): p.requires_grad_(True) if grad_ckpt: - # trade compute for activation memory so a LARGE InfoNCE batch (= more in-batch negatives, the dominant lever) - # fits — a 4B at batch 4 was negative-starved. use_cache off is required; input-require-grads lets grad reach - # checkpointed segments when only LoRA trains (frozen embeddings). + # Trade additional computation for lower activation memory. Caching must be disabled, and input gradients are + # enabled so checkpointed segments remain connected when embeddings are frozen and only LoRA is trainable. base.config.use_cache = False peft.gradient_checkpointing_enable() peft.enable_input_require_grads() @@ -394,8 +394,13 @@ def _load(attn: Optional[str]) -> _Backbone: class LangSetModel(nn.Module): - """LLM backbone (LoRA) + EmitHead. The latent lives in the model's own hidden space; the geometry is defined - by the `target_text` the Trainer contrasts against (see Trainer).""" + """Wrap a language-model backbone with continuous or codebook-based emission heads. + + Depending on its configuration, a model can produce a single embedding, multiple learned latent vectors, + or an autoregressive sequence of named-state emissions. Use :meth:`from_pretrained` or :meth:`from_scratch` + to construct a model, :meth:`encode` for Sentence-Transformer-compatible embeddings, and :meth:`rollout` for + autoregressive codebook emission. + """ def __init__( self, @@ -419,9 +424,9 @@ def __init__( self.embed = backbone.get_input_embeddings() self.h = _cfg_int(backbone.config, "hidden_size") self.vocab_size = _cfg_int(backbone.config, "vocab_size") - # Gemma E-series (3n/4) use Per-Layer Embeddings: each layer mixes in an embedding indexed by TOKEN ID. - # We pass `per_layer_inputs` explicitly (real tokens -> real PLE, synthetic emit/feedback tokens -> zeros) - # so `inputs_embeds` forwards don't crash on the reverse-ID lookup. 0 => not a PLE model (no-op). + # Gemma E-series models use per-layer embeddings indexed by token ID. Pass `per_layer_inputs` explicitly: + # real tokens receive their PLE values, while synthetic emit and feedback positions receive zeros. A zero + # `_ple_dim` identifies a backbone without per-layer embeddings. self._ple_dim = int(getattr(backbone.config, "hidden_size_per_layer_input", 0) or 0) self._n_layers = int(getattr(backbone.config, "num_hidden_layers", 0) or 0) # A text tower (when build_backbone unwrapped to one) always returns `last_hidden_state`, so we don't ask the @@ -449,20 +454,18 @@ def __init__( self.llm_model = llm_model self.max_len = max_len self._lora_top_k = 0 # overwritten by from_pretrained; persisted in config - # pool_mode="last": SKIP the learned emit-query; read the backbone's LAST real-token hidden and project it - # (head.out_proj). Lets a FROZEN strong-embedder backbone (e.g. F2LLM) be specialized by training ONLY the - # projection head -> nothing is backpropped through the layers (no grad_ckpt, huge batch). "" => emit-head - # (default, byte-identical: forward() takes the learned-query path exactly as before). + # `pool_mode="last"` bypasses the learned emit query and projects the final real-token hidden state. When + # the backbone is frozen, this allows head-only training without retaining a backbone autograd graph. The + # empty mode uses the learned-query emission path. self.pool_mode = pool_mode self._frozen_bb = ( False # set by from_pretrained(freeze_backbone=True); gates the no-grad backbone read ) - # RANDOM-INIT bookkeeping (set by from_scratch). A pretrained model rebuilds its backbone from `llm_model`, so - # persistence stores only LoRA; a random-init model has no such source, so save/load must carry the FULL net. + # A pretrained model can rebuild its base weights from `llm_model`, so persistence needs only its adapters. + # A randomly initialized model has no external weight source and must persist the complete backbone. self._pretrained = True - # FULL-FINETUNE bookkeeping (set by from_pretrained(train_base=True)). A pretrained model normally stores only - # LoRA on save (base rebuilds from `llm_model`), but train_base=True trains the WHOLE base — those weights have - # no source to rebuild from, so save/snapshot/load must carry the FULL backbone exactly like a random-init net. + # Full fine-tuning changes base weights that cannot be reconstructed from `llm_model`; those runs persist the + # complete backbone just like randomly initialized models. self._full_ft = False self._tokenizer_id: Optional[str] = ( None # decoupled HF tokenizer id (None => same as llm_model/arch) @@ -470,10 +473,8 @@ def __init__( self._arch_overrides: Optional[dict] = ( None # config shrink applied to the from-scratch backbone ) - # PERSISTED AUXILIARY HEADS (langset.heads): a name->nn.Linear map of PERSISTED heads (transient heads are - # never registered here — they live and die in the trainer). Each carries a metadata spec (read site, loss, - # in/out dims, CE classes) so save_pretrained can serialize it and `head_output(name, ...)` can read it back - # at inference. Empty by default => byte-identical: no extra params, nothing written to disk. + # Persisted auxiliary heads are registered by name with metadata describing their read site, loss, dimensions, + # and optional classes. Transient trainer-only heads are not registered here. self.aux_heads: nn.ModuleDict = nn.ModuleDict() self.aux_head_specs: dict[str, dict[str, Any]] = {} @@ -502,6 +503,16 @@ def from_pretrained( code_tau: float = 0.07, res_dim: int = 0, ) -> "LangSetModel": + """Construct a LangSet model from a pretrained Hugging Face model. + + The backbone is adapted with LoRA by default. Set ``train_base=True`` for full fine-tuning, or combine + ``pool_mode="last"`` with ``freeze_backbone=True`` to train only the output projection over frozen + backbone features. When ``latent_dim`` is omitted, it defaults to the backbone hidden size. + + Set ``code_emit=True`` to enable codebook-based emission and configure ``n_codes`` for the codebook size. + Set ``multi_latent=True`` when the model will be trained with a multi-latent emission strategy. The + returned model is moved to ``device``, or to CUDA when available if no device is specified. + """ from transformers import AutoTokenizer # type: ignore[import-untyped] dev = device or ("cuda" if torch.cuda.is_available() else "cpu") @@ -539,15 +550,14 @@ def from_pretrained( model._lora_top_k = int(lora_top_k) # persisted in config so load() rebuilds same adapters model._full_ft = bool( train_base - ) # train_base trains the whole base -> persist FULL backbone - if pool_mode == "last": # WARM-START the pool head at the base's NATIVE embedding: identity + ) # full fine-tuning requires persisting the complete backbone + if pool_mode == "last": # initialize the projection as an identity mapping torch.nn.init.eye_( model.head.out_proj.weight - ) # out_proj -> head_project == normalize(last-token hidden) == + ) # `head_project` initially preserves the normalized final-token hidden state torch.nn.init.zeros_( model.head.out_proj.bias - ) # the base's own readout (last-token pool) -> starts at base - # quality and REFINES, instead of a random head destroying the base geometry and relearning it worse. + ) if freeze_backbone: # FROZEN base: only the head trains -> backbone read needs no graph for p in model.backbone.parameters(): p.requires_grad_(False) @@ -575,13 +585,14 @@ def from_scratch( code_tau: float = 0.07, res_dim: int = 0, ) -> "LangSetModel": - """RANDOM-INIT control arm (the "does pretraining matter" baseline). `arch` names an HF model whose - ARCHITECTURE is copied, but the weights are NOT loaded — the backbone starts from scratch and trains fully - (no LoRA). The tokenizer is decoupled: pass any HF `tokenizer_id` (default = `arch`) and the fresh embedding - table is sized to it, so there is no baked-in tokenizer. `arch_overrides` shrinks the net, e.g. - `{"num_hidden_layers": 4, "hidden_size": 256, "num_attention_heads": 4, "num_key_value_heads": 4, - "intermediate_size": 1024}`. Everything downstream (emit head, target source, and losses) is - identical to `from_pretrained`; only the source of the backbone weights differs.""" + """Construct a LangSet model with randomly initialized backbone weights. + + ``arch`` identifies a Hugging Face configuration, but its pretrained weights are not loaded. The complete + backbone remains trainable and no LoRA adapters are added. ``tokenizer_id`` defaults to ``arch``, and the + new input embedding table is sized for that tokenizer. Use ``arch_overrides`` to replace configuration + fields such as the number of layers, hidden size, or attention-head count. The returned model is moved to + ``device``, or to CUDA when available if no device is specified. + """ from transformers import AutoTokenizer # type: ignore[import-untyped] dev = device or ("cuda" if torch.cuda.is_available() else "cpu") @@ -678,8 +689,11 @@ def _last_hidden(out: _HiddenOutput) -> torch.Tensor: return h if h is not None else out.hidden_states[-1] def _pool_hidden(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - """Last real-token hidden of the raw text (frozen-backbone read when _frozen_bb). [B, h]. This is the STATIC - feature the frozen-pool fast path caches ONCE; head_project() then turns it into the trainable latent.""" + """Return each row's final real-token hidden state with shape ``[batch, hidden_size]``. + + When the backbone is frozen, the forward pass runs without gradient + tracking so callers can cache these features for projection-head training. + """ if self._frozen_bb: # frozen backbone -> read under no_grad (zero activation memory) with torch.no_grad(): hid = self._last_hidden( @@ -693,11 +707,15 @@ def _pool_hidden(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> return hid[torch.arange(hid.size(0), device=hid.device), last] # [B, h] def head_project(self, feats: torch.Tensor) -> torch.Tensor: - """Project pooled features through the trainable head -> normalized latent. [B, h] -> [B, d].""" + """Project ``[batch, hidden_size]`` pooled features to normalized latent vectors.""" return F.normalize(self.head.out_proj(feats.float()), p=2, dim=-1) def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - """Read input text, emit the latent. Returns [B, d].""" + """Emit latent vectors from a tokenized batch. + + Returns ``[batch, latent_dim]`` when one latent is configured, or + ``[batch, n_latents, latent_dim]`` otherwise. + """ if self.pool_mode == "last": # POOL path: no emit query; last real-token hidden -> out_proj return self.head_project(self._pool_hidden(input_ids, attention_mask)) nl = self.head.n_latents @@ -729,7 +747,13 @@ def encode( show_progress_bar: bool = False, device: Optional[str] = None, ) -> Union[np.ndarray, torch.Tensor]: - """Sentence-Transformer-compatible. This is the method SetFit calls on its body.""" + """Encode one or more sentences using the Sentence-Transformer-compatible interface. + + Inputs are tokenized in batches and truncated to :attr:`max_len`. A single string returns one embedding; + a sequence returns a batch. Results are NumPy arrays by default or CPU tensors when + ``convert_to_numpy=False``. ``show_progress_bar`` and ``device`` are accepted for interface compatibility + but are currently ignored. + """ single = isinstance(sentences, str) texts = [sentences] if single else list(sentences) was_training = self.training @@ -737,10 +761,10 @@ def encode( out: list[torch.Tensor] = [] with ( torch.no_grad() - ): # eval-only: LoRA params require grad, so w/o this every forward builds + ): for i in range( 0, len(texts), batch_size - ): # a throwaway autograd graph -> slower + huge VRAM (caps batch) + ): enc = self.tokenizer( texts[i : i + batch_size], padding=True, @@ -759,25 +783,32 @@ def encode( return emb.numpy() if convert_to_numpy else emb def emit(self, sentences: Union[str, list[str]], **kw: Unknown) -> torch.Tensor: - # kw is a genuine passthrough to encode() -> Unknown (gradual) so it forwards into encode's typed - # params without ANN401 tripping on Any. cast the union return (encode -> ndarray | Tensor). + """Encode text and return the result as a PyTorch tensor. + + Additional keyword arguments are forwarded to :meth:`encode`. This is equivalent to calling + ``encode(..., convert_to_numpy=False)``. + """ + # `kw` is a typed passthrough to `encode`; `Unknown` avoids weakening the public signature to `Any`. return cast("torch.Tensor", self.encode(sentences, convert_to_numpy=False, **kw)) # --- auxiliary supervised heads (langset.heads) -------------------------------------------------------------- def seed_hidden(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: - """Per-sequence backbone hidden for a LEFT-padded batch -> [B, h]. Left padding puts every row's last real - token at the final column, so the pooled/final hidden is simply `hid[:, -1]` (this is the "hidden" read - site a value/time head reads; distinct from `_pool_hidden`, which assumes RIGHT padding). Grad flows into - the backbone/LoRA when they are trainable, so the head can SHAPE the representation.""" + """Return the final-token backbone hidden state for each sequence. + + The batch must be left-padded so every sequence's final real token occupies the last column. The result + has shape ``[batch, hidden_size]`` and remains connected to trainable backbone parameters. + """ hid = self._last_hidden( self._run_backbone(self.embed(input_ids), attention_mask, input_ids, 0) ) return hid[:, -1] def add_aux_head(self, module: nn.Linear, spec: dict[str, Any]) -> None: - """Register a PERSISTED auxiliary head so `save_pretrained` serializes it and `head_output` can read it back. - `spec` is the head's metadata (name / reads / loss / in_dim / out_dim / classes). Called by the trainer for - each persisted Head; the module object is SHARED with the trainer so in-place training keeps it current.""" + """Register an auxiliary linear head and its persistence metadata. + + ``spec`` describes the head name, read site, loss, input and output dimensions, and optional class labels. + Registered heads are serialized by :meth:`save_pretrained` and can be queried with :meth:`head_output`. + """ name = str(spec["name"]) self.aux_heads[name] = module self.aux_head_specs[name] = dict(spec) @@ -790,12 +821,18 @@ def head_output( batch_size: int = 32, reduce: str = "mean", ) -> Union[torch.Tensor, list[torch.Tensor]]: - """Query a PERSISTED auxiliary head at inference — the readout that makes a value/time head useful. - `reads="hidden"` heads read the pooled seed hidden -> one [out_dim] vector per sentence ([N, out_dim]). - `reads="recon"` heads read every emitted latent of the rollout: reduce="mean" -> [N, out_dim] (mean over a - row's emitted latents); reduce="none" -> a per-row list of [Li, out_dim] (the dense per-tick readout). - For a `loss="mse"` value head out_dim=1 (a scalar per state); for a `loss="ce"` head the columns are class - logits (argmax -> `aux_head_specs[name]['classes']`).""" + """Run a persisted auxiliary head on one or more sentences. + + Heads with ``reads="hidden"`` consume one pooled seed representation per sentence. Heads with + ``reads="recon"`` consume every latent produced by :meth:`rollout`; ``reduce="mean"`` averages each + sequence, while ``reduce="none"`` preserves the per-step outputs. MSE heads typically return scalar + values. Classification heads return logits whose columns correspond to + ``aux_head_specs[name]["classes"]``. + + Raises: + KeyError: If ``name`` is not registered. + ValueError: If the reduction or the head's read site is unsupported. + """ if name not in self.aux_heads: raise KeyError(f"no persisted head {name!r}; have {sorted(self.aux_heads)}") if reduce not in ( @@ -851,8 +888,11 @@ def head_output( @torch.no_grad() def generate_text(self, prompt: str, max_new: int = 200) -> str: - """Greedy text generation via the TIED input embedding (the lm_head is stripped). Used to MEASURE whether - [LEARN]/train_base actually injected knowledge — ask the trained model a question and read its answer.""" + """Generate text greedily using the tied input embeddings as the output projection. + + Generation stops at the tokenizer's EOS token or after ``max_new`` tokens. The model is left in evaluation + mode after this call. + """ self.eval() tok, dev = self.tokenizer, self.device msgs = [{"role": "user", "content": prompt}] @@ -900,11 +940,17 @@ def rollout( tuple[torch.Tensor, ...], tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]], ]: - """Autoregressively emit named-state superpositions and feed each one back. + """Autoregressively emit codebook-based latent states. - ``temperature`` reshapes the code distribution before it is committed; zero uses the trained - distribution unchanged. ``return_soft`` returns the committed vectors and per-step code entropy. - QueryBridge is a separate non-autoregressive family and routes inference through its objective. + Each emitted state is fed back into the backbone until its stop logit exceeds ``stop_threshold`` or + ``max_steps`` is reached. This method requires a model constructed with ``code_emit=True``. + ``temperature <= 0`` uses the learned code distribution unchanged; positive values rescale its logits. + Batched outputs zero-pad halted rows, while outputs for a single string are trimmed to their emitted length. + + By default, returns the emitted latents. ``return_lengths`` adds per-row lengths; ``return_confidence`` + returns ``(latents, lengths, confidence)``; and ``return_soft`` returns + ``(latents, lengths, committed_latents, entropy)``. These flags are checked in that precedence order: + ``return_soft``, then ``return_confidence``, then ``return_lengths``. """ if not self.head.code_emit: raise ValueError( @@ -991,24 +1037,21 @@ def rollout_train_state( kv_cache: bool = False, return_emit_hidden: bool = False, ) -> tuple[torch.Tensor, ...]: - """Autoregressive concept/codebook pass. - - Maps targets to nearest code indices for code-classification objectives while preserving the exact target - vectors for teacher forcing. Returns code logits ``[B, L+1, 1, n_codes]``, STOP logits, indices, and exact - target vectors. - - `ss_prob`=0 (default): pure TEACHER FORCING in ONE forward pass — every position predicted from the true - prefix, gradients flow at most ONE hop (byte-identical to before). `ss_prob`>0: SCHEDULED SAMPLING for the - first `train_hops` positions (None = all) — with prob `ss_prob` each of those positions is fed the model's - OWN emitted latent instead of the ground truth, so the emitter learns to consume its own (imperfect) - predictions. This is the exposure-bias fix that makes MULTI-HOP rollout trained rather than emergent. Self- - fed latents are DETACHED (standard scheduled sampling). Cost = train_hops+1 - backbone passes (positions past train_hops are teacher-forced in one pass). - - `ss_mask` (optional [B, H] bool): the PRECOMPUTED per-(row, hop) self-feed decisions. When given, the loop - uses `ss_mask[:, h]` in place of a fresh `torch.rand < ss_prob` draw. This makes the rollout DETERMINISTIC - given the mask, so GradCache's phase-1 (no_grad, full batch) and phase-2 (grad, per chunk) forwards produce - identical `recon` and the cached gradients line up exactly. None = sample as usual (byte-identical).""" + """Run an autoregressive concept/codebook training pass. + + Targets are mapped to nearest code indices for code-classification objectives while their exact vectors are + retained for teacher forcing. Returns code logits with shape ``[B, L + 1, 1, n_codes]``, stop logits, code + indices, and exact target vectors. When ``return_emit_hidden=True``, the corresponding emission hidden states + are appended to the result. + + With ``ss_prob=0``, all positions are predicted from the true prefix in one forward pass. With a positive + probability, the first ``train_hops`` positions may consume detached model emissions instead of ground-truth + latents; later positions remain teacher-forced. This scheduled-sampling path requires multiple backbone + forwards. + + ``ss_mask`` optionally supplies precomputed ``[B, H]`` self-feed decisions. Reusing the same mask makes + GradCache's detached and gradient-enabled passes produce aligned emissions and cached gradients. + """ assert self.head.multi_latent and self.head.code_emit bsz, s_len = input_ids.size(0), input_ids.size(1) n = target_latents.size(1) @@ -1016,7 +1059,7 @@ def rollout_train_state( codes = codes.view(bsz, n, -1) # [B, L, 1] recon = recon.view(bsz, n, -1) # [B, L, d] — clean feedback + recon target H = n if train_hops is None else max(0, min(int(train_hops), n)) - if ss_prob <= 0 or n == 0 or H == 0: # TEACHER-FORCED one-shot (default, fast) + if ss_prob <= 0 or n == 0 or H == 0: # one-pass teacher forcing seed = self.embed(input_ids) fb = self.head.feedback( recon.detach().to(seed.dtype) @@ -1033,7 +1076,7 @@ def rollout_train_state( if return_emit_hidden else (code_lg, stop_lg, codes, recon) ) - # SCHEDULED-SAMPLING multi-hop path + # Scheduled-sampling path: self-fed positions are evaluated serially. dev = recon.device if ( kv_cache @@ -1147,16 +1190,23 @@ def _feed( return code_lg, stop_lg, codes, recon def get_sentence_embedding_dimension(self) -> int: + """Return the configured latent embedding dimension.""" return self.latent_dim def as_sentence_transformer(self) -> SentenceTransformer: - """Wrap as a `sentence_transformers.SentenceTransformer` so it drops into SetFit as `model_body`.""" + """Return a SentenceTransformer wrapper suitable for APIs such as SetFit.""" from langset.st_module import to_sentence_transformer return to_sentence_transformer(self) # ---- persistence (LoRA + head + config; backbone rebuilt from ids) ---- def save_pretrained(self, path: Union[str, Path]) -> None: + """Save the LangSet configuration and trainable state to a directory. + + Pretrained models whose base weights can be reconstructed store the emission head and LoRA adapters. + Randomly initialized or fully fine-tuned models store the complete backbone. Persisted auxiliary heads and + emission-bridge state are included when present. + """ import json p = Path(path) @@ -1228,6 +1278,12 @@ def load( device: Optional[str] = None, attn_implementation: str = "sdpa", ) -> "LangSetModel": + """Load a model saved by :meth:`save_pretrained`. + + Reconstructs the original pretrained or randomly initialized backbone, restores the saved trainable state, + rebuilds persisted auxiliary heads, moves the model to ``device``, and returns it in evaluation mode. + Checkpoints using the removed FSQ emitter are rejected. + """ import json p = Path(path) diff --git a/src/langset/selection.py b/src/langset/selection.py index ace09bf..fbd7ecc 100644 --- a/src/langset/selection.py +++ b/src/langset/selection.py @@ -1,6 +1,12 @@ -"""Validation metrics for early-stop. The hard-won rule: NEVER select on training loss — a contrastive objective -can minimize it by COLLAPSING the geometry. Selection scores held-out input-view <-> target-view retrieval and -held-out reconstruction, with a hard collapse penalty (see Trainer).""" +"""Collapse-aware validation metrics for checkpoint selection. + +Training loss is deliberately not used to select the best model: a contrastive +objective can reduce its loss by collapsing the representation geometry. These +metrics instead evaluate held-out input-to-target retrieval and reconstruction, +then apply a hard penalty when the embedding geometry collapses. The trainer +uses them for early stopping and checkpoint selection so optimization and model +selection measure different things. +""" from __future__ import annotations diff --git a/src/langset/st_module.py b/src/langset/st_module.py index 47f1e97..8b11075 100644 --- a/src/langset/st_module.py +++ b/src/langset/st_module.py @@ -1,8 +1,17 @@ -"""Wrap a trained LangSetModel as a `sentence_transformers.SentenceTransformer` so it drops into SetFit as the -`model_body`: +"""Sentence-Transformers compatibility for trained langset models. + +This adapter wraps a trained single-latent :class:`langset.LangSetModel` as a +``sentence_transformers.SentenceTransformer``-compatible object. It provides +the interface expected by SetFit and similar tooling while keeping langset's +model and latent geometry underneath. + +Typical usage:: body = langset_model.as_sentence_transformer() setfit_model = SetFitModel(model_body=body, labels=[...]) + +Use this integration for embedding-style models; autoregressive multi-latent +rollouts remain on :class:`langset.LangSetModel` itself. """ from __future__ import annotations diff --git a/src/langset/strategies.py b/src/langset/strategies.py index f1ef6e5..d3a62b7 100644 --- a/src/langset/strategies.py +++ b/src/langset/strategies.py @@ -1,13 +1,8 @@ -"""Swappable training strategies for the multi-latent trainer (GoF Strategy pattern). +"""Interchangeable strategies for multi-latent training. -The multi-latent step is assembled from a few interchangeable pieces — the emission objective, the target -source, the aux loss terms, plus small function-strategies (epoch ordering, checkpoint selection, seed -building). Each is a class or callable with a fixed interface and a DEFAULT implementation here that reproduces -the historical behavior byte-for-byte (guarded by tests/test_trainer_multi_characterization.py). - -`TrainingArguments` holds these as INJECTABLE fields (defaults below), so selecting a different behavior is -passing a different implementation — `TrainingArguments(target_source=SIGRegTarget)` — not toggling a flag that -the trainer then branches on. The trainer builds each once and uses it with no per-feature `if`. +A training step combines an emission objective, a target source, optional loss terms, and small callables for +epoch ordering, checkpoint selection, and seed construction. :class:`TrainingArguments` stores the selected +implementations, which are instantiated once and used through shared interfaces. """ from __future__ import annotations @@ -44,13 +39,11 @@ # ---- per-step context handed to the aux loss terms ---------------------------------------------- @dataclass class MultiStepCtx: - """Read-only snapshot of ONE multi-latent training step, handed to every aux `_LossTerm`. Assembled right - after the emission forward; a term pulls the few fields it needs and returns its contribution (or None). + """Inputs shared by auxiliary loss terms for one multi-latent step. - Shape legend used below: B = rows in this batch · L = `lmax` (max emitted items across the batch's rows, - the padded time dim) · N = number of VALID emissions in the batch (= Σ lens_l, since rows differ in length) - · d = latent dim. The canonical flattened view a term works in is `recon[valid]` -> [N, d], - and `flat_texts` / `lens_l` / `bidx` are all aligned to that same row-major order. + ``recon`` and ``target_lat`` have shape ``[B, L, d]`` and ``valid`` marks the ``N`` real emission slots. + Indexing either latent tensor with ``valid`` produces the canonical flattened ``[N, d]`` view. + ``flat_texts``, ``lens_l``, and ``bidx`` follow the same row-major ordering. """ trainer: Trainer # the owning Trainer; read its PER-ROW data (indexed by dataset row id): @@ -84,16 +77,15 @@ class MultiStepCtx: # ---- aux loss terms ----------------------------------------------------------------------------- class _LossTerm: - """Strategy for one weighted, optional term added on top of the base emission loss (one per historical - `if a.lam_x > 0:` block). Terms are built once and iterated each step; each self-skips when inapplicable.""" + """Interface for an optional weighted term added to the base emission loss. + + Terms are built once, evaluated each step, and return ``None`` when their required inputs or weights are absent. + """ key: str = "" # this term's log/agg name (e.g. "loss_multi_nce"); set by each subclass - isolated_backward: bool = ( - False # if True the term is NOT summed into the shared loss; instead the trainer runs - ) - # its forward+backward SEPARATELY, AFTER the main loss.backward() has freed its - # graph, so the two graphs never coexist (peak activation = max, not sum). Grads - # accumulate into .grad before the single opt.step() -> same step, batch unchanged. + # Isolated terms run after the shared loss backward pass, so their graph does not coexist with the main graph. + # Their gradients accumulate before the same optimizer step. + isolated_backward: bool = False def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float]]: """Compute this term for the step described by `c`. Return `(key, raw_unweighted_loss, weight)` — the @@ -103,9 +95,10 @@ def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float def identical_text_mask(c: MultiStepCtx, fn_mask: torch.Tensor) -> None: - """A negative-mask: MUTATES `fn_mask` ([N, N] bool, N = valid emissions) in place, setting [i, j] = True for - pairs that must NOT be treated as negatives of each other. Default policy: two emissions with IDENTICAL - target text share the same true geometry, so they aren't negatives (mirrors the single-latent mask_keys path).""" + """Mask pairs with identical target text from the in-batch negative set. + + Mutates the ``[N, N]`` boolean ``fn_mask`` in place by setting excluded pairs to ``True``. + """ grp: dict[str, list[int]] = {} for ii, tx in enumerate(c.flat_texts): # flat_texts is row-major aligned with recon[valid] grp.setdefault(tx, []).append(ii) @@ -118,9 +111,11 @@ def identical_text_mask(c: MultiStepCtx, fn_mask: torch.Tensor) -> None: class MultiNCETerm(_LossTerm): - """IN-BATCH-NEGATIVE InfoNCE: each emitted recon vs the batch's EMA targets, own target = positive, others - = negatives, minus the `maskers`' false-negatives. On by default (lam_multi_nce). Ported from the - single-latent self-contrastive loss.""" + """Apply in-batch InfoNCE between emitted latents and target-source latents. + + Each emission is paired with its aligned target. Other targets act as negatives unless excluded by a configured + masker. The term is controlled by ``lam_multi_nce`` and is skipped when the target source suppresses NCE. + """ key = "loss_multi_nce" @@ -198,21 +193,21 @@ def contribute(self, c: MultiStepCtx) -> Optional[tuple[str, torch.Tensor, float def build_loss_terms(args: TrainingArguments) -> list[_LossTerm]: - """DEFAULT loss-term set, built once from the args. Fixed order (label -> multi_nce -> hard_neg -> sup) so the - float summation is byte-identical; each term self-skips when its weight/column is absent. Inject - `TrainingArguments(loss_terms=...)` with your own builder (or add terms like CoTGenTerm) to change the set. - NOTE: the phase head is no longer a term here — it is the `phase` instance of the generic AUXILIARY-HEAD plug - (langset.heads), applied inline in `_train_multi` right after this loop (the same summation position the old - `PhaseTerm` held, so lam_phase>0 stays byte-identical). `PhaseTerm` is kept for back-compat injection.""" + """Build the default auxiliary loss terms. + + Returns an in-batch InfoNCE term with identical-text false-negative masking, followed by the + supervised-contrastive term. Each term skips itself when its required weight or dataset labels are absent. + """ return [MultiNCETerm(maskers=[identical_text_mask]), SupConTerm()] class CoTGenTerm(_LossTerm): - """Exp-B: teach the model to GENERATE the row's chain-of-thought from the clean seed (doc=seed -> target= - cot_text) via the tied embedding — the SAME CE machinery the latents use, co-trained in the same step. - `isolated_backward` so its (long) CoT graph never coexists with the latent graph. Pairs with the seed+CoT - conditioning that `cot_seed_texts` applies to the emission forward — inject BOTH (see build_cot_loss_terms). - Self-skips when the batch's rows carry no CoT text (so it's inert if injected without a `cot_text` column).""" + """Train next-token generation of each row's ``cot_text`` from its input seed. + + The term uses an isolated backward pass so its long generation graph does not coexist with the latent-emission + graph. Pair it with :func:`cot_seed_texts` when emissions should also condition on the provided reasoning. + Batches without reasoning text are skipped. + """ key = "loss_cot" isolated_backward = True @@ -236,12 +231,9 @@ def _tokm(texts: list[str], mx: int, side: str) -> tuple[torch.Tensor, torch.Ten ) return e["input_ids"].to(dev), e["attention_mask"].to(dev) - # CoT blocks are long (p50~726, p90~1541 tok) -> keep full a.max_len, don't truncate hard. Pin padding sides - # EXPLICITLY (not the tokenizer's mutable default): the SEED is LEFT-padded so its last real token lands at - # index sd-1 (the position that predicts the first CoT token), and the CoT is RIGHT-padded so its real tokens - # sit adjacent to the seed with pads trailing -> no padding between seed and CoT, and the CE never conditions - # on a pad hidden. (Mirrors the emission forward's left-pad; a right-defaulting tokenizer would otherwise - # silently condition short seeds' CoT on padding.) + # Keep reasoning blocks up to the configured sequence limit. Seeds are left-padded so their final real token + # predicts the first reasoning token; reasoning targets are right-padded so no padding separates seed and + # target tokens, and padding remains excluded from cross-entropy. di, dm = _tokm([self_.input_text[k] for k in c.bidx], a.max_len, "left") ti, tm = _tokm([self_.cot_texts[k] or " " for k in c.bidx], a.max_len, "right") loss_cot = cot_loss( @@ -258,9 +250,10 @@ def _tokm(texts: list[str], mx: int, side: str) -> tuple[torch.Tensor, torch.Ten def build_cot_loss_terms(args: TrainingArguments) -> list[_LossTerm]: - """Exp-B loss set — INJECT via `TrainingArguments(loss_terms=build_cot_loss_terms)` (pair with - `seed_builder=cot_seed_texts`). The DEFAULT terms plus CoTGenTerm (isolated-backward), so the model is - co-trained to generate the row's reasoning while it emits the latents. Needs a `cot_text` dataset column.""" + """Build the default auxiliary terms plus :class:`CoTGenTerm`. + + Use with ``seed_builder=cot_seed_texts`` and a dataset containing ``cot_text``. + """ return [*build_loss_terms(args), CoTGenTerm()] @@ -278,10 +271,12 @@ class EmissionOut: class _EmissionObjective: - """Strategy for turning the seeded forward into emitted latents + the base emission loss. Selected ONCE per run so the step loop has no - emission `if`s. `codebook` is a class flag the free-run rollout reads to pick which emission head to use - (True = autoregressive codebook, False = strategy-owned continuous vectors). All objectives share the __init__ signature - (model, args, dev, trainer) so they are interchangeable as an injected `TrainingArguments.emission`.""" + """Interface for multi-latent emission objectives. + + An objective converts seeded inputs and target latents into gradient-bearing emitted latents and a base loss. + ``codebook`` indicates whether inference uses the model's autoregressive codebook rollout. Implementations share + a common constructor and can be selected through ``TrainingArguments.emission``. + """ codebook: bool = True @@ -322,11 +317,12 @@ def build_targets( d: int, dev: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, list[int], int]: - """Shape the per-row target item lists + their encoded latents into the emission's teacher-forcing tensors - `(target_lat [B,L,d], valid [B,L], lens_l, lmax)`. EXTRACTED verbatim from the trainer's inline block so the - STRATEGY owns target↔slot shaping: the autoregressive default is positional teacher forcing (item i -> slot i); a - future matching family (DETR/parallel-query) overrides this (e.g. all-valid, assignment deferred to its own - Hungarian match inside emit()). Pure tensor assembly — no model forward.""" + """Arrange encoded target items into padded teacher-forcing tensors. + + Returns ``(target_lat, valid, lens_l, lmax)``. ``target_lat`` has shape ``[B, L, d]``, ``valid`` marks real + target slots, ``lens_l`` contains per-row lengths, and ``lmax`` is the padded sequence length. The default + implementation preserves target-item order and performs no model forward. + """ lmax = max(len(x) for x in ent_lists) b = len(ent_lists) target_lat = torch.zeros(b, lmax, d, device=dev) @@ -342,9 +338,11 @@ def build_targets( return target_lat, valid, lens_l, lmax def emit_infer(self, texts: list[str], max_steps: int) -> tuple[torch.Tensor, torch.Tensor]: - """Inference emission: texts -> (lat [B, Lmax, d], lens [B]), zero-padding halted rows. The eval and - `model.rollout` route through here so a non-AR family (parallel-query) can emit in one pass. Default - delegates to the model's autoregressive concept/codebook rollout.""" + """Emit latent sequences for inference. + + Returns zero-padded latents with shape ``[B, L, d]`` and per-row lengths. The default implementation + delegates to the model's autoregressive codebook rollout; non-autoregressive objectives may override it. + """ lat, lens = self.m.rollout( # ty: ignore[invalid-assignment] # return_lengths=True -> (lat, lens) texts, max_steps=max_steps, return_lengths=True ) @@ -353,8 +351,7 @@ def emit_infer(self, texts: list[str], max_steps: int) -> tuple[torch.Tensor, to def z_for_reg( self, em: EmissionOut, target_lat: torch.Tensor, valid: torch.Tensor, lmax: int ) -> tuple[torch.Tensor, torch.Tensor]: - """The (predicted, target) latents a TargetSource.regularizer (e.g. SIGReg) constrains. Default = - the emitted vs target latents directly ([N, d]).""" + """Return the ``[N, d]`` predicted and target representations passed to target regularization.""" return em.recon[valid], target_lat[valid] @@ -414,7 +411,7 @@ def emit( with torch.no_grad(): # the target's law over codes, recovered from the target latent w = F.relu(target_lat.float() @ code.t()) # [b, lmax, n_codes] w = w / w.sum(-1, keepdim=True).clamp_min(1e-9) - # MEMBERSHIP: soft-target CE over the codes ALONE, at real emission slots. + # Apply soft-target cross-entropy over code membership at real emission slots. loss_code = soft_target_cross_entropy_loss( SoftTargetCrossEntropyContext( logits=code_logits[:, :lmax, 0, :].float(), @@ -424,11 +421,8 @@ def emit( ).to_tensor() # Termination is independent of the member distribution, so set width cannot weaken its supervision. loss_stop = stop_loss(StopLossContext(logits=stop_lg, lengths=lens_l)).to_tensor() - # `recon` is the contract every loss term reads as "the model's emission, with gradient" (MultiNCETerm - # and the emitter both build their query from it). The rollout's own recon is the TARGET here -- this - # path is lossless, so it carries no gradient at all, and a term querying it would compare each target - # against itself: near-zero loss, exactly zero gradient, silently inert. Hand back the DIFFERENTIABLE - # committed mixture instead, which is both the true emission and the vector fed back each tick. + # Auxiliary terms consume `recon` as the gradient-bearing model emission. Return the differentiable committed + # mixture rather than the teacher-forcing target, which carries no emitter gradient on this path. mix = F.normalize( code_logits[:, :lmax, 0, :].float().softmax(-1) @ code, dim=-1 ) # [b, lmax, d] @@ -447,8 +441,8 @@ def emit( def z_for_reg( self, em: EmissionOut, target_lat: torch.Tensor, valid: torch.Tensor, lmax: int ) -> tuple[torch.Tensor, torch.Tensor]: - # Regularize the DISTRIBUTIONS, the codebook analogue of codebook's pre-quantization z: predicted law vs target - # law over codes, so a penalty spreads usage across the alphabet instead of collapsing onto a few codes. + # Regularize predicted and target laws over codes so the penalty acts on codebook usage rather than only on + # committed latent vectors. code_logits = em.code_logits assert code_logits is not None code = self.m.head.code @@ -458,22 +452,24 @@ def z_for_reg( class ConceptObjective(_EmissionObjective): - """Emit a state as a superposition over NAMED CONCEPTS, per facet, with a residual for the unnamed. - - The row format is text the whole way down: - - {"input_text": …, "target_text": …, - "concepts": {"vocals": ["yell-singing", "gang-vocals"], "mood": ["angry-but-vulnerable"]}} - - langset discovers each facet's alphabet by scanning the column, so no index is ever written by hand, and - the emission's leading dims are CONSTITUTED from those concepts — project them back on the codebook and you - get what the latent is made of, in your own vocabulary, with no probe. Each facet is normalized separately, - so a wide `vocals` mixture does not make `tempo` look uncertain. - - Reach for this when the twin's geometry is good but WRONG SOMEWHERE. The twin gives a usable space for free - and in-batch negatives separate it, but neither lets you say along WHICH axis. Naming one facet overwrites - that neighbourhood and leaves the rest to the residual — so the normal configuration is a narrow state half - and a wide residual, not the reverse. You are patching a geometry, not adopting an ontology. + """Emit a latent with independently normalized named-concept facets. + + The dataset column selected by ``concept_field`` contains one concept mapping per target step. For example:: + + { + "input_text": "...", + "target_texts": ["..."], + "concepts": [ + { + "vocals": ["yell-singing", "gang-vocals"], + "mood": ["angry-but-vulnerable"], + } + ], + } + + LangSet discovers each facet's alphabet from the dataset and builds a fixed codebook using ``code_source``. + Only facets present in a row are supervised. Each facet is normalized independently, and ``res_dim`` optionally + reserves dimensions for an unnamed residual representation. """ codebook = True @@ -492,8 +488,7 @@ def __init__( assert self.laws is not None, ( "ConceptObjective needs a parsed concepts column — set TrainingArguments.concept_field" ) - # Build the codebook from the DISCOVERED alphabet: every concept becomes a vector via `code_source`, - # laid out per facet. Once, at setup, over the concept names — then frozen for the run. + # Build a fixed code vector for every discovered concept, grouped into facet spans. alpha = getattr(trainer, "concept_alphabet", None) if alpha and not bool(model.head.code.abs().sum()): names_per_facet = [(f, alpha[f]) for f in alpha] @@ -565,7 +560,7 @@ def emit( flat = code_logits[:, :lmax, 0, :].float() # [b, lmax, n_codes] tgt, seen = self._target_law(b, lmax, lens_l, bidx) - # one soft-target CE PER FACET, over that facet's members only + # Normalize and supervise each stated facet independently. losses, per_facet = [], {} for fi, (m_lo, m_hi, _, _) in enumerate(self.spans): sel = seen[:, :, fi] & valid @@ -584,7 +579,7 @@ def emit( loss_stop = stop_loss(StopLossContext(logits=stop_lg, lengths=lens_l)).to_tensor() - # the emission the rest of the trainer sees — WITH gradient, or every aux term silently idles + # Auxiliary terms require the gradient-bearing committed emission rather than teacher-forcing targets. p = m.head.concept_probs(code_logits[:, :lmax]) state = F.normalize(p @ m.head.code, dim=-1) emission = ( @@ -607,28 +602,15 @@ def emit( class StateResidualObjective(_EmissionObjective): - """STATE + RESIDUAL emission: a named half and an unnamed half, in one latent. - - The latent splits at `head.res_dim`. The STATE half is a mixture over a fixed alphabet of members you can - name, trained by soft-target cross-entropy against the members that are actually live this tick. The - RESIDUAL half is trained by whatever the surrounding terms already do (recon to the target, in-batch NCE), - so it carries what the alphabet cannot say. - - Why both. A named alphabet is the strongest signal available -- the emission is CONSTITUTED from the members - rather than having them decoded out of it -- but it is also a ceiling: anything the alphabet cannot name, - the emission cannot hold, and a world model whose ontology has gaps will silently refuse to represent them. - The residual is where the unnamed goes, and how much it carries is measurable: ablate it at eval and the - drop is how incomplete the alphabet was. + """Emit a latent composed of a named codebook state and an unnamed residual. - It also keeps a self-supervised component in the design. The state half's meaning is externally specified - (you chose the members); the residual half's is not, so the latent is not wholly hand-defined. + The leading ``latent_dim - res_dim`` dimensions contain a normalized mixture over the fixed state codebook. + The trailing ``res_dim`` dimensions contain a learned residual. State labels come from the dataset column + selected by ``state_field`` and are represented as a per-step list of active member indices. Target probability + mass is divided evenly among the active members at each step. - Labels come from a per-row column named by `state_field`: a per-tick list of ACTIVE member indices (sparse, - e.g. `[[3, 47], [12], ...]`). Mass is split evenly over each tick's live members, so a tick holding four - members asks the emission to spread over four -- the superposition, written as the target. - - TrainingArguments(emission=StateResidualObjective, state_field="frontier", state_classes=256) - LangSetModel.from_pretrained(..., code_emit=True, n_codes=256, res_dim=64) + Construct the model with ``code_emit=True``, set ``n_codes`` to the state alphabet size, and choose a positive + ``res_dim`` when an unnamed residual is required. """ codebook = True @@ -642,8 +624,7 @@ def __init__( ) self.n_codes = int(model.head.n_codes) self.res_dim = int(model.head.res_dim) - # Build the codebook ONCE, here, from the injected `code_source` — over the alphabet's NAMES, not over - # the training data. Skipped when a codebook was already installed by hand (head.set_code before train). + # Build the fixed codebook from member names unless the caller installed one before training. names = getattr(args, "code_names", None) if names is not None and not bool(model.head.code.abs().sum()): codes = build_codebook( @@ -691,7 +672,7 @@ def emit( ) code = m.head.code - # STATE: soft-target CE over the alphabet, mass split evenly across each tick's live members. + # Split target probability mass evenly across each step's active members. tgt = torch.zeros(b, lmax, self.n_codes, device=dev) has = torch.zeros(b, lmax, dtype=torch.bool, device=dev) assert self.labels is not None @@ -710,14 +691,12 @@ def emit( ) ).to_tensor() - # TERMINATION: its own sigmoid, never folded into the member softmax. Folding is only fair when the - # member target is one-hot; against a 1/k-diffuse target the gradient suppressing STOP scales with - # P(STOP), which weakens as the set widens, so the rollout truncates exactly where sets get wide. + # Termination uses an independent sigmoid. Folding it into a diffuse member distribution would make stop + # supervision depend on the number of active members. loss_stop = stop_loss(StopLossContext(logits=stop_lg, lengths=lens_l)).to_tensor() - # THE EMISSION the rest of the trainer sees: state mixture ++ residual, exactly what feeds back. It must - # be the emission WITH gradient -- every aux term (in-batch NCE, hard negatives) builds its query from - # `recon`, and handing back the stop-grad target instead leaves them running but training nothing. + # Return the same gradient-bearing state-plus-residual emission that is fed back. Auxiliary losses use this + # value as their query and must not receive the detached teacher-forcing target. p = code_logits[:, :lmax, 0, :].float().softmax(-1) state = F.normalize(p @ code, dim=-1) if self.res_dim: @@ -737,41 +716,37 @@ def emit( # ---- target source ------------------------------------------------------------------------------ class _TargetSource: - """Strategy for the TARGET latents the emission trains toward, plus any anti-collapse regularization. - Default = stop-grad EMA twin. Selected ONCE per run. All sources share the __init__ signature - (model, args, tok, dev) so they are interchangeable as an injected `TrainingArguments.target_source`.""" + """Interface for target-latent providers and optional anti-collapse regularization. - suppresses_nce: bool = ( - False # if True the trainer skips the in-batch NCE term (a live-target source that - ) - # already prevents collapse via `regularizer` doesn't need — and fights — it) - wants_regularizer: bool = ( - False # if True the trainer computes objective.z_for_reg and adds `regularizer` to the - ) - # loss; keeps that (non-trivial) work off the default path when False - twin: Optional[LangSetModel] = ( - None # the model the EVAL block encodes its retrieval bank with (the EMA copy for the - ) - # default; the online model itself for a live-target source). Set by subclasses. + Sources share a common constructor and can be selected through ``TrainingArguments.target_source``. + """ + + suppresses_nce: bool = False # skip in-batch NCE when the target source replaces it + wants_regularizer: bool = False # request objective representations for `regularizer` + twin: Optional[LangSetModel] = None # encoder used to build the evaluation retrieval bank def encode(self, texts: list[str]) -> torch.Tensor: - """Emit each text -> [n, d] L2-normalized target latents (no grad for the EMA default). Used both for the - per-step targets and, via MultiStepCtx.target_source, for the hard-negative bank.""" + """Encode texts as ``[n, d]`` L2-normalized target latents. + + The gradient policy is defined by the target source. The trainer uses these values for per-step targets and + for the hard-negative bank. + """ raise NotImplementedError def update(self) -> None: - """Called once AFTER each opt.step(). The EMA default nudges the twin toward the online weights; a - live-target source has nothing to track, so this is a no-op.""" + """Update source state after an optimizer step; the base implementation is a no-op.""" def regularizer(self, z_pred: torch.Tensor, z_tgt: torch.Tensor) -> Optional[torch.Tensor]: - """Optional extra anti-collapse loss on the emitted (`z_pred`) and target (`z_tgt`) latents, added to the - step loss. None for the EMA default (the stop-grad twin is what prevents collapse there).""" + """Return an optional regularization loss for predicted and target representations.""" return None class EMATwinTarget(_TargetSource): - """DEFAULT: a stop-grad EMA copy of the online model supplies the target latents (BYOL/JEPA) so both sides - don't move together and collapse. Byte-identical to the historical inline twin + emit_texts + ema_update.""" + """Produce stop-gradient targets from an exponential-moving-average copy of the online model. + + After each optimizer step, :meth:`update` moves the target model toward the trainable online parameters using + ``ema_m``. The lagged target provides a more stable comparison geometry than using the online model on both sides. + """ suppresses_nce = False @@ -793,9 +768,8 @@ def __init__( ] def encode(self, texts: list[str]) -> torch.Tensor: - # Single-latent emission of each text -> [N, d] normalized, no_grad. Truncated to target_max_len - # (default 64: targets are short descriptors; raise it when a target is a DOCUMENT, e.g. emit_seed's - # phase-0 target is a full abstract). Short future strings are already < 64 so unaffected. + # Encode each text as one normalized target latent without gradients. Inputs are truncated to + # `target_max_len`, which callers should increase for document-length targets. a, tok, dev = self.a, self.tok, self.dev e = tok( texts, padding=True, truncation=True, max_length=a.target_max_len, return_tensors="pt" @@ -813,19 +787,12 @@ def update(self) -> None: class CachedTarget(_TargetSource): - """FROZEN-ENCODER target source with an encode-once cache — the two-stage (V-JEPA) split for the multi-latent - world model. The target geometry is a FIXED encoder, so every text->latent map is CONSTANT for the whole run: - encode each unique text ONCE, memoize, then serve lookups. The step loop stops re-encoding BOTH the per-step - targets (trainer flat_texts) and the batch-pooled hard-negative bank (the emitter) — bringing the 'train only - the head on cached vectors, epochs in seconds' win that langset's single-latent frozen-pool path already has - to the multi-latent rollout. It's the fix for re-encoding fixed data every epoch. - - Encoder source: `args.target_encoder_ckpt` (a saved LangSetModel dir) when set — the intended use, a SEPARATE - already-good geometry (e.g. a trained affordance/embedding model) that the emitter learns to roll FORWARD in, - so no encoder is co-trained at all. Otherwise a frozen snapshot of the online model at init (only meaningful - when it already starts from a good encoder). INJECT via - `TrainingArguments(target_source=CachedTarget, target_encoder_ckpt="path/to/encoder")`. `update()` is a no-op - (the geometry is fixed) and the eval retrieval bank encodes through the same frozen geometry (`twin = enc`).""" + """Produce targets with a frozen encoder and cache embeddings by input text. + + If ``target_encoder_ckpt`` is configured, the encoder is loaded from that LangSet checkpoint. Otherwise, the + source is a frozen copy of the online model taken at initialization. Each unique text is encoded once and + subsequent calls reuse the cached latent. :meth:`update` is a no-op because the target geometry is fixed. + """ suppresses_nce = False @@ -875,11 +842,12 @@ def update(self) -> None: class SIGRegTarget(_TargetSource): - """EMA-free anti-collapse (LeJEPA, arXiv:2511.08544). INJECT via `TrainingArguments(target_source=SIGRegTarget)`. - Targets come from the LIVE model WITH gradient (no stop-grad twin); collapse is prevented by an isotropic-Gaussian - SIGReg penalty on the emitted representation (via `regularizer`) instead of by a twin. So it drops the twin's VRAM + target - forward, the in-batch NCE is suppressed (the regularizer replaces it), and eval encodes with the live model itself. - Reads scalar knobs off args: sigreg_lambda (loss weight, applied in the trainer), sigreg_knots, sigreg_slices.""" + """Use the online model for gradient-bearing targets and apply SIGReg regularization. + + This target source does not maintain an EMA copy. Both predicted and target representations remain connected to + the online model, and independent SIGReg penalties are applied to them. The default in-batch NCE term is + suppressed. Configure the regularizer with ``sigreg_lambda``, ``sigreg_knots``, and ``sigreg_slices``. + """ suppresses_nce = True wants_regularizer = True @@ -896,8 +864,7 @@ def __init__( self.sig_reg = SIGReg(knots=args.sigreg_knots, slices=args.sigreg_slices).to(dev) def encode(self, texts: list[str]) -> torch.Tensor: - # LIVE target WITH gradient (no no_grad, no twin): both the emitted and target latents move, and SIGReg — - # not a stop-grad twin — is what stops them collapsing together. + # Keep live target embeddings connected to the online model; SIGReg supplies the configured regularization. a, tok, dev = self.a, self.tok, self.dev e = tok( texts, padding=True, truncation=True, max_length=a.target_max_len, return_tensors="pt" @@ -911,7 +878,7 @@ def regularizer(self, z_pred: torch.Tensor, z_tgt: torch.Tensor) -> Optional[tor # ---- concepts (the text-in format for a named, superposed state) --------------------------------- -# A row carries a `concepts` column of NAMED FACETS, each holding the concepts that are true of it: +# A row's `concepts` column maps facet names to the concepts that are active: # # "concepts": {"vocals": ["yell-singing", "gang-vocals"], "tempo": {"7": 0.1, "8": 0.9}} # @@ -919,11 +886,10 @@ def regularizer(self, z_pred: torch.Tensor, z_tgt: torch.Tensor) -> Optional[tor # weights, which also encodes a CONTINUOUS value as a mixture over ordered concepts: 7.9 is 0.1 of "7" and 0.9 # of "8", interpolation included, and unlike a regressed scalar you can still read what it says. # -# Everything is text. The alphabet is not configured — it is DISCOVERED by scanning the column, the way a -# tokenizer's vocabulary is, so nobody writes or maintains an index. Multi-latent rows pass a LIST of these -# dicts, one per tick, aligned with `target_texts`. +# The alphabet is discovered by scanning the column. Multi-latent rows provide one mapping per step, aligned with +# `target_texts`. def parse_concepts(raw: object) -> "dict[str, dict[str, float]]": - """One row's (or tick's) concepts -> {facet: {concept: weight}}, weights normalized within each facet.""" + """Parse one concept mapping and normalize positive weights within each facet.""" if raw is None: return {} if not isinstance(raw, dict): @@ -952,8 +918,10 @@ def parse_concepts(raw: object) -> "dict[str, dict[str, float]]": def discover_concept_alphabet(rows_concepts: "list[object]") -> "dict[str, list[str]]": - """Scan the corpus's concept column -> {facet: sorted concept names}. Sorted for determinism, so the same - corpus always yields the same layout and a checkpoint stays readable.""" + """Return sorted concept names for each facet discovered in the dataset. + + Sorting facets and members makes the resulting codebook layout deterministic. + """ seen: dict[str, set[str]] = {} for raw in rows_concepts: ticks = raw if isinstance(raw, (list, tuple)) else [raw] @@ -964,20 +932,14 @@ def discover_concept_alphabet(rows_concepts: "list[object]") -> "dict[str, list[ # ---- code sources (where a named member's vector comes from) ------------------------------------- -# A `code_source` maps the alphabet's member NAMES to their vectors: (names, dim, model) -> [n_members, dim]. -# Called ONCE at setup, over the alphabet (hundreds to a few thousand short strings), never per batch. The -# result is frozen into a buffer for the run — a codebook that re-embeds while the model trains is a learned -# codebook in disguise, and re-opens the collapse problem the fixed codebook was chosen to avoid. -# -# The choice is a real trade, not a default: an orthonormal codebook decodes losslessly by plain matmul (so a -# recall number measures the EMISSION, with no probe as a confound) but its members carry no relation to each -# other; embedded codes put `ph2` near `ph3` and `melanoma` near `breast_neoplasms`, at the cost of exact -# recovery. Benchmarks want the first, real domains usually want the second. +# A `code_source` maps member names to a `[n_members, dim]` tensor during setup. The result is frozen for the run. +# Orthonormal codes support exact linear recovery but encode no member relationships; model-derived codes may +# preserve information from token embeddings but are not orthogonal or exactly invertible. def random_orthonormal_codes(names: list[str], dim: int, model: LangSetModel) -> torch.Tensor: - """Arbitrary but perfectly decodable: a seeded orthonormal frame, one row per member (`C Cᵀ = I`). + """Return a deterministic random orthonormal codebook with one row per member. - Member vectors are unrelated by construction — cell 10 and cell 11 are as orthogonal as cell 10 and cell - 200 — so nothing about the geometry can flatter a result. That is exactly why it belongs on a benchmark.""" + Requires ``len(names) <= dim`` and produces rows satisfying ``C Cᵀ = I``. + """ assert len(names) <= dim, ( f"random_orthonormal_codes needs dim >= n_members for an orthonormal frame; got dim={dim}, " f"n_members={len(names)}. Use model_embedded_codes (no such limit) or widen the state half." @@ -988,12 +950,11 @@ def random_orthonormal_codes(names: list[str], dim: int, model: LangSetModel) -> def model_embedded_codes(names: list[str], dim: int, model: LangSetModel) -> torch.Tensor: - """The base model's own reading of each member's NAME, mean-pooled over its tokens. + """Build code vectors by mean-pooling input embeddings for each member name. - Semantically related members land near each other, which is information the random frame throws away, and - the codebook becomes model-derived rather than hand-designed. Not orthonormal, so a mixture no longer - recovers its members exactly — read out with top-k over member scores rather than expecting a clean inverse. - Unlimited alphabet size (no dim >= n_members constraint).""" + Each vector is truncated or padded to ``dim`` and normalized. The resulting rows are not guaranteed to be + orthogonal or exactly invertible, and the alphabet size is not limited by ``dim``. + """ emb = model.embed.weight # [V, h]; tied to the LM head on most small models out = [] for nm in names: @@ -1006,21 +967,21 @@ def model_embedded_codes(names: list[str], dim: int, model: LangSetModel) -> tor def orthogonalized_codes(names: list[str], dim: int, model: LangSetModel) -> torch.Tensor: - """Embed the names, then orthogonalize — keeps a lossless readout and as much of the semantic arrangement - as an orthonormal frame can hold. The distortion is the price; the members are no longer purely the model's - own vectors.""" + """Embed member names and apply QR decomposition to produce orthonormal rows. + + This enables an orthogonal readout but generally changes pairwise relationships between the name embeddings. + """ e = model_embedded_codes(names, dim, model) q = torch.linalg.qr(e.t().float()).Q[:, : len(names)] return q.t().contiguous() def twin_encoded_codes(names: list[str], dim: int, model: LangSetModel) -> torch.Tensor: - """Encode each member name with the model's own emit path, so the codebook lands in the SAME space as the - training targets. + """Encode member names through the model's emission path and normalize the resulting code vectors. - This is the coherent option: the state mixture and the twin's target then share one geometry, so the state - loss and the recon loss are talking about the same thing, and the residual is precisely the part of the - target the named members cannot span. Requires the emit path to be usable at setup.""" + This places the codebook in the same output space as model-produced target latents. The emission path must be + usable during setup. + """ with torch.no_grad(): z = model.emit(list(names)) # [n_members, latent_dim], the target space itself z = z.float()[:, :dim] if z.size(-1) >= dim else F.pad(z.float(), (0, dim - z.size(-1))) @@ -1056,22 +1017,20 @@ def build_codebook( def multi_epoch_order( tr_idx: list[int], rng_t: torch.Generator, args: TrainingArguments, seeds: list[str] ) -> list[int]: - """DEFAULT epoch ordering: a plain shuffle of the training positions. Inject a different `epoch_order` to - change it.""" + """Return a random permutation of the training positions for one epoch.""" return torch.randperm(len(tr_idx), generator=rng_t).tolist() def multi_select_metric(mode: str, mrr: float, pur: float, ep: int) -> float: - """DEFAULT checkpoint-selection signal from the epoch's metrics. retr_mrr (default) / purity / blend. Inject - a different `selector` to change it (e.g. one that keeps the last epoch).""" + """Return the configured retrieval, purity, or blended checkpoint-selection score.""" return pur if mode == "purity" else (mrr + pur) if mode == "blend" else mrr def last_epoch_selector(mode: str, mrr: float, pur: float, ep: int) -> float: - """SUPERPOSITION selector — INJECT via `TrainingArguments(selector=last_epoch_selector)`. No early-stop signal; - keeps the LAST epoch (returns float(ep)). Under superposition training retr_mrr selects for a collapsed - one-future-per-seed geometry — exactly the wrong target when you WANT the latent to spread over a seed's - alternative futures, so retr_mrr is meant to fall and must not gate selection.""" + """Return the epoch index so each evaluated epoch supersedes earlier checkpoints. + + The attached ``needs_final_epoch`` flag ensures the final epoch is evaluated even when ``eval_every > 1``. + """ return float(ep) @@ -1082,14 +1041,14 @@ def last_epoch_selector(mode: str, mrr: float, pur: float, ep: int) -> float: def multi_seed_texts(trainer: Trainer, seeds: list[str], args: TrainingArguments) -> list[str]: - """DEFAULT texts fed to the EMISSION forward — what the model reads before emitting its latents = the raw - input seeds. Inject a different `seed_builder` to change it (e.g. append per-row CoT); targets/eval keep raw seeds.""" + """Return the raw input seeds used by the default emission forward.""" return seeds def cot_seed_texts(trainer: Trainer, seeds: list[str], args: TrainingArguments) -> list[str]: - """Exp-B seed-builder — INJECT via `TrainingArguments(seed_builder=cot_seed_texts)` (pair with - `loss_terms=build_cot_loss_terms`). Conditions the emission forward on each row's teacher-forced reasoning - (seed + CoT) so the latents are emitted AFTER the reasoning; targets and eval keep the raw seeds, and - CoTGenTerm trains the model to produce that reasoning itself.""" + """Append each row's reasoning text to its seed before emission. + + Pair with :func:`build_cot_loss_terms` so the same reasoning is trained autoregressively. Targets and evaluation + continue to use the raw seeds. + """ return [f"{s}\n\nReasoning:\n{trainer.cot_texts[i]}" for i, s in enumerate(seeds)] diff --git a/src/langset/trainer.py b/src/langset/trainer.py index 0ceb1cf..b1be2f1 100644 --- a/src/langset/trainer.py +++ b/src/langset/trainer.py @@ -1,9 +1,11 @@ -"""Trainer: fit the LLM emitter so `emit(input_text)` lands where `emit(target_text)` does — a native -self-contrastive objective (both views in the model's own space, in-batch negatives). The target text DEFINES -the geometry. Two light aux terms keep it grounded and spread; selection is collapse-aware. +"""Training loops for single- and multi-latent LangSet models. -Dataset rows: `input_text` (what you have at inference) + `target_text` (a description of the same item that -defines where it should land). Pass a `datasets.Dataset` or `list[dict]`; use `column_mapping` to rename. +Single-latent training aligns ``input_text`` embeddings with corresponding ``target_text`` embeddings using a +self-contrastive objective. Multi-latent training consumes ``target_texts`` lists and delegates emission and +target construction to the configured strategy classes. + +Datasets may be ``datasets.Dataset`` instances or lists of dictionaries. Use ``column_mapping`` to map custom +column names to LangSet's canonical fields. """ from __future__ import annotations @@ -50,8 +52,7 @@ def _wandb_config(a: TrainingArguments) -> dict[str, Any]: - """`vars(a)` for wandb, with the injected strategy fields (classes/callables — not JSON-serializable) rendered - as their names, so the config still LOGS which strategy each run used (e.g. target_source="SIGRegTarget").""" + """Return a wandb-safe argument mapping with strategy classes and callables rendered by name.""" cfg = dict(vars(a)) for k, v in cfg.items(): if callable(v): @@ -74,9 +75,12 @@ def _fuse_views( mask_b: torch.Tensor, pad_id: int, ) -> tuple[torch.Tensor, torch.Tensor]: - """FUSE two views for ONE forward: pad both to the batch's common real max-len, stack on the batch dim. - RIGHT-padded -> real tokens lead; the padded columns are masked so each row's emit is IDENTICAL to a separate - per-view forward. Halves kernel launches and grad-ckpt recomputes. Split the output back at row B.""" + """Combine two right-padded views into one backbone forward. + + Both views are resized to the batch's common maximum real length and stacked along the batch dimension. Only + masked padding columns are added or removed, so each row has the same retained-token output as a separate + per-view forward. The caller splits the output at the original batch size. + """ L = int(max(mask_a.sum(dim=1).max().item(), mask_b.sum(dim=1).max().item())) def _fit(x: torch.Tensor, mk: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -95,10 +99,10 @@ def _fit(x: torch.Tensor, mk: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] def _dyn_trim(ids: torch.Tensor, mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """DYNAMIC PADDING: sequences are pre-tokenized/padded to the DATASET max, so short batches still forward at - full width. Trim each batch to its own real max length (RIGHT-padded, so real tokens are the leading columns). - The dropped columns are pure padding the attention mask already zeros -> forward output on the kept tokens is - IDENTICAL. Big saving when doc lengths vary (legal corpora do). Guarded by test_train_identity.""" + """Trim a right-padded batch to its longest non-padding sequence. + + Because only masked padding columns are removed, outputs for retained positions are unchanged. + """ n = int(mask.sum(dim=1).max().item()) if n <= 0 or n >= ids.size(1): return ids, mask @@ -107,9 +111,7 @@ def _dyn_trim(ids: torch.Tensor, mask: torch.Tensor) -> tuple[torch.Tensor, torc # ---- text replay (rehearsal) — shared by the single- and multi-latent learn paths --------------- def _require_emit_rows(is_learn: list[bool], learn_field: Optional[str]) -> None: - """Text replay pulls `learn`-tagged rows OUT of the emit/latent split (they only rehearse text). If EVERY row - is tagged, that split is empty — training would silently no-op (no emit steps) or crash deep in eval. Fail - loudly and early instead.""" + """Require at least one emission row when text-replay rows are separated from the training split.""" if is_learn and all(is_learn): raise ValueError( f"learn_field '{learn_field}' tagged all {len(is_learn)} rows as 'learn' — no rows left for the emit " @@ -120,10 +122,11 @@ def _require_emit_rows(is_learn: list[bool], learn_field: Optional[str]) -> None def _tokenize_replay( tok: PreTrainedTokenizerBase, texts: list[str], max_len: int, side: str, dev: torch.device ) -> tuple[torch.Tensor, torch.Tensor]: - """Tokenize replay text with an EXPLICIT padding side. The doc (the conditioning context) MUST be left-padded - so every row's last real token lands in the final column — that's the position whose hidden predicts the first - target token in `learn_loss`. The target is right-padded (its pad is masked out of the CE). Defaulting the side - (right-pad the doc) silently conditions the replay loss on padding for any row shorter than the batch max.""" + """Tokenize replay text with an explicit padding side. + + Conditioning documents must be left-padded so their final real token predicts the first target token in + ``learn_loss``. Targets are right-padded, and their padding is excluded from cross-entropy. + """ e = tok( texts, padding=True, @@ -136,12 +139,11 @@ def _tokenize_replay( def _snapshot_best(m: LangSetModel) -> dict[str, Any]: - """Snapshot best-so-far weights to restore at the end of training. MIRRORS LangSetModel.save_pretrained's - branch: a PRETRAINED + frozen-base model rebuilds its backbone from `llm_model`, so LoRA-only is enough; a - RANDOM-INIT model (from_scratch) OR a train_base=True full-finetune has no source to rebuild from, so we must - snapshot the FULL backbone — otherwise the best-so-far restore silently keeps the LAST-epoch backbone (a - random-init net has no 'lora' keys, so the old lora-only snapshot restored nothing into it; a full-FT run - trains base weights the lora-only snapshot never captured). Applies to single-latent, multi-latent, and masked.""" + """Snapshot all state required to restore the best epoch. + + Pretrained models whose base weights can be reconstructed store the head and LoRA parameters. Randomly + initialized and fully fine-tuned models store the complete backbone. Persisted auxiliary heads are included. + """ snap: dict[str, Any] = { "head": {k: v.detach().cpu().clone() for k, v in m.head.state_dict().items()} } @@ -162,7 +164,7 @@ def _snapshot_best(m: LangSetModel) -> dict[str, Any]: def _restore_best(m: LangSetModel, best_state: dict[str, Any]) -> None: - """Restore a _snapshot_best() payload. Back-compat: a legacy resume checkpoint carries only 'lora'.""" + """Restore a snapshot created by :func:`_snapshot_best`, including legacy LoRA-only snapshots.""" m.head.load_state_dict(best_state["head"]) if "backbone" in best_state: # random-init: full backbone m.backbone.load_state_dict(best_state["backbone"], strict=False) @@ -173,34 +175,45 @@ def _restore_best(m: LangSetModel, best_state: dict[str, Any]) -> None: # ---- single-latent step engines ----------------------------------------------------------------- -# The real axis of the single-latent path is "where do pred/target/hard-neg features come from, and how is the -# step run": the LIVE backbone (default) vs. FROZEN-POOL cached vectors (pool_mode="last" + frozen backbone). Both -# obey ONE interface; train() picks the engine ONCE, so the epoch loop and eval block have no per-feature `if`s. +# Single-latent steps obtain prediction, target, and hard-negative features either from the trainable backbone or +# from cached vectors produced by a frozen backbone. Both implementations expose the same step interface. class _StepEngine: - """Strategy for producing the contrastive features of one step. `supports_recon` gates the recon aux — the - frozen-pool engine has no backbone in the loop, so recon (which decodes via the backbone) is unavailable there.""" + """Interface for producing contrastive features for one single-latent step. + + ``supports_recon`` indicates whether the backbone is available for reconstruction loss. + """ supports_recon: bool = True def precompute(self) -> None: - """Called ONCE before the epoch loop (frozen-pool caches all view features here; backbone impl = no-op).""" + """Prepare reusable features before the epoch loop; the live-backbone implementation is a no-op.""" def featurize( self, idx: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: - """-> (pred [B,d] grad flows, target [B,d] honoring stop_grad_target, hard_neg [Hn,d] or None).""" + """Build the input, target, and optional hard-negative features for a batch. + + Returns ``(prediction, target, hard_negatives)``. Prediction and target have shape + ``[batch, latent_dim]``. Hard negatives are either ``None`` or a tensor of negative features. Whether + gradients flow through the target is controlled by ``stop_grad_target``. + """ raise NotImplementedError def val_embeddings(self, val_idx: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """-> (emit_in, emit_tg) numpy for the eval MRR/retrieval block.""" + """Return NumPy input and target embeddings for retrieval evaluation.""" raise NotImplementedError class BackboneStepEngine(_StepEngine): - """DEFAULT: runs the LIVE backbone. Owns the fuse_views / stop_grad_target / _dyn_trim variations internally - (they are "how to run the backbone", cohesive here). Byte-identical to the historical default path.""" + """Produce contrastive features by running the trainable backbone for each step. + + This is the default single-latent engine. It computes input, target, and + optional hard-negative features and supports reconstruction auxiliaries; + runtime options such as view fusion and target stop-gradient are handled + inside :meth:`featurize`. + """ supports_recon = True @@ -271,9 +284,12 @@ def val_embeddings(self, val_idx: np.ndarray) -> tuple[np.ndarray, np.ndarray]: class FrozenPoolStepEngine(_StepEngine): - """FROZEN-POOL: backbone frozen + pool_mode="last" -> features are STATIC. Encode every view ONCE in precompute() - and train only the head on the cached vectors (no backbone in the step loop) => epochs run in seconds. No recon - (it needs the frozen-out backbone), so select by retrieval MRR.""" + """Cache frozen-backbone features and train only the projection head. + + This engine is used when ``pool_mode="last"`` and the backbone is frozen. It precomputes input, target, and + optional hard-negative features before the epoch loop. Reconstruction loss is unavailable because the + backbone does not participate in training. + """ supports_recon = False @@ -362,6 +378,17 @@ def val_embeddings(self, val_idx: np.ndarray) -> tuple[np.ndarray, np.ndarray]: class Trainer: + """Train a :class:`LangSetModel` on text-to-latent alignment data. + + The trainer selects the single- or multi-latent path from the model configuration. Single-latent rows require + ``input_text`` and ``target_text``; multi-latent rows require ``input_text`` and a non-empty ``target_texts`` + list. Optional columns configure hard negatives, false-negative masks, concepts, states, auxiliary heads, and + text replay according to :class:`TrainingArguments`. + + ``eval_dataset`` is currently reserved and ignored; validation rows are split from ``train_dataset``. + ``on_checkpoint``, when provided, is called after a durable checkpoint is written. + """ + def __init__( self, model: LangSetModel, @@ -373,17 +400,14 @@ def __init__( ) -> None: self.model = model self.args = args - # if set, the best-so-far model is written to output_dir on every improvement and this is called after - # (e.g. modal Volume.commit) so another process can eval the live best checkpoint mid-training. + # When set, this callback runs after a best-so-far checkpoint is written, allowing an external system to + # publish or evaluate the durable checkpoint. self.on_checkpoint = on_checkpoint - # ONE switch routes the whole trainer: a multi_latent model emits a VARIABLE-LENGTH latent set, so it reads - # a `target_texts` (list[str] per row) column and runs the selected emission strategy; otherwise the single-latent - # self-contrastive path (byte-for-byte unchanged) reads a scalar `target_text` column. + # A multi-latent model reads a per-row `target_texts` list and uses the selected emission strategy. The + # single-latent path reads one `target_text` value per row. self.multi_latent = bool(model.head.multi_latent) - # ROLLOUT MUST BE TRAINED. A multi_latent model emits its latent set AUTOREGRESSIVELY at inference (rollout() - # feeds each emitted latent back), so training PURELY teacher-forced (ss_prob=0) is exposure-biased and the - # rollout is never actually trained — you cannot roll out what you did not train. Enforce it against the - # ss_prob sentinel: unset (None) -> 0.25 (scheduled sampling on); an EXPLICIT ss_prob=0 -> hard error. + # Multi-latent inference feeds emitted states back autoregressively. Require positive scheduled sampling so + # training exposes the model to its own prior emissions. An unspecified probability defaults to 0.25. if self.multi_latent: if args.emission is None: raise ValueError( @@ -456,10 +480,8 @@ def __init__( f"multi_latent Trainer needs a 'target_texts' column of non-empty lists; row {i} = {v!r}" ) self.target_texts.append([str(x) for x in v]) - # optional Exp-B CoT: a per-row reasoning STRING the model learns to GENERATE before the latents and - # conditions the emission on (via seed_builder=cot_seed_texts + loss_terms=build_cot_loss_terms). Absent - # column -> empty strings, so the default (non-CoT) strategies stay byte-identical: multi_seed_texts uses - # raw seeds and CoTGenTerm — even if injected — self-skips on all-empty reasoning. + # Optional per-row reasoning text used by the CoT loss and seed strategies. Missing values become empty + # strings; the default strategies ignore them, and CoTGenTerm skips batches without reasoning text. cot_key = inv.get( "cot_text", "cot_text" ) # honor column_mapping (a renamed reasoning column) @@ -486,8 +508,7 @@ def __init__( ] for v in raw_hn ] - # optional MULTI-latent supervised-contrastive: a per-row LIST of group labels aligned 1:1 with - # target_texts (each emitted item's stage/group). Shapes emissions into separate regions (weight lam_sup). + # Optional supervised-contrastive labels, aligned one-to-one with each row's target texts. self.sup_labels: Optional[list[list[str]]] = None sup_field = getattr(args, "sup_field", None) if sup_field is not None: @@ -495,10 +516,8 @@ def __init__( self.sup_labels = [ [str(x) for x in (v if isinstance(v, (list, tuple)) else [v])] for v in raw_sup ] - # CONCEPTS (ConceptObjective): a per-row dict of named facets -> the concepts true of each, e.g. - # {"vocals": ["yell-singing", "gang-vocals"], "tempo": {"7": 0.1, "8": 0.9}}. Multi-latent rows pass - # a LIST of those dicts, one per tick. The alphabet is DISCOVERED here by scanning the column — no - # index is ever authored — then frozen into the head's codebook by the injected `code_source`. + # ConceptObjective accepts one facet mapping per target step. Discover a deterministic alphabet from + # those mappings, then let the configured code source construct the fixed codebook. self.concept_alphabet: Optional[dict[str, list[str]]] = None self.concept_laws: Optional[list[list[dict[int, dict[int, float]]]]] = None concept_field = getattr(args, "concept_field", None) @@ -528,9 +547,7 @@ def __init__( f"{sizes}", flush=True, ) - # STATE MEMBER labels (StateResidualObjective): same sparse per-tick shape as the frontier column, but - # these are not decoded OUT of the emission by an aux head — they CONSTITUTE its named half, so the - # objective reads them directly rather than through a probe. + # StateResidualObjective reads per-step member indices directly as supervision for its named component. self.state_labels: Optional[list[list[list[int]]]] = None state_field = getattr(args, "state_field", None) if state_field is not None: @@ -539,17 +556,12 @@ def __init__( [[int(i) for i in tick] for tick in row[: args.max_target_items]] for row in raw_st ] - # optional PLUGGABLE AUXILIARY HEADS (langset.heads): each user Head names a `target` dataset column. Read - # the RAW per-row values here (a list per row for reads="recon", a scalar for reads="hidden"); resolution - # to class maps / float tensors happens per-loss-kind in _train_multi. The phase shim reuses sup_labels, - # so it is NOT read here. [] = off (byte-identical). Reject on the single-latent path (heads are - # multi-latent only — the phase-head lineage they generalize lives only here). + # Read each auxiliary head's target column. Runtime resolution later converts these raw row values into + # class IDs or numeric tensors according to the configured loss. self.head_cols: dict[str, list[object]] = {} for h in getattr(args, "heads", []): self.head_cols[h.name] = list(cols[inv.get(h.target, h.target)]) - # TEXT REPLAY tag (multi-latent): rows marked "learn" rehearse the backbone's plain next-token ability - # (interleaved at `learn_ratio` in _train_multi). Must be set HERE — this branch returns before the - # single-latent is_learn setup below. learn_field unset / learn_ratio=0 -> all-False (feature off). + # Mark rows used for multi-latent text replay. This branch returns before the single-latent replay setup. self.is_learn: list[bool] = [False] * len(self.input_text) learn_field = getattr(args, "learn_field", None) if learn_field is not None and args.learn_ratio > 0: @@ -609,6 +621,12 @@ def __init__( print(f"[langset] {len(self.input_text)} rows{jepa}{masked}{hn}{lr}", flush=True) def train(self) -> LangSetModel: + """Train the configured model and return the best restored model. + + Dispatches to the multi-latent strategy loop when the model uses a multi-latent head; otherwise runs + single-latent contrastive training. Checkpointing, early stopping, text replay, and auxiliary losses are + controlled by :class:`TrainingArguments`. + """ if self.multi_latent: return self._train_multi() a, m = self.args, self.model @@ -617,9 +635,8 @@ def train(self) -> LangSetModel: rng = np.random.default_rng(a.seed) tok = m.tokenizer - # Convert all text to tokenized tensors def tok_to(texts: list[str], mx: int) -> tuple[torch.Tensor, torch.Tensor]: - """Tokenize text to input_ids and attention_mask tensors""" + """Tokenize text into input-ID and attention-mask tensors on the model device.""" e = tok(texts, padding=True, truncation=True, max_length=mx, return_tensors="pt") return e["input_ids"].to(dev), e["attention_mask"].to(dev) @@ -710,8 +727,7 @@ def tok_to(texts: list[str], mx: int) -> tuple[torch.Tensor, torch.Tensor]: ) def save_resume(next_ep: int) -> None: - """Atomically persist FULL training state (weights+opt+connector+epoch+best+rng) so a preempt/retry resumes - from the last epoch boundary instead of ep0. tmp+rename => a mid-write preempt cannot corrupt the good file.""" + """Save model, optimizer, selection, and RNG state through an atomic file replacement.""" if not a.resume_dir: return d = Path(a.resume_dir) @@ -736,11 +752,10 @@ def save_resume(next_ep: int) -> None: ) tmp.replace(d / "resume.pt") if self.on_checkpoint is not None: - self.on_checkpoint() # e.g. Volume.commit() -> durable across preempt + self.on_checkpoint() # Notify the caller after the checkpoint becomes durable. - # ---- pick the step engine ONCE (see the _StepEngine classes above): FROZEN-POOL (backbone frozen + - # pool_mode="last" -> features static, cached once, head-only training in seconds) vs the LIVE backbone - # (default). The epoch loop + eval below then run ONE path with no per-feature `if`s. ---- + # Select the feature engine once. Frozen last-token pooling caches backbone features; other configurations + # compute features through the backbone during each step. if getattr(m, "pool_mode", "") == "last" and getattr(m, "_frozen_bb", False): engine: _StepEngine = FrozenPoolStepEngine( m, a, ids, mask, t2_ids, t2_mask, hn_ids, hn_mask @@ -759,12 +774,10 @@ def save_resume(next_ep: int) -> None: self.input_text, self.target_text, ) - engine.precompute() # no-op for the backbone; frozen-pool encodes+caches every view here + engine.precompute() # The frozen-pool engine caches all view features here. - # OPTIONAL step-level diagnostics (env-gated, OFF by default = byte-identical). Mirrors the multi-latent path's - # profiler harness for the single-latent trainer: LANGSET_PROFILE_STEPS=N captures a torch profiler over the - # first N steps then dumps a CUDA-time table and EXITS; LANGSET_MEM_PRINT=N prints the MEASURED VRAM peak for - # the first N steps (for right-sizing batch/seq-len instead of guessing at OOMs). + # Environment-gated diagnostics. LANGSET_PROFILE_STEPS profiles the requested number of steps and exits; + # LANGSET_MEM_PRINT reports CUDA memory usage for the requested number of steps. import os as _os import time as _time @@ -794,12 +807,13 @@ def save_resume(next_ep: int) -> None: ) def _grad_cache_step(idx: torch.Tensor) -> float: - """GradCache: EXACT full-batch contrastive gradient with peak activation = one gc_chunk (Gao et al. - 2021). Lets `batch_size` (in-batch negatives) grow far past what one graph fits. Phase 1 encodes the - whole batch in no_grad chunks (only [B,d] embeddings survive), builds the loss on the full batch and - caches d(loss)/d(embedding); Phase 2 re-forwards each chunk WITH grad and injects the cached grads via - autograd.backward -> exact param grad. emb_slot heads get their grad in phase 1 (they read only the - cached embeddings); the backbone gets its grad in phase 2. Requires dropout==0.""" + """Compute a full-batch contrastive gradient with chunk-sized backbone activation memory. + + The first pass computes detached batch embeddings and their loss gradients. Each chunk is then + recomputed with autograd enabled, and the cached embedding gradients are propagated through the + backbone. Embedding-slot heads receive gradients from the first pass. Deterministic forwards, including + zero dropout, are required. + """ ch = a.gc_chunk or a.batch_size chunks = [idx[j : j + ch] for j in range(0, len(idx), ch)] preds: list[torch.Tensor] = [] @@ -950,10 +964,9 @@ def _grad_cache_step(idx: torch.Tensor) -> float: connector=connector, ) ).to_tensor() - ) # SKIP so recon's fp32 full-vocab ([B,S,vocab]) projection graph is NOT built every step - # (that graph, not the stripped lm_head, OOM'd a 0.6B at 84GB without grad_ckpt). Also lets - # frozen-pool run. Direct path only: grad_cache asserts lam_recon==0 (embedding-only caching - # cannot hold the per-token backbone graph recon needs). Default 0.3 -> unchanged. + ) + # Building reconstruction only when enabled avoids its full-vocabulary projection graph. + # GradCache does not support this per-token backbone graph and requires `lam_recon == 0`. opt.zero_grad() loss.backward() opt.step() @@ -1001,10 +1014,9 @@ def _grad_cache_step(idx: torch.Tensor) -> float: ) # engine owns HOW val embeddings are produced mrr = selection.retrieval_mrr(emit_in, emit_tg)["mrr"] collapse = selection.collapse_score(emit_in) - if ( - not engine.supports_recon or a.lam_recon == 0.0 - ): # recon not the objective (frozen-pool OR lam_recon=0) -> - recon_val = 0.0 # select by retrieval MRR, and skip the wasteful recon-val + if not engine.supports_recon or a.lam_recon == 0.0: + # Without a reconstruction objective, select by collapse-penalized retrieval and skip reconstruction. + recon_val = 0.0 sel_score = mrr - _COLLAPSE_PENALTY * max( 0.0, collapse - _COLLAPSE_FLOOR ) # fp32 vocab projection @@ -1123,22 +1135,12 @@ def _train_multi(self) -> LangSetModel: perm[cut:].tolist() or perm[:1].tolist() ) # never-empty val (a tiny smoke can fill train) - # Target source (stop-grad EMA twin by default): supplies the target latents so both sides don't move - # together and collapse. Selected ONCE (see the _TargetSource seam). - target_source: _TargetSource = a.target_source( - m, a, tok, dev - ) # target strategy (INJECTED), built ONCE - - # PLUGGABLE AUXILIARY HEADS (langset.heads) — the generalization of the phase head. Each Head hangs a small - # supervised head off the model at a chosen READ SITE ("recon" = per emitted latent, like the phase head | - # "hidden" = pooled per-sequence backbone hidden) with a chosen LOSS ("ce"/"mse"/custom callable) and - # LIFECYCLE (transient shaping-gradient, or PERSISTED + queryable at inference). `lam_phase` is folded in as - # the `phase` SHIM (a transient recon+CE head over sup_labels), built FIRST so its nn.Linear draws RNG at the - # exact point the old inline phase head did -> lam_phase>0 is byte-identical (guarded by the multi-latent - # golden). User heads follow. Each RtHead owns its nn.Linear (trainer-owned, added to the optimizer below); - # a PERSISTED head is ALSO registered on the model (shared module) so save_pretrained serializes it and - # `head_output()` can query it at inference. phase_head/phase_ids stay None/{} for MultiStepCtx back-compat - # (PhaseTerm was dropped from the default loss terms; phase is applied inline like the other heads). + # Build the configured target source once for the training run. + target_source: _TargetSource = a.target_source(m, a, tok, dev) + + # Resolve auxiliary heads before constructing the optimizer. The phase compatibility option becomes a + # transient classification head; user-configured heads follow. Persisted heads share their module with the + # model so save_pretrained() and head_output() can retain and query them. phase_head: Optional[torch.nn.Module] = None phase_ids: dict[str, int] = {} _spec_values: list[tuple[Head, list[object]]] = [] @@ -1199,9 +1201,10 @@ def _train_multi(self) -> LangSetModel: @torch.no_grad() def evaluate() -> dict[str, float]: - """Free-roll each val seed -> emitted latents; decode each by nearest-neighbor against an EMA-emitted bank - of the val `target_texts`. Reports (a) retrieval MRR vs the chain's OWN targets and (b) a NON-COLLAPSE - diversity count = distinct nearest-bank items produced (state must not mean-collapse to one mode).""" + """Evaluate free-running validation emissions against the target-source retrieval bank. + + Reports retrieval MRR against each row's targets and the number of distinct nearest-bank items emitted. + """ m.eval() import time as _et @@ -1320,8 +1323,7 @@ def evaluate() -> dict[str, float]: ) def save_resume(next_ep: int) -> None: - """Atomically persist FULL multi-latent training state so a preempt/retry resumes from the last epoch - boundary instead of ep0. tmp+rename => a mid-write preempt cannot corrupt the good file.""" + """Save multi-latent training and RNG state through an atomic file replacement.""" if not a.resume_dir: return d = Path(a.resume_dir) @@ -1388,12 +1390,10 @@ def _rf( return _rfn(name) return _nullctx() - loss_terms = a.loss_terms(a) # aux separation/shaping terms (INJECTED), built ONCE + loss_terms = a.loss_terms(a) # Build the configured auxiliary loss terms once. - # optional TEXT REPLAY (multi-latent). Rows tagged `learn` (via `learn_field`) rehearse the backbone's plain - # text ability with a next-token CE on (input_text -> target_texts[0]), interleaved with the latent objective - # so multi-latent co-training doesn't erode the LM. Ported from the single-latent learn path; projection via - # the tied input embedding (no lm_head). learn_ratio=0 / no learn rows = off (byte-identical to before). + # Rows marked for text replay rehearse next-token prediction from input_text to the first target text. The + # replay loss uses the tied input embedding because the language-model output head is not retained. learn_pool = [ i for i in range(len(seeds)) if getattr(self, "is_learn", [False] * len(seeds))[i] ] @@ -1442,11 +1442,13 @@ def _multi_grad_cache_step( flat_texts: list[str], agg: dict[str, float], ) -> float: - """Multi-latent GradCache. Phase 1 rolls out the FULL batch under no_grad (a SHARED ss_mask makes the - scheduled-sampling rollout deterministic), runs the cross-batch recon-pure terms on the full-batch - recon, and caches d(loss)/d(recon). Phase 2 re-rolls each gc_chunk WITH grad, backprops the (row-weighted) - base loss, and injects the cached recon-grads -> the cross-batch term is EXACT full-batch; the per-row - base loss is accumulated (pragmatic, see grad_cache docs). Peak activation = one chunk.""" + """Compute multi-latent cross-batch gradients with chunk-sized activation memory. + + The detached full-batch rollout uses a shared scheduled-sampling mask, evaluates cross-batch terms, and + caches gradients with respect to emitted latents. Each chunk is then rolled out again with autograd + enabled to propagate the cached gradients and its row-weighted base loss. Cross-batch terms retain their + full-batch semantics; base losses are accumulated per row. + """ ss_prob = a.ss_prob assert ( ss_prob is not None @@ -1629,9 +1631,8 @@ def _multi_grad_cache_step( _k, _raw, _w = contrib loss = loss + _w * _raw agg[_k] = agg.get(_k, 0.0) + float(_raw.detach()) - for h in rt_heads: # PLUGGABLE AUX HEADS (langset.heads): phase shim + user heads. Applied HERE (right - # after the loss terms, before sigreg) — the exact summation slot the old PhaseTerm held, so a - # phase-only run is byte-identical. Each self-skips (loss_on -> None) when nothing supervises it. + # Apply resolved auxiliary heads after the general loss terms and before target regularization. + for h in rt_heads: if ( h.spec.reads == "recon" ): # per emitted latent (grad shapes the emission geometry, like the phase head) diff --git a/src/langset/training_args.py b/src/langset/training_args.py index ab8f5fd..9158c37 100644 --- a/src/langset/training_args.py +++ b/src/langset/training_args.py @@ -1,4 +1,16 @@ -"""Training configuration for langset.""" +"""Declarative training configuration for langset's trainer and strategies. + +:class:`TrainingArguments` collects the knobs that control tokenization, +optimization, latent emission, auxiliary objectives, validation, checkpointing, +and strategy injection. It is intentionally configuration-only: the trainer +consumes these values, while the interchangeable strategy classes and loss +functions implement the behavior. + +Single-latent and multi-latent runs share the same configuration object. For +multi-latent training, fields such as ``emission``, ``target_source``, +``loss_terms``, and ``seed_builder`` select implementations rather than adding +feature-specific branches to the trainer. +""" from __future__ import annotations diff --git a/uv.lock b/uv.lock index ce07e31..5337913 100644 --- a/uv.lock +++ b/uv.lock @@ -755,6 +755,7 @@ dependencies = [ [package.optional-dependencies] dev = [ + { name = "pdoc" }, { name = "pytest" }, { name = "ruff" }, { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -764,6 +765,9 @@ dev = [ { name = "scipy", version = "1.18.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "ty" }, ] +docs = [ + { name = "pdoc" }, +] probes = [ { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "scikit-learn", version = "1.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -786,6 +790,8 @@ wandb = [ requires-dist = [ { name = "datasets", specifier = ">=2.15.0" }, { name = "numpy", specifier = ">=1.26" }, + { name = "pdoc", marker = "extra == 'dev'", specifier = ">=15" }, + { name = "pdoc", marker = "extra == 'docs'", specifier = ">=15" }, { name = "peft", specifier = ">=0.13.2" }, { name = "pytest", marker = "extra == 'dev'" }, { name = "ruff", marker = "extra == 'dev'" }, @@ -803,7 +809,16 @@ requires-dist = [ { name = "ty", marker = "extra == 'dev'", specifier = ">=0.0.61" }, { name = "wandb", marker = "extra == 'wandb'", specifier = ">=0.18" }, ] -provides-extras = ["setfit", "wandb", "probes", "dev"] +provides-extras = ["setfit", "wandb", "probes", "docs", "dev"] + +[[package]] +name = "markdown2" +version = "2.5.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e4/ae/07d4a5fcaa5509221287d289323d75ac8eda5a5a4ac9de2accf7bbcc2b88/markdown2-2.5.5.tar.gz", hash = "sha256:001547e68f6e7fcf0f1cb83f7e82f48aa7d48b2c6a321f0cd20a853a8a2d1664", size = 157249, upload-time = "2026-03-02T20:46:53.411Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/af/4b3891eb0a49d6cfd5cbf3e9bf514c943afc2b0f13e2c57cc57cd88ecc21/markdown2-2.5.5-py3-none-any.whl", hash = "sha256:be798587e09d1f52d2e4d96a649c4b82a778c75f9929aad52a2c95747fa26941", size = 56250, upload-time = "2026-03-02T20:46:52.032Z" }, +] [[package]] name = "markupsafe" @@ -1561,6 +1576,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0f/54/68a0978d1ef8502b8492099beaa6e7a0c1b32e3b5d4f677f5810cb08711c/pandas-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:b2c95f8bfc1ee412bf482605d7bfd30c12d1d26bd59fdd91efeef1d4718decb1", size = 9466464, upload-time = "2026-05-11T18:54:22.754Z" }, ] +[[package]] +name = "pdoc" +version = "16.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jinja2" }, + { name = "markdown2" }, + { name = "markupsafe" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/fe/ab3f34a5fb08c6b698439a2c2643caf8fef0d61a86dd3fdcd5501c670ab8/pdoc-16.0.0.tar.gz", hash = "sha256:fdadc40cc717ec53919e3cd720390d4e3bcd40405cb51c4918c119447f913514", size = 111890, upload-time = "2025-10-27T16:02:16.345Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/a1/56a17b7f9e18c2bb8df73f3833345d97083b344708b97bab148fdd7e0b82/pdoc-16.0.0-py3-none-any.whl", hash = "sha256:070b51de2743b9b1a4e0ab193a06c9e6c12cf4151cf9137656eebb16e8556628", size = 100014, upload-time = "2025-10-27T16:02:15.007Z" }, +] + [[package]] name = "peft" version = "0.19.1" From 3299e81870d016f26c6c8a6935118e5282549ecf Mon Sep 17 00:00:00 2001 From: Stephen Solka Date: Tue, 18 Aug 2026 14:06:26 -0400 Subject: [PATCH 2/3] Format modeling documentation --- src/langset/modeling.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/langset/modeling.py b/src/langset/modeling.py index 3b3d34c..40a9dd1 100644 --- a/src/langset/modeling.py +++ b/src/langset/modeling.py @@ -555,9 +555,7 @@ def from_pretrained( torch.nn.init.eye_( model.head.out_proj.weight ) # `head_project` initially preserves the normalized final-token hidden state - torch.nn.init.zeros_( - model.head.out_proj.bias - ) + torch.nn.init.zeros_(model.head.out_proj.bias) if freeze_backbone: # FROZEN base: only the head trains -> backbone read needs no graph for p in model.backbone.parameters(): p.requires_grad_(False) @@ -759,12 +757,8 @@ def encode( was_training = self.training self.eval() out: list[torch.Tensor] = [] - with ( - torch.no_grad() - ): - for i in range( - 0, len(texts), batch_size - ): + with torch.no_grad(): + for i in range(0, len(texts), batch_size): enc = self.tokenizer( texts[i : i + batch_size], padding=True, From a89c0510c5a4bef37afc9e0db1da4bc2f0c755bf Mon Sep 17 00:00:00 2001 From: Stephen Solka Date: Tue, 18 Aug 2026 14:13:20 -0400 Subject: [PATCH 3/3] Address API documentation review --- src/langset/data.py | 4 ++-- src/langset/loss.py | 32 ++++++++++++++++---------------- src/langset/modeling.py | 16 ++++++++-------- src/langset/selection.py | 14 ++++++-------- src/langset/st_module.py | 4 ++-- src/langset/strategies.py | 12 ++++++------ src/langset/trainer.py | 13 +++++++------ src/langset/training_args.py | 2 +- 8 files changed, 48 insertions(+), 49 deletions(-) diff --git a/src/langset/data.py b/src/langset/data.py index cfbe66f..1db192a 100644 --- a/src/langset/data.py +++ b/src/langset/data.py @@ -9,9 +9,9 @@ interest and cannot be trivially regenerated from ``input_text``; otherwise the objective mostly distills a text encoder. -Pass a ``datasets.Dataset`` or ``list[dict]`` to :class:`langset.Trainer`. Use +Pass a ``datasets.Dataset`` or ``list[dict]`` to `langset.Trainer`. Use ``column_mapping`` when an existing dataset uses different column names, or -use :func:`from_records` to project arbitrary records onto the basic contract. +use `from_records` to project arbitrary records onto the basic contract. """ from __future__ import annotations diff --git a/src/langset/loss.py b/src/langset/loss.py index 76eabff..77a0def 100644 --- a/src/langset/loss.py +++ b/src/langset/loss.py @@ -2,22 +2,22 @@ This module separates the mathematical objectives from the training loop. Each loss function accepts a typed context dataclass containing the tensors -and configuration it needs, and returns a :class:`Loss`. The trainer can then +and configuration it needs, and returns a `Loss`. The trainer can then combine the result with other objectives and call ``to_tensor()`` for the scalar used in backpropagation. The objectives fall into four groups: -- **Latent grounding:** :func:`recon_loss` teaches an emitted latent to +- **Latent grounding:** `recon_loss` teaches an emitted latent to reconstruct its target text. -- **Representation alignment:** :func:`sl_loss` and :func:`info_nce_loss` +- **Representation alignment:** `sl_loss` and `info_nce_loss` align paired embeddings and contrastive candidates. The supervised contrastive objective groups emitted embeddings by labels. -- **Multi-latent supervision:** :func:`soft_target_cross_entropy_loss` - trains distributions over concepts or states, while :func:`stop_loss` +- **Multi-latent supervision:** `soft_target_cross_entropy_loss` + trains distributions over concepts or states, while `stop_loss` teaches autoregressive emissions when to terminate. -- **Text-generation support:** :func:`learn_loss` rehearses document-to-target - associations, and :func:`cot_loss` trains optional chain-of-thought +- **Text-generation support:** `learn_loss` rehearses document-to-target + associations, and `cot_loss` trains optional chain-of-thought generation. The context classes are intentionally explicit: they document tensor shapes, @@ -80,7 +80,7 @@ class TrainingContext: @dataclass class ReconstructionLossContext(TrainingContext): - """Inputs to :func:`recon_loss`. + """Inputs to `recon_loss`. ``latent`` contains one latent vector per selected row. ``tr_ids`` and ``tr_mask`` contain tokenized target text for every dataset row; ``rows`` @@ -166,7 +166,7 @@ def recon_loss(ctx: ReconstructionLossContext) -> Loss: @dataclass class LearnLossContext(TrainingContext): - """Inputs to :func:`learn_loss` for a batch from the replay pool.""" + """Inputs to `learn_loss` for a batch from the replay pool.""" pos: torch.Tensor """Indices selecting replay examples from the tokenized replay pool.""" @@ -272,7 +272,7 @@ def sl_loss( ``mask_keys`` exclude known false negatives from that anchor's denominator. This function prepares those single-latent candidates, positive indices, and exclusion mask, then delegates the primary contrastive calculation to - :func:`info_nce_loss`. + `info_nce_loss`. The contrastive term has an implicit weight of ``1.0``. When ``args.lam_uniform > 0`` and the batch contains multiple rows, the returned @@ -337,7 +337,7 @@ def sl_loss( # ---- multi-latent loss kernels ------------------------------------------------------------------ @dataclass class InfoNCELossContext: - """Inputs to :func:`info_nce_loss`. + """Inputs to `info_nce_loss`. The candidate at ``positive_indices[row]`` is the positive for ``anchors[row]``. All other unmasked candidates are negatives. @@ -389,7 +389,7 @@ def info_nce_loss(ctx: InfoNCELossContext) -> Loss: @dataclass class SoftTargetCrossEntropyContext: - """Inputs to :func:`soft_target_cross_entropy_loss`.""" + """Inputs to `soft_target_cross_entropy_loss`.""" logits: torch.Tensor """Unnormalized class scores, shaped ``[..., class_count]``.""" @@ -430,7 +430,7 @@ def soft_target_cross_entropy_loss(ctx: SoftTargetCrossEntropyContext) -> Loss: @dataclass class StopLossContext: - """Inputs to :func:`stop_loss` for an autoregressive multi-latent emission.""" + """Inputs to `stop_loss` for an autoregressive multi-latent emission.""" logits: torch.Tensor """One stop logit per row and emission position, shaped ``[batch_size, max_items + 1]``.""" @@ -465,7 +465,7 @@ def stop_loss(ctx: StopLossContext) -> Loss: @dataclass class SupervisedContrastiveLossContext: - """Inputs to :func:`supervised_contrastive_loss`.""" + """Inputs to `supervised_contrastive_loss`.""" embeddings: torch.Tensor """Emitted embeddings, shaped ``[item_count, embedding_dim]``.""" @@ -538,7 +538,7 @@ def supervised_contrastive_loss(ctx: SupervisedContrastiveLossContext) -> Loss: @dataclass class CoTLossContext(TrainingContext): - """Inputs to :func:`cot_loss` after the strategy tokenizes seed and reasoning text.""" + """Inputs to `cot_loss` after the strategy tokenizes seed and reasoning text.""" seed_ids: torch.Tensor """Left-padded seed token IDs, shaped ``[batch_size, seed_length]``.""" @@ -590,7 +590,7 @@ def cot_loss(ctx: CoTLossContext) -> Loss: @dataclass class BridgeLossContext: - """Inputs to :func:`bridge_loss` after Hungarian matching prepares bridge targets.""" + """Inputs to `bridge_loss` after Hungarian matching prepares bridge targets.""" matched_predictions: torch.Tensor """Hungarian-matched query embeddings, shaped ``[match_count, embedding_dim]``.""" diff --git a/src/langset/modeling.py b/src/langset/modeling.py index 40a9dd1..ac7f2fb 100644 --- a/src/langset/modeling.py +++ b/src/langset/modeling.py @@ -1,6 +1,6 @@ """Model and emission-head implementations for LangSet. -:class:`LangSetModel` combines a language-model backbone with a learned emission interface. The emission head can +`LangSetModel` combines a language-model backbone with a learned emission interface. The emission head can project backbone hidden states into a single embedding, multiple continuous latents, or autoregressive codebook states. Optional LoRA adaptation supports parameter-efficient fine-tuning. @@ -397,8 +397,8 @@ class LangSetModel(nn.Module): """Wrap a language-model backbone with continuous or codebook-based emission heads. Depending on its configuration, a model can produce a single embedding, multiple learned latent vectors, - or an autoregressive sequence of named-state emissions. Use :meth:`from_pretrained` or :meth:`from_scratch` - to construct a model, :meth:`encode` for Sentence-Transformer-compatible embeddings, and :meth:`rollout` for + or an autoregressive sequence of named-state emissions. Use `from_pretrained` or `from_scratch` + to construct a model, `encode` for Sentence-Transformer-compatible embeddings, and `rollout` for autoregressive codebook emission. """ @@ -747,7 +747,7 @@ def encode( ) -> Union[np.ndarray, torch.Tensor]: """Encode one or more sentences using the Sentence-Transformer-compatible interface. - Inputs are tokenized in batches and truncated to :attr:`max_len`. A single string returns one embedding; + Inputs are tokenized in batches and truncated to `max_len`. A single string returns one embedding; a sequence returns a batch. Results are NumPy arrays by default or CPU tensors when ``convert_to_numpy=False``. ``show_progress_bar`` and ``device`` are accepted for interface compatibility but are currently ignored. @@ -779,7 +779,7 @@ def encode( def emit(self, sentences: Union[str, list[str]], **kw: Unknown) -> torch.Tensor: """Encode text and return the result as a PyTorch tensor. - Additional keyword arguments are forwarded to :meth:`encode`. This is equivalent to calling + Additional keyword arguments are forwarded to `encode`. This is equivalent to calling ``encode(..., convert_to_numpy=False)``. """ # `kw` is a typed passthrough to `encode`; `Unknown` avoids weakening the public signature to `Any`. @@ -801,7 +801,7 @@ def add_aux_head(self, module: nn.Linear, spec: dict[str, Any]) -> None: """Register an auxiliary linear head and its persistence metadata. ``spec`` describes the head name, read site, loss, input and output dimensions, and optional class labels. - Registered heads are serialized by :meth:`save_pretrained` and can be queried with :meth:`head_output`. + Registered heads are serialized by `save_pretrained` and can be queried with `head_output`. """ name = str(spec["name"]) self.aux_heads[name] = module @@ -818,7 +818,7 @@ def head_output( """Run a persisted auxiliary head on one or more sentences. Heads with ``reads="hidden"`` consume one pooled seed representation per sentence. Heads with - ``reads="recon"`` consume every latent produced by :meth:`rollout`; ``reduce="mean"`` averages each + ``reads="recon"`` consume every latent produced by `rollout`; ``reduce="mean"`` averages each sequence, while ``reduce="none"`` preserves the per-step outputs. MSE heads typically return scalar values. Classification heads return logits whose columns correspond to ``aux_head_specs[name]["classes"]``. @@ -1272,7 +1272,7 @@ def load( device: Optional[str] = None, attn_implementation: str = "sdpa", ) -> "LangSetModel": - """Load a model saved by :meth:`save_pretrained`. + """Load a model saved by `save_pretrained`. Reconstructs the original pretrained or randomly initialized backbone, restores the saved trainable state, rebuilds persisted auxiliary heads, moves the model to ``device``, and returns it in evaluation mode. diff --git a/src/langset/selection.py b/src/langset/selection.py index fbd7ecc..edd0e88 100644 --- a/src/langset/selection.py +++ b/src/langset/selection.py @@ -1,11 +1,9 @@ -"""Collapse-aware validation metrics for checkpoint selection. - -Training loss is deliberately not used to select the best model: a contrastive -objective can reduce its loss by collapsing the representation geometry. These -metrics instead evaluate held-out input-to-target retrieval and reconstruction, -then apply a hard penalty when the embedding geometry collapses. The trainer -uses them for early stopping and checkpoint selection so optimization and model -selection measure different things. +"""Validation metrics for retrieval quality and representation collapse. + +This module provides held-out retrieval MRR, mean-cosine collapse scoring, and +k-nearest-neighbor purity. The trainer combines these metrics with +reconstruction evaluation and a collapse penalty when selecting checkpoints; +the metric functions here do not apply that policy themselves. """ from __future__ import annotations diff --git a/src/langset/st_module.py b/src/langset/st_module.py index 8b11075..891b03a 100644 --- a/src/langset/st_module.py +++ b/src/langset/st_module.py @@ -1,6 +1,6 @@ """Sentence-Transformers compatibility for trained langset models. -This adapter wraps a trained single-latent :class:`langset.LangSetModel` as a +This adapter wraps a trained single-latent `langset.LangSetModel` as a ``sentence_transformers.SentenceTransformer``-compatible object. It provides the interface expected by SetFit and similar tooling while keeping langset's model and latent geometry underneath. @@ -11,7 +11,7 @@ setfit_model = SetFitModel(model_body=body, labels=[...]) Use this integration for embedding-style models; autoregressive multi-latent -rollouts remain on :class:`langset.LangSetModel` itself. +rollouts remain on `langset.LangSetModel` itself. """ from __future__ import annotations diff --git a/src/langset/strategies.py b/src/langset/strategies.py index d3a62b7..107d669 100644 --- a/src/langset/strategies.py +++ b/src/langset/strategies.py @@ -1,7 +1,7 @@ """Interchangeable strategies for multi-latent training. A training step combines an emission objective, a target source, optional loss terms, and small callables for -epoch ordering, checkpoint selection, and seed construction. :class:`TrainingArguments` stores the selected +epoch ordering, checkpoint selection, and seed construction. `TrainingArguments` stores the selected implementations, which are instantiated once and used through shared interfaces. """ @@ -205,7 +205,7 @@ class CoTGenTerm(_LossTerm): """Train next-token generation of each row's ``cot_text`` from its input seed. The term uses an isolated backward pass so its long generation graph does not coexist with the latent-emission - graph. Pair it with :func:`cot_seed_texts` when emissions should also condition on the provided reasoning. + graph. Pair it with `cot_seed_texts` when emissions should also condition on the provided reasoning. Batches without reasoning text are skipped. """ @@ -250,7 +250,7 @@ def _tokm(texts: list[str], mx: int, side: str) -> tuple[torch.Tensor, torch.Ten def build_cot_loss_terms(args: TrainingArguments) -> list[_LossTerm]: - """Build the default auxiliary terms plus :class:`CoTGenTerm`. + """Build the default auxiliary terms plus `CoTGenTerm`. Use with ``seed_builder=cot_seed_texts`` and a dataset containing ``cot_text``. """ @@ -744,7 +744,7 @@ def regularizer(self, z_pred: torch.Tensor, z_tgt: torch.Tensor) -> Optional[tor class EMATwinTarget(_TargetSource): """Produce stop-gradient targets from an exponential-moving-average copy of the online model. - After each optimizer step, :meth:`update` moves the target model toward the trainable online parameters using + After each optimizer step, `update` moves the target model toward the trainable online parameters using ``ema_m``. The lagged target provides a more stable comparison geometry than using the online model on both sides. """ @@ -791,7 +791,7 @@ class CachedTarget(_TargetSource): If ``target_encoder_ckpt`` is configured, the encoder is loaded from that LangSet checkpoint. Otherwise, the source is a frozen copy of the online model taken at initialization. Each unique text is encoded once and - subsequent calls reuse the cached latent. :meth:`update` is a no-op because the target geometry is fixed. + subsequent calls reuse the cached latent. `update` is a no-op because the target geometry is fixed. """ suppresses_nce = False @@ -1048,7 +1048,7 @@ def multi_seed_texts(trainer: Trainer, seeds: list[str], args: TrainingArguments def cot_seed_texts(trainer: Trainer, seeds: list[str], args: TrainingArguments) -> list[str]: """Append each row's reasoning text to its seed before emission. - Pair with :func:`build_cot_loss_terms` so the same reasoning is trained autoregressively. Targets and evaluation + Pair with `build_cot_loss_terms` so the same reasoning is trained autoregressively. Targets and evaluation continue to use the raw seeds. """ return [f"{s}\n\nReasoning:\n{trainer.cot_texts[i]}" for i, s in enumerate(seeds)] diff --git a/src/langset/trainer.py b/src/langset/trainer.py index b1be2f1..1c4c08c 100644 --- a/src/langset/trainer.py +++ b/src/langset/trainer.py @@ -164,7 +164,7 @@ def _snapshot_best(m: LangSetModel) -> dict[str, Any]: def _restore_best(m: LangSetModel, best_state: dict[str, Any]) -> None: - """Restore a snapshot created by :func:`_snapshot_best`, including legacy LoRA-only snapshots.""" + """Restore a snapshot created by `_snapshot_best`, including legacy LoRA-only snapshots.""" m.head.load_state_dict(best_state["head"]) if "backbone" in best_state: # random-init: full backbone m.backbone.load_state_dict(best_state["backbone"], strict=False) @@ -212,7 +212,7 @@ class BackboneStepEngine(_StepEngine): This is the default single-latent engine. It computes input, target, and optional hard-negative features and supports reconstruction auxiliaries; runtime options such as view fusion and target stop-gradient are handled - inside :meth:`featurize`. + inside `featurize`. """ supports_recon = True @@ -378,15 +378,16 @@ def val_embeddings(self, val_idx: np.ndarray) -> tuple[np.ndarray, np.ndarray]: class Trainer: - """Train a :class:`LangSetModel` on text-to-latent alignment data. + """Train a `LangSetModel` on text-to-latent alignment data. The trainer selects the single- or multi-latent path from the model configuration. Single-latent rows require ``input_text`` and ``target_text``; multi-latent rows require ``input_text`` and a non-empty ``target_texts`` list. Optional columns configure hard negatives, false-negative masks, concepts, states, auxiliary heads, and - text replay according to :class:`TrainingArguments`. + text replay according to `TrainingArguments`. ``eval_dataset`` is currently reserved and ignored; validation rows are split from ``train_dataset``. - ``on_checkpoint``, when provided, is called after a durable checkpoint is written. + ``on_checkpoint``, when provided, is called after the checkpoint is written locally so callers can perform + external persistence or synchronization. """ def __init__( @@ -625,7 +626,7 @@ def train(self) -> LangSetModel: Dispatches to the multi-latent strategy loop when the model uses a multi-latent head; otherwise runs single-latent contrastive training. Checkpointing, early stopping, text replay, and auxiliary losses are - controlled by :class:`TrainingArguments`. + controlled by `TrainingArguments`. """ if self.multi_latent: return self._train_multi() diff --git a/src/langset/training_args.py b/src/langset/training_args.py index 9158c37..b14a5b3 100644 --- a/src/langset/training_args.py +++ b/src/langset/training_args.py @@ -1,6 +1,6 @@ """Declarative training configuration for langset's trainer and strategies. -:class:`TrainingArguments` collects the knobs that control tokenization, +`TrainingArguments` collects the knobs that control tokenization, optimization, latent emission, auxiliary objectives, validation, checkpointing, and strategy injection. It is intentionally configuration-only: the trainer consumes these values, while the interchangeable strategy classes and loss