vllm_omni.model_executor.custom_process_mixin ¶

CustomProcessMixin ¶

Mixin class for all stages in the Omni model.

postprocess ¶

postprocess(model_output, **info_dict: object)

Postprocess the model output. Returns the postprocessed model output and the save dictionary. Args: model_output: The model output to postprocess.

preprocess ¶

preprocess(
    input_ids: Tensor,
    input_embeds: Tensor,
    **input_dict: object,
) -> tuple[Tensor, Tensor, dict]

Process the input_ids and input_embeds for the given input_dict. Returns the processed input_ids, input_embeds, and the input_dict. If the stage don't applicable, return the original input_ids, input_embeds, and an empty dict.

set_custom_postprocess ¶

set_custom_postprocess(postprocess_fn: Callable) -> None

Set a postprocess function for the stage. Args: postprocess_fn: The postprocess function to register.

set_custom_preprocess ¶

set_custom_preprocess(preprocess_fn: Callable) -> None

Set a preprocess function for the stage. Args: preprocess_fn: The preprocess function to register.