code: def inference(
      self,
      input: torch.Tensor,  # (batch, len_q) int32
      input_sub: torch.Tensor,  # (batch, len_q) int32
      position: torch.Tensor,  # (batch, len_q)  int32
      context: torch.Tensor,  # (batch, len_q) bool
      sample_ids: torch.Tensor,  # (batch, len_q) int32
      num_segments: torch.Tensor,  # (batch, len_q) int32
      segment: torch.Tensor,  # (batch, len_q) int32
      segment_rel_offset: torch.Tensor,  # (batch, len_q) int32
      segment_rel: torch.Tensor,  # (batch, num_segment_bucket) int32
      ext_table_ids: torch.Tensor,  # (ext_table_size) int32
      ext_table_sub: torch.Tensor,  # (ext_table_size) int32
      past_key_values: Optional[CPMBeeInferenceState] = None,
  ) -> Tuple[torch.Tensor, torch.Tensor, CPMBeeInferenceState]:
      with torch.no_grad():
          if past_key_values is None:
              present_position = position
              present_context = context
              present_sample_ids = sample_ids
              present_num_segments = num_segments
              present_segments = segment
              present_buffer = None
          else:
              present_position = torch.cat(
                  [past_key_values["buffer_position"], position], dim=-1
              )
              present_context = torch.cat(
                  [past_key_values["buffer_context"], context], dim=-1
              )
              present_sample_ids = torch.cat(
                  [past_key_values["buffer_sample_ids"], sample_ids],
                  dim=-1,
              )
              present_num_segments = torch.cat(
                  [
                      past_key_values["buffer_num_segments"],
                      num_segments,
                  ],
                  dim=-1,
              )
              present_segments = torch.cat(
                  [past_key_values["buffer_segments"], segment], dim=-1
              )
              present_buffer = past_key_values["buffer"]
  
          batch = input.size(0)
          len_q = input.size(1)
          len_buffer = present_position.size(1)
  
          segment_rel_2d = torch.masked_fill(
              segment[:, :, None] * num_segments[:, :, None]
              + present_segments[:, None, :]
              + segment_rel_offset[:, :, None],
              ~(
                  (
                      sample_ids[:, :, None]
                      == present_sample_ids[:, None, :]
                  )
              ),  # not in the same sample
              0,  # avoid torch.gather overflow
          ).view(batch, len_q * len_buffer)
  
          segment_bucket = torch.gather(
              input=segment_rel,
              dim=1,
              index=segment_rel_2d.long(),
          ).view(batch, len_q, len_buffer)
  
          segment_bucket.masked_fill_(
              ~(
                  (
                      sample_ids[:, :, None]
                      == present_sample_ids[:, None, :]
                  )
              ),  # not in the same span or sample
              1,  # bucket is used for in-context samples
          )
  
          # directional mask
          directional_mask_2d = (
              present_position[:, None, :] <= position[:, :, None]
          )
          # sample mask
          sample_mask_2d = (sample_ids[:, :, None] == 0) | (
              sample_ids[:, :, None] == present_sample_ids[:, None, :]
          )
          # context mask
          attention_mask = present_context[:, None, :] | (
              context[:, :, None].logical_not()
              & directional_mask_2d.view(batch, len_q, len_buffer)
          )
          # span mask
          attention_mask = attention_mask & sample_mask_2d
          # length mask
          mask_1d = present_num_segments != 0
          attention_mask = (
              mask_1d.view(batch, 1, len_buffer) & attention_mask
          )
  
          hidden_states = self.input_embedding(input, input_sub)
  
          position_bias = self.position_bias(
              position, present_position, segment_bucket
          )
          hidden_states, present_key_values = self.encoder(
              hidden_states,
              attention_mask,
              position_bias,
              True,
              present_buffer,
          )
          ext_table = self.input_embedding(ext_table_ids, ext_table_sub)
          logits = self.input_embedding.projection(
              hidden_states, ext_table
          )
  
          return (
              logits,
              hidden_states,
              {
                  "buffer_position": present_position,
                  "buffer_context": present_context,
                  "buffer_sample_ids": present_sample_ids,
                  "buffer_num_segments": present_num_segments,
                  "buffer_segments": present_segments,
                  "buffer": present_key_values,
              },
          )
  
file_path: DeepKE/example/llm/CPM-Bee/src/cpm_live/models/bee_torch.py
imports: from typing import Optional, Tuple
  import torch
  from ..native_layers import Encoder, EmbeddingExt, BucketPositionBias
  from .bee import CPMBeeConfig, CPMBeeInferenceState
label_desc: Use image node_id 3 for calling the CPMBeeTorch obj's underlying member method code with example usage: obj.inference(input, input_sub, position, context, sample_ids, num_segments, segment, segment_rel_offset, segment_rel, ext_table_ids, ext_table_sub, past_key_values) and returns: logits,  hidden_states,  dict