nemo_gym.token_id_capture.records

Define training-token records extracted from served responses.

A TokenEntry contains one model call’s training data. It stores the exact prompt token ids. It stores generated token ids and their log probabilities. Evaluation uses a separate ModelCallRecord. Evaluation records do not carry token arrays. Both records share a model_call_id.

Module Contents

Classes

Name	Description
`ParentResolutionStatus`	Describe whether a model call has a proven captured predecessor.
`TokenEntry`	Store one model call’s content and token metadata.

Functions

Name	Description
`compute_digest`	Digest of an exact token sequence.
`cumulative_tokens`	The full sequence a child of this call must start with.
`encode_token_ids`	Stable, length-delimited, big-endian encoding of a token sequence.
`extract_token_fields`	Pull the token-id fields off a served response, or `None` if absent.
`response_to_output_items`	Normalize a served response to a list of content-bearing Responses output items.
`stamp_lineage`	Fill token lineage and the request-time parent decision.
`strip_token_fields`	Drop the token arrays from output items, keeping the content.

Data

DIGEST_VERSION

TOKEN_ENTRY_MIN_SCHEMA_VERSION

TOKEN_ENTRY_RECORD_SCHEMA_VERSION

TOKEN_FIELDS

_DIGEST_DOMAIN

_EMPTY_DIGEST

API

class nemo_gym.token_id_capture.records.ParentResolutionStatus

Bases: enum.Enum

Describe whether a model call has a proven captured predecessor.

RESOLVED

= 'resolved'

ROOT

= 'root'

UNRESOLVED

= 'unresolved'

class nemo_gym.token_id_capture.records.TokenEntry()

Bases: BaseModel

Store one model call’s content and token metadata.

The rollout id identifies the training sample. The model call id joins evaluation context. output_items preserves assistant text and tool calls. Text-based penalties require that content. Token arrays are stored once at the top level. token_item_index identifies their original output item. A trajectory builder can restore chain-correct token fields there.

continuation_context_digest

str = ''

continuation_context_len

int = 0

continuation_fingerprint

str = ''

created_at

float = 0.0

cum_len

int | None = None

digest

str | None = None

fingerprint_version

int | None = None

generation_log_probs

list[float]

generation_token_ids

list[int]

model

str = ''

model_call_id

str

model_config

= ConfigDict(extra='allow')

output_items

list[dict] = Field(default_factory=list)

parent_call_id

str | None = None

parent_resolution

ParentResolutionStatus | None = None

parent_resolution_reason

str = ''

prefix_requested

bool = False

prefix_supplied

bool = False

prompt_is_delta

bool = False

prompt_token_ids

list[int]

response_id

str | None = None

rollout_id

str

routed_experts

Any | None = None

schema_version

int = TOKEN_ENTRY_RECORD_SCHEMA_VERSION

token_item_index

int | None = None

nemo_gym.token_id_capture.records.TokenEntry._refuse_a_newer_record() -> 'TokenEntry'

Accept older records and reject newer records.

Missing older fields use their defaults. Unknown newer fields may change token semantics. Rejecting them prevents silent training corruption.

nemo_gym.token_id_capture.records.compute_digest(
    token_ids: list[int]
) -> str

Digest of an exact token sequence.

The builder verifies a claimed parent by hashing the corresponding prompt prefix. A mismatch quarantines the call. This prevents stale or interleaved records from merging silently.

nemo_gym.token_id_capture.records.cumulative_tokens(
    entry: nemo_gym.token_id_capture.records.TokenEntry
) -> list[int]

The full sequence a child of this call must start with.

Delta records require parent-chain reconstruction.

nemo_gym.token_id_capture.records.encode_token_ids(
    token_ids: list[int]
) -> bytes

Stable, length-delimited, big-endian encoding of a token sequence.

Independent implementations can hash the same bytes. One vectorized pack replaces the per-token loop. The byte layout remains identical to the original encoding.

nemo_gym.token_id_capture.records.extract_token_fields(
    response_json: dict
) -> dict | None

Pull the token-id fields off a served response, or None if absent.

Handle Responses output items and Chat Completions messages. Exactly one item may carry token metadata. Return None when no item carries token ids.

nemo_gym.token_id_capture.records.response_to_output_items(
    payload: dict
) -> list[dict]

Normalize a served response to a list of content-bearing Responses output items.

Responses payloads already carry output. Chat payloads carry choices[*].message. Wrap each assistant message as a Responses message item.

nemo_gym.token_id_capture.records.stamp_lineage(
    entry: nemo_gym.token_id_capture.records.TokenEntry,
    parent_call_id: str | None,
    parent_resolution: nemo_gym.token_id_capture.records.ParentResolutionStatus | None = None,
    cumulative: list[int] | None = None
) -> nemo_gym.token_id_capture.records.TokenEntry

Fill token lineage and the request-time parent decision.

cum_len and digest describe the full sequence. A delta entry must pass cumulative explicitly. parent_resolution=None preserves records built by compatibility callers.

nemo_gym.token_id_capture.records.strip_token_fields(
    items: list[dict]
) -> tuple[list[dict], int | None]

Drop the token arrays from output items, keeping the content.

Return the stripped items and the index of their token-bearing item. Capture requires exactly one token-bearing item. The arrays are held once on the entry. Storing them again per item would roughly double the record size.

nemo_gym.token_id_capture.records.DIGEST_VERSION = 1

nemo_gym.token_id_capture.records.TOKEN_ENTRY_MIN_SCHEMA_VERSION = 1

nemo_gym.token_id_capture.records.TOKEN_ENTRY_RECORD_SCHEMA_VERSION = 1

nemo_gym.token_id_capture.records.TOKEN_FIELDS = ('prompt_token_ids', 'generation_token_ids', 'generation_log_probs', 'routed_exp...

nemo_gym.token_id_capture.records._DIGEST_DOMAIN = b'nemo-gym-tokens'

nemo_gym.token_id_capture.records._EMPTY_DIGEST = hashlib.sha256(_DIGEST_DOMAIN).hexdigest()