nemo_gym.token_id_capture.builder

Build trainable trajectories from a frozen token capture.

The builder consumes TokenEntry records from a TokenCaptureSnapshot. The snapshot is frozen before the consumer passes its entries to the builder.

per_request creates one training sequence per call. It does not infer relationships between calls. It can return multiple trajectories.

prefix_merging chains calls by token-prefix relationships. Each call uses the earlier call with the longest matching cumulative sequence as its parent. The cumulative sequence contains the prompt and generation. This strategy rebuilds an append-only rollout as one chain. A rewritten or compacted prompt starts a new root. Identical candidate parents are ambiguous. The builder quarantines the ambiguous subtree.

Both strategies are independent of capture order. prefix_merging processes entries by increasing prompt length. This order comes from the tokens. A parent’s prompt is shorter than its child’s prompt.

Loss masks follow token provenance. Policy-generated tokens have a mask value of 1 and retain their captured log probabilities. Prompt tokens have a mask value of 0.

Module Contents

Classes

Name	Description
`BuildNotes`	Describe what the build kept, dropped, or could not resolve.
`BuildOutput`	-
`Chain`	-
`ChainLink`	-
`_Node`	-
`_PrefixIndex`	Index cumulative token sequences for linear-time parent lookup.
`_PrefixTrieNode`	-

Functions

Name	Description
`_infer_parent`	Infer the parent from the longest cumulative prefix.
`assert_prefix_contiguity`	Require each generated item to extend all preceding tokens.
`per_request`	-
`prefix_merging`	-
`project_chain_to_output_items`	Project a chain into Responses output items with contiguous prompts.
`project_main_chain_response`	Rebuild the main chain as a Responses object whose output items are contiguous.
`run_builder`	Chain frozen snapshot entries with the named strategy.

Data

_BUILDERS

API

class nemo_gym.token_id_capture.builder.BuildNotes(
    builder: str,
    roots: int = 0,
    chains: int = 0,
    generated_tokens_captured: int = 0,
    generated_tokens_delivered: int = 0,
    delivered_fraction: float = 0.0,
    unresolved_retries: list[str] = list(),
    empty_generation_calls: list[str] = list()
)

Dataclass

Describe what the build kept, dropped, or could not resolve.

The consumer converts these fields into run metrics. Typed fields prevent renamed keys from appearing as zero values on dashboards.

builder

str

chains

int = 0

delivered_fraction

float = 0.0

empty_generation_calls

list[str] = field(default_factory=list)

generated_tokens_captured

int = 0

generated_tokens_delivered

int = 0

roots

int = 0

unresolved_retries

list[str] = field(default_factory=list)

class nemo_gym.token_id_capture.builder.BuildOutput(
    chains: list[nemo_gym.token_id_capture.builder.Chain],
    quarantined: list[str] = list(),
    notes: nemo_gym.token_id_capture.builder.BuildNotes = (lambda: BuildNotes(builder...
)

Dataclass

chains

list[Chain]

notes

BuildNotes

quarantined

list[str] = field(default_factory=list)

class nemo_gym.token_id_capture.builder.Chain(
    chain_id: str,
    links: list[nemo_gym.token_id_capture.builder.ChainLink] = list(),
    root_prompt: list[int] = list()
)

Dataclass

chain_id

str

links

list[ChainLink] = field(default_factory=list)

root_prompt

list[int] = field(default_factory=list)

nemo_gym.token_id_capture.builder.Chain.validate() -> None

Require one log probability for each generated token.

A trainer cannot use a chain with mismatched token and log-probability counts.

class nemo_gym.token_id_capture.builder.ChainLink(
    entry: nemo_gym.token_id_capture.records.TokenEntry,
    interstitial: list[int]
)

Dataclass

entry

TokenEntry

interstitial

list[int]

class nemo_gym.token_id_capture.builder._Node(
    entry: nemo_gym.token_id_capture.records.TokenEntry,
    cumulative: list[int],
    parent: '_Node | None' = None,
    children: list['_Node'] = list(),
    quarantined: bool = False
)

Dataclass

children

list['_Node'] = field(default_factory=list)

cumulative

list[int]

entry

TokenEntry

parent

'_Node | None' = None

quarantined

bool = False

class nemo_gym.token_id_capture.builder._PrefixIndex()

Index cumulative token sequences for linear-time parent lookup.

_root

= _PrefixTrieNode()

nemo_gym.token_id_capture.builder._PrefixIndex.add(
    candidate: nemo_gym.token_id_capture.builder._Node
) -> None

nemo_gym.token_id_capture.builder._PrefixIndex.infer_parent(
    prompt: list[int]
) -> tuple[nemo_gym.token_id_capture.builder._Node | None, bool]

class nemo_gym.token_id_capture.builder._PrefixTrieNode(
    children: dict[int, '_PrefixTrieNode'] = dict(),
    candidates: list[nemo_gym.token_id_capture.builder._Node] = list()
)

Dataclass

candidates

list[_Node] = field(default_factory=list)

children

dict[int, '_PrefixTrieNode'] = field(default_factory=dict)

nemo_gym.token_id_capture.builder._infer_parent(
    prompt: list[int],
    index: nemo_gym.token_id_capture.builder._PrefixIndex
) -> tuple['_Node | None', bool]

Infer the parent from the longest cumulative prefix.

This is the fallback when no verified parent link exists. Identical cumulative sequences are ambiguous. The caller quarantines an ambiguous subtree.

nemo_gym.token_id_capture.builder.assert_prefix_contiguity(
    response: dict
) -> None

Require each generated item to extend all preceding tokens.

The preceding tokens include the prompt and all prior generations. Raise AssertionError when the response is not contiguous.

nemo_gym.token_id_capture.builder.per_request(
    entries: list[nemo_gym.token_id_capture.records.TokenEntry]
) -> nemo_gym.token_id_capture.builder.BuildOutput

nemo_gym.token_id_capture.builder.prefix_merging(
    entries: list[nemo_gym.token_id_capture.records.TokenEntry]
) -> nemo_gym.token_id_capture.builder.BuildOutput

nemo_gym.token_id_capture.builder.project_chain_to_output_items(
    chain: nemo_gym.token_id_capture.builder.Chain
) -> list[dict]

Project a chain into Responses output items with contiguous prompts.

Preserve captured assistant text and tool calls. Put the contiguous prompt on the item that carries each generation. Each generated item’s prompt extends the previous generated item. Preserve text for downstream scoring. Create a token-only item only when a call has no captured content items.

nemo_gym.token_id_capture.builder.project_main_chain_response(
    rollout_id: str,
    out: nemo_gym.token_id_capture.builder.BuildOutput,
    model: str = ''
) -> dict

Rebuild the main chain as a Responses object whose output items are contiguous.

The result is a Gym-native Responses payload. It contains object: "response", output items, and usage. Token fields describe one unbroken sequence across the rollout. The sequence combines items from multiple model calls.

nemo_gym.token_id_capture.builder.run_builder(
    entries: list[nemo_gym.token_id_capture.records.TokenEntry],
    builder: str = 'prefix_merging'
) -> nemo_gym.token_id_capture.builder.BuildOutput

Chain frozen snapshot entries with the named strategy.

nemo_gym.token_id_capture.builder._BUILDERS: dict[str, Callable[[list[TokenEntry]], BuildOutput]] = {'per_request': per_request, 'prefix_merging': prefix_merging}