nemo_automodel.components.datasets.zephon.output

AutoModel batch representations after Zephon collation and attention policy.

Module Contents

Functions

Name	Description
`convert_batch_output`	Describe packed documents using the selected AutoModel input format.

API

nemo_automodel.components.datasets.zephon.output.convert_batch_output(
    batch: nemo_automodel.components.datasets.zephon.attention.TrainingBatch,
    output_format: nemo_automodel.components.datasets.packing.BatchOutputFormat
) -> nemo_automodel.components.datasets.zephon.attention.TrainingBatch

Describe packed documents using the selected AutoModel input format.

Zephon owns packing, padding and shifting. THD adds the native length-based schema; AutoModel’s existing preparation flattens tensors later, after this map. No token is moved or re-tokenized, and positions are never reconstructed from the length metadata. Xdoc policy has already selected the attention IDs.

Parameters:

batch

TrainingBatch

Owned mapping of CPU token, label, padding and document tensors of shape [batch, sequence]. Positions have shape [batch, sequence] or [axes, batch, sequence]. Image tensors retain their processor layouts. Positive attention_document_ids delimit contiguous subdocuments; zero marks exclusions. Padding must be a trailing suffix in each packed document. Each packed document must contain a real token. THD updates mapping entries in place without mutating existing tensors.

output_format

BatchOutputFormat

“batched” preserves the mapping. “thd” selects the native seq_lens input route; it does not select a packing algorithm.

Returns: TrainingBatch

The same mapping. THD adds int64 seq_lens/seq_lens_padded [batch, max_subdocuments] (-1000 for unused

Raises:

ValueError: For unknown formats, non-suffix padding, noncontiguous subdocuments, empty documents or processor-excluded real tokens. Those exclusions require the batched mask/metadata representation.