from typing import (Any, Dict, Optional)

from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
from examples.cambricon_custom_func.context_parallel.mlu_hijack.worker.model_runner_base import vllm__worker__model_runner_base___init_sampling_metadata_from_tensor_dict
from vllm.worker.model_runner_base import _init_attn_metadata_from_tensor_dict

@classmethod
def vllm__worker__model_runner__ModelInputForGPUWithSamplingMetadata__from_broadcasted_tensor_dict(
    cls,
    tensor_dict: Dict[str, Any],
    attn_backend: Optional["AttentionBackend"] = None,
) -> "ModelInputForGPUWithSamplingMetadata":
    '''
    ==========================
    Modify by Context Parallel
    ==========================
    @brief: force apply hijacked function.
    '''
    tensor_dict = vllm__worker__model_runner_base___init_sampling_metadata_from_tensor_dict(tensor_dict)
    '''
    =======================
    End of Context Parallel
    =======================
    '''
    if attn_backend is not None:
        tensor_dict = _init_attn_metadata_from_tensor_dict(
            attn_backend, tensor_dict)
    return cls(**tensor_dict)

MluHijackObject.apply_hijack(
    ModelInputForGPUWithSamplingMetadata,
    ModelInputForGPUWithSamplingMetadata.from_broadcasted_tensor_dict,
    vllm__worker__model_runner__ModelInputForGPUWithSamplingMetadata__from_broadcasted_tensor_dict
)
