nemo_gym.token_id_capture.adapters.vllm

Dependency-light extraction adapter for vLLM chat completions.

Module Contents

Classes

Name	Description
`VLLMCaptureAdapter`	Translate vLLM request/response payloads at the framework boundary.

Functions

Name	Description
`_message`	-
`_single_choice`	-
`extract_generation_token_info`	Read exact generation IDs/log probabilities from supported vLLM shapes.

Data

PREFIX_IDS_FIELD

PROMPT_IDS_FIELD

ROUTED_EXPERTS_FIELD

API

class nemo_gym.token_id_capture.adapters.vllm.VLLMCaptureAdapter()

Translate vLLM request/response payloads at the framework boundary.

nemo_gym.token_id_capture.adapters.vllm.VLLMCaptureAdapter.enter_prefix(
    request_payload: dict[str, typing.Any],
    prefix_ids: list[int]
) -> dict[str, typing.Any]

nemo_gym.token_id_capture.adapters.vllm.VLLMCaptureAdapter.extract_extras(
    response_payload: dict[str, typing.Any]
) -> dict[str, typing.Any] | None

nemo_gym.token_id_capture.adapters.vllm.VLLMCaptureAdapter.extract_generation(
    response_payload: dict[str, typing.Any]
) -> tuple[list[int], list[float]]

nemo_gym.token_id_capture.adapters.vllm.VLLMCaptureAdapter.extract_prompt_ids(
    response_payload: dict[str, typing.Any]
) -> list[int]

nemo_gym.token_id_capture.adapters.vllm._message(
    choice: dict[str, typing.Any]
) -> dict[str, typing.Any]

nemo_gym.token_id_capture.adapters.vllm._single_choice(
    response_payload: dict[str, typing.Any]
) -> dict[str, typing.Any]

nemo_gym.token_id_capture.adapters.vllm.extract_generation_token_info(
    choice: dict[str, typing.Any]
) -> tuple[list[int], list[float]]

Read exact generation IDs/log probabilities from supported vLLM shapes.

nemo_gym.token_id_capture.adapters.vllm.PREFIX_IDS_FIELD = 'required_prefix_token_ids'

nemo_gym.token_id_capture.adapters.vllm.PROMPT_IDS_FIELD = 'prompt_token_ids'

nemo_gym.token_id_capture.adapters.vllm.ROUTED_EXPERTS_FIELD = 'routed_experts'