nemo_gym.models.token_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
`_engine_refusal_ended_the_chain`	Whether `refusal` was the harness’s next request after the delivered terminal call and did not fit.
`_materialize_delta_prompts`	Rebuild full prompts for delta records by walking their parent chains.
`_resolve_parent`	Find this call’s parent.
`_same_full_sequence`	Whether two calls’ full token sequences, prompt and generation together, are identical.
`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.models.token_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.models.token_capture.builder.BuildOutput(
    chains: list[nemo_gym.models.token_capture.builder.Chain],
    quarantined: list[str] = list(),
    notes: nemo_gym.models.token_capture.builder.BuildNotes = (lambda: BuildNotes(builder...,
    entries: dict[str, nemo_gym.models.token_capture.records.TokenEntry] = dict()
)

Dataclass

chains

list[Chain]

entries

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

notes

BuildNotes

quarantined

list[str] = field(default_factory=list)

class nemo_gym.models.token_capture.builder.Chain(
    chain_id: str,
    links: list[nemo_gym.models.token_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.models.token_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.models.token_capture.builder.ChainLink(
    entry: nemo_gym.models.token_capture.records.TokenEntry,
    interstitial: list[int]
)

Dataclass

entry

TokenEntry

interstitial

list[int]

class nemo_gym.models.token_capture.builder._Node(
    entry: nemo_gym.models.token_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.models.token_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.models.token_capture.builder._NullPrefixIndex.add(
    candidate: '_Node'
) -> None

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

class nemo_gym.models.token_capture.builder._PrefixIndex()

Index cumulative token sequences for linear-time parent lookup.

_root

= _PrefixTrieNode()

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

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

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

Dataclass

candidates

list[_Node] = field(default_factory=list)

children

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

nemo_gym.models.token_capture.builder._engine_refusal_ended_the_chain(
    refusal: nemo_gym.models.token_capture.records.EngineRefusalRecord,
    terminal: nemo_gym.models.token_capture.records.TokenEntry,
    out: nemo_gym.models.token_capture.builder.BuildOutput
) -> bool

Whether refusal was the harness’s next request after the delivered terminal call and did not fit.

Two conditions. The refusal carries the context-window code: an output-token parameter the engine rejected (code is None) says nothing about the conversation’s length, so it is reported in metrics["engine_refusals"] but sets no verdict. And the refused request continued the delivered terminal call: its request-time lineage names that call, or a call whose full token sequence is identical to it (a retry twin), and no twin of the terminal call was served at or after the moment the refusal was recorded. A twin served at or after the refusal is a successful re-serving of the call the refused request continued, so the harness recovered from the refusal and the retained twin’s own ending stands. The ordering reads the whole twin group, whichever twin attribution retained, so the verdict does not depend on which twin drew the smaller call id when attribution could only name them by content. A refusal on another branch of the rollout, one that continued an earlier call the harness later moved past, and one whose parent could not be resolved leave the verdict alone.

Known limitation: an auxiliary request the harness sends after the session completed (a summary or title call, or an end-of-session compaction) can carry the terminal call’s whole history, so it continues the terminal call in lineage too. The capture cannot tell such a request from the harness’s next turn: a context-window refusal of it marks the rollout incomplete although the harness finished on its own, and delivery then replaces a harness completed status with that verdict.

The twin comparison reads the calls’ token sequences from out.entries on demand, so a build without refusals compares nothing.

nemo_gym.models.token_capture.builder._materialize_delta_prompts(
    entries: list[nemo_gym.models.token_capture.records.TokenEntry]
) -> tuple[list[nemo_gym.models.token_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.models.token_capture.builder._resolve_parent(
    node: '_Node',
    by_call_id: dict[str, '_Node'],
    prefix_index: nemo_gym.models.token_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.models.token_capture.builder._same_full_sequence(
    a: nemo_gym.models.token_capture.records.TokenEntry,
    b: nemo_gym.models.token_capture.records.TokenEntry
) -> bool

Whether two calls’ full token sequences, prompt and generation together, are identical.

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

nemo_gym.models.token_capture.builder.project_chain_to_output_items(
    chain: nemo_gym.models.token_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.models.token_capture.builder.project_main_chain_response(
    rollout_id: str,
    out: nemo_gym.models.token_capture.builder.BuildOutput,
    model: str = '',
    finish_reason_resolved: bool = True,
    engine_refusals: tuple[nemo_gym.models.token_capture.records.EngineRefusalRecord, ...] = ()
) -> 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.

When the delivered chain’s terminal call stopped for length or content_filter, the payload carries status: "incomplete" with the matching incomplete_details.reason (max_output_tokens or content_filter), the verdict the Responses converter gives a chat completion that ended the same way. The finish reason is recorded by the local capture path (capture_tokens); under worker custody the ledger row carries it and staging.rebuild.LinearizedRow.incomplete_reason applies the same mapping. finish_reason_resolved=False withholds the verdict: terminal attribution passes it when the terminal call was chosen among token-identical retries that recorded different finish reasons, so the status would otherwise depend on which twin drew the smaller call id. The withholding covers the finish-reason verdict only; a refusal that ended the chain still sets the status, because the refusal record names the call it continued and its verdict never depended on the twins’ finish reasons.

engine_refusals are the frozen snapshot’s engine refusals. A refusal of the request that continued the delivered terminal call, refused for the context window, also makes the payload incomplete with reason max_output_tokens: the harness’s next request did not fit, so the rollout ended at the model’s limit although its last served call completed (_engine_refusal_ended_the_chain states the two conditions). A refusal the harness recovered from, with a call served after it (a token-identical retry of the terminal call that was served at or after the refusal included), a refusal on another branch of the rollout, a refusal of an output-token parameter, and one whose parent is unknown do not change the verdict. A refusal and a served length stop give the same status; the consumer’s metrics["engine_refusals"] tells them apart.

nemo_gym.models.token_capture.builder.run_builder(
    entries: list[nemo_gym.models.token_capture.records.TokenEntry],
    builder: str = 'prefix_merging',
    terminal_call_id: str | None = None
) -> nemo_gym.models.token_capture.builder.BuildOutput

Chain frozen snapshot entries with the named strategy.

terminal_call_id anchors chain selection for prefix_merging.

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