class_docstr: Transformers handler class for sequence, token classification and question answering.
class_docstr_tok: ["Transformers","handler","class","for","sequence",",","token","classification","and","question","answering","."]
code: def _rearrange_kv_cache_for_decode(self, ids: List[str]) -> None:
      req_id_to_batch_idx = {
          req_id: idx
          for idx, req_id in enumerate(self.batch_idx_to_req_ids)
          if req_id is not None
      }
      decode_indices = [req_id_to_batch_idx[req_id] for req_id in ids]
  
      prefill_ids = list(
          set(
              [
                  req_id
                  for req_id in self.batch_idx_to_req_ids
                  if req_id is not None
              ]
          )
          - set(ids)
      )
      prefill_indices = [
          req_id_to_batch_idx[req_id] for req_id in prefill_ids
      ]
  
      none_indices = [
          idx
          for idx, req_id in enumerate(self.batch_idx_to_req_ids)
          if req_id is None
      ]
  
      new_arrangement = torch.tensor(
          decode_indices + prefill_indices + none_indices,
          dtype=torch.long,
          device=self.device,
      )
      for l in self.model.layers:
          l.attention.cache_k = l.attention.cache_k[
              new_arrangement, ...
          ]
          l.attention.cache_v = l.attention.cache_v[
              new_arrangement, ...
          ]
      self.batch_idx_to_req_ids = [
          self.batch_idx_to_req_ids[idx]
          for idx in decode_indices + prefill_indices + none_indices
      ]
  
imports: import json
  import logging
  import os
  import sys
  import time
  from abc import ABC
  from typing import List
  import packaging.version
  import requests
  import torch
  import transformers
  from ts.torch_handler.base_handler import BaseHandler
  from generate import sample_top_p
  from llama2 import Llama
label_desc: Use image node_id 11 for calling the LlamaHandler obj's underlying member method code with example usage: obj._rearrange_kv_cache_for_decode(ids) without return types