code: def inference(
      self,
      input: torch.Tensor,  # (batch, seqlen)
      length: torch.Tensor,  # (batch)
      context: torch.Tensor,  # (batch, seqlen)
      position: torch.Tensor,  # (batch, seqlen)
      segment: torch.Tensor,  # (batch, seqlen)
      span: torch.Tensor,  # (batch, seqlen)
      past_key_values=None,  # num_layers * 2 * (batch, num_heads, seqlen, dim_head)
  ):
      batch = input.size(0)
  
      if past_key_values is None:
          past_length = 0
          past_key_values = tuple([None] * self.encoder.num_layers)
          input_prompt = input[:, : self.prompt_length].contiguous()
          input_ids = input[:, self.prompt_length :].contiguous()
  
          prompt_states = self.prompt_embedding(input_prompt)
          hidden_states = self.input_embedding(input_ids)
          segment_states = self.segment_embedding(segment)
          hidden_states = (
              torch.cat([prompt_states, hidden_states], 1)
              + segment_states
          )
  
      else:
          past_length = past_key_values[0][0].size(-2)
          segment_states = self.segment_embedding(segment)
          hidden_states = (
              self.input_embedding(input) + segment_states[:, -1:, :]
          )
  
      seqlen = past_length + input.size(1)
  
      with torch.no_grad():
          device = input.device
          directional_mask_2d = torch.arange(
              seqlen, device=device
          ) <= torch.arange(seqlen, device=device).view(-1, 1)
          attention_mask = context[:, None, :] | (
              context[:, :, None].logical_not()
              & directional_mask_2d.view(1, seqlen, seqlen)
          )
          attention_mask = attention_mask & (
              span[:, None, :] == span[:, :, None]
          )
          # mask for left paddding
          mask_1d = (
              torch.tensor(list(range(seqlen))[::-1], device=device)[
                  None, :
              ].repeat(batch, 1)
              < length[:, None]
          )
          attention_mask = (
              mask_1d.view(batch, seqlen, 1)
              & mask_1d.view(batch, 1, seqlen)
              & attention_mask
          )
  
      position_bias = self.position_bias(
          position, position, segment, segment
      )
  
      attention_mask = attention_mask[:, past_length:, :]
      position_bias = position_bias[:, :, past_length:, :]
  
      hidden_states, present_key_values = self.encoder(
          hidden_states,
          attention_mask,
          position_bias,
          True,
          past_key_values,
      )
      logits = self.input_embedding.projection(hidden_states)
      return logits, hidden_states, present_key_values
  
file_path: DeepKE/example/llm/CPM-Bee/src/cpm_live/models/ant.py
imports: from typing import List, Optional, Tuple
  import torch
  from ..utils import Config
  from ..layers import Encoder, Embedding, SegmentPositionEmbedding
  import bmtrain
label_desc: Use image node_id 3 for calling the CPMAnt obj's underlying member method code with example usage: obj.inference(input, length, context, position, segment, span, past_key_values) and returns: logits,  hidden_states,  present_key_values