Text Model API

The model orchestrates two-stream losses, cache refresh, outcome monitoring and durable epoch health, with objective/checkpoint contract req11-v1.

Model

The text stage's Lightning module: the shared encoder, trained on Req 11's three terms over a live code cache and selected on the outcome panel's validation MRR (spec 4.1-4.4; D6).

  • SharedEncoder: one LoRA-tuned backbone over field-marked channels, masked fusion, one affine map to dimension d and the geometry head (Req 12), which gives each text its point, its radius r and its direction û.
  • A step reads two streams (spec 4.3): a chunk of the codes, as anchors, and a chunk of the task queries. Every candidate comes from the code cache, each code's (r, û) at its last refresh, with the step's anchors replaced by their live points.
  • Req 11's terms (spec 4.1): the task term over each query's candidates and the code-code listwise term over each anchor's J_a, both under the arm's own distance, and, in the hyperbolic arm only, the radial term (Req 12); under moe only, the experts' load balancing.
  • The cache is refreshed at fit start and at each epoch's end. The end-of-epoch refresh feeds the outcome monitor's read, whose MRR is logged as val/outcome_mrr and steps the plateau.

The model is decomposed into functional mixins: - LossMixin: the experts' load-balancing term (moe only) - LoggingMixin: the epoch's health logs (P20) - OptimizerMixin: AdamW, the warmup, the plateau and the logit scales' clamp (P16)

NAICSContrastiveModel

Bases: LossMixin, LoggingMixin, OptimizerMixin, LightningModule

The text stage's model: the shared encoder, trained on Req 11's three terms over a live code cache and selected on the outcome panel's validation MRR (spec 4.1-4.4; D6).

Supervision comes only from one validated supervision bundle: its codebook, its tree metric D* and its unary pairs (CodeTargets), held as buffers in codebook order and never saved.

Parameters:

Name Type Description Default
base_model_name str

HuggingFace model name for the backbone

 'sentence-transformers/all-MiniLM-L6-v2'
lora_r int

LoRA rank

 8
lora_alpha int

LoRA alpha scaling factor

 16
lora_dropout float

LoRA dropout rate

 0.1
fusion str

Channel fusion: masked_mean (default), attention or moe. The load-balancing term exists only under moe (R11)

 'masked_mean'
dimension int

Embedding dimension, one of 8, 16 or 32: the width of the one Linear(hidden → d) before the geometry head

 16
geometry str

The geometry arm (Req 12): euclidean, spherical or hyperbolic (the default), which picks the head; the radial term applies under hyperbolic only

 'hyperbolic'
num_experts int

Number of MoE experts (moe only)

 4
top_k int

Number of experts each row is routed to (moe only)

 2
moe_hidden_dim int

Hidden dimension of the experts (moe only)

 1024
radius_bound float

R, the hyperbolic head's bound on every radius: r = R · tanh(‖v‖ / R) (spec 4.2); the flat heads do not read it

 8.0
code_code_weight float

w_c, the code-code term's weight in the total (spec 4.1)

 1.0
radial_weight float

w_r, the radial term's weight in the total

 1.0
target_temperature float

τ_t, the temperature of the code-code target softmax(−D* / τ_t)

 1.0
radial_step float

ρ, the radius from one level to the next: the radial target is ρ · (λ − 1)

 1.0
logit_scale_init float

Where both logit scales s = exp(θ) start

 1.0
logit_scale_range Tuple[float, float]

(low, high), the range both logit scales are clamped to

 (0.01, 100.0)
learning_rate float

AdamW's base learning rate

 0.0001
weight_decay float

AdamW's weight decay, on every parameter but the logit scales

 0.01
warmup_epochs int

W, the epochs of the linear warmup; 0 for none

 1
lr_plateau_factor float

The factor the plateau cuts the learning rate by

 0.5
lr_plateau_patience int

The epochs without a better val/outcome_mrr before a cut

 2
load_balancing_coef float

The load-balancing term's coefficient (moe only)

 0.01
seed int

