add source files

configuration_cpmbee.py

@@ -0,0 +1,132 @@
+# coding=utf-8
+# Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" CpmBee model configuration"""
+
+from typing import List, Optional, Tuple, Union
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+CPMBEE_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "openbmb/cpm-bee-10b": "https://huggingface.co/openbmb/cpm-bee-10b/resolve/main/config.json",
+    "openbmb/cpm-bee-5b": "https://huggingface.co/openbmb/cpm-bee-5b/resolve/main/config.json",
+    "openbmb/cpm-bee-2b": "https://huggingface.co/openbmb/cpm-bee-2b/resolve/main/config.json",
+    "openbmb/cpm-bee-1b": "https://huggingface.co/openbmb/cpm-bee-1b/resolve/main/config.json",
+    # See all CpmBee models at https://huggingface.co/models?filter=cpmbee
+}
+
+
+class CpmBeeConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`CpmBeeModel`]. It is used to instbeeiate an
+    CPMBee model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the CPMBee
+    [openbmb/cpm-bee-10b](https://huggingface.co/openbmb/cpm-bee-10b) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 30720):
+            Vocabulary size of the CPMBee model. Defines the number of different tokens that can be represented by the
+            `input` passed when calling [`CpmBeeModel`].
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the encoder layers.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads in the Transformer encoder.
+        dim_head (`int`, *optional*, defaults to 128):
+            Dimension of attention heads for each attention layer in the Transformer encoder.
+        dim_ff (`int`, *optional*, defaults to 10240):
+            Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 48):
+            Number of layers of the Transformer encoder.
+        dropout_p (`float`, *optional*, defaults to 0.1):
+            The dropout probabilitiy for all fully connected layers in the embeddings, encoder.
+        position_bias_num_buckets (`int`, *optional*, defaults to 512):
+            The number of position_bias buckets.
+        position_bias_num_segment_buckets (`int`, *optional*, defaults to 32):
+            The number of segment buckets.
+        position_bias_max_distance (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        eps (`float`, *optional*, defaults to 1e-6):
+            The epsilon used by the layer normalization layers.
+        init_std (`float`, *optional*, defaults to 1.0):
+            Initialize parameters with std = init_std.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether to use cache.
+        distance_scale (`float` or `int`, *optional*, defaults to 16):
+            Scale the rotary embedding.
+        mask_modules (`list` or `tuple`, *optional*, defaults to None):
+            Decides which feedforward block or attention block is pruned.
+        half (`bool`, *optional*, defaults to `False`):
+            Decides the model parameters are half-precision or not.
+
+    Example:
+
+    ```python
+    >>> from transformers import CpmBeeModel, CpmBeeConfig
+
+    >>> # Initializing a CPMBee cpm-bee-10b style configuration
+    >>> configuration = CpmBeeConfig()
+
+    >>> # Initializing a model from the cpm-bee-10b style configuration
+    >>> model = CpmBeeModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+    model_type = "cpmbee"
+
+    def __init__(
+        self,
+        vocab_size: int = 30720,
+        hidden_size: int = 4096,
+        num_attention_heads: int = 64,
+        dim_head: int = 64,
+        dim_ff: int = 10240,
+        num_hidden_layers: int = 32,
+        dropout_p: int = 0.0,
+        position_bias_num_buckets: int = 256,
+        position_bias_num_segment_buckets: int = 32,
+        position_bias_max_distance: int = 2048,
+        eps: int = 1e-6,
+        init_std: float = 1.0,
+        use_cache: bool = True,
+        distance_scale: Union[int, float] = 16,
+        mask_modules: Optional[Union[List, Tuple]] = None,
+        half: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.position_bias_num_segment_buckets = position_bias_num_segment_buckets
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.dim_head = dim_head
+        self.dim_ff = dim_ff
+        self.num_hidden_layers = num_hidden_layers
+        self.position_bias_num_buckets = position_bias_num_buckets
+        self.position_bias_max_distance = position_bias_max_distance
+        self.dropout_p = dropout_p
+        self.eps = eps
+        self.use_cache = use_cache
+        self.vocab_size = vocab_size
+        self.init_std = init_std
+        self.distance_scale = distance_scale
+        self.half = half
+        self.mask_modules = mask_modules

modeling_cpmbee.py

@@ -0,0 +1,1944 @@
+# coding=utf-8
+# Copyright 2022 The OpenBMB Team The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch CpmBee model."""
+import copy
+import math
+from collections import UserDict
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from ...generation.beam_search import BeamHypotheses, BeamSearchScorer
+from ...generation.streamers import BaseStreamer
+from ...generation.utils import (
+    GenerationConfig,
+    LogitsProcessorList,
+    StoppingCriteriaList,
+    dist,
+    inspect,
+    is_deepspeed_zero3_enabled,
+    warnings,
+)
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
+from .configuration_cpmbee import CpmBeeConfig
+from .tokenization_cpmbee import CpmBeeTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "openbmb/cpm-bee-10b"
+_CONFIG_FOR_DOC = "CpmBeeConfig"
+
+CPMBEE_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "openbmb/cpm-bee-10b",
+    "openbmb/cpm-bee-5b",
+    "openbmb/cpm-bee-2b",
+    "openbmb/cpm-bee-1b",
+    # See all CPMBee models at https://huggingface.co/models?filter=cpmbee
+]
+
+
+class CpmBeeLinear(nn.Linear):
+    def __init__(self, dim_in, dim_out, dtype):
+        """
+        Construct a linear for CPMBee. It contains a scale operation.
+        """
+        super().__init__(dim_in, dim_out, bias=False)
+        self.dim_in = self.in_features = dim_in
+        self.dim_out = self.out_features = dim_out
+
+        self.weight = torch.nn.parameter.Parameter(torch.empty((dim_out, dim_in), dtype=dtype))
+
+    def forward(self, x: torch.Tensor):
+        """
+        Args:
+            x (`torch.Tensor` of shape `(batch, seq_len, dim_in)`): The input of linear layer
+        Returns:
+            `torch.Tensor` of shape `(batch, seq_len, dim_out)`: The output of the linear transform y.
+        """
+        x = nn.functional.linear(x, self.weight)
+        x = x / math.sqrt(self.dim_in)
+        return x
+
+
+class CpmBeeLayerNorm(nn.Module):
+    """
+    We use Root Mean Square (RMS) Layer Normalization, please see https://arxiv.org/abs/1910.07467 for details."
+    """
+
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__()
+
+        self.eps = config.eps
+        self.dim_norm = config.hidden_size
+        self.weight = nn.Parameter(torch.empty(config.hidden_size, dtype=config.torch_dtype))
+
+    def forward(self, hidden_states: torch.Tensor):
+        """
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
+        """
+        if hidden_states.size(-1) != self.dim_norm:
+            raise AssertionError("hidden_states.size(-1) != self.dim_norm")
+        old_dtype = hidden_states.dtype
+        variance = hidden_states.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
+        hidden_states = (hidden_states * torch.rsqrt(variance + self.eps)).to(old_dtype) * self.weight
+        return hidden_states
+
+
+class CpmBeeAttention(nn.Module):
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__()
+        self.dim_model = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.dim_head = config.dim_head
+
+        self.project_q = CpmBeeLinear(self.dim_model, self.num_heads * self.dim_head, dtype=config.torch_dtype)
+        self.project_k = CpmBeeLinear(self.dim_model, self.num_heads * self.dim_head, dtype=config.torch_dtype)
+        self.project_v = CpmBeeLinear(self.dim_model, self.num_heads * self.dim_head, dtype=config.torch_dtype)
+
+        self.attention_out = CpmBeeLinear(self.num_heads * self.dim_head, self.dim_model, dtype=config.torch_dtype)
+
+        self.softmax = torch.nn.Softmax(dim=-1)
+
+        if config.dropout_p is not None:
+            self.dropout = torch.nn.Dropout(p=config.dropout_p)
+        else:
+            self.dropout = None
+
+    def forward(
+        self,
+        hidden_q: torch.Tensor,
+        hidden_kv: torch.Tensor,
+        attention_mask: torch.BoolTensor,
+        position_bias: torch.Tensor,
+        output_attentions: Optional[bool] = False,
+        past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: Optional[bool] = None,
+    ):
+        """
+        Args:
+            hidden_q (`torch.Tensor`):
+                Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
+            hidden_kv (`torch.Tensor` of shape `(batch, len_k, dim_model)`)):
+                Tensor *key_value* and *query* of shape `(batch, len_k, dim_model)`
+            attention_mask (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
+                Avoid invalid areas to participate in the calculation of self-attention.
+            position_bias (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
+                Provide positional information to self-attention block.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers.
+            past_key_values (`Tuple[torch.Tensor, torch.Tensor]`, *optional*):
+                Cached past key and value projection states.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+        """
+        batch_size = hidden_q.size(0)
+        len_q = hidden_q.size(1)
+        len_k = hidden_kv.size(1)
+
+        query = self.project_q(hidden_q)
+        key = self.project_k(hidden_kv)
+        value = self.project_v(hidden_kv)
+
+        query = query.view(batch_size, len_q, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
+        key = key.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
+        value = value.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
+
+        if past_key_values is not None:
+            key = torch.cat([past_key_values[0], key], dim=-2)
+            value = torch.cat([past_key_values[1], value], dim=-2)
+            len_k = key.size(-2)
+
+        # (batch_size, num_heads, len_q, dim_head) @ (batch_size, num_heads, dim_head, len_k) -> (batch_size, num_heads, len_q, len_k)
+        score = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.dim_head)
+        score = score + position_bias
+
+        score = torch.masked_fill(
+            score,
+            attention_mask.view(batch_size, 1, len_q, len_k) == torch.tensor(False),
+            torch.scalar_tensor(float("-inf"), device=score.device, dtype=score.dtype),
+        )
+        score = self.softmax(score)
+
+        score = torch.masked_fill(
+            score,
+            attention_mask.view(batch_size, 1, len_q, len_k) == torch.tensor(False),
+            torch.scalar_tensor(0, device=score.device, dtype=score.dtype),
+        )
+        if output_attentions:
+            attn_weights = score
+        else:
+            attn_weights = None
+
+        if self.dropout is not None:
+            score = self.dropout(score)
+
+        # (batch_size, num_heads, len_q, len_k) @ (batch_size, num_heads, len_k, dim_head) -> (batch_size, num_heads, len_q, dim_head)
+        score = torch.matmul(score, value)
+
+        score = score.view(batch_size, self.num_heads, len_q, self.dim_head).permute(0, 2, 1, 3)
+        score = score.contiguous().view(batch_size, len_q, self.num_heads * self.dim_head)
+
+        score = self.attention_out(score)
+
+        past_key_values = None
+        if use_cache:
+            past_key_values = (key, value)
+
+        return score, attn_weights, past_key_values
+
+
+class CpmBeeSelfAttentionBlock(nn.Module):
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__()
+        self.layernorm_before_attention = CpmBeeLayerNorm(config)
+        self.self_attention = CpmBeeAttention(config)
+        if config.dropout_p:
+            self.dropout = torch.nn.Dropout(config.dropout_p)
+        else:
+            self.dropout = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_bias: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+        past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: Optional[bool] = None,
+    ):
+        """
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch, len_seq, dim_model)`):
+                Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
+            attention_mask (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
+                Avoid invalid areas to participate in the calculation of self-attention.
+            position_bias (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
+                Provide positional information to self-attention block.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers.
+            past_key_values (`Tuple(torch.FloatTensor)`, *optional*):
+                Cached past key and value projection states.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+        """
+        outputs = self.layernorm_before_attention(hidden_states)
+        outputs = self.self_attention(
+            outputs, outputs, attention_mask, position_bias, output_attentions, past_key_values, use_cache
+        )
+
+        outputs, attn_weights, current_key_value = outputs
+
+        if self.dropout is not None:
+            outputs = self.dropout(outputs)
+        hidden_states = (hidden_states + outputs) / 1.05
+
+        return hidden_states, attn_weights, current_key_value
+
+
+class CpmBeeDenseGatedACT(nn.Module):
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__()
+        self.w_0 = CpmBeeLinear(config.hidden_size, config.dim_ff, dtype=config.torch_dtype)
+        self.w_1 = CpmBeeLinear(config.hidden_size, config.dim_ff, dtype=config.torch_dtype)
+        self.act = torch.nn.GELU()
+
+    def forward(self, hidden_states: torch.Tensor):
+        """Transform an input tensor from one feature space to another via a nonlinear operation
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
+        """
+        gate_score = self.act(self.w_0(hidden_states))
+        hidden_states = self.w_1(hidden_states)
+
+        hidden_states = gate_score * hidden_states
+        return hidden_states
+
+
+class CpmBeeFeedForward(nn.Module):
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__()
+        self.w_in = CpmBeeDenseGatedACT(config)
+        if config.dropout_p is not None:
+            self.dropout = torch.nn.Dropout(config.dropout_p)
+        else:
+            self.dropout = None
+
+        self.w_out = CpmBeeLinear(config.dim_ff, config.hidden_size, dtype=config.torch_dtype)
+
+    def forward(self, hidden_states: torch.Tensor):
+        """
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
+        """
+        hidden_states = self.w_in(hidden_states)
+
+        if self.dropout is not None:
+            hidden_states = self.dropout(hidden_states)
+
+        hidden_states = self.w_out(hidden_states)
+
+        return hidden_states
+
+
+class CpmBeeFFNBlock(nn.Module):
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__()
+        self.layernorm_before_ffn = CpmBeeLayerNorm(config)
+        self.ffn = CpmBeeFeedForward(config)
+        if config.dropout_p:
+            self.dropout = torch.nn.Dropout(config.dropout_p)
+        else:
+            self.dropout = None
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+    ):
+        """
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch, len_seq, dim_model)`):
+                Hidden states before feed forward layer.
+        """
+        ln_outputs = self.layernorm_before_ffn(hidden_states)
+        outputs = self.ffn(ln_outputs)
+        if self.dropout is not None:
+            outputs = self.dropout(outputs)
+        hidden_states = (hidden_states + outputs) / 1.05
+        return hidden_states
+
+
+class CpmBeeTransformerBlock(nn.Module):
+    def __init__(self, config: CpmBeeConfig, mask_att: bool = False, mask_ffn: bool = False):
+        super().__init__()
+        self.mask_att = mask_att
+        self.mask_ffn = mask_ffn
+
+        if not self.mask_att:
+            self.self_att = CpmBeeSelfAttentionBlock(config)
+        if not self.mask_ffn:
+            self.ffn = CpmBeeFFNBlock(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_bias: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+        past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: Optional[bool] = None,
+    ):
+        """
+        Args:
+            hidden_states (`torch.Tensor`):
+                Input to the layer of shape `(batch, seq_len, dim_model)`
+            attention_mask (`torch.Tensor`):
+                Avoid invalid areas to participate in the calculation of shape `(batch, seq_len, seq_len)`
+            position_bias (`torch.Tensor`):
+                Provides position information to attention mechanism of shape `(num_heads, seq_len, seq_len)`
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers.
+            past_key_values (`Tuple[torch.Tensor, torch.Tensor])`, *optional*):
+                Cached past key and value projection states
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+        """
+        if not self.mask_att:
+            hidden_states = self.self_att(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_bias=position_bias,
+                output_attentions=output_attentions,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+
+            hidden_states, attn_weights, current_key_value = hidden_states
+        else:
+            attn_weights, current_key_value = None, (None, None)
+
+        if not self.mask_ffn:
+            hidden_states = self.ffn(hidden_states)
+
+        return hidden_states, attn_weights, current_key_value
+
+
+class CpmBeeEncoder(nn.Module):
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__()
+        self.num_layers = config.num_hidden_layers
+        if config.mask_modules is not None:
+            assert len(config.mask_modules) == self.num_layers, "The total number of masks should equal to num_layers"
+            for mask_module in config.mask_modules:
+                assert len(mask_module) == 2, "For encoder, each mask should be (mask_att, mask_ffn)"
+        else:
+            config.mask_modules = [(False, False)] * self.num_layers
+
+        self.layers = nn.ModuleList(
+            [
+                CpmBeeTransformerBlock(
+                    config, mask_att=config.mask_modules[ith][0], mask_ffn=config.mask_modules[ith][1]
+                )
+                for ith in range(self.num_layers)
+            ]
+        )
+
+        self.output_layernorm = CpmBeeLayerNorm(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_bias: torch.Tensor,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: Optional[bool] = None,
+    ):
+        """
+        Args:
+            hidden_states (`torch.Tensor`):
+                Input to the layer of shape `(batch, seq_len, dim_model)`
+            attention_mask (`torch.Tensor`):
+                Avoid invalid areas to participate in the calculation of shape `(batch, seq_len, seq_len)`
+            position_bias (`torch.Tensor`):
+                Provides position information to attention mechanism of shape `(num_heads, seq_len, seq_len)`
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers.
+            past_key_values (`Tuple[torch.Tensor, torch.Tensor])`, *optional*):
+                Cached past key and value projection states
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+        """
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        current_key_values = () if use_cache else None
+
+        for i, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            layer_outputs = layer(
+                hidden_states,
+                attention_mask,
+                position_bias,
+                output_attentions=output_attentions,
+                past_key_values=past_key_values[i] if past_key_values else None,
+                use_cache=use_cache,
+            )
+            hidden_states, attn_weights, current_key_value = layer_outputs
+            if output_attentions:
+                all_self_attns += (attn_weights,)
+            if current_key_value is not None:
+                current_key_values = current_key_values + (current_key_value,)
+
+        hidden_states = self.output_layernorm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        return hidden_states, current_key_values, all_hidden_states, all_self_attns
+
+
+class CpmBeeBucketPositionBias(nn.Module):
+    def __init__(self, config: CpmBeeConfig) -> None:
+        super().__init__()
+
+        self.num_heads = config.num_attention_heads
+        self.num_buckets = config.position_bias_num_buckets
+        self.num_segment_bucket = config.position_bias_num_segment_buckets
+        self.max_distance = config.position_bias_max_distance
+
+        self.relative_attention_bias = nn.Parameter(
+            torch.empty(
+                config.position_bias_num_buckets + config.position_bias_num_segment_buckets,
+                config.num_attention_heads,
+                dtype=config.torch_dtype,
+            ),
+        )
+
+    def forward(self, query_pos: torch.Tensor, key_pos: torch.Tensor, rel_buckets: torch.Tensor):
+        with torch.no_grad():
+            batch = key_pos.size(0)
+            keylen = key_pos.size(1)
+            querylen = query_pos.size(1)
+
+            if key_pos.size(0) != query_pos.size(0):
+                raise AssertionError(
+                    f"key_pos.size(0) should be equal to query_pos.size(0), but got {key_pos.size(0)} and {query_pos.size(0)}!"
+                )
+            if rel_buckets.size(0) != batch:
+                raise AssertionError(
+                    f"rel_buckets.size(0) should be equal to batch, but got {rel_buckets.size(0)} and {batch}!"
+                )
+            if rel_buckets.size(1) != querylen:
+                raise AssertionError(
+                    f"rel_buckets.size(1) should be equal to querylen, but got {rel_buckets.size(1)} and {querylen}!"
+                )
+            if rel_buckets.size(2) != keylen:
+                raise AssertionError(
+                    f"rel_buckets.size(2) should be equal to keylen, but got {rel_buckets.size(2)} and {keylen}!"
+                )
+
+            relative_position_bucket = rel_buckets - 1 + self.num_buckets
+
+            inner_segment_bucket = self._position_bucket(
+                key_pos[..., None, :] - query_pos[..., :, None],
+                num_buckets=self.num_buckets,
+                max_distance=self.max_distance,
+            )
+            relative_position_bucket = torch.where(
+                rel_buckets == 0,
+                inner_segment_bucket,
+                relative_position_bucket,
+            )
+
+        embeds = nn.functional.embedding(relative_position_bucket, self.relative_attention_bias)
+        embeds = embeds.permute(0, 3, 1, 2).contiguous()
+        return embeds
+
+    def _position_bucket(self, relative_position, num_buckets=32, max_distance=128):
+        relative_buckets = 0
+        num_buckets //= 2
+        relative_buckets = (relative_position > 0).to(torch.int32) * num_buckets
+        relative_position = torch.abs(relative_position)
+        max_exact = num_buckets // 2
+        is_small = relative_position < max_exact
+        relative_postion_if_large = max_exact + (
+            torch.log(relative_position.float() / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).to(torch.int32)
+        relative_postion_if_large = torch.min(
+            relative_postion_if_large,
+            torch.full_like(relative_postion_if_large, num_buckets - 1),
+        )
+        relative_buckets += torch.where(is_small, relative_position.to(torch.int32), relative_postion_if_large)
+        return relative_buckets
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->CPMBee
+class CpmBeeOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class CpmBeeRotaryEmbedding(nn.Module):
+    """
+    RotaryEmbedding embeds the unk token and special token. It will embeds the "...<mask>...<mask>...<unk>...<unk>..."
+    to "...<mask_0>...<mask_1>...<unk_0>...<unk_1>..."" to help model to specify different special tokens and unk
+    tokens.
+    """
+
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, config.hidden_size, 2, dtype=torch.float32) / config.hidden_size))
+        self.distance_scale = config.distance_scale
+        self.dtype = config.torch_dtype
+        self.inv_freq = inv_freq.to(config.torch_dtype)
+
+    def forward(self, x: torch.Tensor, x_pos: torch.Tensor):
+        inv_freq = self.inv_freq.to(device=x.device, dtype=self.dtype)
+
+        x_pos = x_pos * self.distance_scale
+        freqs = x_pos[..., None].to(self.dtype) * inv_freq[None, :]  # (..., dim/2)
+
+        emb = torch.cat((freqs, freqs), dim=-1)  # (..., dim)
+        emb_cos = emb.cos()  # (..., dim)
+        emb_sin = emb.sin()  # (..., dim)
+
+        rotate_x = torch.cat([-x[..., x.size(-1) // 2 :], x[..., : x.size(-1) // 2]], dim=-1)  # (..., dim)
+
+        return x * emb_cos + rotate_x * emb_sin
+
+
+class CpmBeeEmbeddingExt(nn.Embedding):
+    """
+    Contains a RotaryEmbedding.
+    """
+
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__(config.vocab_size, config.hidden_size, dtype=config.torch_dtype)
+        self.dim_model = config.hidden_size
+        self.rotary_emb = CpmBeeRotaryEmbedding(config)
+
+    def forward(self, ids: torch.Tensor, ids_sub: torch.Tensor):
+        embeds = super().forward(ids) / math.sqrt(self.dim_model)
+        return self.rotary_emb(embeds, ids_sub)
+
+    def projection(self, x: torch.Tensor, ext_table: Optional[torch.Tensor] = None):
+        logits = nn.functional.linear(x / math.sqrt(self.dim_model), self.weight)
+        if ext_table is not None:
+            logits_ext = nn.functional.linear(x, ext_table)
+            logits = torch.cat([logits, logits_ext], dim=-1)
+        return logits
+
+
+class CpmBeePreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = CpmBeeConfig
+    base_model_prefix = "cpmbee"
+    supports_gradient_checkpointing = True
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.init_std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        # still needed
+        elif isinstance(module, CpmBeeEmbeddingExt):
+            module.weight.data.normal_(mean=0.0, std=self.config.init_std)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, CpmBeeLayerNorm):
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, CpmBeeBucketPositionBias):
+            module.relative_attention_bias.data.normal_(mean=0.0, std=self.config.init_std)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, CpmBeeEncoder):
+            module.gradient_checkpointing = value
+
+
+CPMBEE_START_DOCSTRING = r"""
+    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
+    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
+    behavior.
+
+    Parameters
+        config ([`~CpmBeeConfig`]): Model configuration class with all the parameters of the
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+CPMBEE_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
+            Indices of input sequence tokens in the vocabulary.
+
+            Indices can be obtained using [`CPMBeeTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        input_id_sub (`torch.Tensor` of shape `(batch_size, seq_len)`):
+            Subscription of input sequence tokens in the vocabulary.
+
+            Subscription of normal text will be zero while the special tokens of each group will be the 0, 1, 2, ...
+            <ans_0>, <ans_1>, <ans_2> ... belongs to group <ans>. <mask_0>, <mask_1>, <mask_2> ... belongs to group
+            <mask>.
+        position (`torch.Tensor` of shape `(batch_size, seq_len)`):
+            The position of input sequence tokens in the vocabulary for each segment. if segment1 is 0, 1, 2 and
+            segment2 is 0, 1, 2, 3, the position will be 0, 1, 2, 0, 1, 2, 3
+        context (`torch.Tensor` of shape `(batch_size, seq_len)`):
+            Whether this token id is context or not. If is context, the value is 1. If not, the value is 0. If a token
+            id is context, it does not need to be predicted.
+        sample_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
+            Give a sample id to every token id. The token ids with same sample ids belongs to the same sample.
+        num_segments (`torch.Tensor` of shape `(batch_size, seq_len)`):
+            Total number of segments in the current input.
+        segment (`torch.Tensor` of shape `(batch_size, seq_len)`):
+            Give a segment id to every token id. The token ids with same segment ids belongs to the same sample.
+
+            Generally, a string key or value in input data will be a segment. For example, input {"input": "hello, ",
+            "<ans>": ""}, the segments includes: "input", "hello, ", "<ans>" and "".
+        segment_rel_offset (`torch.Tensor` of shape `(batch_size, seq_len)`):
+            The offset of segment rel.
+        segment_rel (`torch.Tensor` of shape `(batch_size, seq_len)`):
+            The segment relevance. A relative implementation of measuring the importance of segments.
+        past_states (`Dict[str, Union[torch.Tensor, List]]`):
+            Store the history information including position, context, sample_ids, num_segments, segment and
+            past_key_values.
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers.
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            A dummy arguments for CPMBee. The `past_states` contains pre-computed hidden-states (key and values in the
+            self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values` input) and
+            other history arguments to speed up sequential decoding.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare CPMBee Model outputting raw hidden-states without any specific head on top.",
+    CPMBEE_START_DOCSTRING,
+)
+class CpmBeeModel(CpmBeePreTrainedModel):
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__(config)
+        if config.half:
+            config.torch_dtype = torch.half
+        else:
+            config.torch_dtype = torch.float
+        self.encoder = CpmBeeEncoder(config)
+        self.input_embedding = CpmBeeEmbeddingExt(config)
+        self.position_bias = CpmBeeBucketPositionBias(config)
+        self.vocab_size = config.vocab_size
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.input_embedding
+
+    def set_input_embeddings(self, embeddings, **kwargs):
+        self.input_embedding = embeddings
+
+    @add_start_docstrings_to_model_forward(CPMBEE_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=BaseModelOutputWithPast,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        input_id_sub: Optional[torch.Tensor] = None,
+        position: Optional[torch.Tensor] = None,
+        context: Optional[torch.Tensor] = None,
+        sample_ids: Optional[torch.Tensor] = None,
+        num_segments: Optional[torch.Tensor] = None,
+        segment: Optional[torch.Tensor] = None,
+        segment_rel_offset: Optional[torch.Tensor] = None,
+        segment_rel: Optional[torch.Tensor] = None,
+        past_states: Optional[Dict] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        past_key_values: Optional[List] = None,
+        use_cache: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        # dummy setting for common tests
+        if input_id_sub is None:
+            dtype, device = input_ids.dtype, input_ids.device
+            batch, seq_length = input_ids.size()
+            segment = torch.where(input_ids != 0, 2, 0).to(dtype=dtype, device=device)
+            context = torch.full((batch, seq_length), 1, dtype=dtype, device=device)
+            position = torch.arange(seq_length, dtype=dtype, device=device).repeat(batch, 1)
+            input_id_sub = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
+            segment_rel_offset = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
+            segment_rel = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
+            num_segments = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
+            sample_ids = torch.zeros_like(input_ids)
+
+        with torch.no_grad():
+            if past_states 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_states["buffer_position"], position], dim=-1)
+                present_context = torch.cat([past_states["buffer_context"], context], dim=-1)
+                present_sample_ids = torch.cat([past_states["buffer_sample_ids"], sample_ids], dim=-1)
+                present_num_segments = torch.cat([past_states["buffer_num_segments"], num_segments], dim=-1)
+                present_segments = torch.cat([past_states["buffer_segments"], segment], dim=-1)
+                present_buffer = past_states["buffer"]
+
+            batch = input_ids.size(0)
+            len_q = input_ids.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_ids, input_id_sub)
+        position_bias = self.position_bias(position, present_position, segment_bucket)
+        hidden_states, present_key_values, all_hidden_states, all_attentions = self.encoder(
+            hidden_states,
+            attention_mask,
+            position_bias,
+            output_attentions,
+            output_hidden_states,
+            present_buffer,
+            use_cache,
+        )
+
+        if not return_dict:
+            return tuple(
+                v for v in [hidden_states, present_key_values, all_hidden_states, all_attentions] if v is not None
+            )
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=present_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+class CpmBeeBeamHypotheses(BeamHypotheses):
+    def __init__(self, num_beams: int, length_penalty: float, early_stopping: bool, max_length: Optional[int] = None):
+        """
+        Override BeamHypotheses for CpmBee. The hyp to add is list but not tensor.
+        """
+        super().__init__(num_beams, length_penalty, early_stopping, max_length)
+
+    def add(self, hyp: List, sum_logprobs: float, beam_indices: Optional[torch.LongTensor] = None):
+        """
+        Add a new hypothesis to the list.
+        """
+        score = sum_logprobs / (len(hyp) ** self.length_penalty)
+        if len(self) < self.num_beams or score > self.worst_score:
+            self.beams.append((score, hyp, beam_indices))
+            if len(self) > self.num_beams:
+                sorted_next_scores = sorted([(s, idx) for idx, (s, _, _) in enumerate(self.beams)])
+                del self.beams[sorted_next_scores[0][1]]
+                self.worst_score = sorted_next_scores[1][0]
+            else:
+                self.worst_score = min(score, self.worst_score)
+
+
+class CpmBeeBeamSearchScorer(BeamSearchScorer):
+    """
+    Override BeamSearchScorer for CPMBee to support:
+    1. Replace beam_tokens by beam_states, containing `idx`, `ans`, `nx_token_id`...
+    2. The `process` will update the beam_states
+    3. The `finalize` will just return the best hypotheses as a list.
+    """
+
+    def __init__(
+        self,
+        batch_size: int,
+        num_beams: int,
+        device: torch.device,
+        length_penalty: Optional[float] = 1.0,
+        do_early_stopping: Optional[Union[bool, str]] = False,
+        num_beam_hyps_to_keep: Optional[int] = 1,
+        num_beam_groups: Optional[int] = 1,
+        max_length: Optional[int] = None,
+        **model_kwargs,
+    ):
+        self.num_beams = num_beams
+        self.device = device
+        self.length_penalty = length_penalty
+        self.do_early_stopping = do_early_stopping
+        self.num_beam_hyps_to_keep = num_beam_hyps_to_keep
+        self.num_beam_groups = num_beam_groups
+        self.group_size = self.num_beams // self.num_beam_groups
+
+        self._is_init = False
+        self._beam_hyps = [
+            CpmBeeBeamHypotheses(
+                num_beams=self.num_beams,
+                length_penalty=self.length_penalty,
+                early_stopping=self.do_early_stopping,
+                max_length=max_length,
+            )
+            for _ in range(batch_size)
+        ]
+        self._done = torch.tensor([False for _ in range(batch_size)], dtype=torch.bool, device=self.device)
+
+        self.beam_states = []
+        for sent_id in range(batch_size):
+            instance_beam_states = []
+
+            for _ in range(self.num_beams):
+                instance_beam_states.append(
+                    {
+                        "idx": 0,
+                        "ans": [],
+                        "nx_token_id": 6,
+                        "nx_token_sub": 0,
+                        "nx_segment_id": model_kwargs["other_info"][sent_id]["predict_segments"][0][0],
+                        "nx_position": 0,
+                    }
+                )
+            self.beam_states.append(instance_beam_states)
+
+    def process(
+        self,
+        batch_size: int,
+        cur_len: int,
+        _next_scores: torch.FloatTensor,
+        next_scores: torch.FloatTensor,
+        next_tokens: torch.LongTensor,
+        vocab_size: Optional[int] = None,
+        pad_token_id: Optional[int] = None,
+        bos_token_id: Optional[int] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        max_length: Optional[int] = None,
+        ext_table_sub_cpu: Optional[torch.Tensor] = None,
+        ext_table_ids_cpu: Optional[torch.Tensor] = None,
+        **model_kwargs,
+    ) -> Tuple[torch.Tensor]:
+        next_beam_state = []
+        for sent_id in range(batch_size):
+            self._done[sent_id] = self._done[sent_id] or self._beam_hyps[sent_id].is_done(
+                next_scores[sent_id].max().item(), cur_len
+            )
+            if self._done[sent_id]:
+                next_beam_state.append(
+                    [
+                        (
+                            {
+                                "idx": 0,
+                                "ans": [],
+                                "nx_token_id": pad_token_id,
+                                "nx_token_sub": 0,
+                                "nx_segment_id": 0,
+                                "nx_position": 0,
+                            },
+                            0,
+                            0,
+                        )
+                    ]
+                    * self.num_beams
+                )
+                continue
+
+            next_instance_beam_states = []
+
+            for idx, value in zip(next_tokens[sent_id], next_scores[sent_id]):
+                beam_id = torch.div(idx, _next_scores.size(-1), rounding_mode="floor").item()
+                word_id = (idx % _next_scores.size(-1)).item()
+
+                curr_info = self.beam_states[sent_id][beam_id]
+                if (
+                    word_id == eos_token_id
+                    and (curr_info["idx"] + 1 == len(model_kwargs["other_info"][sent_id]["predict_segments"]))
+                ) or cur_len == max_length:
+                    self._beam_hyps[sent_id].add(
+                        self.beam_states[sent_id][beam_id]["ans"]
+                        + [
+                            (
+                                word_id,
+                                model_kwargs["other_info"][sent_id]["predict_segments"][curr_info["idx"]][1],
+                            )
+                        ],
+                        value.item(),
+                    )
+                elif word_id == eos_token_id:
+                    next_instance_beam_states.append(
+                        (
+                            {
+                                "idx": curr_info["idx"] + 1,
+                                "ans": curr_info["ans"]
+                                + [
+                                    (
+                                        word_id,
+                                        model_kwargs["other_info"][sent_id]["predict_segments"][curr_info["idx"]][1],
+                                    )
+                                ],
+                                "nx_token_id": bos_token_id,
+                                "nx_token_sub": 0,
+                                "nx_segment_id": model_kwargs["other_info"][sent_id]["predict_segments"][
+                                    curr_info["idx"] + 1
+                                ][0],
+                                "nx_position": 0,
+                            },
+                            value.item(),
+                            sent_id * self.num_beams + beam_id,
+                        )
+                    )
+
+                else:
+                    raw_word_id = word_id
+                    word_id_sub = 0
+                    if word_id >= vocab_size:
+                        word_id -= vocab_size
+                        word_id_sub = int(ext_table_sub_cpu[word_id].item())
+                        word_id = int(ext_table_ids_cpu[word_id].item())
+
+                    next_instance_beam_states.append(
+                        (
+                            {
+                                "idx": curr_info["idx"],
+                                "ans": curr_info["ans"]
+                                + [
+                                    (
+                                        raw_word_id,
+                                        model_kwargs["other_info"][sent_id]["predict_segments"][curr_info["idx"]][1],
+                                    )
+                                ],
+                                "nx_token_id": word_id,
+                                "nx_token_sub": word_id_sub,
+                                "nx_segment_id": curr_info["nx_segment_id"],
+                                "nx_position": curr_info["nx_position"] + 1,
+                            },
+                            value.item(),
+                            sent_id * self.num_beams + beam_id,
+                        )
+                    )
+
+                if len(next_instance_beam_states) == self.num_beams:
+                    break
+            assert len(next_instance_beam_states) == 0 if cur_len == max_length else self.num_beams
+            next_beam_state.append(next_instance_beam_states)
+
+        if cur_len == max_length:
+            return None
+
+        beam_reorder_idx = []
+        beam_new_scores = []
+        beam_states = []
+        for sent_id in range(batch_size):
+            instance_beam_states = []
+            for beam_id in range(self.num_beams):
+                state, value, beam_idx = next_beam_state[sent_id][beam_id]
+                beam_reorder_idx.append(beam_idx)
+                beam_new_scores.append(value)
+                instance_beam_states.append(state)
+            beam_states.append(instance_beam_states)
+        self.beam_states = beam_states
+
+        return UserDict(
+            {
+                "next_beam_scores": torch.tensor(beam_new_scores, device=self.device).view(-1),
+                "next_beam_states": beam_states,
+                "next_beam_indices": torch.tensor(beam_reorder_idx, dtype=torch.int32, device=self.device).view(-1),
+            }
+        )
+
+    def finalize(self) -> Tuple[torch.LongTensor]:
+        results = []
+        for _, hypotheses in enumerate(self._beam_hyps):
+            best_hyp = max(hypotheses.beams, key=lambda x: x[0])[1]
+            results.append(best_hyp)
+        return results
+
+    @staticmethod
+    def apply_repetition_penalty(
+        logits,
+        batch_size,
+        num_beams,
+        prev_output_tokens,
+        repetition_penalty,
+        start_idx=None,
+        end_idx=None,
+        window_size=None,
+    ):
+        # only conduct repetition penalty for the output
+        assert repetition_penalty >= 1, "repetition penalty coefficient should >= 1"
+        # repetition penalty (from CTRL paper https://arxiv.org/abs/1909.05858)
+        for i in range(batch_size * num_beams):
+            if start_idx is None or end_idx is None:
+                output_tokens = prev_output_tokens[i].tolist()
+            else:
+                if end_idx >= start_idx:
+                    if window_size:
+                        output_tokens = prev_output_tokens[i][
+                            max(start_idx, end_idx + 1 - window_size) : end_idx + 1
+                        ].tolist()
+                    else:
+                        output_tokens = prev_output_tokens[i][start_idx : end_idx + 1].tolist()
+                else:
+                    output_tokens = []
+            for previous_token in set(output_tokens):
+                # if score < 0 then repetition penalty has to
+                # multiplied to reduce the previous token probability
+                if logits[i, previous_token] < 0:
+                    logits[i, previous_token] *= repetition_penalty
+                else:
+                    logits[i, previous_token] /= repetition_penalty
+
+
+@add_start_docstrings(
+    """
+    The CPMBee Model with a language modeling head on top (linear layer with weights tied to the input embeddings).
+    """,
+    CPMBEE_START_DOCSTRING,
+)
+class CpmBeeForCausalLM(CpmBeePreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"lm_head.weight"]
+
+    def __init__(self, config: CpmBeeConfig):
+        super().__init__(config)
+        self.cpmbee = CpmBeeModel(config)
+
+        # lm_head.weight is tied to cpmbee.input_embedding.weight
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(CPMBEE_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=CausalLMOutputWithPast,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        input_id_sub: Optional[torch.Tensor] = None,
+        position: Optional[torch.Tensor] = None,
+        context: Optional[torch.Tensor] = None,
+        sample_ids: Optional[torch.Tensor] = None,
+        num_segments: Optional[torch.Tensor] = None,
+        segment: Optional[torch.Tensor] = None,
+        segment_rel_offset: Optional[torch.Tensor] = None,
+        segment_rel: Optional[torch.Tensor] = None,
+        past_states: Optional[Dict] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        past_key_values: Optional[List] = None,
+        use_cache: Optional[bool] = None,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+        ext_table_ids: Optional[torch.Tensor] = None,  # (ext_table_size) int32
+        ext_table_sub: Optional[torch.Tensor] = None,  # (ext_table_size) int32
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            input_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                Indices of input sequence tokens in the vocabulary.
+
+                Indices can be obtained using [`CPMBeeTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+                [`PreTrainedTokenizer.__call__`] for details.
+
+                [What are input IDs?](../glossary#input-ids)
+            input_id_sub (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                Subscription of input sequence tokens in the vocabulary.
+
+                Subscription of normal text will be zero while the special tokens of each group will be the 0, 1, 2,
+                ... <ans_0>, <ans_1>, <ans_2> ... belongs to group <ans>. <mask_0>, <mask_1>, <mask_2> ... belongs to
+                group <mask>.
+            position (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                The position of input sequence tokens in the vocabulary for each segment. if segment1 is 0, 1, 2 and
+                segment2 is 0, 1, 2, 3, the position will be 0, 1, 2, 0, 1, 2, 3
+            context (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                Whether this token id is context or not. If is context, the value is 1. If not, the value is 0. If a
+                token id is context, it does not need to be predicted.
+            sample_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                Give a sample id to every token id. The token ids with same sample ids belongs to the same sample.
+            num_segments (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                Total number of segments in the current input.
+            segment (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                Give a segment id to every token id. The token ids with same segment ids belongs to the same sample.
+
+                Generally, a string key or value in input data will be a segment. For example, input {"input": "hello,
+                ", "<ans>": ""}, the segments includes: "input", "hello, ", "<ans>" and "".
+            segment_rel_offset (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                The offset of segment rel.
+            segment_rel (`torch.Tensor` of shape `(batch_size, seq_len)`):
+                The segment relevance. A relative implementation of measuring the importance of segments.
+            past_states (`Dict[str, Union[torch.Tensor, List]]`):
+                Store the history information including position, context, sample_ids, num_segments, segment and
+                past_key_values.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers.
+            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+                A dummy arguments for CPMBee. The `past_states` contains pre-computed hidden-states (key and values in
+                the self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values`
+                input) and other history arguments to speed up sequential decoding.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+            ext_table_ids (`torch.Tensor`, *optional*):
+                ext_table ids for embedding projection.
+            ext_table_sub (`torch.Tensor`, *optional*):
+                ext_table subscriptions for embedding projection.
+
+        Example:
+
+        Text Generation with CpmBeeForCausalLM.
+        ```python
+        >>> from transformers import CpmBeeTokenizer, CpmBeeForCausalLM
+
+        >>> texts = {"input": "今天天气不错，", "<ans>": ""}
+        >>> model = CpmBeeForCausalLM.from_pretrained("openbmb/cpm-bee-10b")
+        >>> tokenizer = CPMBeeTokenizer.from_pretrained("openbmb/cpm-bee-10b")
+        >>> output_texts = model.generate({"input": "今天天气不错，", "<ans>": ""}, tokenizer)
+        >>> print(output_texts)
+        {'input': '今天天气不错，', '<ans>': '适合睡觉。'}
+        ```
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        model_output = self.cpmbee(
+            input_ids,
+            input_id_sub,
+            position,
+            context,
+            sample_ids,
+            num_segments,
+            segment,
+            segment_rel_offset,
+            segment_rel,
+            past_states,
+            output_attentions,
+            output_hidden_states,
+            past_key_values,
+            use_cache,
+            return_dict,
+        )
+        hidden_states = model_output.last_hidden_state if return_dict else model_output[0]
+
+        if ext_table_ids is not None:
+            ext_table = self.cpmbee.input_embedding(ext_table_ids, ext_table_sub)
+        else:
+            ext_table = None
+        logits = self.cpmbee.input_embedding.projection(hidden_states, ext_table)
+
+        loss = None
+        if labels is not None:
+            loss_func = nn.CrossEntropyLoss()
+            loss = loss_func(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + model_output[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=model_output.past_key_values,
+            hidden_states=model_output.hidden_states,
+            attentions=model_output.attentions,
+        )
+
+    def get_input_embeddings(self):
+        return self.cpmbee.input_embedding
+
+    def set_input_embeddings(self, embeddings):
+        self.cpmbee.input_embedding = embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.Tensor,
+        batch_size: int,
+        beam_scorer: CpmBeeBeamSearchScorer = None,
+        input_id_subs: Optional[torch.Tensor] = None,
+        input_pos: Optional[torch.Tensor] = None,
+        segment_ids: Optional[torch.Tensor] = None,
+        batch_ext_table_ids: Optional[torch.Tensor] = None,
+        batch_ext_table_sub: Optional[torch.Tensor] = None,
+        other_info: Optional[Dict] = None,
+        **model_kwargs,
+    ):
+        """
+        Choose the current input according to beam states.
+        """
+        # init preparation
+        context = model_kwargs.get("context")
+        sample_ids = model_kwargs.get("sample_ids")
+        segment_rel_offset = model_kwargs.get("segment_rel_offset")
+        num_segments = model_kwargs.get("num_segments")
+        segment_rel = model_kwargs.get("segment_rel")
+        past_states = model_kwargs.get("past_states", None)
+        past_key_values = model_kwargs.get("past_key_values", None)
+        _input_ids = input_ids
+
+        # update input in generation
+        if beam_scorer is not None:
+            tmp_input = []
+            tmp_input_sub = []
+            tmp_position = []
+            tmp_segment = []
+            for sent_id in range(batch_size):
+                for beam_id in range(beam_scorer.num_beams):
+                    tmp_input.append(beam_scorer.beam_states[sent_id][beam_id]["nx_token_id"])
+                    tmp_input_sub.append(beam_scorer.beam_states[sent_id][beam_id]["nx_token_sub"])
+                    tmp_position.append(beam_scorer.beam_states[sent_id][beam_id]["nx_position"])
+                    tmp_segment.append(beam_scorer.beam_states[sent_id][beam_id]["nx_segment_id"])
+
+            model_kwargs["input_id_subs"] = input_id_subs = torch.tensor(
+                tmp_input_sub, dtype=torch.int32, device=self.device
+            ).view(batch_size * beam_scorer.num_beams, 1)
+            model_kwargs["input_pos"] = input_pos = torch.tensor(
+                tmp_position, dtype=torch.int32, device=self.device
+            ).view(batch_size * beam_scorer.num_beams, 1)
+            model_kwargs["segment_ids"] = segment_ids = torch.tensor(
+                tmp_segment, dtype=torch.int32, device=self.device
+            ).view(batch_size * beam_scorer.num_beams, 1)
+            input_ids = torch.cat(
+                [
+                    input_ids,
+                    torch.tensor(tmp_input, dtype=torch.int32, device=self.device).view(
+                        batch_size * beam_scorer.num_beams, 1
+                    ),
+                ],
+                dim=-1,
+            )
+            _input_ids = input_ids[:, -1:]
+
+        return {
+            "input_ids": _input_ids,
+            "input_id_sub": input_id_subs,
+            "position": input_pos,
+            "context": context,
+            "sample_ids": sample_ids,
+            "segment_rel_offset": segment_rel_offset,
+            "segment": segment_ids,
+            "num_segments": num_segments,
+            "segment_rel": segment_rel,
+            "use_cache": True,
+            "past_key_values": past_key_values,
+            "ext_table_ids": batch_ext_table_ids,
+            "ext_table_sub": batch_ext_table_sub,
+            "past_states": past_states,
+        }, input_ids
+
+    def _update_model_kwargs_for_generation(
+        self,
+        outputs: ModelOutput,
+        model_inputs=None,
+        **model_kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Concatenate the history input and current input.
+        """
+
+        old_past_states = model_kwargs["past_states"]
+        model_kwargs["past_states"] = {
+            "buffer_position": torch.cat([old_past_states["buffer_position"], model_inputs["position"]], dim=-1),
+            "buffer_context": torch.cat([old_past_states["buffer_context"], model_inputs["context"]], dim=-1),
+            "buffer_sample_ids": torch.cat([old_past_states["buffer_sample_ids"], model_inputs["sample_ids"]], dim=-1),
+            "buffer_num_segments": torch.cat(
+                [old_past_states["buffer_num_segments"], model_inputs["num_segments"]], dim=-1
+            ),
+            "buffer_segments": torch.cat([old_past_states["buffer_segments"], model_inputs["segment"]], dim=-1),
+            "buffer": outputs.past_key_values,
+        }
+
+        return model_kwargs
+
+    def _reorder_cache(self, past_key_values: Dict, beam_idx: torch.Tensor):
+        beam_idx = beam_idx.tolist()
+        for kw in past_key_values.keys():
+            if kw == "buffer":
+                buf_list = past_key_values[kw]
+                nw_buf_list = []
+                for buf in buf_list:
+                    if buf == (None, None):
+                        nw_buf_list.append((None, None))
+                    else:
+                        k_buf, v_buf = buf
+                        nw_buf_list.append((k_buf[beam_idx, :], v_buf[beam_idx, :]))
+                past_key_values[kw] = nw_buf_list
+            else:
+                past_key_values[kw] = past_key_values[kw][beam_idx, :]
+
+        return past_key_values
+
+    @staticmethod
+    def _expand_inputs_for_generation(
+        expand_size: int = 1,
+        is_encoder_decoder: bool = False,
+        input_ids: Optional[torch.LongTensor] = None,
+        **model_kwargs,
+    ) -> Tuple[torch.LongTensor, Dict[str, Any]]:
+        """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]"""
+
+        # do not expand ext_table_ids and ext_table_sub
+        def _expand_dict_for_generation(dict_to_expand):
+            for key in dict_to_expand:
+                if (
+                    dict_to_expand[key] is not None
+                    and isinstance(dict_to_expand[key], torch.Tensor)
+                    and "ext_table" not in key
+                ):
+                    dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+            return dict_to_expand
+
+        if input_ids is not None:
+            input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+
+        model_kwargs = _expand_dict_for_generation(model_kwargs)
+
+        if is_encoder_decoder:
+            if model_kwargs.get("encoder_outputs") is None:
+                raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+            model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+
+        return input_ids, model_kwargs
+
+    def adjust_logits_during_generation(
+        self,
+        logits: torch.FloatTensor,
+        batch_size: int,
+        beam_size: int,
+        vocab_size: int,
+        ext_table_ids: torch.Tensor,
+        **model_kwargs,
+    ) -> torch.FloatTensor:
+        """
+        Implement in subclasses of [`PreTrainedModel`] for custom behavior to adjust the logits in the generate method.
+        """
+        for sent_id in range(batch_size):
+            if 1 not in model_kwargs["other_info"][sent_id]["ext_table"]:
+                # unk is not allowed, mask unk
+                logits[sent_id * beam_size : (sent_id + 1) * beam_size, 1] = -10000
+            ext_ids = set()
+            for v in model_kwargs["other_info"][sent_id]["ext_table"].keys():
+                ext_ids.add(v)
+            for ext_id in range(vocab_size, vocab_size + ext_table_ids.size(0)):
+                if ext_id not in ext_ids:
+                    logits[sent_id * beam_size : (sent_id + 1) * beam_size, ext_id] = -10000
+        return logits
+
+    def beam_search(
+        self,
+        input_ids: torch.LongTensor,
+        beam_scorer: CpmBeeBeamSearchScorer,
+        repetition_penalty: Optional[float] = 1.0,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        max_length: Optional[int] = None,
+        pad_token_id: Optional[int] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        bos_token_id: Optional[Union[int, List[int]]] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_scores: Optional[bool] = None,
+        return_dict_in_generate: Optional[bool] = None,
+        synced_gpus: bool = False,
+        **model_kwargs,
+    ) -> List:
+        """
+        Override the beam_search for CPMBee.
+        """
+        # init values
+        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+        pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
+        eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
+        bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id
+        max_length = max_length if max_length is not None else self.generation_config.max_length
+        output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.generation_config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
+        )
+        return_dict_in_generate = (
+            return_dict_in_generate
+            if return_dict_in_generate is not None
+            else self.generation_config.return_dict_in_generate
+        )
+
+        batch_size = len(beam_scorer._beam_hyps)
+        num_beams = beam_scorer.num_beams
+
+        batch_beam_size, cur_len = input_ids.shape
+
+        if num_beams * batch_size != batch_beam_size:
+            raise ValueError(
+                f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
+            )
+
+        # init attention / hidden states / scores tuples
+        scores = () if (return_dict_in_generate and output_scores) else None
+        beam_indices = (
+            tuple(() for _ in range(batch_beam_size)) if (return_dict_in_generate and output_scores) else None
+        )
+        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+        cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+
+        # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens
+        # of the first beam are considered to avoid sampling the exact same tokens across all beams.
+        beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=self.device)
+        beam_scores[:, 1:] = -1e9
+        beam_scores = beam_scores.view((batch_size * num_beams,))
+
+        this_peer_finished = False  # used by synced_gpus only
+
+        # init inference
+        model_inputs, input_ids = self.prepare_inputs_for_generation(input_ids, batch_size, **model_kwargs)
+        pred_start_index = input_ids.size(-1)
+        outputs = self(
+            **model_inputs,
+            return_dict=True,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+
+        # update model_kwargs
+        model_kwargs["past_states"] = {
+            "buffer_position": model_inputs["position"],
+            "buffer_context": model_inputs["context"],
+            "buffer_sample_ids": model_inputs["sample_ids"],
+            "buffer_num_segments": model_inputs["num_segments"],
+            "buffer_segments": model_inputs["segment"],
+            "buffer": outputs.past_key_values,
+        }
+        model_kwargs["context"] = torch.ones(batch_beam_size, dtype=torch.bool, device=self.device).view(
+            batch_beam_size, 1
+        )
+        model_kwargs["sample_ids"] = torch.zeros(batch_beam_size, dtype=torch.int32, device=self.device).view(
+            batch_beam_size, 1
+        )
+        model_kwargs["num_segments"] = model_kwargs["num_segments"][:, -1:]
+        model_kwargs["segment_rel_offset"] = model_kwargs["segment_rel_offset"][:, -1:]
+        model_kwargs["past_key_values"] = outputs.past_key_values
+
+        ext_table_ids_cpu = model_inputs["ext_table_ids"].cpu()
+        ext_table_sub_cpu = model_inputs["ext_table_sub"].cpu()
+
+        cur_len = 0
+        while True:
+            if synced_gpus:
+                # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
+                # The following logic allows an early break if all peers finished generating their sequence
+                this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)
+                # send 0.0 if we finished, 1.0 otherwise
+                dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
+                # did all peers finish? the reduced sum will be 0.0 then
+                if this_peer_finished_flag.item() == 0.0:
+                    break
+
+            model_inputs, input_ids = self.prepare_inputs_for_generation(
+                input_ids, batch_size, beam_scorer, **model_kwargs
+            )
+
+            outputs = self(
+                **model_inputs,
+                return_dict=True,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+            )
+
+            next_token_logits = outputs.logits[:, -1, :]
+
+            if all(beam_scorer._done):
+                break
+            # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id`
+            # cannot be generated both before and after the `nn.functional.log_softmax` operation.
+            vocab_size = next_token_logits.shape[-1]
+            next_token_logits = self.adjust_logits_during_generation(
+                next_token_logits, batch_size, num_beams, vocab_size, ext_table_ids_cpu, **model_kwargs
+            )
+
+            # repetition_penalty
+            beam_scorer.apply_repetition_penalty(
+                next_token_logits,
+                batch_size,
+                num_beams,
+                model_inputs["input_ids"],
+                repetition_penalty,
+                pred_start_index,
+                model_inputs["input_ids"].size(-1) - 1,
+                None,
+            )
+
+            _next_token_scores = nn.functional.log_softmax(
+                next_token_logits, dim=-1
+            )  # (batch_size * num_beams, vocab_size)
+
+            next_token_scores_processed = logits_processor(input_ids, _next_token_scores)
+            # next_token_scores_processed = _next_token_scores
+            next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(_next_token_scores)
+
+            # Store scores, attentions and hidden_states when required
+            if return_dict_in_generate:
+                if output_scores:
+                    scores += (next_token_scores_processed,)
+                if output_attentions:
+                    decoder_attentions += (
+                        (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
+                    )
+                    if self.config.is_encoder_decoder:
+                        cross_attentions += (outputs.cross_attentions,)
+
+                if output_hidden_states:
+                    decoder_hidden_states += (
+                        (outputs.decoder_hidden_states,)
+                        if self.config.is_encoder_decoder
+                        else (outputs.hidden_states,)
+                    )
+
+            # reshape for beam search
+            vocab_size = next_token_scores.shape[-1]
+            next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)
+
+            # Sample 2 next tokens for each beam (so we have some spare tokens and match output of beam search)
+            next_token_scores, next_tokens = torch.topk(
+                next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True
+            )
+
+            beam_outputs = beam_scorer.process(
+                batch_size,
+                cur_len,
+                _next_token_scores,
+                next_token_scores,
+                next_tokens,
+                vocab_size=vocab_size,
+                pad_token_id=pad_token_id,
+                bos_token_id=bos_token_id,
+                eos_token_id=eos_token_id,
+                max_length=max_length,
+                ext_table_ids_cpu=ext_table_ids_cpu,
+                ext_table_sub_cpu=ext_table_sub_cpu,
+                **model_kwargs,
+            )
+            if beam_outputs is None:
+                break
+            beam_idx = beam_outputs["next_beam_indices"]
+            beam_scores = beam_outputs["next_beam_scores"]
+
+            input_ids = input_ids[beam_idx.tolist(), :]
+            model_kwargs = self._update_model_kwargs_for_generation(outputs, model_inputs, **model_kwargs)
+            if model_kwargs["past_states"] is not None:
+                model_kwargs["past_states"] = self._reorder_cache(model_kwargs["past_states"], beam_idx)
+
+            if return_dict_in_generate and output_scores:
+                beam_indices = tuple((beam_indices[beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices))))
+
+            cur_len += 1
+
+            if beam_scorer.is_done or cur_len == max_length + 1:
+                if not synced_gpus:
+                    break
+                else:
+                    this_peer_finished = True
+
+        sequence_outputs = beam_scorer.finalize()
+
+        return sequence_outputs
+
+    def _generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        repetition_penalty: Optional[float] = 1.0,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        stopping_criteria: Optional[StoppingCriteriaList] = None,
+        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
+        synced_gpus: Optional[bool] = None,
+        streamer: Optional["BaseStreamer"] = None,
+        **kwargs,
+    ) -> List:
+        r"""
+        The generation of CPMBee.
+        1. It will use beam search as generation strategy.
+        2. It will use CpmBeeBeamSearchScorer as the beamsearch scorer.
+        """
+        if synced_gpus is None:
+            if is_deepspeed_zero3_enabled() and dist.get_world_size() > 1:
+                synced_gpus = True
+            else:
+                synced_gpus = False
+
+        # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
+        self._validate_model_class()
+
+        # priority: `generation_config` argument > `model.generation_config` (the default generation config)
+        if generation_config is None:
+            # legacy: users may modify the model configuration to control generation -- update the generation config
+            # model attribute accordingly, if it was created from the model config
+            if self.generation_config._from_model_config:
+                new_generation_config = GenerationConfig.from_model_config(self.config)
+                if new_generation_config != self.generation_config:
+                    warnings.warn(
+                        "You have modified the pretrained model configuration to control generation. This is a"
+                        " deprecated strategy to control generation and will be removed soon, in a future version."
+                        " Please use a generation configuration file (see"
+                        " https://huggingface.co/docs/transformers/main_classes/text_generation)"
+                    )
+                    self.generation_config = new_generation_config
+            generation_config = self.generation_config
+
+        generation_config = copy.deepcopy(generation_config)
+        model_kwargs = generation_config.update(**kwargs)  # All unused kwargs must be model kwargs
+        generation_config.validate()
+        self._validate_model_kwargs(model_kwargs.copy())
+
+        # 2. Set generation parameters if not already defined
+        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+        if generation_config.pad_token_id is None and generation_config.eos_token_id is not None:
+            if model_kwargs.get("attention_mask", None) is None:
+                logger.warning(
+                    "The attention mask and the pad token id were not set. As a consequence, you may observe "
+                    "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results."
+                )
+            eos_token_id = generation_config.eos_token_id
+            if isinstance(eos_token_id, list):
+                eos_token_id = eos_token_id[0]
+            logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{eos_token_id} for open-end generation.")
+            generation_config.pad_token_id = eos_token_id
+
+        # 3. Define model inputs
+        # inputs_tensor has to be defined
+        # model_input_name is defined if model-specific keyword input is passed
+        # otherwise model_input_name is None
+        # all model-specific keyword inputs are removed from `model_kwargs`
+        inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
+            inputs, generation_config.bos_token_id, model_kwargs
+        )
+        batch_size = inputs_tensor.shape[0]
+
+        # 4. Define other model kwargs
+        model_kwargs["output_attentions"] = generation_config.output_attentions
+        model_kwargs["output_hidden_states"] = generation_config.output_hidden_states
+        model_kwargs["use_cache"] = generation_config.use_cache
+
+        accepts_attention_mask = "attention_mask" in set(inspect.signature(self.forward).parameters.keys())
+        requires_attention_mask = "encoder_outputs" not in model_kwargs
+
+        if model_kwargs.get("attention_mask", None) is None and requires_attention_mask and accepts_attention_mask:
+            model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
+                inputs_tensor, generation_config.pad_token_id, generation_config.eos_token_id
+            )
+
+        # decoder-only models should use left-padding for generation
+        if not self.config.is_encoder_decoder:
+            # If `input_ids` was given, check if the last id in any sequence is `pad_token_id`
+            # Note: If using, `inputs_embeds` this check does not work, because we want to be more hands-off.
+            if (
+                generation_config.pad_token_id is not None
+                and len(inputs_tensor.shape) == 2
+                and torch.sum(inputs_tensor[:, -1] == generation_config.pad_token_id) > 0
+            ):
+                logger.warning(
+                    "A decoder-only architecture is being used, but right-padding was detected! For correct "
+                    "generation results, please set `padding_side='left'` when initializing the tokenizer."
+                )
+
+        # 5. Prepare `input_ids` which will be used for auto-regressive generation
+        input_ids = inputs_tensor if model_input_name == "input_ids" else model_kwargs.pop("input_ids")
+
+        if streamer is not None:
+            streamer.put(input_ids.cpu())
+
+        # 6. Prepare `max_length` depending on other stopping criteria.
+        input_ids_seq_length = input_ids.shape[-1]
+        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+        if has_default_max_length and generation_config.max_new_tokens is None:
+            warnings.warn(
+                f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
+                "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
+                " recommend using `max_new_tokens` to control the maximum length of the generation.",
+                UserWarning,
+            )
+        elif generation_config.max_new_tokens is not None:
+            if not has_default_max_length:
+                logger.warning(
+                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
+                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
+                    "Please refer to the documentation for more information. "
+                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
+                )
+            generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
+
+        if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length:
+            raise ValueError(
+                f"Unfeasible length constraints: the minimum length ({generation_config.min_length}) is larger than"
+                f" the maximum length ({generation_config.max_length})"
+            )
+        if input_ids_seq_length >= generation_config.max_length:
+            input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
+            logger.warning(
+                f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
+                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
+                " increasing `max_new_tokens`."
+            )
+
+        if streamer is not None and (generation_config.num_beams > 1):
+            raise ValueError(
+                "`streamer` cannot be used with beam search (yet!). Make sure that `num_beams` is set to 1."
+            )
+
+        if self.device.type != input_ids.device.type:
+            warnings.warn(
+                "You are calling .generate() with the `input_ids` being on a device type different"
+                f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
+                f" is on {self.device.type}. You may experience unexpected behaviors or slower generation."
+                " Please make sure that you have put `input_ids` to the"
+                f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"
+                " running `.generate()`.",
+                UserWarning,
+            )
+
+        # 7. prepare distribution pre_processing samplers
+        logits_processor = self._get_logits_processor(
+            generation_config=generation_config,
+            input_ids_seq_length=input_ids_seq_length,
+            encoder_input_ids=inputs_tensor,
+            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+            logits_processor=logits_processor,
+        )
+
+        # 8. prepare beam search scorer
+        beam_scorer = CpmBeeBeamSearchScorer(
+            batch_size=batch_size,
+            num_beams=generation_config.num_beams,
+            device=inputs_tensor.device,
+            length_penalty=generation_config.length_penalty,
+            do_early_stopping=generation_config.early_stopping,
+            num_beam_hyps_to_keep=generation_config.num_return_sequences,
+            max_length=generation_config.max_length,
+            **kwargs,
+        )
+        # 9. interleave input_ids with `num_beams` additional sequences per batch
+        input_ids, model_kwargs = self._expand_inputs_for_generation(
+            input_ids=input_ids,
+            expand_size=generation_config.num_beams,
+            is_encoder_decoder=self.config.is_encoder_decoder,
+            **model_kwargs,
+        )
+        # 10. run beam search
+        return self.beam_search(
+            input_ids,
+            beam_scorer,
+            repetition_penalty=repetition_penalty,
+            logits_processor=logits_processor,
+            pad_token_id=generation_config.pad_token_id,
+            eos_token_id=generation_config.eos_token_id,
+            output_scores=generation_config.output_scores,
+            return_dict_in_generate=generation_config.return_dict_in_generate,
+            synced_gpus=synced_gpus,
+            **model_kwargs,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        data_list: Union[Dict, List[Dict]],
+        tokenizer: CpmBeeTokenizer,
+        generation_config=None,
+        **kwargs,
+    ):
+        """
+        Override the generate for CPMBee. It will accept dict or list(dict) as input and returns dict or list(dict)
+        with `<ans>` filled.
+
+        Parameters:
+            data_list (`dict` or `list(dict)`):
+                The sequence used as a prompt for the generation or as model inputs to the encoder. If dict, data_list
+                will be wrapped as a list.
+            tokenizer: (`CpmBeeTokenizer`):
+                The tokenizer.
+            generation_config (`~generation.GenerationConfig`, *optional*):
+                The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+                passed to generate matching the attributes of `generation_config` will override them. If
+                `generation_config` is not provided, the default will be used, which had the following loading
+                priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+                configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+                default values, whose documentation should be checked to parameterize generation.
+        """
+        if isinstance(data_list, dict):
+            data_list = [data_list]
+        input_encoded = tokenizer(data_list, return_tensors="pt", padding=True, device=self.device)
+        input_encoded.update(kwargs)
+        input_encoded["generation_config"] = generation_config
+
+        decode_res = self._generate(**input_encoded)
+
+        for sent_id, result in enumerate(decode_res):
+            ans_result_map: Dict[int, List[int]] = {}
+            for raw_word_id, ans_id in result:
+                if ans_id not in ans_result_map:
+                    ans_result_map[ans_id] = []
+                ans_result_map[ans_id].append(raw_word_id)
+
+            answer_placeholders = input_encoded["other_info"][sent_id]["answer_placeholders"]
+            ext_table = input_encoded["other_info"][sent_id]["ext_table"]
+            data = data_list[sent_id]
+            for ans_id, token_ids in ans_result_map.items():
+                if token_ids[-1] == tokenizer.eos_token_id:
+                    token_ids = token_ids[:-1]
+                text = tokenizer.decode(token_ids, ext_table)
+                path = answer_placeholders[ans_id - 1]
+
+                if len(path) > 0:
+                    p = data["<ans>"]
+                    for part in path[:-1]:
+                        p = p[part]
+                    p[path[-1]] = text
+                else:
+                    data["<ans>"] = text
+            for ans_id in range(len(answer_placeholders)):
+                if (ans_id + 1) not in ans_result_map:
+                    path = answer_placeholders[ans_id]
+                    p = data["<ans>"]
+                    for part in path[:-1]:
+                        p = p[part]
+                    p[path[-1]] = None
+        return data_list

test_modeling_cpmbee.py

@@ -0,0 +1,183 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Testing suite for the PyTorch CpmBee model. """
+
+
+import unittest
+
+from transformers.testing_utils import is_torch_available, require_torch, tooslow
+
+from ...generation.test_utils import torch_device
+from ...test_configuration_common import ConfigTester
+from ...test_modeling_common import ModelTesterMixin, ids_tensor
+from ...test_pipeline_mixin import PipelineTesterMixin
+
+
+if is_torch_available():
+    import torch
+
+    from transformers import (
+        CpmBeeConfig,
+        CpmBeeForCausalLM,
+        CpmBeeModel,
+        CpmBeeTokenizer,
+    )
+
+
+@require_torch
+class CpmBeeModelTester:
+    def __init__(
+        self,
+        parent,
+        batch_size=2,
+        seq_length=8,
+        is_training=True,
+        use_token_type_ids=False,
+        use_input_mask=False,
+        use_labels=False,
+        use_mc_token_ids=False,
+        vocab_size=99,
+        hidden_size=32,
+        num_hidden_layers=3,
+        num_attention_heads=4,
+        intermediate_size=37,
+        num_buckets=32,
+        max_distance=128,
+        position_bias_num_segment_buckets=32,
+        init_std=1.0,
+        return_dict=True,
+    ):
+        self.parent = parent
+        self.batch_size = batch_size
+        self.seq_length = seq_length
+        self.is_training = is_training
+        self.use_token_type_ids = use_token_type_ids
+        self.use_input_mask = use_input_mask
+        self.use_labels = use_labels
+        self.use_mc_token_ids = use_mc_token_ids
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.position_bias_num_segment_buckets = position_bias_num_segment_buckets
+        self.init_std = init_std
+        self.return_dict = return_dict
+
+    def prepare_config_and_inputs(self):
+        input_ids = {}
+        input_ids["input_ids"] = ids_tensor([self.batch_size, self.seq_length], self.vocab_size).type(torch.int32)
+        input_ids["use_cache"] = False
+
+        config = self.get_config()
+
+        return (config, input_ids)
+
+    def get_config(self):
+        return CpmBeeConfig(
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            dim_ff=self.intermediate_size,
+            position_bias_num_buckets=self.num_buckets,
+            position_bias_max_distance=self.max_distance,
+            position_bias_num_segment_buckets=self.position_bias_num_segment_buckets,
+            use_cache=True,
+            init_std=self.init_std,
+            return_dict=self.return_dict,
+        )
+
+    def create_and_check_cpmbee_model(self, config, input_ids, *args):
+        model = CpmBeeModel(config=config)
+        model.to(torch_device)
+        model.eval()
+
+        hidden_states = model(**input_ids).last_hidden_state
+
+        self.parent.assertEqual(hidden_states.shape, (self.batch_size, self.seq_length, config.hidden_size))
+
+    def create_and_check_lm_head_model(self, config, input_ids, *args):
+        model = CpmBeeForCausalLM(config)
+        model.to(torch_device)
+        input_ids["input_ids"] = input_ids["input_ids"].to(torch_device)
+        model.eval()
+
+        model_output = model(**input_ids)
+        self.parent.assertEqual(
+            model_output.logits.shape,
+            (self.batch_size, self.seq_length, config.vocab_size),
+        )
+
+    def prepare_config_and_inputs_for_common(self):
+        config, inputs_dict = self.prepare_config_and_inputs()
+        return config, inputs_dict
+
+
+@require_torch
+class CpmBeeModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
+    all_model_classes = (CpmBeeModel, CpmBeeForCausalLM) if is_torch_available() else ()
+    pipeline_model_mapping = (
+        {"feature-extraction": CpmBeeModel, "text-generation": CpmBeeForCausalLM} if is_torch_available() else {}
+    )
+
+    test_pruning = False
+    test_missing_keys = False
+    test_mismatched_shapes = False
+    test_head_masking = False
+    test_resize_embeddings = False
+
+    def setUp(self):
+        self.model_tester = CpmBeeModelTester(self)
+        self.config_tester = ConfigTester(self, config_class=CpmBeeConfig)
+
+    def test_config(self):
+        self.config_tester.create_and_test_config_common_properties()
+        self.config_tester.create_and_test_config_to_json_string()
+        self.config_tester.create_and_test_config_to_json_file()
+        self.config_tester.create_and_test_config_from_and_save_pretrained()
+        self.config_tester.check_config_can_be_init_without_params()
+        self.config_tester.check_config_arguments_init()
+
+    def test_inputs_embeds(self):
+        unittest.skip("CPMBee doesn't support input_embeds.")(self.test_inputs_embeds)
+
+    def test_retain_grad_hidden_states_attentions(self):
+        unittest.skip(
+            "CPMBee doesn't support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
+                 So is attentions. We strongly recommand you use loss to tune model."
+        )(self.test_retain_grad_hidden_states_attentions)
+
+    def test_cpmbee_model(self):
+        config, inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_cpmbee_model(config, inputs)
+
+    def test_cpmbee_lm_head_model(self):
+        config, inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_lm_head_model(config, inputs)
+
+
+@require_torch
+class CpmBeeForCausalLMlIntegrationTest(unittest.TestCase):
+    @tooslow
+    def test_simple_generation(self):
+        texts = {"input": "今天天气不错，", "<ans>": ""}
+        model = CpmBeeForCausalLM.from_pretrained("openbmb/cpm-bee-10b")
+        tokenizer = CpmBeeTokenizer.from_pretrained("openbmb/cpm-bee-10b")
+        output_texts = model.generate(texts, tokenizer)
+        expected_output = {"input": "今天天气不错，", "<ans>": "适合睡觉。"}
+        self.assertEqual(expected_output["<ans>"], output_texts["<ans>"])

test_tokenization_cpmbee.py

@@ -0,0 +1,187 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Testing suite for the PyTorch CpmBee tokenizer. """
+
+import os
+import unittest
+
+from transformers.models.cpmbee.tokenization_cpmbee import VOCAB_FILES_NAMES, CpmBeeTokenizer
+from transformers.tokenization_utils import AddedToken
+
+from ...test_tokenization_common import TokenizerTesterMixin
+
+
+class CPMBeeTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
+    tokenizer_class = CpmBeeTokenizer
+    test_rust_tokenizer = False
+
+    def setUp(self):
+        super().setUp()
+
+        vocab_tokens = [
+            "<d>",
+            "</d>",
+            "<s>",
+            "</s>",
+            "</_>",
+            "<unk>",
+            "<pad>",
+            "<mask>",
+            "</n>",
+            "我",
+            "是",
+            "C",
+            "P",
+            "M",
+            "B",
+            "e",
+            "e",
+        ]
+        self.vocab_file = os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["vocab_file"])
+        vocab_tokens = list(set(vocab_tokens))
+        with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
+            vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
+
+    # override test_add_tokens_tokenizer because <...> is special token in CpmBeeTokenizer.
+    def test_add_tokens_tokenizer(self):
+        tokenizers = self.get_tokenizers(do_lower_case=False)
+        for tokenizer in tokenizers:
+            with self.subTest(f"{tokenizer.__class__.__name__}"):
+                vocab_size = tokenizer.vocab_size
+                all_size = len(tokenizer)
+
+                self.assertNotEqual(vocab_size, 0)
+
+                # We usually have added tokens from the start in tests because our vocab fixtures are
+                # smaller than the original vocabs - let's not assert this
+                # self.assertEqual(vocab_size, all_size)
+
+                new_toks = ["aaaaa bbbbbb", "cccccccccdddddddd"]
+                added_toks = tokenizer.add_tokens(new_toks)
+                vocab_size_2 = tokenizer.vocab_size
+                all_size_2 = len(tokenizer)
+
+                self.assertNotEqual(vocab_size_2, 0)
+                self.assertEqual(vocab_size, vocab_size_2)
+                self.assertEqual(added_toks, len(new_toks))
+                self.assertEqual(all_size_2, all_size + len(new_toks))
+
+                tokens = tokenizer.encode("aaaaa bbbbbb low cccccccccdddddddd l", add_special_tokens=False)
+
+                self.assertGreaterEqual(len(tokens), 4)
+                self.assertGreater(tokens[0], tokenizer.vocab_size - 1)
+                self.assertGreater(tokens[-2], tokenizer.vocab_size - 1)
+
+                new_toks_2 = {"eos_token": ">>>>|||<||<<|<<", "pad_token": "<<<<<|||;;;||;"}
+                added_toks_2 = tokenizer.add_special_tokens(new_toks_2)
+                vocab_size_3 = tokenizer.vocab_size
+                all_size_3 = len(tokenizer)
+
+                self.assertNotEqual(vocab_size_3, 0)
+                self.assertEqual(vocab_size, vocab_size_3)
+                self.assertEqual(added_toks_2, len(new_toks_2))
+                self.assertEqual(all_size_3, all_size_2 + len(new_toks_2))
+
+                tokens = tokenizer.encode(
+                    ">>>>|||<||<<|<< aaaaabbbbbb low cccccccccdddddddd <<<<<|||;;;||; l", add_special_tokens=False
+                )
+
+                self.assertGreaterEqual(len(tokens), 6)
+                self.assertGreater(tokens[0], tokenizer.vocab_size - 1)
+                self.assertGreater(tokens[0], tokens[1])
+                self.assertGreater(tokens[-2], tokenizer.vocab_size - 1)
+                self.assertGreater(tokens[-2], tokens[-3])
+                self.assertEqual(tokens[0], tokenizer.eos_token_id)
+                self.assertEqual(tokens[-2], tokenizer.pad_token_id)
+
+    def test_added_tokens_do_lower_case(self):
+        tokenizers = self.get_tokenizers(do_lower_case=True)
+        for tokenizer in tokenizers:
+            with self.subTest(f"{tokenizer.__class__.__name__}"):
+                if not hasattr(tokenizer, "do_lower_case") or not tokenizer.do_lower_case:
+                    continue
+
+                special_token = tokenizer.all_special_tokens[0]
+
+                text = special_token + " aaaaa bbbbbb low cccccccccdddddddd l " + special_token
+                text2 = special_token + " AAAAA BBBBBB low CCCCCCCCCDDDDDDDD l " + special_token
+
+                toks_before_adding = tokenizer.tokenize(text)  # toks before adding new_toks
+
+                new_toks = ["aaaaa bbbbbb", "cccccccccdddddddd", "AAAAA BBBBBB", "CCCCCCCCCDDDDDDDD"]
+                added = tokenizer.add_tokens([AddedToken(tok, lstrip=True, rstrip=True) for tok in new_toks])
+
+                toks_after_adding = tokenizer.tokenize(text)
+                toks_after_adding2 = tokenizer.tokenize(text2)
+
+                # Rust tokenizers dont't lowercase added tokens at the time calling `tokenizer.add_tokens`,
+                # while python tokenizers do, so new_toks 0 and 2 would be treated as the same, so do new_toks 1 and 3.
+                self.assertIn(added, [2, 4])
+
+                self.assertListEqual(toks_after_adding, toks_after_adding2)
+                self.assertTrue(
+                    len(toks_before_adding) > len(toks_after_adding),  # toks_before_adding should be longer
+                )
+
+                # Check that none of the special tokens are lowercased
+                sequence_with_special_tokens = "A " + " yEs ".join(tokenizer.all_special_tokens) + " B"
+                # Convert the tokenized list to str as some special tokens are tokenized like normal tokens
+                # which have a prefix spacee e.g. the mask token of Albert, and cannot match the original
+                # special tokens exactly.
+                tokenized_sequence = "".join(tokenizer.tokenize(sequence_with_special_tokens))
+
+                for special_token in tokenizer.all_special_tokens:
+                    self.assertTrue(special_token in tokenized_sequence)
+
+        tokenizers = self.get_tokenizers(do_lower_case=True)
+        for tokenizer in tokenizers:
+            with self.subTest(f"{tokenizer.__class__.__name__}"):
+                if hasattr(tokenizer, "do_lower_case") and tokenizer.do_lower_case:
+                    continue
+
+                special_token = tokenizer.all_special_tokens[0]
+
+                text = special_token + " aaaaa bbbbbb low cccccccccdddddddd l " + special_token
+                text2 = special_token + " AAAAA BBBBBB low CCCCCCCCCDDDDDDDD l " + special_token
+
+                toks_before_adding = tokenizer.tokenize(text)  # toks before adding new_toks
+
+                new_toks = ["aaaaa bbbbbb", "cccccccccdddddddd", "AAAAA BBBBBB", "CCCCCCCCCDDDDDDDD"]
+                added = tokenizer.add_tokens([AddedToken(tok, lstrip=True, rstrip=True) for tok in new_toks])
+                self.assertIn(added, [2, 4])
+
+                toks_after_adding = tokenizer.tokenize(text)
+                toks_after_adding2 = tokenizer.tokenize(text2)
+
+                self.assertEqual(len(toks_after_adding), len(toks_after_adding2))  # Length should still be the same
+                self.assertNotEqual(
+                    toks_after_adding[1], toks_after_adding2[1]
+                )  # But at least the first non-special tokens should differ
+                self.assertTrue(
+                    len(toks_before_adding) > len(toks_after_adding),  # toks_before_adding should be longer
+                )
+
+    def test_pre_tokenization(self):
+        tokenizer = CpmBeeTokenizer.from_pretrained("openbmb/cpm-bee-10b")
+        texts = {"input": "你好，", "<ans>": ""}
+        tokens = tokenizer(texts)
+        tokens = tokens["input_ids"][0]
+
+        input_tokens = [6, 8, 7, 6, 65678, 7, 6, 10273, 246, 7, 6, 9, 7]
+        self.assertListEqual(tokens, input_tokens)
+
+        normalized_text = "<s><root></s><s>input</s><s>你好，</s><s><ans></s>"
+        reconstructed_text = tokenizer.decode(tokens)
+        self.assertEqual(reconstructed_text, normalized_text)

tokenization_cpmbee.py

@@ -0,0 +1,868 @@
+# coding=utf-8
+# Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes for CpmBee."""
+import json
+import os
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+from typing_extensions import TypedDict
+
+from ...tokenization_utils import PaddingStrategy, PreTrainedTokenizer, TensorType
+from ...tokenization_utils_base import AddedToken, BatchEncoding, TextInput, TruncationStrategy
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "openbmb/cpm-bee-10b": "https://huggingface.co/openbmb/cpm-bee-10b/blob/main/vocab.txt",
+        "openbmb/cpm-bee-5b": "https://huggingface.co/openbmb/cpm-bee-5b/blob/main/vocab.txt",
+        "openbmb/cpm-bee-2b": "https://huggingface.co/openbmb/cpm-bee-2b/blob/main/vocab.txt",
+        "openbmb/cpm-bee-1b": "https://huggingface.co/openbmb/cpm-bee-1b/blob/main/vocab.txt",
+    },
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "openbmb/cpm-bee-10b": 4096,
+    "openbmb/cpm-bee-5b": 4096,
+    "openbmb/cpm-bee-2b": 4096,
+    "openbmb/cpm-bee-1b": 4096,
+}
+
+
+class _PrevExtTableStates(TypedDict):
+    ext_table: Dict[int, str]
+    token_id_table: Dict[str, Dict[int, int]]
+
+
+CPMBeeInputType = Union[str, Dict[str, "CPMBeeInputType"]]
+
+
+def rel_to_bucket(n_up: int, n_down: int, max_depth: int = 8):
+    ret = n_up * max_depth + n_down
+    if ret == 0:
+        return ret
+    else:
+        # bucket 1 is reserved for incontext samples
+        return ret + 1
+
+
+class _DictTree(TypedDict):
+    value: str
+    children: List["_DictTree"]
+    depth: int
+    segment_id: int
+    need_predict: bool
+
+
+class CpmBeeTokenizer(PreTrainedTokenizer):
+    """
+    Construct a CPMBee tokenizer.
+
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+        bos_token (`str`, *optional*, defaults to `"<s>"`):
+            The beginning of sequence token.
+        eos_token (`str`, *optional*, defaults to `"</s>"`):
+            The end of sequence token.
+        line_token (`str`, *optional*, defaults to `"\n"`):
+            The line token.
+        space_token (`str`, *optional*, defaults to `" "`):
+            The space token.
+        unk_token (`str`, *optional*, defaults to `"<unk>"`):
+            The unknown token.
+        mask_token (`str`, *optional*, defaults to `"<mask>"`):
+            The mask token.
+        pad_token (`str`, *optional*, defaults to `"<pad>"`):
+            The token used for padding.
+        padding_side (`str`, *optional*, defaults to `"left"`):
+            The padding side. CPM-Bee will use left padding by default.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names: List[str] = [
+        "input_ids",
+        "attention_mask",
+        "input_id_sub",
+        "position",
+        "context",
+        "sample_ids",
+        "num_segments",
+        "segment",
+        "segment_rel_offset",
+        "segment_rel",
+    ]
+    add_prefix_space = False
+
+    def __init__(
+        self,
+        vocab_file,
+        bos_token="<s>",
+        eos_token="</s>",
+        line_token="\n",
+        space_token=" ",
+        unk_token="<unk>",
+        mask_token="<mask>",
+        pad_token="<pad>",
+        padding_side="left",
+        **kwargs,
+    ):
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            line_token=line_token,
+            space_token=space_token,
+            unk_token=unk_token,
+            mask_token=mask_token,
+            pad_token=pad_token,
+            padding_side=padding_side,
+            **kwargs,
+        )
+
+        self.encoder: Dict[str, int] = {}
+
+        with open(vocab_file, "r", encoding="utf-8") as reader:
+            for token in reader.readlines():
+                token = token.rstrip("\n")
+                if len(token) == 0:
+                    continue
+                self.encoder[token] = len(self.encoder)
+
+        self.encoder[" "] = self.encoder["</_>"]
+        self.encoder["\n"] = self.encoder["</n>"]
+        del self.encoder["</_>"]
+        del self.encoder["</n>"]
+
+        self.decoder = {v: k for k, v in self.encoder.items()}
+
+        self._max_word_len = max([len(x) for x in self.encoder.keys()])
+        self.cpmbee_special_tokens = {k: v for k, v in self.encoder.items() if k.startswith("<") and k.endswith(">")}
+
+        self.ext_table: Dict[int, str] = {}
+        self.ext_table_rev: Dict[str, int] = {}
+
+        self.token_id_table: Dict[str, Dict[int, int]] = {}
+        self.ext_special_tokens = []
+
+        self.ext_args_for_model = [
+            "input_id_subs",
+            "input_pos",
+            "context",
+            "segment_ids",
+            "segment_rel_offset",
+            "segment_rel",
+            "sample_ids",
+            "num_segments",
+            "predict_segments",
+            "answer_placeholders",
+            "ext_table",
+            "token_id_table",
+        ]
+
+    @property
+    def bod_token_id(self):
+        return self.encoder[self.bod_token]
+
+    @property
+    def eod_token_id(self):
+        return self.encoder[self.eod_token]
+
+    @property
+    def newline_id(self):
+        return self.encoder[self.line_token]
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.encoder)
+
+    def __len__(self):
+        """
+        Size of the full vocabulary with the added tokens.
+        """
+        return self.vocab_size + len(self.added_tokens_encoder)
+
+    def get_vocab(self):
+        return dict(self.encoder, **self.added_tokens_encoder)
+
+    def get_piece(self, text: str) -> str:
+        """
+        Match with maximum length.
+        """
+        len_text = len(text)
+        for i in range(len(text)):
+            sub = text[: len_text - i]
+            if (sub in self.encoder) or (sub in self.added_tokens_encoder):
+                return sub
+        return text[0]
+
+    def tokenize(self, text: TextInput, **kwargs) -> List[str]:
+        r"""
+        Override the `tokenize` to meet the needs of CPMBee:
+        1. Mark the special token with `<` and `>`. The `<>` will be ignored.
+        2. Split sentences by the marked special tokens.
+        3. Record the marked special token by `ext_table` and `ext_table_rev`.
+        4. Tokenize the sentence without special tokens.
+        """
+        for_cpmbee = kwargs.get("for_cpmbee", False)
+        all_special_tokens_extended = {
+            str(t): t for t in self.all_special_tokens_extended if isinstance(t, AddedToken)
+        }
+
+        sentence_split = [""]
+        is_special_token = False
+        for i, c in enumerate(text):
+            if is_special_token:
+                if c == "<":
+                    tail = sentence_split.pop(-1)
+                    sentence_split[-1] += tail
+                    sentence_split.append(c)
+                elif c == ">":
+                    # end of special token
+                    sentence_split[-1] += c
+                    if sentence_split[-1] == "<>":
+                        continue
+                    is_special_token = False
+                    sentence_split.append("")
+                else:
+                    sentence_split[-1] += c
+            else:
+                if c == "<":
+                    is_special_token = True
+                    sentence_split.append(c)
+                else:
+                    sentence_split[-1] += c
+        if is_special_token:
+            tail = sentence_split.pop(-1)
+            sentence_split[-1] += tail
+
+        output_tokens = []
+        for i, part in enumerate(sentence_split):
+            if (i & 1) == 1:
+                # special token
+                output_tokens.append(part)
+                if for_cpmbee and (part not in self.encoder) and (part not in self.ext_table_rev):
+                    self.ext_table_rev[part] = len(self.ext_table_rev) + self.vocab_size
+                    self.ext_table[self.ext_table_rev[part]] = part
+            else:
+                output_tokens.extend(self._tokenize(part, for_cpmbee=for_cpmbee))
+
+        # drop spaces
+        for i, token in enumerate(output_tokens):
+            if token in self.added_tokens_encoder:
+                token = all_special_tokens_extended.get(token, None)
+                left = output_tokens[i - 1] if i > 0 else None
+                right = output_tokens[i + 1] if i < len(output_tokens) - 1 else None
+                if isinstance(token, AddedToken):
+                    if token.rstrip and right:
+                        # A bit counter-intuitive but we strip the left of the string
+                        # since tok_extended.rstrip means the special token is eating all white spaces on its right
+                        output_tokens[i + 1] = right.lstrip()
+                    # Strip white spaces on the left
+                    if token.lstrip and left:
+                        output_tokens[i - 1] = left.rstrip()  # Opposite here
+                else:
+                    if right:
+                        output_tokens[i + 1] = right.lstrip()
+                    if left:
+                        output_tokens[i - 1] = left.rstrip()
+
+        skipped_tokens = []
+        for token in output_tokens:
+            if not token:
+                continue
+            else:
+                skipped_tokens.append(token)
+
+        return skipped_tokens
+
+    def _tokenize(self, text, **kwargs):
+        """
+        Converts a string in a sequence of tokens (string), using the tokenizer. Split in words for word-based
+        vocabulary.
+
+        Do NOT take care of added tokens. Record the unk tokens and special tokens in `ext_table` and `ext_table_rev`.
+        """
+        for_cpmbee = kwargs.get("for_cpmbee", False)
+        output_tokens = []
+
+        part_st = 0
+        last_unk = None
+        while part_st < len(text):
+            piece = self.get_piece(text[part_st:])
+            if piece in self.encoder or self.added_tokens_encoder:
+                if last_unk is None:
+                    output_tokens.append(piece)
+                else:
+                    if for_cpmbee and (last_unk not in self.ext_table_rev):
+                        self.ext_table_rev[last_unk] = len(self.ext_table_rev) + self.vocab_size
+                        self.ext_table[self.ext_table_rev[last_unk]] = last_unk
+                    output_tokens.append(last_unk)
+                    output_tokens.append(piece)
+                    last_unk = None
+            else:
+                if last_unk is None:
+                    last_unk = piece
+                else:
+                    last_unk += piece
+            part_st += len(piece)
+        if last_unk is not None:
+            # part end with UNK
+            if for_cpmbee and (last_unk not in self.ext_table_rev):
+                self.ext_table_rev[last_unk] = len(self.ext_table_rev) + self.vocab_size
+                self.ext_table[self.ext_table_rev[last_unk]] = last_unk
+            output_tokens.append(last_unk)
+
+        return output_tokens
+
+    def check(self, token):
+        return token in self.encoder
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        return "".join(tokens)
+
+    def _convert_token_to_id(self, token: str):
+        """Converts a token (str) in an id using the vocab and ext_table."""
+        if token in self.encoder:
+            return self.encoder.get(token)
+        elif token in self.ext_table_rev:
+            return self.ext_table_rev[token]
+        elif token in self.added_tokens_encoder:
+            return self.added_tokens_encoder[token]
+        else:
+            return self.unk_token_id
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab and ext_table."""
+        if index in self.ext_table:
+            return self.ext_table[index]
+        elif index in self.added_tokens_decoder:
+            return self.added_tokens_decoder[index]
+        else:
+            if index >= 0:
+                return self.decoder[index]
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if os.path.isdir(save_directory):
+            vocab_file = os.path.join(
+                save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+            )
+        else:
+            vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
+        index = 0
+        self.encoder["</n>"] = self.encoder["\n"]
+        del self.encoder["\n"]
+        self.encoder["</_>"] = self.encoder[" "]
+        del self.encoder[" "]
+        with open(vocab_file, "w", encoding="utf-8") as writer:
+            for token, token_index in sorted(self.encoder.items(), key=lambda x: x[1]):
+                if index != token_index:
+                    logger.warning(
+                        f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
+                        " Please check that the vocabulary is not corrupted!"
+                    )
+                    index = token_index
+                writer.write(token + "\n")
+                index += 1
+        return (vocab_file,)
+
+    def __call__(self, text, *args, **kwargs):
+        r"""
+        CPMBee `call` method will use `_tokenize_cpmbee` when the input type is dict.
+        """
+        if isinstance(text, dict):
+            return self._batch_tokenize_cpmbee([text], *args, **kwargs)
+        elif isinstance(text, (list, tuple)):
+            if isinstance(text[0], dict):
+                return self._batch_tokenize_cpmbee(text, *args, **kwargs)
+            else:
+                return super().__call__(text, *args, **kwargs)
+        else:
+            return super().__call__(text, *args, **kwargs)
+
+    # 分词
+    def _tokenize_cpmbee(self, data: TextInput, *args, **kwargs) -> List[str]:
+        """
+        A tokenize method to process dict data. Exclusive for CPMBee.
+        """
+        if isinstance(data, str):
+            data = json.loads(data)
+        if not isinstance(data, Dict):
+            raise TypeError(
+                "CpmBeeTokenizer input data should be dict or str in dict format, but got {}".format(type(data))
+            )
+
+        # 1. prepare answer placeholder
+        answer_placeholders = []
+
+        def _put_placeholder(data: Any, path: List[str] = []):
+            if isinstance(data, dict):
+                ret = {}
+                for k, v in data.items():
+                    ret[k] = _put_placeholder(v, path + [k])
+                return ret
+            else:
+                answer_placeholders.append(path)
+                return "<ans_{}>".format(len(answer_placeholders))
+
+        data["<ans>"] = _put_placeholder(data["<ans>"])
+
+        (
+            input_ids,
+            input_id_subs,
+            context,
+            segment_ids,
+            segment_rel,
+            n_segments,
+            table_states,
+        ) = self.convert_data_to_id(data, shuffle_answer=False, max_depth=8)
+
+        # <ans> mapping from sub to id
+        sub_ans_map: Dict[int, int] = {}
+        for fake_id, token_sub in table_states["token_id_table"]["<ans>"].items():
+            token = table_states["ext_table"][fake_id]
+            if token.startswith("<ans_") and token.endswith(">"):
+                ans_id = int(token[5:-1])
+                sub_ans_map[token_sub] = ans_id
+
+        tmp_input_ids = []
+        tmp_input_sub = []
+        tmp_input_seg = []
+
+        # get predict segments
+        predict_segments: List[Tuple[int, int]] = []
+        for i in range(input_ids.shape[0]):
+            if context[i] == 0:
+                if input_ids[i] == self.encoder["<ans>"]:
+                    # is ans
+                    # (segment_id, ans_id)
+                    predict_segments.append((segment_ids[i], sub_ans_map[input_id_subs[i]]))
+            else:
+                tmp_input_ids.append(input_ids[i])
+                tmp_input_sub.append(input_id_subs[i])
+                tmp_input_seg.append(segment_ids[i])
+
+        if len(predict_segments) == 0:
+            raise ValueError("No answer to predict")
+
+        input_ids = np.array(tmp_input_ids, dtype=np.int32)  # all context
+        input_id_subs = np.array(tmp_input_sub, dtype=np.int32)  # [0, 0, 0, 0, 1, 0, 0, 2, 0, ...]
+        context = np.full_like(tmp_input_ids, 1, dtype=np.int8)  # [1, 1, 1, ...]
+        segment_ids = np.array(tmp_input_seg, dtype=np.int32)  # [0, 0, 0, 1, 1, 1, 2, 2, 2, 2, ...]
+        sample_ids = np.zeros(input_ids.shape, dtype=np.int32)  # [0, 0, 0, 0, ...]
+        segment_rel_offset = np.zeros(input_ids.shape, dtype=np.int32)  # [0, 0, 0, ...]
+        num_segments = np.full(input_ids.shape, n_segments, dtype=np.int32)  # [n_seg, n_seg, n_seg, ...]
+        input_pos = np.arange(input_ids.shape[0], dtype=np.int32)  # [0, 1, 2, 3, 4, ...]
+
+        return (
+            self.prepare_for_model(
+                input_ids.tolist(),
+                input_id_subs=input_id_subs.tolist(),
+                input_pos=input_pos.tolist(),
+                context=context.tolist(),
+                segment_ids=segment_ids.tolist(),
+                segment_rel_offset=segment_rel_offset.tolist(),
+                segment_rel=segment_rel.tolist(),
+                sample_ids=sample_ids.tolist(),
+                num_segments=num_segments.tolist(),
+                **kwargs,
+            ),
+            predict_segments,
+            answer_placeholders,
+            table_states["ext_table"],
+            table_states["token_id_table"],
+        )
+
+    def _batch_tokenize_cpmbee(self, data_lst, *args, **kwargs):
+        """
+        Batched _token_cpmbee.
+        """
+        device = kwargs.get("device", "cpu")
+        return_tensors = kwargs.get("return_tensors", None)
+        batch_outputs = {}
+        segment_rel_pack = []
+        other_info = []
+
+        batch_ext_table_map: Dict[Tuple[int, int], int] = {}
+        batch_ext_table_ids: List[int] = []
+        batch_ext_table_sub: List[int] = []
+
+        for data in data_lst:
+            self.ext_table = {}
+            self.ext_table_rev = {}
+            self.token_id_table = {}
+            (outputs, predict_segments, answer_placeholders, ext_table, token_id_table) = self._tokenize_cpmbee(
+                data,
+                truncation=None,
+                padding=PaddingStrategy.DO_NOT_PAD.value,
+                max_length=None,
+                pad_to_multiple_of=None,
+                return_attention_mask=False,
+                return_tensors=None,
+            )
+            rev_ext_table = {}
+            for token, mp in token_id_table.items():
+                if token == "<ans>":
+                    continue
+                token_id = self.encoder[token]
+                for fake_id, token_sub in mp.items():
+                    if token_sub > 0:
+                        if (token_id, token_sub) not in batch_ext_table_map:
+                            batch_ext_table_map[(token_id, token_sub)] = len(batch_ext_table_ids) + self.vocab_size
+                            batch_ext_table_ids.append(token_id)
+                            batch_ext_table_sub.append(token_sub)
+                        rev_ext_table[batch_ext_table_map[(token_id, token_sub)]] = ext_table[fake_id]
+                    else:
+                        rev_ext_table[token_id] = ext_table[fake_id]
+
+            segment_rel_pack.append(np.array(outputs.pop("segment_rel")))
+            other_info.append(
+                {
+                    "predict_segments": predict_segments,
+                    "answer_placeholders": answer_placeholders,
+                    "ext_table": rev_ext_table,
+                }
+            )
+
+            for key, value in outputs.items():
+                if key not in batch_outputs:
+                    batch_outputs[key] = []
+                batch_outputs[key].append(value)
+
+        max_length = max([len(item) for item in batch_outputs[self.model_input_names[0]]])
+        batch_size = len(batch_outputs[self.model_input_names[0]])
+        for i in range(batch_size):
+            inputs = {k: v[i] for k, v in batch_outputs.items()}
+
+            for k, v in inputs.items():
+                required_input = v
+
+                needs_to_be_padded = len(required_input) != max_length
+
+                if needs_to_be_padded:
+                    difference = max_length - len(required_input)
+                    batch_outputs[k][i] = [self.pad_token_id] * difference + required_input
+
+        max_num_rels = 0
+        for rel in segment_rel_pack:
+            max_num_rels = max(max_num_rels, rel.shape[0])
+        padded_rels = np.zeros((len(segment_rel_pack), max_num_rels), dtype=np.int32)
+        for i, rel in enumerate(segment_rel_pack):
+            padded_rels[i, : rel.shape[0]] = rel
+        batch_outputs["segment_rel"] = padded_rels
+        batch_outputs["batch_ext_table_ids"] = np.array(batch_ext_table_ids, dtype=np.int32)
+        batch_outputs["batch_ext_table_sub"] = np.array(batch_ext_table_sub, dtype=np.int32)
+        batch_outputs = BatchEncoding(batch_outputs, tensor_type=return_tensors)
+        if return_tensors == "pt":
+            batch_outputs = batch_outputs.to(device=device)
+        batch_outputs["other_info"] = other_info
+
+        return batch_outputs
+
+    def convert_data_to_id(
+        self,
+        data: Any,
+        prev_ext_states: Optional[_PrevExtTableStates] = None,
+        shuffle_answer: bool = True,
+        max_depth: int = 8,
+    ):
+        """
+        Parse a dict to data ids. Exclusive for CPMBee. It will
+        1. parse the dict to segments and get segment_rel, which for calculating of position_bias.
+        2. tokenize every segment.
+        """
+        root: _DictTree = {
+            "value": "<root>",
+            "children": [],
+            "depth": 0,
+            "segment_id": 0,
+            "need_predict": False,
+        }
+
+        segments = [root]
+
+        def _build_dict_tree(data: CPMBeeInputType, depth: int, need_predict: bool) -> List[_DictTree]:
+            if isinstance(data, dict):
+                ret_list: List[_DictTree] = []
+                curr_items = list(data.items())
+                if need_predict and shuffle_answer:
+                    access_idx = np.arange(len(curr_items))
+                    np.random.shuffle(access_idx)
+                    curr_items = [curr_items[idx] for idx in access_idx]
+                for k, v in curr_items:
+                    child_info: _DictTree = {
+                        "value": k,
+                        "children": [],
+                        "depth": depth,
+                        "segment_id": len(segments),
+                        "need_predict": False,  # only leaves are contexts
+                    }
+                    segments.append(child_info)
+                    child_info["children"] = _build_dict_tree(
+                        v, depth + 1, need_predict or (depth == 1 and k == "<ans>")
+                    )  # elements in <root>.<ans>
+
+                    ret_list.append(child_info)
+                return ret_list
+            else:
+                assert isinstance(data, str), "Invalid data {}".format(data)
+                ret: _DictTree = {
+                    "value": data,
+                    "children": [],
+                    "depth": depth,
+                    "segment_id": len(segments),
+                    "need_predict": need_predict,
+                }
+                segments.append(ret)
+                return [ret]
+
+        root["children"] = _build_dict_tree(data, 1, False)
+
+        num_segments = len(segments)
+        segment_rel = np.zeros((num_segments * num_segments,), dtype=np.int32)
+
+        def _build_segment_rel(node: _DictTree) -> List[Tuple[int, int]]:
+            ret: List[Tuple[int, int]] = [(node["segment_id"], node["depth"])]
+            for child in node["children"]:
+                sub = _build_segment_rel(child)
+                for seg_id_1, depth_1 in sub:
+                    for seg_id_2, depth_2 in ret:
+                        n_up = min(depth_1 - node["depth"], max_depth - 1)
+                        n_down = min(depth_2 - node["depth"], max_depth - 1)
+                        segment_rel[seg_id_1 * num_segments + seg_id_2] = rel_to_bucket(
+                            n_up, n_down, max_depth=max_depth
+                        )
+                        segment_rel[seg_id_2 * num_segments + seg_id_1] = rel_to_bucket(
+                            n_down, n_up, max_depth=max_depth
+                        )
+                ret.extend(sub)
+            return ret
+
+        _build_segment_rel(root)
+
+        input_ids: List[int] = []
+        input_id_subs: List[int] = []
+        segment_bound: List[Tuple[int, int]] = []
+
+        if prev_ext_states is not None:
+            self.ext_table = prev_ext_states["ext_table"]
+            self.token_id_table = prev_ext_states["token_id_table"]
+
+        for seg in segments:
+            # tokenize
+            tokens = self.convert_tokens_to_ids(self.tokenize(seg["value"], for_cpmbee=True))
+
+            token_id_subs = []
+            reid_token_ids = []
+            for idx in tokens:
+                if idx in self.ext_table:
+                    # unk or special token
+                    token = self.ext_table[idx]
+                    if token.startswith("<") and token.endswith(">"):
+                        # special token
+                        if "_" in token:
+                            token_name = token[1:-1].split("_", maxsplit=1)[0]
+                        else:
+                            token_name = token[1:-1]
+                        token_name = "<{}>".format(token_name)
+                    else:
+                        token_name = "<unk>"
+
+                    if token_name not in self.token_id_table:
+                        self.token_id_table[token_name] = {}
+                    if idx not in self.token_id_table[token_name]:
+                        self.token_id_table[token_name][idx] = len(self.token_id_table[token_name])
+                    if token_name not in self.encoder:
+                        raise ValueError("Invalid token {}".format(token))
+                    reid_token_ids.append(self.encoder[token_name])
+                    token_id_subs.append(self.token_id_table[token_name][idx])
+                else:
+                    reid_token_ids.append(idx)
+                    token_id_subs.append(0)
+            tokens = [self.bos_token_id] + reid_token_ids
+            token_id_subs = [0] + token_id_subs
+            # eos_id 表示 no need_predict
+            if not seg["need_predict"]:  # eos
+                tokens = tokens + [self.eos_token_id]
+                token_id_subs = token_id_subs + [0]
+            else:
+                # no eos
+                pass
+            begin = len(input_ids)
+            input_ids.extend(tokens)
+            input_id_subs.extend(token_id_subs)
+            end = len(input_ids)
+            segment_bound.append((begin, end))
+
+        ids = np.array(input_ids, dtype=np.int32)
+        id_subs = np.array(input_id_subs, dtype=np.int32)
+        segs = np.zeros((ids.shape[0],), dtype=np.int32)  # 按segment_bound对seg编号
+        context = np.zeros((ids.shape[0],), dtype=np.int8)
+        for i, (begin, end) in enumerate(segment_bound):
+            if not segments[i]["need_predict"]:
+                context[begin:end] = 1
+            segs[begin:end] = i
+
+        curr_ext_table_states: _PrevExtTableStates = {
+            "ext_table": self.ext_table,
+            "token_id_table": self.token_id_table,
+        }
+        return ids, id_subs, context, segs, segment_rel, num_segments, curr_ext_table_states
+
+    def prepare_for_model(
+        self,
+        ids: List[int],
+        pair_ids: Optional[List[int]] = None,
+        add_special_tokens: bool = True,
+        padding: Union[bool, str, PaddingStrategy] = False,
+        truncation: Union[bool, str, TruncationStrategy] = None,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_token_type_ids: Optional[bool] = None,
+        return_attention_mask: Optional[bool] = None,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+        prepend_batch_axis: bool = False,
+        **kwargs,
+    ) -> BatchEncoding:
+        """
+        Prepares a sequence of input id, or a pair of sequences of inputs ids so that it can be used by the model. It
+        adds special tokens, truncates sequences if overflowing while taking into account the special tokens and
+        manages a moving window (with user defined stride) for overflowing tokens. Please Note, for *pair_ids*
+        different than `None` and *truncation_strategy = longest_first* or `True`, it is not possible to return
+        overflowing tokens. Such a combination of arguments will raise an error.
+
+        Args:
+            ids (`List[int]`):
+                Tokenized input ids of the first sequence. Can be obtained from a string by chaining the `tokenize` and
+                `convert_tokens_to_ids` methods.
+            pair_ids (`List[int]`, *optional*):
+                Tokenized input ids of the second sequence. Can be obtained from a string by chaining the `tokenize`
+                and `convert_tokens_to_ids` methods.
+        """
+
+        # Backward compatibility for 'truncation_strategy', 'pad_to_max_length'
+        padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            pad_to_multiple_of=pad_to_multiple_of,
+            verbose=verbose,
+            **kwargs,
+        )
+
+        pair = bool(pair_ids is not None)
+        len_ids = len(ids)
+        len_pair_ids = len(pair_ids) if pair else 0
+
+        if return_token_type_ids and not add_special_tokens:
+            raise ValueError(
+                "Asking to return token_type_ids while setting add_special_tokens to False "
+                "results in an undefined behavior. Please set add_special_tokens to True or "
+                "set return_token_type_ids to None."
+            )
+
+        if (
+            return_overflowing_tokens
+            and truncation_strategy == TruncationStrategy.LONGEST_FIRST
+            and pair_ids is not None
+        ):
+            raise ValueError(
+                "Not possible to return overflowing tokens for pair of sequences with the "
+                "`longest_first`. Please select another truncation strategy than `longest_first`, "
+                "for instance `only_second` or `only_first`."
+            )
+
+        # Load from model defaults
+        if return_token_type_ids is None:
+            return_token_type_ids = "token_type_ids" in self.model_input_names
+        if return_attention_mask is None:
+            return_attention_mask = "attention_mask" in self.model_input_names
+
+        encoded_inputs = {}
+
+        # Compute the total size of the returned encodings
+        total_len = len_ids + len_pair_ids + (self.num_special_tokens_to_add(pair=pair) if add_special_tokens else 0)
+
+        # Truncation: Handle max sequence length
+        overflowing_tokens = []
+        if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE and max_length and total_len > max_length:
+            ids, pair_ids, overflowing_tokens = self.truncate_sequences(
+                ids,
+                pair_ids=pair_ids,
+                num_tokens_to_remove=total_len - max_length,
+                truncation_strategy=truncation_strategy,
+                stride=stride,
+            )
+
+        if return_overflowing_tokens:
+            encoded_inputs["overflowing_tokens"] = overflowing_tokens
+            encoded_inputs["num_truncated_tokens"] = total_len - max_length
+
+        # Add special tokens
+        if add_special_tokens:
+            sequence = self.build_inputs_with_special_tokens(ids, pair_ids)
+            token_type_ids = self.create_token_type_ids_from_sequences(ids, pair_ids)
+        else:
+            sequence = ids + pair_ids if pair else ids
+            token_type_ids = [0] * len(ids) + ([0] * len(pair_ids) if pair else [])
+
+        # Build output dictionary
+        encoded_inputs["input_ids"] = sequence
+        if return_token_type_ids:
+            encoded_inputs["token_type_ids"] = token_type_ids
+        if return_special_tokens_mask:
+            if add_special_tokens:
+                encoded_inputs["special_tokens_mask"] = self.get_special_tokens_mask(ids, pair_ids)
+            else:
+                encoded_inputs["special_tokens_mask"] = [0] * len(sequence)
+
+        # Check lengths
+        self._eventual_warn_about_too_long_sequence(encoded_inputs["input_ids"], max_length, verbose)
+
+        # Padding
+        if padding_strategy != PaddingStrategy.DO_NOT_PAD or return_attention_mask:
+            encoded_inputs = self.pad(
+                encoded_inputs,
+                max_length=max_length,
+                padding=padding_strategy.value,
+                pad_to_multiple_of=pad_to_multiple_of,
+                return_attention_mask=return_attention_mask,
+            )
+
+        if return_length:
+            encoded_inputs["length"] = len(encoded_inputs["input_ids"])
+
+        # for CPMBee, encode all the model arguments
+        for arg in self.ext_args_for_model:
+            v = kwargs.get(arg, None)
+            if v is not None:
+                encoded_inputs[arg] = v
+
+        batch_outputs = BatchEncoding(
+            encoded_inputs, tensor_type=return_tensors, prepend_batch_axis=prepend_batch_axis
+        )
+
+        return batch_outputs