code: def _prepare_model_inputs(self, ids):
      lengths = list(
          torch.sum(
              self.context.cache[i]["encoded"]["attention_mask"], dim=1
          ).item()
          for i in ids
      )
      max_len = max(lengths)
  
      device = next(iter(self.model.parameters())).device
      dtype = torch.int64
      config = {"device": device, "dtype": dtype}
  
      input_ids = []
      attention_mask = []
      kv_cache = {}
      for req_id, seq_len in zip(ids, lengths):
          input_ids.append(
              self.context.cache[req_id]["encoded"]["input_ids"]
          )
          attention_mask.append(
              self.context.cache[req_id]["encoded"]["attention_mask"]
          )
  
          for layer_idx, layer_kv in enumerate(
              self.context.cache[req_id]["encoded"]["past_key_values"]
          ):
              k, v = layer_kv
              kv_cache[layer_idx] = kv_cache.get(layer_idx, {})
              kv_cache[layer_idx][0] = kv_cache.get(layer_idx, {}).get(
                  0, []
              ) + [k]
              kv_cache[layer_idx][1] = kv_cache.get(layer_idx, {}).get(
                  1, []
              ) + [v]
          padded_len = input_ids[-1].size()[-1]
          if padded_len < max_len:
              # Apply padding to input_ids, attention_mask and past_key_values
              n = max_len - seq_len
              input_ids[-1] = torch.concat(
                  (
                      self.tokenizer.pad_token_id
                      + torch.zeros((1, n), **config),
                      input_ids[-1],
                  ),
                  dim=1,
              )
              attention_mask[-1] = torch.concat(
                  (torch.zeros((1, n), **config), attention_mask[-1]),
                  dim=1,
              )
  
              size_delta = list(kv_cache[0][0][-1].size())
              size_delta[2] = n
              dtype = kv_cache[0][0][-1].dtype
              for layer_idx in range(len(kv_cache)):
                  kv_cache[layer_idx][0][-1] = torch.concat(
                      (
                          torch.zeros(size_delta, **config),
                          kv_cache[layer_idx][0][-1],
                      ),
                      dim=2,
                  )
                  kv_cache[layer_idx][1][-1] = torch.concat(
                      (
                          torch.zeros(size_delta, **config),
                          kv_cache[layer_idx][1][-1],
                      ),
                      dim=2,
                  )
  
          elif padded_len > max_len:
              # Truncate padding from input_ids, attention_mask and past_key_values
              input_ids[-1] = input_ids[-1][:, -max_len:]
              attention_mask[-1] = attention_mask[-1][:, -max_len:]
  
              for layer_idx in range(len(kv_cache)):
                  kv_cache[layer_idx][0][-1] = kv_cache[layer_idx][0][
                      -1
                  ][:, :, (-max_len + 1) :, :]
                  kv_cache[layer_idx][1][-1] = kv_cache[layer_idx][1][
                      -1
                  ][:, :, (-max_len + 1) :, :]
          del self.context.cache[req_id]["encoded"]["past_key_values"]
  
      for layer_idx in range(len(kv_cache)):
          kv_cache[layer_idx][0] = torch.concat(
              kv_cache[layer_idx][0], dim=0
          )
          kv_cache[layer_idx][1] = torch.concat(
              kv_cache[layer_idx][1], dim=0
          )
  
      kv_cache = tuple(
          (kv_cache[layer_idx][0], kv_cache[layer_idx][1])
          for layer_idx in range(len(kv_cache))
      )
  
      encoded = {
          "input_ids": torch.concat(input_ids, dim=0),
          "attention_mask": torch.concat(attention_mask, dim=0),
          "past_key_values": kv_cache,
      }
      return encoded
  
file_path: serve/test/pytest/test_data/streaming/stream_handler.py
imports: import json
  import logging
  import types
  import torch
  from transformers import AutoModelForCausalLM, AutoTokenizer
  from ts.torch_handler.base_handler import BaseHandler
label_desc: Use image node_id 7 for calling the StreamingHandler obj's underlying member method code with example usage: obj._prepare_model_inputs(ids) and returns: encoded