The run's seed, which every monitor read names

 0
run_settings Optional[Dict[str, Any]]

The run's free settings (utils/training.run_settings), saved so an exact resume under other settings can be refused; None outside train

 None
supervision_manifest_path Optional[str]

Manifest of the validated supervision bundle (required: it is the one authority for the codes, D* and the unary pairs)

 None
supervision_contract_version str

Expected supervision contract version

 CONTRACT_VERSION
summaries Optional[str]

The sha256 of the window-fitting summaries the token cache applied, or None for a backbone with no pin; recorded in the checkpoint contract

 None
checkpoint_contract Optional[CheckpointContract]

Optional runtime contract; must match the loaded bundle

 None
supervision_bundle Optional[ValidatedSupervisionBundle]

Optional already-validated bundle for supervision_manifest_path

 None
monitor Optional[OutcomeMonitor]

The run's OutcomeMonitor, which reads the validation split at each epoch's end; with None, no epoch reads it, logs val/outcome_mrr or steps the plateau

 None

checkpoint_contract, supervision_bundle and monitor are not saved in the hyperparameters.

Raises:

Type Description
ValueError

If the fusion, dimension or geometry is unknown, a setting is out of its range, the manifest is missing, a pre-validated bundle is not the configured one, or the runtime contract is not the bundle's.

compute_losses(batch)

One step's losses over its two streams (spec 4.1, 4.3).

The step's anchors and queries go through the encoder, and nothing else does. Every candidate is the code cache's point, with the anchors' rows replaced by their live points, so gradient reaches the codes through the anchors alone, as anchors and as candidates. The distances and the terms run in float32 with autocast off: only the backbone runs in reduced precision (spec 4.2). Every distance is the arm's own (head.pair_distance), and the radial term exists only in the hyperbolic arm (Req 12).

Parameters:

Name Type Description Default
batch Dict[str, Dict[str, Any]]

One step, as StepDataset builds it: codes (inputs, ids and levels) and queries (inputs, levels, targets and negatives).

 required

Returns:

Type Description
StepLosses

The step's losses.

Raises:

Type Description
RuntimeError

If no refresh has built the code cache.

forward(channel_inputs)

Forward pass through the shared encoder.

Parameters:

Name Type Description Default
channel_inputs Dict[str, Dict[str, Tensor]]

Per field, tokenized text and a boolean present, as stack_text_inputs builds them: a code batch's four channels, or query

 required

Returns:

Type Description
Dict[str, Tensor]

Dictionary containing:
Dict[str, Tensor]
  • embedding: The arm's points, Lorentz (batch_size, dimension + 1) under hyperbolic
Dict[str, Tensor]
  • tangent: The coordinates the export writes (batch_size, dimension)
Dict[str, Tensor]
  • radius, direction: Each point's r (batch_size,) and û (batch_size, dimension)
Dict[str, Tensor]
  • gate_probs, top_k_indices: The experts' gates, under moe fusion only

on_load_checkpoint(checkpoint)

Refuse a checkpoint of any other objective, supervision contract or encoder architecture; then stash its training run and epoch for on_train_start.

Runs for Lightning exact resume and load_from_checkpoint before the state dict loads, so a checkpoint of another objective, such as every one saved before Stage 7, or of the four-copy encoder meets a refusal that cites D2, never a key mismatch. It reads and writes no file, since load_from_checkpoint runs it too.

on_save_checkpoint(checkpoint)

Record the supervision contract this checkpoint belongs to, Req 11's objective and the encoder architecture included (spec 4.5), and the training run's id.

on_train_epoch_end()

End the epoch (P15, P18): refresh the cache; then, with a monitor, read the validation split on it, log the MRR as val/outcome_mrr, step the plateau on it and append the read to the run's records; then log the health values and write the epoch summary (P20).

The MRR is one float64 value: the one logged, the one the plateau steps on and the one the record holds. This hook runs before ModelCheckpoint's, so the epoch's checkpoint holds the stepped plateau, and the run's records hold the epoch's read.

