Add pdoc API reference

.gitignore

@@ -5,6 +5,7 @@ __pycache__/
 .eggs/
 build/
 dist/
+docs/
 
 # environments
 .venv/

.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"
+  }
+]

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`

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"

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"]

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

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="<col>"` (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,22 +127,23 @@ 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
 
     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])

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 `langset.Trainer`. Use
+``column_mapping`` when an existing dataset uses different column names, or
+use `from_records` to project arbitrary records onto the basic contract.
 """
 
 from __future__ import annotations