`nemo_rl.models.generation.dynamo.token_wrapper`#

OpenAI-compatible token wrapper for NeMo-Gym traffic to Dynamo.

Module Contents#

Classes#

DynamoTokenWrapperServer

Small HTTP server that supplies tokenized chat prompts to Dynamo.

Functions#

`_coerce_token_id_list`
`_coerce_logprob_list`
`_strip_gym_token_metadata`
`_chat_template_kwargs`
`_request_add_generation_prompt`
`_apply_chat_template`
`_render_prompt_token_ids`
`_render_prompt_token_ids_with_optional_prefix`
`_latest_tokenized_assistant_index`
`_derive_required_prefix_token_ids`	Return the exact prompt and generation IDs from the latest model turn.
`_normalize_tool_arguments_for_template`	Make OpenAI tool calls renderable by model chat templates.
`_validate_engine_data`
`_inject_gym_token_metadata`	Expose Dynamo engine metadata on the message fields consumed by Gym.
`prepare_dynamo_chat_completion_request`	Prepare a NeMo-Gym chat-completion request for Dynamo token input.

Data#

`_GYM_TOKEN_METADATA_FIELDS`
`_TOOL_ARGUMENT_MAPPING_ERROR`

API#

nemo_rl.models.generation.dynamo.token_wrapper._GYM_TOKEN_METADATA_FIELDS#: (‘prompt_token_ids’, ‘generation_token_ids’, ‘generation_log_probs’)

nemo_rl.models.generation.dynamo.token_wrapper._TOOL_ARGUMENT_MAPPING_ERROR#: ‘Can only get item pairs from a mapping.’

nemo_rl.models.generation.dynamo.token_wrapper._coerce_token_id_list(value: Any, field_name: str) → list[int]#

nemo_rl.models.generation.dynamo.token_wrapper._coerce_logprob_list( value: Any, field_name: str, ) → list[float]#

nemo_rl.models.generation.dynamo.token_wrapper._strip_gym_token_metadata( messages: list[Any], ) → list[Any]#

nemo_rl.models.generation.dynamo.token_wrapper._chat_template_kwargs( request_body: dict[str, Any], tokenizer_chat_template_kwargs: Optional[dict[str, Any]], ) → dict[str, Any]#

nemo_rl.models.generation.dynamo.token_wrapper._request_add_generation_prompt( request_body: dict[str, Any], ) → bool#

nemo_rl.models.generation.dynamo.token_wrapper._apply_chat_template( *, tokenizer: Any, request_body: dict[str, Any], messages: list[Any], tokenizer_chat_template_kwargs: Optional[dict[str, Any]], exclude_tools_when_tool_choice_none: bool, add_generation_prompt: bool, tokenize: bool, ) → Any#

nemo_rl.models.generation.dynamo.token_wrapper._render_prompt_token_ids( *, tokenizer: Any, request_body: dict[str, Any], messages: list[Any], tokenizer_chat_template_kwargs: Optional[dict[str, Any]], exclude_tools_when_tool_choice_none: bool, add_generation_prompt: bool, ) → list[int]#

nemo_rl.models.generation.dynamo.token_wrapper._render_prompt_token_ids_with_optional_prefix( *, tokenizer: Any, request_body: dict[str, Any], messages: list[Any], tokenizer_chat_template_kwargs: Optional[dict[str, Any]], exclude_tools_when_tool_choice_none: bool, add_generation_prompt: bool, assistant_index: int | None, ) → tuple[list[int], list[int] | None]#

nemo_rl.models.generation.dynamo.token_wrapper._latest_tokenized_assistant_index( messages: list[Any], ) → Optional[int]#

nemo_rl.models.generation.dynamo.token_wrapper._derive_required_prefix_token_ids( messages: list[Any], ) → list[int] | None#: Return the exact prompt and generation IDs from the latest model turn.

nemo_rl.models.generation.dynamo.token_wrapper._normalize_tool_arguments_for_template( messages: list[Any], *, before_index: int, ) → None#

Make OpenAI tool calls renderable by model chat templates.

OpenAI chat messages carry function.arguments as a JSON string, while some model templates iterate those arguments as a mapping. Normalize only the local template copy; the request forwarded to Dynamo retains its original OpenAI payload.

nemo_rl.models.generation.dynamo.token_wrapper._validate_engine_data( response_body: dict[str, Any], ) → tuple[list[int], list[int], list[float]]#

nemo_rl.models.generation.dynamo.token_wrapper._inject_gym_token_metadata( response_body: dict[str, Any], ) → None#: Expose Dynamo engine metadata on the message fields consumed by Gym.

nemo_rl.models.generation.dynamo.token_wrapper.prepare_dynamo_chat_completion_request( request_body: dict[str, Any], *, tokenizer: Any, tokenizer_chat_template_kwargs: Optional[dict[str, Any]] = None, exclude_tools_when_tool_choice_none: bool, ) → dict[str, Any]#: Prepare a NeMo-Gym chat-completion request for Dynamo token input.

class nemo_rl.models.generation.dynamo.token_wrapper.DynamoTokenWrapperServer( *, dynamo_frontend_base_url: str, tokenizer: Any, tokenizer_chat_template_kwargs: Optional[dict[str, Any]], exclude_tools_when_tool_choice_none: bool, request_timeout_s: Optional[float], )#