on_train_start()

Start the training run: keep the restored run's id or mint one, ready the monitor's records and epoch summary, and build the code cache (spec 4.3, 4.4).

On exact resume, on_load_checkpoint has stashed the checkpoint's training run and epoch: the run keeps its id, and the monitor keeps its records through that epoch. The cache is never saved, so a resume rebuilds it from the restored weights.

refresh_code_cache(code_rows)

Re-encode every code into code_cache (spec 4.3, P14): in eval mode, without gradient, in float32 with autocast off, in codebook order. The training flags are put back.

Parameters:

Name Type Description Default
code_rows Sequence[Mapping[str, Any]]

Every code's cached token rows, in codebook order.

 required

Returns:

Type Description
CodeCache

The new cache.

training_step(batch, batch_idx)

One step over the two streams: its losses, kept for the epoch's health logs.

Parameters:

Name Type Description Default
batch Dict[str, Dict[str, Any]]

One step, as StepDataset builds it

 required
batch_idx int

The step's index in the epoch

 required

Returns:

Type Description
Tensor

The step's total loss.

StepLosses

Bases: NamedTuple

One step's losses (spec 4.1).

Attributes:

Name Type Description
total Tensor

What the step optimizes: L = L_task + w_c · L_cc, plus w_r · L_rad in the hyperbolic arm, plus the load-balancing term times its coefficient under moe.
task Tensor

L_task, the task term.
code_code Tensor

L_cc, the code-code listwise term.
radial Optional[Tensor]

L_rad, the radial term; None unless the arm is hyperbolic (Req 12).
load_balancing Optional[Tensor]

The experts' load-balancing term, before its coefficient; None unless the fusion is moe.
anchor_radius Tensor

The anchors' live radii r_a, (A,): in the hyperbolic arm, every term's gradient reaches the radius through them (Verification "Radius").

Active Mixins

Mixins module for the NAICSContrastiveModel.

The model's three mixins:

  • LossMixin: the experts' load-balancing term, under moe only (R11)
  • LoggingMixin: the epoch's health logs (P20)
  • OptimizerMixin: AdamW, the warmup, the plateau on the monitor's MRR and the logit scales' clamp (P16)

LoggingMixin

Mixin providing the health logs.

This mixin expects the following attributes on the class: - logit_scale_task, logit_scale_code: the two LogitScale modules - code_cache: the CodeCache of the last refresh, or None - code_levels: each code's level, in codebook order (a buffer) - log: LightningModule.log

LossMixin

Mixin providing the MoE load-balancing term and its utilization logs.

This mixin expects the following attributes on the class: - device: torch.device - trainer: PyTorch Lightning trainer (for distributed logging) - logger: PyTorch Lightning logger (for histograms) - global_step: int training step counter

OptimizerMixin

Mixin providing the optimizer, its schedule and the logit scales' clamp.

This mixin expects the following attributes on the class: - hparams: learning_rate, weight_decay, warmup_epochs, lr_plateau_factor and lr_plateau_patience - logit_scale_task, logit_scale_code: the two LogitScale modules - trainer: inside a fit, its global_step and num_training_batches

configure_optimizers()

AdamW over two groups and the plateau, registered non-strict on val/outcome_mrr.

Returns:

Type Description
Dict[str, Any]

The optimizer and the plateau's scheduler config.

lr_scheduler_step(scheduler, metric)

Lightning steps no plateau: on_train_epoch_end steps it with the monitor's MRR, before ModelCheckpoint saves the epoch (P16). Any other scheduler steps as Lightning's.

optimizer_step(epoch, batch_idx, optimizer, optimizer_closure=None)

One optimizer step, at the warmup's rate while it lasts, then the scales' clamp (P16).

While global_step < W · S, for W warmup epochs of S steps, every group's rate is base · (global_step + 1) / (W · S), so the warmup's last step runs at the base rate. After it, the rate is the plateau's alone.

Raises:

Type Description
ValueError

If the warmup has no step count to run over: the train loader is unsized.