def get_merged_sample_with_past_kv_inputs(
        self,
        model_name: str,
        seq_len: int,
        past_seq_len: int,
        use_fp16: bool = False,
        trust_remote_code=None,
    ):
        """Get inputs for forward pass with past_key_values.

        This is for the "merged" model which can be used for both prompt generation and token generation.
        For prompt generation, past_seq_len = 0 and seq_len >= 1. past_kv is a list of tuples of empty tensors.
        For token generation, past_seq_len >= 1 and seq_len = 1.

        Shape of outputs:
            input_ids: (seq_len)
            attention_mask: (past_seq_len + seq_len)
            position_ids: (seq_len)
            past_key: (num_heads, past_seq_len, head_size)
            past_value: (num_heads, past_seq_len, head_size)
        """
        # Note: No need for separate function for legacy prompt and token generation
        # prompt generation (get_sample_inputs):
        #   past_seq_len = 0, seq_len >= 1, use_gqa = False, use_fp16 = False
        #   and remove past_kv from the output
        # token generation (get_sample_with_past_kv_inputs):
        #   past_seq_len >= 1, seq_len = 1, use_gqa = False, use_fp16 = False
        # By using a single function with no default values, we can avoid confusion and are deliberate about the sizes
        # can instead write dummy input functions like 'get_merged_decoder_with_past_dummy_inputs' if needed

        # Using Namespace class to access dict items like class attributes
        model_attributes = get_model_config(model_name, trust_remote_code=trust_remote_code).to_dict()
        world_size = model_attributes.get("world_size", 1)
        vocab_size = model_attributes.get("vocab_size", 50256)
        input_ids = torch.randint(low=0, high=vocab_size, size=(seq_len,), dtype=torch.int64)
        attention_mask = torch.ones(past_seq_len + seq_len, dtype=torch.int64)
        position_ids = self.get_position_ids(attention_mask, past_seq_len=past_seq_len)
        position_ids = position_ids.to(torch.int64)
        past_kv = self.get_past_kv_inputs(model_attributes, past_seq_len, use_fp16=use_fp16, world_size=world_size)

        return (input_ids, attention_mask, position_ids, past_kv)
Get inputs for forward pass with past_key_values.

        This is for the "merged" model which can be used for both prompt generation and token generation.
        For prompt generation, past_seq_len = 0 and seq_len >= 1. past_kv is a list of tuples of empty tensors.
        For token generation, past_seq_len >= 1 and seq_len = 1.

        Shape of outputs:
            input_ids: (seq_len)
            attention_mask: (past_seq_len + seq_len)
            position_ids: (seq_len)
            past_key: (num_heads, past_seq_len, head_size)
            past_value: (num_heads, past_seq_len, head_size)
        
olive/data/component/dataset.py