nemo_gym.token_id_capture.builder

Build trainable trajectories from a frozen token capture.

The consumer freezes a TokenCaptureSnapshot before passing its TokenEntry records here. The builder can return multiple trajectories when a rollout contains independent call chains.

Each current record states whether the call starts a root, continues a verified parent, or has unresolved ancestry. The builder verifies every resolved parent against the child’s prompt tokens before joining the calls. An unresolved call begins a separate fragment. The builder never infers a parent across that boundary.

prefix_merging uses token-prefix matching only when a verified parent is absent from the frozen snapshot. For example, capture can filter a parent that generated no tokens. In that case, prefix matching may reconnect the child to a verified surviving ancestor.

The build does not depend on capture order. prefix_merging processes entries by increasing prompt length because a parent’s prompt is shorter than its child’s.

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`	-
`_NullPrefixIndex`	Avoid building a token-prefix index when every parent is present.
`_PrefixIndex`	Index cumulative token sequences for linear-time parent lookup.
`_PrefixTrieNode`	-

Functions

Name	Description
`_materialize_delta_prompts`	Rebuild full prompts for delta records by walking their parent chains.
`_resolve_parent`	Find this call’s parent.
`assert_prefix_contiguity`	Require each generated item to extend all preceding tokens.
`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(),
    terminal_call_id: str | None = None,
    terminal_chain: str = '',
    empty_generation_calls: list[str] = list(),
    parent_link_failures: dict[str, int] = dict(),
    unresolved_parent_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

parent_link_failures

dict[str, int] = field(default_factory=dict)

roots

int = 0

terminal_call_id

str | None = None

terminal_chain

str = ''

unresolved_parent_calls

list[str] = field(default_factory=list)

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,
    unresolved_boundary: bool = False
)

Dataclass

children

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

cumulative

list[int]

entry

TokenEntry

parent

'_Node | None' = None

quarantined

bool = False

unresolved_boundary

bool = False

class nemo_gym.token_id_capture.builder._NullPrefixIndex()

Avoid building a token-prefix index when every parent is present.

Every supported entry carries a request-time parent decision. Only missing-parent recovery needs the trie.

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

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

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._materialize_delta_prompts(
    entries: list[nemo_gym.token_id_capture.records.TokenEntry]
) -> tuple[list[nemo_gym.token_id_capture.records.TokenEntry], list[str]]

Rebuild full prompts for delta records by walking their parent chains.

Return full-prompt entries and call ids with broken chains.

nemo_gym.token_id_capture.builder._resolve_parent(
    node: '_Node',
    by_call_id: dict[str, '_Node'],
    prefix_index: nemo_gym.token_id_capture.builder._PrefixIndex
) -> tuple['_Node | None', bool, str | None]

Find this call’s parent.

New records preserve the request-time parent decision. Token-prefix matching recovers a resolved link whose parent is absent from this build. This can happen when the parent had an empty generation and was filtered out. note reports why the recorded link was not used as recorded.

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.prefix_merging(
    entries: list[nemo_gym.token_id_capture.records.TokenEntry],
    terminal_call_id: str | None = None
) -> 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',
    terminal_call_id: str | None = None
) -> nemo_gym.token_id_capture.builder.BuildOutput

Chain frozen snapshot entries with the named strategy.

terminal_call_id anchors chain selection for prefix_merging.

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