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
`TokenEntry`	Store one model call’s content and token metadata.

Functions

Name	Description
`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.
`strip_token_fields`	Drop the token arrays from output items, keeping the content.

Data

TOKEN_ENTRY_RECORD_SCHEMA_VERSION

TOKEN_FIELDS

API

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.

created_at

float = 0.0

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)

prompt_token_ids

list[int]

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.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.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.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...