DataModule API¶

Bases: LightningDataModule

The text stage's data: two streams over every code and every task query (Req 10; spec 4.3).

prepare_data builds the tokenization cache, or keeps one built from the same inputs. setup builds the StepDataset once: every code's cached token rows in codebook order, and the bundle's task queries, each tokenized once as a query: text with the cache's tokenizer at its window. code_rows holds those rows for the model's code cache, which encodes them at fit start and at each epoch's end.

The one loader is train_dataloader (P12). There is no validation loader, so Lightning runs no validation loop: validation is the outcome monitor (spec 4.4). Construction reads nothing; the bundle is loaded, and refused, in prepare_data and setup.

Parameters:

Name	Type	Description	Default
token_config	TokenizationConfig	The tokenization cache to build and read: the descriptions, the tokenizer, the window and the file. train passes code_token_config's, the cache the export and the reads load.	required
seed	int	The run's seed, from which each epoch's two permutations are drawn.	required
queries_per_step	int	The most task queries one step reads.	DEFAULT_QUERIES_PER_STEP
supervision_manifest_path	Optional[str]	The bundle's manifest, read when no bundle is given.	None
supervision_contract_version	str	The contract the manifest must carry.	CONTRACT_VERSION
supervision_bundle	Optional[ValidatedSupervisionBundle]	The validated bundle, when the caller has loaded it.	None

Raises:

Type	Description
ValueError	If queries_per_step is below 1.

Bases: Dataset

The steps of one epoch over two streams: every code once as an anchor, and every task query once (Req 10; spec 4.3).

Each epoch permutes the codes and the queries (epoch_permutation, under the run's seed and the epoch) and cuts both orders into the same S even chunks (even_chunks), with S from steps_per_epoch: step s reads code chunk s and query chunk s. A step carries only its anchors' token rows and its queries'. Every candidate a term scores comes from the model's code cache, so no step carries a candidate pool.

The epoch comes only from set_epoch, which TrainDatasetEpochCallback calls at each epoch start, and its two permutations are drawn at its first step and kept for the others. set_epoch reaches only the main process's copy of the dataset, so the steps are read there, as DataLoader(dataset, batch_size=None, shuffle=False, num_workers=0) reads them.

A step is a dict of two dicts:

• codes: inputs, the anchors' token rows through stack_text_inputs; and ids and levels, int64 (A,).
• queries: inputs, the queries' token rows under the field query; levels, int64 (Q,); and targets and negatives, bool (Q, N) masks over the codes of each query's T and N.

Parameters:

Name	Type	Description	Default
code_rows	Sequence[Mapping[str, Any]]	Every code's token row, in codebook order, as the tokenization cache holds it.	required
code_levels	Sequence[int]	Every code's level, in codebook order.	required
queries	TokenizedQueries	The tokenized task queries.	required
n_codes	int	N, the number of codes.	required
seed	int	The run's seed.	required
queries_per_step	int	The most queries one step reads.	required

Raises:

Type	Description
ValueError	If the code rows or levels are not one per code; if there is no task query, or there are more steps than codes, since every step needs at least one anchor; or if a query names a code id outside the codes, has no target, or has a target at another level.

The task queries, each tokenized once as a query: text, with its level and code ids.

Built by tokenize_task_queries. Every field is in query order.

Attributes:

Name	Type	Description
texts	Tuple[str, ...]	Each query's text.
tokens	Tuple[Dict[str, Any], ...]	Each query's query token row, as tokenize_field returns it.
levels	Tuple[int, ...]	Each query's level.
target_ids	Tuple[Tuple[int, ...], ...]	The code ids of each query's targets, T.
negative_ids	Tuple[Tuple[int, ...], ...]	The code ids of each query's forced negatives, N.

Bases: Callback

Set the step dataset's epoch to the trainer's at each epoch start (P27).

Register it on every Trainer that fits a NAICSDataModule: Lightning dispatches on_train_epoch_start to callbacks and the LightningModule, never to a LightningDataModule, and the step dataset refuses a step read before any epoch is set. Lightning never prefetches from a loader whose length it knows, so the epoch is set before the epoch's first step is read.