jeffreygo commited on
Commit
83e21b6
1 Parent(s): b197fa6

add source files

Browse files
configuration_cpmbee.py ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """ CpmBee model configuration"""
16
+
17
+ from typing import List, Optional, Tuple, Union
18
+
19
+ from ...configuration_utils import PretrainedConfig
20
+ from ...utils import logging
21
+
22
+
23
+ logger = logging.get_logger(__name__)
24
+
25
+ CPMBEE_PRETRAINED_CONFIG_ARCHIVE_MAP = {
26
+ "openbmb/cpm-bee-10b": "https://huggingface.co/openbmb/cpm-bee-10b/resolve/main/config.json",
27
+ "openbmb/cpm-bee-5b": "https://huggingface.co/openbmb/cpm-bee-5b/resolve/main/config.json",
28
+ "openbmb/cpm-bee-2b": "https://huggingface.co/openbmb/cpm-bee-2b/resolve/main/config.json",
29
+ "openbmb/cpm-bee-1b": "https://huggingface.co/openbmb/cpm-bee-1b/resolve/main/config.json",
30
+ # See all CpmBee models at https://huggingface.co/models?filter=cpmbee
31
+ }
32
+
33
+
34
+ class CpmBeeConfig(PretrainedConfig):
35
+ r"""
36
+ This is the configuration class to store the configuration of a [`CpmBeeModel`]. It is used to instbeeiate an
37
+ CPMBee model according to the specified arguments, defining the model architecture. Instantiating a configuration
38
+ with the defaults will yield a similar configuration to that of the CPMBee
39
+ [openbmb/cpm-bee-10b](https://huggingface.co/openbmb/cpm-bee-10b) architecture.
40
+
41
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
42
+ documentation from [`PretrainedConfig`] for more information.
43
+
44
+ Args:
45
+ vocab_size (`int`, *optional*, defaults to 30720):
46
+ Vocabulary size of the CPMBee model. Defines the number of different tokens that can be represented by the
47
+ `input` passed when calling [`CpmBeeModel`].
48
+ hidden_size (`int`, *optional*, defaults to 4096):
49
+ Dimension of the encoder layers.
50
+ num_attention_heads (`int`, *optional*, defaults to 32):
51
+ Number of attention heads in the Transformer encoder.
52
+ dim_head (`int`, *optional*, defaults to 128):
53
+ Dimension of attention heads for each attention layer in the Transformer encoder.
54
+ dim_ff (`int`, *optional*, defaults to 10240):
55
+ Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
56
+ num_hidden_layers (`int`, *optional*, defaults to 48):
57
+ Number of layers of the Transformer encoder.
58
+ dropout_p (`float`, *optional*, defaults to 0.1):
59
+ The dropout probabilitiy for all fully connected layers in the embeddings, encoder.
60
+ position_bias_num_buckets (`int`, *optional*, defaults to 512):
61
+ The number of position_bias buckets.
62
+ position_bias_num_segment_buckets (`int`, *optional*, defaults to 32):
63
+ The number of segment buckets.
64
+ position_bias_max_distance (`int`, *optional*, defaults to 2048):
65
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
66
+ just in case (e.g., 512 or 1024 or 2048).
67
+ eps (`float`, *optional*, defaults to 1e-6):
68
+ The epsilon used by the layer normalization layers.
69
+ init_std (`float`, *optional*, defaults to 1.0):
70
+ Initialize parameters with std = init_std.
71
+ use_cache (`bool`, *optional*, defaults to `True`):
72
+ Whether to use cache.
73
+ distance_scale (`float` or `int`, *optional*, defaults to 16):
74
+ Scale the rotary embedding.
75
+ mask_modules (`list` or `tuple`, *optional*, defaults to None):
76
+ Decides which feedforward block or attention block is pruned.
77
+ half (`bool`, *optional*, defaults to `False`):
78
+ Decides the model parameters are half-precision or not.
79
+
80
+ Example:
81
+
82
+ ```python
83
+ >>> from transformers import CpmBeeModel, CpmBeeConfig
84
+
85
+ >>> # Initializing a CPMBee cpm-bee-10b style configuration
86
+ >>> configuration = CpmBeeConfig()
87
+
88
+ >>> # Initializing a model from the cpm-bee-10b style configuration
89
+ >>> model = CpmBeeModel(configuration)
90
+
91
+ >>> # Accessing the model configuration
92
+ >>> configuration = model.config
93
+ ```"""
94
+ model_type = "cpmbee"
95
+
96
+ def __init__(
97
+ self,
98
+ vocab_size: int = 30720,
99
+ hidden_size: int = 4096,
100
+ num_attention_heads: int = 64,
101
+ dim_head: int = 64,
102
+ dim_ff: int = 10240,
103
+ num_hidden_layers: int = 32,
104
+ dropout_p: int = 0.0,
105
+ position_bias_num_buckets: int = 256,
106
+ position_bias_num_segment_buckets: int = 32,
107
+ position_bias_max_distance: int = 2048,
108
+ eps: int = 1e-6,
109
+ init_std: float = 1.0,
110
+ use_cache: bool = True,
111
+ distance_scale: Union[int, float] = 16,
112
+ mask_modules: Optional[Union[List, Tuple]] = None,
113
+ half: bool = False,
114
+ **kwargs,
115
+ ):
116
+ super().__init__(**kwargs)
117
+ self.position_bias_num_segment_buckets = position_bias_num_segment_buckets
118
+ self.hidden_size = hidden_size
119
+ self.num_attention_heads = num_attention_heads
120
+ self.dim_head = dim_head
121
+ self.dim_ff = dim_ff
122
+ self.num_hidden_layers = num_hidden_layers
123
+ self.position_bias_num_buckets = position_bias_num_buckets
124
+ self.position_bias_max_distance = position_bias_max_distance
125
+ self.dropout_p = dropout_p
126
+ self.eps = eps
127
+ self.use_cache = use_cache
128
+ self.vocab_size = vocab_size
129
+ self.init_std = init_std
130
+ self.distance_scale = distance_scale
131
+ self.half = half
132
+ self.mask_modules = mask_modules
modeling_cpmbee.py ADDED
@@ -0,0 +1,1944 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 The OpenBMB Team The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """ PyTorch CpmBee model."""
16
+ import copy
17
+ import math
18
+ from collections import UserDict
19
+ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
20
+
21
+ import torch
22
+ import torch.nn as nn
23
+
24
+ from ...generation.beam_search import BeamHypotheses, BeamSearchScorer
25
+ from ...generation.streamers import BaseStreamer
26
+ from ...generation.utils import (
27
+ GenerationConfig,
28
+ LogitsProcessorList,
29
+ StoppingCriteriaList,
30
+ dist,
31
+ inspect,
32
+ is_deepspeed_zero3_enabled,
33
+ warnings,
34
+ )
35
+ from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, ModelOutput
36
+ from ...modeling_utils import PreTrainedModel
37
+ from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
38
+ from .configuration_cpmbee import CpmBeeConfig
39
+ from .tokenization_cpmbee import CpmBeeTokenizer
40
+
41
+
42
+ logger = logging.get_logger(__name__)
43
+
44
+ _CHECKPOINT_FOR_DOC = "openbmb/cpm-bee-10b"
45
+ _CONFIG_FOR_DOC = "CpmBeeConfig"
46
+
47
+ CPMBEE_PRETRAINED_MODEL_ARCHIVE_LIST = [
48
+ "openbmb/cpm-bee-10b",
49
+ "openbmb/cpm-bee-5b",
50
+ "openbmb/cpm-bee-2b",
51
+ "openbmb/cpm-bee-1b",
52
+ # See all CPMBee models at https://huggingface.co/models?filter=cpmbee
53
+ ]
54
+
55
+
56
+ class CpmBeeLinear(nn.Linear):
57
+ def __init__(self, dim_in, dim_out, dtype):
58
+ """
59
+ Construct a linear for CPMBee. It contains a scale operation.
60
+ """
61
+ super().__init__(dim_in, dim_out, bias=False)
62
+ self.dim_in = self.in_features = dim_in
63
+ self.dim_out = self.out_features = dim_out
64
+
65
+ self.weight = torch.nn.parameter.Parameter(torch.empty((dim_out, dim_in), dtype=dtype))
66
+
67
+ def forward(self, x: torch.Tensor):
68
+ """
69
+ Args:
70
+ x (`torch.Tensor` of shape `(batch, seq_len, dim_in)`): The input of linear layer
71
+ Returns:
72
+ `torch.Tensor` of shape `(batch, seq_len, dim_out)`: The output of the linear transform y.
73
+ """
74
+ x = nn.functional.linear(x, self.weight)
75
+ x = x / math.sqrt(self.dim_in)
76
+ return x
77
+
78
+
79
+ class CpmBeeLayerNorm(nn.Module):
80
+ """
81
+ We use Root Mean Square (RMS) Layer Normalization, please see https://arxiv.org/abs/1910.07467 for details."
82
+ """
83
+
84
+ def __init__(self, config: CpmBeeConfig):
85
+ super().__init__()
86
+
87
+ self.eps = config.eps
88
+ self.dim_norm = config.hidden_size
89
+ self.weight = nn.Parameter(torch.empty(config.hidden_size, dtype=config.torch_dtype))
90
+
91
+ def forward(self, hidden_states: torch.Tensor):
92
+ """
93
+ Args:
94
+ hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
95
+ """
96
+ if hidden_states.size(-1) != self.dim_norm:
97
+ raise AssertionError("hidden_states.size(-1) != self.dim_norm")
98
+ old_dtype = hidden_states.dtype
99
+ variance = hidden_states.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
100
+ hidden_states = (hidden_states * torch.rsqrt(variance + self.eps)).to(old_dtype) * self.weight
101
+ return hidden_states
102
+
103
+
104
+ class CpmBeeAttention(nn.Module):
105
+ def __init__(self, config: CpmBeeConfig):
106
+ super().__init__()
107
+ self.dim_model = config.hidden_size
108
+ self.num_heads = config.num_attention_heads
109
+ self.dim_head = config.dim_head
110
+
111
+ self.project_q = CpmBeeLinear(self.dim_model, self.num_heads * self.dim_head, dtype=config.torch_dtype)
112
+ self.project_k = CpmBeeLinear(self.dim_model, self.num_heads * self.dim_head, dtype=config.torch_dtype)
113
+ self.project_v = CpmBeeLinear(self.dim_model, self.num_heads * self.dim_head, dtype=config.torch_dtype)
114
+
115
+ self.attention_out = CpmBeeLinear(self.num_heads * self.dim_head, self.dim_model, dtype=config.torch_dtype)
116
+
117
+ self.softmax = torch.nn.Softmax(dim=-1)
118
+
119
+ if config.dropout_p is not None:
120
+ self.dropout = torch.nn.Dropout(p=config.dropout_p)
121
+ else:
122
+ self.dropout = None
123
+
124
+ def forward(
125
+ self,
126
+ hidden_q: torch.Tensor,
127
+ hidden_kv: torch.Tensor,
128
+ attention_mask: torch.BoolTensor,
129
+ position_bias: torch.Tensor,
130
+ output_attentions: Optional[bool] = False,
131
+ past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
132
+ use_cache: Optional[bool] = None,
133
+ ):
134
+ """
135
+ Args:
136
+ hidden_q (`torch.Tensor`):
137
+ Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
138
+ hidden_kv (`torch.Tensor` of shape `(batch, len_k, dim_model)`)):
139
+ Tensor *key_value* and *query* of shape `(batch, len_k, dim_model)`
140
+ attention_mask (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
141
+ Avoid invalid areas to participate in the calculation of self-attention.
142
+ position_bias (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
143
+ Provide positional information to self-attention block.
144
+ output_attentions (`bool`, *optional*):
145
+ Whether or not to return the attentions tensors of all attention layers.
146
+ past_key_values (`Tuple[torch.Tensor, torch.Tensor]`, *optional*):
147
+ Cached past key and value projection states.
148
+ use_cache (`bool`, *optional*):
149
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
150
+ (see `past_key_values`).
151
+ """
152
+ batch_size = hidden_q.size(0)
153
+ len_q = hidden_q.size(1)
154
+ len_k = hidden_kv.size(1)
155
+
156
+ query = self.project_q(hidden_q)
157
+ key = self.project_k(hidden_kv)
158
+ value = self.project_v(hidden_kv)
159
+
160
+ query = query.view(batch_size, len_q, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
161
+ key = key.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
162
+ value = value.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
163
+
164
+ if past_key_values is not None:
165
+ key = torch.cat([past_key_values[0], key], dim=-2)
166
+ value = torch.cat([past_key_values[1], value], dim=-2)
167
+ len_k = key.size(-2)
168
+
169
+ # (batch_size, num_heads, len_q, dim_head) @ (batch_size, num_heads, dim_head, len_k) -> (batch_size, num_heads, len_q, len_k)
170
+ score = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.dim_head)
171
+ score = score + position_bias
172
+
173
+ score = torch.masked_fill(
174
+ score,
175
+ attention_mask.view(batch_size, 1, len_q, len_k) == torch.tensor(False),
176
+ torch.scalar_tensor(float("-inf"), device=score.device, dtype=score.dtype),
177
+ )
178
+ score = self.softmax(score)
179
+
180
+ score = torch.masked_fill(
181
+ score,
182
+ attention_mask.view(batch_size, 1, len_q, len_k) == torch.tensor(False),
183
+ torch.scalar_tensor(0, device=score.device, dtype=score.dtype),
184
+ )
185
+ if output_attentions:
186
+ attn_weights = score
187
+ else:
188
+ attn_weights = None
189
+
190
+ if self.dropout is not None:
191
+ score = self.dropout(score)
192
+
193
+ # (batch_size, num_heads, len_q, len_k) @ (batch_size, num_heads, len_k, dim_head) -> (batch_size, num_heads, len_q, dim_head)
194
+ score = torch.matmul(score, value)
195
+
196
+ score = score.view(batch_size, self.num_heads, len_q, self.dim_head).permute(0, 2, 1, 3)
197
+ score = score.contiguous().view(batch_size, len_q, self.num_heads * self.dim_head)
198
+
199
+ score = self.attention_out(score)
200
+
201
+ past_key_values = None
202
+ if use_cache:
203
+ past_key_values = (key, value)
204
+
205
+ return score, attn_weights, past_key_values
206
+
207
+
208
+ class CpmBeeSelfAttentionBlock(nn.Module):
209
+ def __init__(self, config: CpmBeeConfig):
210
+ super().__init__()
211
+ self.layernorm_before_attention = CpmBeeLayerNorm(config)
212
+ self.self_attention = CpmBeeAttention(config)
213
+ if config.dropout_p:
214
+ self.dropout = torch.nn.Dropout(config.dropout_p)
215
+ else:
216
+ self.dropout = None
217
+
218
+ def forward(
219
+ self,
220
+ hidden_states: torch.Tensor,
221
+ attention_mask: torch.Tensor,
222
+ position_bias: Optional[torch.Tensor] = None,
223
+ output_attentions: Optional[bool] = False,
224
+ past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
225
+ use_cache: Optional[bool] = None,
226
+ ):
227
+ """
228
+ Args:
229
+ hidden_states (`torch.Tensor` of shape `(batch, len_seq, dim_model)`):
230
+ Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
231
+ attention_mask (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
232
+ Avoid invalid areas to participate in the calculation of self-attention.
233
+ position_bias (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
234
+ Provide positional information to self-attention block.
235
+ output_attentions (`bool`, *optional*):
236
+ Whether or not to return the attentions tensors of all attention layers.
237
+ past_key_values (`Tuple(torch.FloatTensor)`, *optional*):
238
+ Cached past key and value projection states.
239
+ use_cache (`bool`, *optional*):
240
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
241
+ (see `past_key_values`).
242
+ """
243
+ outputs = self.layernorm_before_attention(hidden_states)
244
+ outputs = self.self_attention(
245
+ outputs, outputs, attention_mask, position_bias, output_attentions, past_key_values, use_cache
246
+ )
247
+
248
+ outputs, attn_weights, current_key_value = outputs
249
+
250
+ if self.dropout is not None:
251
+ outputs = self.dropout(outputs)
252
+ hidden_states = (hidden_states + outputs) / 1.05
253
+
254
+ return hidden_states, attn_weights, current_key_value
255
+
256
+
257
+ class CpmBeeDenseGatedACT(nn.Module):
258
+ def __init__(self, config: CpmBeeConfig):
259
+ super().__init__()
260
+ self.w_0 = CpmBeeLinear(config.hidden_size, config.dim_ff, dtype=config.torch_dtype)
261
+ self.w_1 = CpmBeeLinear(config.hidden_size, config.dim_ff, dtype=config.torch_dtype)
262
+ self.act = torch.nn.GELU()
263
+
264
+ def forward(self, hidden_states: torch.Tensor):
265
+ """Transform an input tensor from one feature space to another via a nonlinear operation
266
+
267
+ Args:
268
+ hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
269
+ """
270
+ gate_score = self.act(self.w_0(hidden_states))
271
+ hidden_states = self.w_1(hidden_states)
272
+
273
+ hidden_states = gate_score * hidden_states
274
+ return hidden_states
275
+
276
+
277
+ class CpmBeeFeedForward(nn.Module):
278
+ def __init__(self, config: CpmBeeConfig):
279
+ super().__init__()
280
+ self.w_in = CpmBeeDenseGatedACT(config)
281
+ if config.dropout_p is not None:
282
+ self.dropout = torch.nn.Dropout(config.dropout_p)
283
+ else:
284
+ self.dropout = None
285
+
286
+ self.w_out = CpmBeeLinear(config.dim_ff, config.hidden_size, dtype=config.torch_dtype)
287
+
288
+ def forward(self, hidden_states: torch.Tensor):
289
+ """
290
+ Args:
291
+ hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
292
+ """
293
+ hidden_states = self.w_in(hidden_states)
294
+
295
+ if self.dropout is not None:
296
+ hidden_states = self.dropout(hidden_states)
297
+
298
+ hidden_states = self.w_out(hidden_states)
299
+
300
+ return hidden_states
301
+
302
+
303
+ class CpmBeeFFNBlock(nn.Module):
304
+ def __init__(self, config: CpmBeeConfig):
305
+ super().__init__()
306
+ self.layernorm_before_ffn = CpmBeeLayerNorm(config)
307
+ self.ffn = CpmBeeFeedForward(config)
308
+ if config.dropout_p:
309
+ self.dropout = torch.nn.Dropout(config.dropout_p)
310
+ else:
311
+ self.dropout = None
312
+
313
+ def forward(
314
+ self,
315
+ hidden_states: torch.Tensor,
316
+ ):
317
+ """
318
+ Args:
319
+ hidden_states (`torch.Tensor` of shape `(batch, len_seq, dim_model)`):
320
+ Hidden states before feed forward layer.
321
+ """
322
+ ln_outputs = self.layernorm_before_ffn(hidden_states)
323
+ outputs = self.ffn(ln_outputs)
324
+ if self.dropout is not None:
325
+ outputs = self.dropout(outputs)
326
+ hidden_states = (hidden_states + outputs) / 1.05
327
+ return hidden_states
328
+
329
+
330
+ class CpmBeeTransformerBlock(nn.Module):
331
+ def __init__(self, config: CpmBeeConfig, mask_att: bool = False, mask_ffn: bool = False):
332
+ super().__init__()
333
+ self.mask_att = mask_att
334
+ self.mask_ffn = mask_ffn
335
+
336
+ if not self.mask_att:
337
+ self.self_att = CpmBeeSelfAttentionBlock(config)
338
+ if not self.mask_ffn:
339
+ self.ffn = CpmBeeFFNBlock(config)
340
+
341
+ def forward(
342
+ self,
343
+ hidden_states: torch.Tensor,
344
+ attention_mask: torch.Tensor,
345
+ position_bias: Optional[torch.Tensor] = None,
346
+ output_attentions: Optional[bool] = False,
347
+ past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
348
+ use_cache: Optional[bool] = None,
349
+ ):
350
+ """
351
+ Args:
352
+ hidden_states (`torch.Tensor`):
353
+ Input to the layer of shape `(batch, seq_len, dim_model)`
354
+ attention_mask (`torch.Tensor`):
355
+ Avoid invalid areas to participate in the calculation of shape `(batch, seq_len, seq_len)`
356
+ position_bias (`torch.Tensor`):
357
+ Provides position information to attention mechanism of shape `(num_heads, seq_len, seq_len)`
358
+ output_attentions (`bool`, *optional*):
359
+ Whether or not to return the attentions tensors of all attention layers.
360
+ past_key_values (`Tuple[torch.Tensor, torch.Tensor])`, *optional*):
361
+ Cached past key and value projection states
362
+ use_cache (`bool`, *optional*):
363
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
364
+ (see `past_key_values`).
365
+ """
366
+ if not self.mask_att:
367
+ hidden_states = self.self_att(
368
+ hidden_states,
369
+ attention_mask=attention_mask,
370
+ position_bias=position_bias,
371
+ output_attentions=output_attentions,
372
+ past_key_values=past_key_values,
373
+ use_cache=use_cache,
374
+ )
375
+
376
+ hidden_states, attn_weights, current_key_value = hidden_states
377
+ else:
378
+ attn_weights, current_key_value = None, (None, None)
379
+
380
+ if not self.mask_ffn:
381
+ hidden_states = self.ffn(hidden_states)
382
+
383
+ return hidden_states, attn_weights, current_key_value
384
+
385
+
386
+ class CpmBeeEncoder(nn.Module):
387
+ def __init__(self, config: CpmBeeConfig):
388
+ super().__init__()
389
+ self.num_layers = config.num_hidden_layers
390
+ if config.mask_modules is not None:
391
+ assert len(config.mask_modules) == self.num_layers, "The total number of masks should equal to num_layers"
392
+ for mask_module in config.mask_modules:
393
+ assert len(mask_module) == 2, "For encoder, each mask should be (mask_att, mask_ffn)"
394
+ else:
395
+ config.mask_modules = [(False, False)] * self.num_layers
396
+
397
+ self.layers = nn.ModuleList(
398
+ [
399
+ CpmBeeTransformerBlock(
400
+ config, mask_att=config.mask_modules[ith][0], mask_ffn=config.mask_modules[ith][1]
401
+ )
402
+ for ith in range(self.num_layers)
403
+ ]
404
+ )
405
+
406
+ self.output_layernorm = CpmBeeLayerNorm(config)
407
+
408
+ def forward(
409
+ self,
410
+ hidden_states: torch.Tensor,
411
+ attention_mask: torch.Tensor,
412
+ position_bias: torch.Tensor,
413
+ output_attentions: Optional[bool] = None,
414
+ output_hidden_states: Optional[bool] = None,
415
+ past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
416
+ use_cache: Optional[bool] = None,
417
+ ):
418
+ """
419
+ Args:
420
+ hidden_states (`torch.Tensor`):
421
+ Input to the layer of shape `(batch, seq_len, dim_model)`
422
+ attention_mask (`torch.Tensor`):
423
+ Avoid invalid areas to participate in the calculation of shape `(batch, seq_len, seq_len)`
424
+ position_bias (`torch.Tensor`):
425
+ Provides position information to attention mechanism of shape `(num_heads, seq_len, seq_len)`
426
+ output_attentions (`bool`, *optional*):
427
+ Whether or not to return the attentions tensors of all attention layers.
428
+ output_hidden_states (`bool`, *optional*):
429
+ Whether or not to return the hidden states of all layers.
430
+ past_key_values (`Tuple[torch.Tensor, torch.Tensor])`, *optional*):
431
+ Cached past key and value projection states
432
+ use_cache (`bool`, *optional*):
433
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
434
+ (see `past_key_values`).
435
+ """
436
+ all_hidden_states = () if output_hidden_states else None
437
+ all_self_attns = () if output_attentions else None
438
+ current_key_values = () if use_cache else None
439
+
440
+ for i, layer in enumerate(self.layers):
441
+ if output_hidden_states:
442
+ all_hidden_states += (hidden_states,)
443
+ layer_outputs = layer(
444
+ hidden_states,
445
+ attention_mask,
446
+ position_bias,
447
+ output_attentions=output_attentions,
448
+ past_key_values=past_key_values[i] if past_key_values else None,
449
+ use_cache=use_cache,
450
+ )
451
+ hidden_states, attn_weights, current_key_value = layer_outputs
452
+ if output_attentions:
453
+ all_self_attns += (attn_weights,)
454
+ if current_key_value is not None:
455
+ current_key_values = current_key_values + (current_key_value,)
456
+
457
+ hidden_states = self.output_layernorm(hidden_states)
458
+
459
+ if output_hidden_states:
460
+ all_hidden_states += (hidden_states,)
461
+
462
+ return hidden_states, current_key_values, all_hidden_states, all_self_attns
463
+
464
+
465
+ class CpmBeeBucketPositionBias(nn.Module):
466
+ def __init__(self, config: CpmBeeConfig) -> None:
467
+ super().__init__()
468
+
469
+ self.num_heads = config.num_attention_heads
470
+ self.num_buckets = config.position_bias_num_buckets
471
+ self.num_segment_bucket = config.position_bias_num_segment_buckets
472
+ self.max_distance = config.position_bias_max_distance
473
+
474
+ self.relative_attention_bias = nn.Parameter(
475
+ torch.empty(
476
+ config.position_bias_num_buckets + config.position_bias_num_segment_buckets,
477
+ config.num_attention_heads,
478
+ dtype=config.torch_dtype,
479
+ ),
480
+ )
481
+
482
+ def forward(self, query_pos: torch.Tensor, key_pos: torch.Tensor, rel_buckets: torch.Tensor):
483
+ with torch.no_grad():
484
+ batch = key_pos.size(0)
485
+ keylen = key_pos.size(1)
486
+ querylen = query_pos.size(1)
487
+
488
+ if key_pos.size(0) != query_pos.size(0):
489
+ raise AssertionError(
490
+ f"key_pos.size(0) should be equal to query_pos.size(0), but got {key_pos.size(0)} and {query_pos.size(0)}!"
491
+ )
492
+ if rel_buckets.size(0) != batch:
493
+ raise AssertionError(
494
+ f"rel_buckets.size(0) should be equal to batch, but got {rel_buckets.size(0)} and {batch}!"
495
+ )
496
+ if rel_buckets.size(1) != querylen:
497
+ raise AssertionError(
498
+ f"rel_buckets.size(1) should be equal to querylen, but got {rel_buckets.size(1)} and {querylen}!"
499
+ )
500
+ if rel_buckets.size(2) != keylen:
501
+ raise AssertionError(
502
+ f"rel_buckets.size(2) should be equal to keylen, but got {rel_buckets.size(2)} and {keylen}!"
503
+ )
504
+
505
+ relative_position_bucket = rel_buckets - 1 + self.num_buckets
506
+
507
+ inner_segment_bucket = self._position_bucket(
508
+ key_pos[..., None, :] - query_pos[..., :, None],
509
+ num_buckets=self.num_buckets,
510
+ max_distance=self.max_distance,
511
+ )
512
+ relative_position_bucket = torch.where(
513
+ rel_buckets == 0,
514
+ inner_segment_bucket,
515
+ relative_position_bucket,
516
+ )
517
+
518
+ embeds = nn.functional.embedding(relative_position_bucket, self.relative_attention_bias)
519
+ embeds = embeds.permute(0, 3, 1, 2).contiguous()
520
+ return embeds
521
+
522
+ def _position_bucket(self, relative_position, num_buckets=32, max_distance=128):
523
+ relative_buckets = 0
524
+ num_buckets //= 2
525
+ relative_buckets = (relative_position > 0).to(torch.int32) * num_buckets
526
+ relative_position = torch.abs(relative_position)
527
+ max_exact = num_buckets // 2
528
+ is_small = relative_position < max_exact
529
+ relative_postion_if_large = max_exact + (
530
+ torch.log(relative_position.float() / max_exact)
531
+ / math.log(max_distance / max_exact)
532
+ * (num_buckets - max_exact)
533
+ ).to(torch.int32)
534
+ relative_postion_if_large = torch.min(
535
+ relative_postion_if_large,
536
+ torch.full_like(relative_postion_if_large, num_buckets - 1),
537
+ )
538
+ relative_buckets += torch.where(is_small, relative_position.to(torch.int32), relative_postion_if_large)
539
+ return relative_buckets
540
+
541
+
542
+ # Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->CPMBee
543
+ class CpmBeeOutput(nn.Module):
544
+ def __init__(self, config):
545
+ super().__init__()
546
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
547
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
548
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
549
+
550
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
551
+ hidden_states = self.dense(hidden_states)
552
+ hidden_states = self.dropout(hidden_states)
553
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
554
+ return hidden_states
555
+
556
+
557
+ class CpmBeeRotaryEmbedding(nn.Module):
558
+ """
559
+ RotaryEmbedding embeds the unk token and special token. It will embeds the "...<mask>...<mask>...<unk>...<unk>..."
560
+ to "...<mask_0>...<mask_1>...<unk_0>...<unk_1>..."" to help model to specify different special tokens and unk
561
+ tokens.
562
+ """
563
+
564
+ def __init__(self, config: CpmBeeConfig):
565
+ super().__init__()
566
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, config.hidden_size, 2, dtype=torch.float32) / config.hidden_size))
567
+ self.distance_scale = config.distance_scale
568
+ self.dtype = config.torch_dtype
569
+ self.inv_freq = inv_freq.to(config.torch_dtype)
570
+
571
+ def forward(self, x: torch.Tensor, x_pos: torch.Tensor):
572
+ inv_freq = self.inv_freq.to(device=x.device, dtype=self.dtype)
573
+
574
+ x_pos = x_pos * self.distance_scale
575
+ freqs = x_pos[..., None].to(self.dtype) * inv_freq[None, :] # (..., dim/2)
576
+
577
+ emb = torch.cat((freqs, freqs), dim=-1) # (..., dim)
578
+ emb_cos = emb.cos() # (..., dim)
579
+ emb_sin = emb.sin() # (..., dim)
580
+
581
+ rotate_x = torch.cat([-x[..., x.size(-1) // 2 :], x[..., : x.size(-1) // 2]], dim=-1) # (..., dim)
582
+
583
+ return x * emb_cos + rotate_x * emb_sin
584
+
585
+
586
+ class CpmBeeEmbeddingExt(nn.Embedding):
587
+ """
588
+ Contains a RotaryEmbedding.
589
+ """
590
+
591
+ def __init__(self, config: CpmBeeConfig):
592
+ super().__init__(config.vocab_size, config.hidden_size, dtype=config.torch_dtype)
593
+ self.dim_model = config.hidden_size
594
+ self.rotary_emb = CpmBeeRotaryEmbedding(config)
595
+
596
+ def forward(self, ids: torch.Tensor, ids_sub: torch.Tensor):
597
+ embeds = super().forward(ids) / math.sqrt(self.dim_model)
598
+ return self.rotary_emb(embeds, ids_sub)
599
+
600
+ def projection(self, x: torch.Tensor, ext_table: Optional[torch.Tensor] = None):
601
+ logits = nn.functional.linear(x / math.sqrt(self.dim_model), self.weight)
602
+ if ext_table is not None:
603
+ logits_ext = nn.functional.linear(x, ext_table)
604
+ logits = torch.cat([logits, logits_ext], dim=-1)
605
+ return logits
606
+
607
+
608
+ class CpmBeePreTrainedModel(PreTrainedModel):
609
+ """
610
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
611
+ models.
612
+ """
613
+
614
+ config_class = CpmBeeConfig
615
+ base_model_prefix = "cpmbee"
616
+ supports_gradient_checkpointing = True
617
+ _keys_to_ignore_on_load_missing = [r"position_ids"]
618
+
619
+ def _init_weights(self, module):
620
+ """Initialize the weights"""
621
+ if isinstance(module, nn.Linear):
622
+ module.weight.data.normal_(mean=0.0, std=self.config.init_std)
623
+ if module.bias is not None:
624
+ module.bias.data.zero_()
625
+ # still needed
626
+ elif isinstance(module, CpmBeeEmbeddingExt):
627
+ module.weight.data.normal_(mean=0.0, std=self.config.init_std)
628
+ elif isinstance(module, nn.LayerNorm):
629
+ module.bias.data.zero_()
630
+ module.weight.data.fill_(1.0)
631
+ elif isinstance(module, CpmBeeLayerNorm):
632
+ module.weight.data.fill_(1.0)
633
+ elif isinstance(module, CpmBeeBucketPositionBias):
634
+ module.relative_attention_bias.data.normal_(mean=0.0, std=self.config.init_std)
635
+
636
+ def _set_gradient_checkpointing(self, module, value=False):
637
+ if isinstance(module, CpmBeeEncoder):
638
+ module.gradient_checkpointing = value
639
+
640
+
641
+ CPMBEE_START_DOCSTRING = r"""
642
+ This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
643
+ it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
644
+ behavior.
645
+
646
+ Parameters
647
+ config ([`~CpmBeeConfig`]): Model configuration class with all the parameters of the
648
+ Initializing with a config file does not load the weights associated with the model, only the
649
+ configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
650
+ """
651
+
652
+ CPMBEE_INPUTS_DOCSTRING = r"""
653
+ Args:
654
+ input_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
655
+ Indices of input sequence tokens in the vocabulary.
656
+
657
+ Indices can be obtained using [`CPMBeeTokenizer`]. See [`PreTrainedTokenizer.encode`] and
658
+ [`PreTrainedTokenizer.__call__`] for details.
659
+
660
+ [What are input IDs?](../glossary#input-ids)
661
+ input_id_sub (`torch.Tensor` of shape `(batch_size, seq_len)`):
662
+ Subscription of input sequence tokens in the vocabulary.
663
+
664
+ Subscription of normal text will be zero while the special tokens of each group will be the 0, 1, 2, ...
665
+ <ans_0>, <ans_1>, <ans_2> ... belongs to group <ans>. <mask_0>, <mask_1>, <mask_2> ... belongs to group
666
+ <mask>.
667
+ position (`torch.Tensor` of shape `(batch_size, seq_len)`):
668
+ The position of input sequence tokens in the vocabulary for each segment. if segment1 is 0, 1, 2 and
669
+ segment2 is 0, 1, 2, 3, the position will be 0, 1, 2, 0, 1, 2, 3
670
+ context (`torch.Tensor` of shape `(batch_size, seq_len)`):
671
+ Whether this token id is context or not. If is context, the value is 1. If not, the value is 0. If a token
672
+ id is context, it does not need to be predicted.
673
+ sample_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
674
+ Give a sample id to every token id. The token ids with same sample ids belongs to the same sample.
675
+ num_segments (`torch.Tensor` of shape `(batch_size, seq_len)`):
676
+ Total number of segments in the current input.
677
+ segment (`torch.Tensor` of shape `(batch_size, seq_len)`):
678
+ Give a segment id to every token id. The token ids with same segment ids belongs to the same sample.
679
+
680
+ Generally, a string key or value in input data will be a segment. For example, input {"input": "hello, ",
681
+ "<ans>": ""}, the segments includes: "input", "hello, ", "<ans>" and "".
682
+ segment_rel_offset (`torch.Tensor` of shape `(batch_size, seq_len)`):
683
+ The offset of segment rel.
684
+ segment_rel (`torch.Tensor` of shape `(batch_size, seq_len)`):
685
+ The segment relevance. A relative implementation of measuring the importance of segments.
686
+ past_states (`Dict[str, Union[torch.Tensor, List]]`):
687
+ Store the history information including position, context, sample_ids, num_segments, segment and
688
+ past_key_values.
689
+ output_attentions (`bool`, *optional*):
690
+ Whether or not to return the attentions tensors of all attention layers.
691
+ output_hidden_states (`bool`, *optional*):
692
+ Whether or not to return the hidden states of all layers.
693
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
694
+ A dummy arguments for CPMBee. The `past_states` contains pre-computed hidden-states (key and values in the
695
+ self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values` input) and
696
+ other history arguments to speed up sequential decoding.
697
+ use_cache (`bool`, *optional*):
698
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
699
+ `past_key_values`).
700
+ labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
701
+ Labels for computing the masked language modeling loss.
702
+ return_dict (`bool`, *optional*):
703
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
704
+ """
705
+
706
+
707
+ @add_start_docstrings(
708
+ "The bare CPMBee Model outputting raw hidden-states without any specific head on top.",
709
+ CPMBEE_START_DOCSTRING,
710
+ )
711
+ class CpmBeeModel(CpmBeePreTrainedModel):
712
+ def __init__(self, config: CpmBeeConfig):
713
+ super().__init__(config)
714
+ if config.half:
715
+ config.torch_dtype = torch.half
716
+ else:
717
+ config.torch_dtype = torch.float
718
+ self.encoder = CpmBeeEncoder(config)
719
+ self.input_embedding = CpmBeeEmbeddingExt(config)
720
+ self.position_bias = CpmBeeBucketPositionBias(config)
721
+ self.vocab_size = config.vocab_size
722
+ self.post_init()
723
+
724
+ def get_input_embeddings(self):
725
+ return self.input_embedding
726
+
727
+ def set_input_embeddings(self, embeddings, **kwargs):
728
+ self.input_embedding = embeddings
729
+
730
+ @add_start_docstrings_to_model_forward(CPMBEE_INPUTS_DOCSTRING)
731
+ @add_code_sample_docstrings(
732
+ checkpoint=_CHECKPOINT_FOR_DOC,
733
+ output_type=BaseModelOutputWithPast,
734
+ config_class=_CONFIG_FOR_DOC,
735
+ )
736
+ def forward(
737
+ self,
738
+ input_ids: torch.Tensor,
739
+ input_id_sub: Optional[torch.Tensor] = None,
740
+ position: Optional[torch.Tensor] = None,
741
+ context: Optional[torch.Tensor] = None,
742
+ sample_ids: Optional[torch.Tensor] = None,
743
+ num_segments: Optional[torch.Tensor] = None,
744
+ segment: Optional[torch.Tensor] = None,
745
+ segment_rel_offset: Optional[torch.Tensor] = None,
746
+ segment_rel: Optional[torch.Tensor] = None,
747
+ past_states: Optional[Dict] = None,
748
+ output_attentions: Optional[bool] = None,
749
+ output_hidden_states: Optional[bool] = None,
750
+ past_key_values: Optional[List] = None,
751
+ use_cache: Optional[bool] = None,
752
+ return_dict: Optional[bool] = None,
753
+ **kwargs,
754
+ ):
755
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
756
+ output_hidden_states = (
757
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
758
+ )
759
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
760
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
761
+
762
+ # dummy setting for common tests
763
+ if input_id_sub is None:
764
+ dtype, device = input_ids.dtype, input_ids.device
765
+ batch, seq_length = input_ids.size()
766
+ segment = torch.where(input_ids != 0, 2, 0).to(dtype=dtype, device=device)
767
+ context = torch.full((batch, seq_length), 1, dtype=dtype, device=device)
768
+ position = torch.arange(seq_length, dtype=dtype, device=device).repeat(batch, 1)
769
+ input_id_sub = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
770
+ segment_rel_offset = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
771
+ segment_rel = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
772
+ num_segments = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
773
+ sample_ids = torch.zeros_like(input_ids)
774
+
775
+ with torch.no_grad():
776
+ if past_states is None:
777
+ present_position = position
778
+ present_context = context
779
+ present_sample_ids = sample_ids
780
+ present_num_segments = num_segments
781
+ present_segments = segment
782
+ present_buffer = None
783
+ else:
784
+ present_position = torch.cat([past_states["buffer_position"], position], dim=-1)
785
+ present_context = torch.cat([past_states["buffer_context"], context], dim=-1)
786
+ present_sample_ids = torch.cat([past_states["buffer_sample_ids"], sample_ids], dim=-1)
787
+ present_num_segments = torch.cat([past_states["buffer_num_segments"], num_segments], dim=-1)
788
+ present_segments = torch.cat([past_states["buffer_segments"], segment], dim=-1)
789
+ present_buffer = past_states["buffer"]
790
+
791
+ batch = input_ids.size(0)
792
+ len_q = input_ids.size(1)
793
+ len_buffer = present_position.size(1)
794
+
795
+ segment_rel_2d = torch.masked_fill(
796
+ segment[:, :, None] * num_segments[:, :, None]
797
+ + present_segments[:, None, :]
798
+ + segment_rel_offset[:, :, None],
799
+ ~((sample_ids[:, :, None] == present_sample_ids[:, None, :])), # not in the same sample
800
+ 0, # avoid torch.gather overflow
801
+ ).view(batch, len_q * len_buffer)
802
+
803
+ segment_bucket = torch.gather(
804
+ input=segment_rel,
805
+ dim=1,
806
+ index=segment_rel_2d.long(),
807
+ ).view(batch, len_q, len_buffer)
808
+
809
+ segment_bucket.masked_fill_(
810
+ ~((sample_ids[:, :, None] == present_sample_ids[:, None, :])), # not in the same span or sample
811
+ 1, # bucket is used for in-context samples
812
+ )
813
+
814
+ # directional mask
815
+ directional_mask_2d = present_position[:, None, :] <= position[:, :, None]
816
+ # sample mask
817
+ sample_mask_2d = (sample_ids[:, :, None] == 0) | (sample_ids[:, :, None] == present_sample_ids[:, None, :])
818
+ # context mask
819
+ attention_mask = present_context[:, None, :] | (
820
+ context[:, :, None].logical_not() & directional_mask_2d.view(batch, len_q, len_buffer)
821
+ )
822
+ # span mask
823
+ attention_mask = attention_mask & sample_mask_2d
824
+ # length mask
825
+ mask_1d = present_num_segments != 0
826
+ attention_mask = mask_1d.view(batch, 1, len_buffer) & attention_mask
827
+
828
+ hidden_states = self.input_embedding(input_ids, input_id_sub)
829
+ position_bias = self.position_bias(position, present_position, segment_bucket)
830
+ hidden_states, present_key_values, all_hidden_states, all_attentions = self.encoder(
831
+ hidden_states,
832
+ attention_mask,
833
+ position_bias,
834
+ output_attentions,
835
+ output_hidden_states,
836
+ present_buffer,
837
+ use_cache,
838
+ )
839
+
840
+ if not return_dict:
841
+ return tuple(
842
+ v for v in [hidden_states, present_key_values, all_hidden_states, all_attentions] if v is not None
843
+ )
844
+
845
+ return BaseModelOutputWithPast(
846
+ last_hidden_state=hidden_states,
847
+ past_key_values=present_key_values,
848
+ hidden_states=all_hidden_states,
849
+ attentions=all_attentions,
850
+ )
851
+
852
+
853
+ class CpmBeeBeamHypotheses(BeamHypotheses):
854
+ def __init__(self, num_beams: int, length_penalty: float, early_stopping: bool, max_length: Optional[int] = None):
855
+ """
856
+ Override BeamHypotheses for CpmBee. The hyp to add is list but not tensor.
857
+ """
858
+ super().__init__(num_beams, length_penalty, early_stopping, max_length)
859
+
860
+ def add(self, hyp: List, sum_logprobs: float, beam_indices: Optional[torch.LongTensor] = None):
861
+ """
862
+ Add a new hypothesis to the list.
863
+ """
864
+ score = sum_logprobs / (len(hyp) ** self.length_penalty)
865
+ if len(self) < self.num_beams or score > self.worst_score:
866
+ self.beams.append((score, hyp, beam_indices))
867
+ if len(self) > self.num_beams:
868
+ sorted_next_scores = sorted([(s, idx) for idx, (s, _, _) in enumerate(self.beams)])
869
+ del self.beams[sorted_next_scores[0][1]]
870
+ self.worst_score = sorted_next_scores[1][0]
871
+ else:
872
+ self.worst_score = min(score, self.worst_score)
873
+
874
+
875
+ class CpmBeeBeamSearchScorer(BeamSearchScorer):
876
+ """
877
+ Override BeamSearchScorer for CPMBee to support:
878
+ 1. Replace beam_tokens by beam_states, containing `idx`, `ans`, `nx_token_id`...
879
+ 2. The `process` will update the beam_states
880
+ 3. The `finalize` will just return the best hypotheses as a list.
881
+ """
882
+
883
+ def __init__(
884
+ self,
885
+ batch_size: int,
886
+ num_beams: int,
887
+ device: torch.device,
888
+ length_penalty: Optional[float] = 1.0,
889
+ do_early_stopping: Optional[Union[bool, str]] = False,
890
+ num_beam_hyps_to_keep: Optional[int] = 1,
891
+ num_beam_groups: Optional[int] = 1,
892
+ max_length: Optional[int] = None,
893
+ **model_kwargs,
894
+ ):
895
+ self.num_beams = num_beams
896
+ self.device = device
897
+ self.length_penalty = length_penalty
898
+ self.do_early_stopping = do_early_stopping
899
+ self.num_beam_hyps_to_keep = num_beam_hyps_to_keep
900
+ self.num_beam_groups = num_beam_groups
901
+ self.group_size = self.num_beams // self.num_beam_groups
902
+
903
+ self._is_init = False
904
+ self._beam_hyps = [
905
+ CpmBeeBeamHypotheses(
906
+ num_beams=self.num_beams,
907
+ length_penalty=self.length_penalty,
908
+ early_stopping=self.do_early_stopping,
909
+ max_length=max_length,
910
+ )
911
+ for _ in range(batch_size)
912
+ ]
913
+ self._done = torch.tensor([False for _ in range(batch_size)], dtype=torch.bool, device=self.device)
914
+
915
+ self.beam_states = []
916
+ for sent_id in range(batch_size):
917
+ instance_beam_states = []
918
+
919
+ for _ in range(self.num_beams):
920
+ instance_beam_states.append(
921
+ {
922
+ "idx": 0,
923
+ "ans": [],
924
+ "nx_token_id": 6,
925
+ "nx_token_sub": 0,
926
+ "nx_segment_id": model_kwargs["other_info"][sent_id]["predict_segments"][0][0],
927
+ "nx_position": 0,
928
+ }
929
+ )
930
+ self.beam_states.append(instance_beam_states)
931
+
932
+ def process(
933
+ self,
934
+ batch_size: int,
935
+ cur_len: int,
936
+ _next_scores: torch.FloatTensor,
937
+ next_scores: torch.FloatTensor,
938
+ next_tokens: torch.LongTensor,
939
+ vocab_size: Optional[int] = None,
940
+ pad_token_id: Optional[int] = None,
941
+ bos_token_id: Optional[int] = None,
942
+ eos_token_id: Optional[Union[int, List[int]]] = None,
943
+ max_length: Optional[int] = None,
944
+ ext_table_sub_cpu: Optional[torch.Tensor] = None,
945
+ ext_table_ids_cpu: Optional[torch.Tensor] = None,
946
+ **model_kwargs,
947
+ ) -> Tuple[torch.Tensor]:
948
+ next_beam_state = []
949
+ for sent_id in range(batch_size):
950
+ self._done[sent_id] = self._done[sent_id] or self._beam_hyps[sent_id].is_done(
951
+ next_scores[sent_id].max().item(), cur_len
952
+ )
953
+ if self._done[sent_id]:
954
+ next_beam_state.append(
955
+ [
956
+ (
957
+ {
958
+ "idx": 0,
959
+ "ans": [],
960
+ "nx_token_id": pad_token_id,
961
+ "nx_token_sub": 0,
962
+ "nx_segment_id": 0,
963
+ "nx_position": 0,
964
+ },
965
+ 0,
966
+ 0,
967
+ )
968
+ ]
969
+ * self.num_beams
970
+ )
971
+ continue
972
+
973
+ next_instance_beam_states = []
974
+
975
+ for idx, value in zip(next_tokens[sent_id], next_scores[sent_id]):
976
+ beam_id = torch.div(idx, _next_scores.size(-1), rounding_mode="floor").item()
977
+ word_id = (idx % _next_scores.size(-1)).item()
978
+
979
+ curr_info = self.beam_states[sent_id][beam_id]
980
+ if (
981
+ word_id == eos_token_id
982
+ and (curr_info["idx"] + 1 == len(model_kwargs["other_info"][sent_id]["predict_segments"]))
983
+ ) or cur_len == max_length:
984
+ self._beam_hyps[sent_id].add(
985
+ self.beam_states[sent_id][beam_id]["ans"]
986
+ + [
987
+ (
988
+ word_id,
989
+ model_kwargs["other_info"][sent_id]["predict_segments"][curr_info["idx"]][1],
990
+ )
991
+ ],
992
+ value.item(),
993
+ )
994
+ elif word_id == eos_token_id:
995
+ next_instance_beam_states.append(
996
+ (
997
+ {
998
+ "idx": curr_info["idx"] + 1,
999
+ "ans": curr_info["ans"]
1000
+ + [
1001
+ (
1002
+ word_id,
1003
+ model_kwargs["other_info"][sent_id]["predict_segments"][curr_info["idx"]][1],
1004
+ )
1005
+ ],
1006
+ "nx_token_id": bos_token_id,
1007
+ "nx_token_sub": 0,
1008
+ "nx_segment_id": model_kwargs["other_info"][sent_id]["predict_segments"][
1009
+ curr_info["idx"] + 1
1010
+ ][0],
1011
+ "nx_position": 0,
1012
+ },
1013
+ value.item(),
1014
+ sent_id * self.num_beams + beam_id,
1015
+ )
1016
+ )
1017
+
1018
+ else:
1019
+ raw_word_id = word_id
1020
+ word_id_sub = 0
1021
+ if word_id >= vocab_size:
1022
+ word_id -= vocab_size
1023
+ word_id_sub = int(ext_table_sub_cpu[word_id].item())
1024
+ word_id = int(ext_table_ids_cpu[word_id].item())
1025
+
1026
+ next_instance_beam_states.append(
1027
+ (
1028
+ {
1029
+ "idx": curr_info["idx"],
1030
+ "ans": curr_info["ans"]
1031
+ + [
1032
+ (
1033
+ raw_word_id,
1034
+ model_kwargs["other_info"][sent_id]["predict_segments"][curr_info["idx"]][1],
1035
+ )
1036
+ ],
1037
+ "nx_token_id": word_id,
1038
+ "nx_token_sub": word_id_sub,
1039
+ "nx_segment_id": curr_info["nx_segment_id"],
1040
+ "nx_position": curr_info["nx_position"] + 1,
1041
+ },
1042
+ value.item(),
1043
+ sent_id * self.num_beams + beam_id,
1044
+ )
1045
+ )
1046
+
1047
+ if len(next_instance_beam_states) == self.num_beams:
1048
+ break
1049
+ assert len(next_instance_beam_states) == 0 if cur_len == max_length else self.num_beams
1050
+ next_beam_state.append(next_instance_beam_states)
1051
+
1052
+ if cur_len == max_length:
1053
+ return None
1054
+
1055
+ beam_reorder_idx = []
1056
+ beam_new_scores = []
1057
+ beam_states = []
1058
+ for sent_id in range(batch_size):
1059
+ instance_beam_states = []
1060
+ for beam_id in range(self.num_beams):
1061
+ state, value, beam_idx = next_beam_state[sent_id][beam_id]
1062
+ beam_reorder_idx.append(beam_idx)
1063
+ beam_new_scores.append(value)
1064
+ instance_beam_states.append(state)
1065
+ beam_states.append(instance_beam_states)
1066
+ self.beam_states = beam_states
1067
+
1068
+ return UserDict(
1069
+ {
1070
+ "next_beam_scores": torch.tensor(beam_new_scores, device=self.device).view(-1),
1071
+ "next_beam_states": beam_states,
1072
+ "next_beam_indices": torch.tensor(beam_reorder_idx, dtype=torch.int32, device=self.device).view(-1),
1073
+ }
1074
+ )
1075
+
1076
+ def finalize(self) -> Tuple[torch.LongTensor]:
1077
+ results = []
1078
+ for _, hypotheses in enumerate(self._beam_hyps):
1079
+ best_hyp = max(hypotheses.beams, key=lambda x: x[0])[1]
1080
+ results.append(best_hyp)
1081
+ return results
1082
+
1083
+ @staticmethod
1084
+ def apply_repetition_penalty(
1085
+ logits,
1086
+ batch_size,
1087
+ num_beams,
1088
+ prev_output_tokens,
1089
+ repetition_penalty,
1090
+ start_idx=None,
1091
+ end_idx=None,
1092
+ window_size=None,
1093
+ ):
1094
+ # only conduct repetition penalty for the output
1095
+ assert repetition_penalty >= 1, "repetition penalty coefficient should >= 1"
1096
+ # repetition penalty (from CTRL paper https://arxiv.org/abs/1909.05858)
1097
+ for i in range(batch_size * num_beams):
1098
+ if start_idx is None or end_idx is None:
1099
+ output_tokens = prev_output_tokens[i].tolist()
1100
+ else:
1101
+ if end_idx >= start_idx:
1102
+ if window_size:
1103
+ output_tokens = prev_output_tokens[i][
1104
+ max(start_idx, end_idx + 1 - window_size) : end_idx + 1
1105
+ ].tolist()
1106
+ else:
1107
+ output_tokens = prev_output_tokens[i][start_idx : end_idx + 1].tolist()
1108
+ else:
1109
+ output_tokens = []
1110
+ for previous_token in set(output_tokens):
1111
+ # if score < 0 then repetition penalty has to
1112
+ # multiplied to reduce the previous token probability
1113
+ if logits[i, previous_token] < 0:
1114
+ logits[i, previous_token] *= repetition_penalty
1115
+ else:
1116
+ logits[i, previous_token] /= repetition_penalty
1117
+
1118
+
1119
+ @add_start_docstrings(
1120
+ """
1121
+ The CPMBee Model with a language modeling head on top (linear layer with weights tied to the input embeddings).
1122
+ """,
1123
+ CPMBEE_START_DOCSTRING,
1124
+ )
1125
+ class CpmBeeForCausalLM(CpmBeePreTrainedModel):
1126
+ _keys_to_ignore_on_load_missing = [r"lm_head.weight"]
1127
+
1128
+ def __init__(self, config: CpmBeeConfig):
1129
+ super().__init__(config)
1130
+ self.cpmbee = CpmBeeModel(config)
1131
+
1132
+ # lm_head.weight is tied to cpmbee.input_embedding.weight
1133
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1134
+ self.post_init()
1135
+
1136
+ @add_start_docstrings_to_model_forward(CPMBEE_INPUTS_DOCSTRING)
1137
+ @add_code_sample_docstrings(
1138
+ checkpoint=_CHECKPOINT_FOR_DOC,
1139
+ output_type=CausalLMOutputWithPast,
1140
+ config_class=_CONFIG_FOR_DOC,
1141
+ )
1142
+ def forward(
1143
+ self,
1144
+ input_ids: Optional[torch.Tensor] = None,
1145
+ input_id_sub: Optional[torch.Tensor] = None,
1146
+ position: Optional[torch.Tensor] = None,
1147
+ context: Optional[torch.Tensor] = None,
1148
+ sample_ids: Optional[torch.Tensor] = None,
1149
+ num_segments: Optional[torch.Tensor] = None,
1150
+ segment: Optional[torch.Tensor] = None,
1151
+ segment_rel_offset: Optional[torch.Tensor] = None,
1152
+ segment_rel: Optional[torch.Tensor] = None,
1153
+ past_states: Optional[Dict] = None,
1154
+ output_attentions: Optional[bool] = None,
1155
+ output_hidden_states: Optional[bool] = None,
1156
+ past_key_values: Optional[List] = None,
1157
+ use_cache: Optional[bool] = None,
1158
+ labels: Optional[torch.Tensor] = None,
1159
+ return_dict: Optional[bool] = None,
1160
+ ext_table_ids: Optional[torch.Tensor] = None, # (ext_table_size) int32
1161
+ ext_table_sub: Optional[torch.Tensor] = None, # (ext_table_size) int32
1162
+ **kwargs,
1163
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1164
+ r"""
1165
+ Args:
1166
+ input_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
1167
+ Indices of input sequence tokens in the vocabulary.
1168
+
1169
+ Indices can be obtained using [`CPMBeeTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1170
+ [`PreTrainedTokenizer.__call__`] for details.
1171
+
1172
+ [What are input IDs?](../glossary#input-ids)
1173
+ input_id_sub (`torch.Tensor` of shape `(batch_size, seq_len)`):
1174
+ Subscription of input sequence tokens in the vocabulary.
1175
+
1176
+ Subscription of normal text will be zero while the special tokens of each group will be the 0, 1, 2,
1177
+ ... <ans_0>, <ans_1>, <ans_2> ... belongs to group <ans>. <mask_0>, <mask_1>, <mask_2> ... belongs to
1178
+ group <mask>.
1179
+ position (`torch.Tensor` of shape `(batch_size, seq_len)`):
1180
+ The position of input sequence tokens in the vocabulary for each segment. if segment1 is 0, 1, 2 and
1181
+ segment2 is 0, 1, 2, 3, the position will be 0, 1, 2, 0, 1, 2, 3
1182
+ context (`torch.Tensor` of shape `(batch_size, seq_len)`):
1183
+ Whether this token id is context or not. If is context, the value is 1. If not, the value is 0. If a
1184
+ token id is context, it does not need to be predicted.
1185
+ sample_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
1186
+ Give a sample id to every token id. The token ids with same sample ids belongs to the same sample.
1187
+ num_segments (`torch.Tensor` of shape `(batch_size, seq_len)`):
1188
+ Total number of segments in the current input.
1189
+ segment (`torch.Tensor` of shape `(batch_size, seq_len)`):
1190
+ Give a segment id to every token id. The token ids with same segment ids belongs to the same sample.
1191
+
1192
+ Generally, a string key or value in input data will be a segment. For example, input {"input": "hello,
1193
+ ", "<ans>": ""}, the segments includes: "input", "hello, ", "<ans>" and "".
1194
+ segment_rel_offset (`torch.Tensor` of shape `(batch_size, seq_len)`):
1195
+ The offset of segment rel.
1196
+ segment_rel (`torch.Tensor` of shape `(batch_size, seq_len)`):
1197
+ The segment relevance. A relative implementation of measuring the importance of segments.
1198
+ past_states (`Dict[str, Union[torch.Tensor, List]]`):
1199
+ Store the history information including position, context, sample_ids, num_segments, segment and
1200
+ past_key_values.
1201
+ output_attentions (`bool`, *optional*):
1202
+ Whether or not to return the attentions tensors of all attention layers.
1203
+ output_hidden_states (`bool`, *optional*):
1204
+ Whether or not to return the hidden states of all layers.
1205
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
1206
+ A dummy arguments for CPMBee. The `past_states` contains pre-computed hidden-states (key and values in
1207
+ the self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values`
1208
+ input) and other history arguments to speed up sequential decoding.
1209
+ use_cache (`bool`, *optional*):
1210
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
1211
+ (see `past_key_values`).
1212
+ labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
1213
+ Labels for computing the masked language modeling loss.
1214
+ return_dict (`bool`, *optional*):
1215
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
1216
+ ext_table_ids (`torch.Tensor`, *optional*):
1217
+ ext_table ids for embedding projection.
1218
+ ext_table_sub (`torch.Tensor`, *optional*):
1219
+ ext_table subscriptions for embedding projection.
1220
+
1221
+ Example:
1222
+
1223
+ Text Generation with CpmBeeForCausalLM.
1224
+ ```python
1225
+ >>> from transformers import CpmBeeTokenizer, CpmBeeForCausalLM
1226
+
1227
+ >>> texts = {"input": "今天天气不错,", "<ans>": ""}
1228
+ >>> model = CpmBeeForCausalLM.from_pretrained("openbmb/cpm-bee-10b")
1229
+ >>> tokenizer = CPMBeeTokenizer.from_pretrained("openbmb/cpm-bee-10b")
1230
+ >>> output_texts = model.generate({"input": "今天天气不错,", "<ans>": ""}, tokenizer)
1231
+ >>> print(output_texts)
1232
+ {'input': '今天天气不错,', '<ans>': '适合睡觉。'}
1233
+ ```
1234
+ """
1235
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1236
+
1237
+ model_output = self.cpmbee(
1238
+ input_ids,
1239
+ input_id_sub,
1240
+ position,
1241
+ context,
1242
+ sample_ids,
1243
+ num_segments,
1244
+ segment,
1245
+ segment_rel_offset,
1246
+ segment_rel,
1247
+ past_states,
1248
+ output_attentions,
1249
+ output_hidden_states,
1250
+ past_key_values,
1251
+ use_cache,
1252
+ return_dict,
1253
+ )
1254
+ hidden_states = model_output.last_hidden_state if return_dict else model_output[0]
1255
+
1256
+ if ext_table_ids is not None:
1257
+ ext_table = self.cpmbee.input_embedding(ext_table_ids, ext_table_sub)
1258
+ else:
1259
+ ext_table = None
1260
+ logits = self.cpmbee.input_embedding.projection(hidden_states, ext_table)
1261
+
1262
+ loss = None
1263
+ if labels is not None:
1264
+ loss_func = nn.CrossEntropyLoss()
1265
+ loss = loss_func(logits.view(-1, logits.size(-1)), labels.view(-1))
1266
+
1267
+ if not return_dict:
1268
+ output = (logits,) + model_output[1:]
1269
+ return ((loss,) + output) if loss is not None else output
1270
+
1271
+ return CausalLMOutputWithPast(
1272
+ loss=loss,
1273
+ logits=logits,
1274
+ past_key_values=model_output.past_key_values,
1275
+ hidden_states=model_output.hidden_states,
1276
+ attentions=model_output.attentions,
1277
+ )
1278
+
1279
+ def get_input_embeddings(self):
1280
+ return self.cpmbee.input_embedding
1281
+
1282
+ def set_input_embeddings(self, embeddings):
1283
+ self.cpmbee.input_embedding = embeddings
1284
+
1285
+ def get_output_embeddings(self):
1286
+ return self.lm_head
1287
+
1288
+ def set_output_embeddings(self, new_embeddings):
1289
+ self.lm_head = new_embeddings
1290
+
1291
+ def prepare_inputs_for_generation(
1292
+ self,
1293
+ input_ids: torch.Tensor,
1294
+ batch_size: int,
1295
+ beam_scorer: CpmBeeBeamSearchScorer = None,
1296
+ input_id_subs: Optional[torch.Tensor] = None,
1297
+ input_pos: Optional[torch.Tensor] = None,
1298
+ segment_ids: Optional[torch.Tensor] = None,
1299
+ batch_ext_table_ids: Optional[torch.Tensor] = None,
1300
+ batch_ext_table_sub: Optional[torch.Tensor] = None,
1301
+ other_info: Optional[Dict] = None,
1302
+ **model_kwargs,
1303
+ ):
1304
+ """
1305
+ Choose the current input according to beam states.
1306
+ """
1307
+ # init preparation
1308
+ context = model_kwargs.get("context")
1309
+ sample_ids = model_kwargs.get("sample_ids")
1310
+ segment_rel_offset = model_kwargs.get("segment_rel_offset")
1311
+ num_segments = model_kwargs.get("num_segments")
1312
+ segment_rel = model_kwargs.get("segment_rel")
1313
+ past_states = model_kwargs.get("past_states", None)
1314
+ past_key_values = model_kwargs.get("past_key_values", None)
1315
+ _input_ids = input_ids
1316
+
1317
+ # update input in generation
1318
+ if beam_scorer is not None:
1319
+ tmp_input = []
1320
+ tmp_input_sub = []
1321
+ tmp_position = []
1322
+ tmp_segment = []
1323
+ for sent_id in range(batch_size):
1324
+ for beam_id in range(beam_scorer.num_beams):
1325
+ tmp_input.append(beam_scorer.beam_states[sent_id][beam_id]["nx_token_id"])
1326
+ tmp_input_sub.append(beam_scorer.beam_states[sent_id][beam_id]["nx_token_sub"])
1327
+ tmp_position.append(beam_scorer.beam_states[sent_id][beam_id]["nx_position"])
1328
+ tmp_segment.append(beam_scorer.beam_states[sent_id][beam_id]["nx_segment_id"])
1329
+
1330
+ model_kwargs["input_id_subs"] = input_id_subs = torch.tensor(
1331
+ tmp_input_sub, dtype=torch.int32, device=self.device
1332
+ ).view(batch_size * beam_scorer.num_beams, 1)
1333
+ model_kwargs["input_pos"] = input_pos = torch.tensor(
1334
+ tmp_position, dtype=torch.int32, device=self.device
1335
+ ).view(batch_size * beam_scorer.num_beams, 1)
1336
+ model_kwargs["segment_ids"] = segment_ids = torch.tensor(
1337
+ tmp_segment, dtype=torch.int32, device=self.device
1338
+ ).view(batch_size * beam_scorer.num_beams, 1)
1339
+ input_ids = torch.cat(
1340
+ [
1341
+ input_ids,
1342
+ torch.tensor(tmp_input, dtype=torch.int32, device=self.device).view(
1343
+ batch_size * beam_scorer.num_beams, 1
1344
+ ),
1345
+ ],
1346
+ dim=-1,
1347
+ )
1348
+ _input_ids = input_ids[:, -1:]
1349
+
1350
+ return {
1351
+ "input_ids": _input_ids,
1352
+ "input_id_sub": input_id_subs,
1353
+ "position": input_pos,
1354
+ "context": context,
1355
+ "sample_ids": sample_ids,
1356
+ "segment_rel_offset": segment_rel_offset,
1357
+ "segment": segment_ids,
1358
+ "num_segments": num_segments,
1359
+ "segment_rel": segment_rel,
1360
+ "use_cache": True,
1361
+ "past_key_values": past_key_values,
1362
+ "ext_table_ids": batch_ext_table_ids,
1363
+ "ext_table_sub": batch_ext_table_sub,
1364
+ "past_states": past_states,
1365
+ }, input_ids
1366
+
1367
+ def _update_model_kwargs_for_generation(
1368
+ self,
1369
+ outputs: ModelOutput,
1370
+ model_inputs=None,
1371
+ **model_kwargs,
1372
+ ) -> Dict[str, Any]:
1373
+ """
1374
+ Concatenate the history input and current input.
1375
+ """
1376
+
1377
+ old_past_states = model_kwargs["past_states"]
1378
+ model_kwargs["past_states"] = {
1379
+ "buffer_position": torch.cat([old_past_states["buffer_position"], model_inputs["position"]], dim=-1),
1380
+ "buffer_context": torch.cat([old_past_states["buffer_context"], model_inputs["context"]], dim=-1),
1381
+ "buffer_sample_ids": torch.cat([old_past_states["buffer_sample_ids"], model_inputs["sample_ids"]], dim=-1),
1382
+ "buffer_num_segments": torch.cat(
1383
+ [old_past_states["buffer_num_segments"], model_inputs["num_segments"]], dim=-1
1384
+ ),
1385
+ "buffer_segments": torch.cat([old_past_states["buffer_segments"], model_inputs["segment"]], dim=-1),
1386
+ "buffer": outputs.past_key_values,
1387
+ }
1388
+
1389
+ return model_kwargs
1390
+
1391
+ def _reorder_cache(self, past_key_values: Dict, beam_idx: torch.Tensor):
1392
+ beam_idx = beam_idx.tolist()
1393
+ for kw in past_key_values.keys():
1394
+ if kw == "buffer":
1395
+ buf_list = past_key_values[kw]
1396
+ nw_buf_list = []
1397
+ for buf in buf_list:
1398
+ if buf == (None, None):
1399
+ nw_buf_list.append((None, None))
1400
+ else:
1401
+ k_buf, v_buf = buf
1402
+ nw_buf_list.append((k_buf[beam_idx, :], v_buf[beam_idx, :]))
1403
+ past_key_values[kw] = nw_buf_list
1404
+ else:
1405
+ past_key_values[kw] = past_key_values[kw][beam_idx, :]
1406
+
1407
+ return past_key_values
1408
+
1409
+ @staticmethod
1410
+ def _expand_inputs_for_generation(
1411
+ expand_size: int = 1,
1412
+ is_encoder_decoder: bool = False,
1413
+ input_ids: Optional[torch.LongTensor] = None,
1414
+ **model_kwargs,
1415
+ ) -> Tuple[torch.LongTensor, Dict[str, Any]]:
1416
+ """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]"""
1417
+
1418
+ # do not expand ext_table_ids and ext_table_sub
1419
+ def _expand_dict_for_generation(dict_to_expand):
1420
+ for key in dict_to_expand:
1421
+ if (
1422
+ dict_to_expand[key] is not None
1423
+ and isinstance(dict_to_expand[key], torch.Tensor)
1424
+ and "ext_table" not in key
1425
+ ):
1426
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
1427
+ return dict_to_expand
1428
+
1429
+ if input_ids is not None:
1430
+ input_ids = input_ids.repeat_interleave(expand_size, dim=0)
1431
+
1432
+ model_kwargs = _expand_dict_for_generation(model_kwargs)
1433
+
1434
+ if is_encoder_decoder:
1435
+ if model_kwargs.get("encoder_outputs") is None:
1436
+ raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
1437
+ model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
1438
+
1439
+ return input_ids, model_kwargs
1440
+
1441
+ def adjust_logits_during_generation(
1442
+ self,
1443
+ logits: torch.FloatTensor,
1444
+ batch_size: int,
1445
+ beam_size: int,
1446
+ vocab_size: int,
1447
+ ext_table_ids: torch.Tensor,
1448
+ **model_kwargs,
1449
+ ) -> torch.FloatTensor:
1450
+ """
1451
+ Implement in subclasses of [`PreTrainedModel`] for custom behavior to adjust the logits in the generate method.
1452
+ """
1453
+ for sent_id in range(batch_size):
1454
+ if 1 not in model_kwargs["other_info"][sent_id]["ext_table"]:
1455
+ # unk is not allowed, mask unk
1456
+ logits[sent_id * beam_size : (sent_id + 1) * beam_size, 1] = -10000
1457
+ ext_ids = set()
1458
+ for v in model_kwargs["other_info"][sent_id]["ext_table"].keys():
1459
+ ext_ids.add(v)
1460
+ for ext_id in range(vocab_size, vocab_size + ext_table_ids.size(0)):
1461
+ if ext_id not in ext_ids:
1462
+ logits[sent_id * beam_size : (sent_id + 1) * beam_size, ext_id] = -10000
1463
+ return logits
1464
+
1465
+ def beam_search(
1466
+ self,
1467
+ input_ids: torch.LongTensor,
1468
+ beam_scorer: CpmBeeBeamSearchScorer,
1469
+ repetition_penalty: Optional[float] = 1.0,
1470
+ logits_processor: Optional[LogitsProcessorList] = None,
1471
+ max_length: Optional[int] = None,
1472
+ pad_token_id: Optional[int] = None,
1473
+ eos_token_id: Optional[Union[int, List[int]]] = None,
1474
+ bos_token_id: Optional[Union[int, List[int]]] = None,
1475
+ output_attentions: Optional[bool] = None,
1476
+ output_hidden_states: Optional[bool] = None,
1477
+ output_scores: Optional[bool] = None,
1478
+ return_dict_in_generate: Optional[bool] = None,
1479
+ synced_gpus: bool = False,
1480
+ **model_kwargs,
1481
+ ) -> List:
1482
+ """
1483
+ Override the beam_search for CPMBee.
1484
+ """
1485
+ # init values
1486
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
1487
+ pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
1488
+ eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
1489
+ bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id
1490
+ max_length = max_length if max_length is not None else self.generation_config.max_length
1491
+ output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
1492
+ output_attentions = (
1493
+ output_attentions if output_attentions is not None else self.generation_config.output_attentions
1494
+ )
1495
+ output_hidden_states = (
1496
+ output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
1497
+ )
1498
+ return_dict_in_generate = (
1499
+ return_dict_in_generate
1500
+ if return_dict_in_generate is not None
1501
+ else self.generation_config.return_dict_in_generate
1502
+ )
1503
+
1504
+ batch_size = len(beam_scorer._beam_hyps)
1505
+ num_beams = beam_scorer.num_beams
1506
+
1507
+ batch_beam_size, cur_len = input_ids.shape
1508
+
1509
+ if num_beams * batch_size != batch_beam_size:
1510
+ raise ValueError(
1511
+ f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}."
1512
+ )
1513
+
1514
+ # init attention / hidden states / scores tuples
1515
+ scores = () if (return_dict_in_generate and output_scores) else None
1516
+ beam_indices = (
1517
+ tuple(() for _ in range(batch_beam_size)) if (return_dict_in_generate and output_scores) else None
1518
+ )
1519
+ decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
1520
+ cross_attentions = () if (return_dict_in_generate and output_attentions) else None
1521
+ decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
1522
+
1523
+ # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens
1524
+ # of the first beam are considered to avoid sampling the exact same tokens across all beams.
1525
+ beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=self.device)
1526
+ beam_scores[:, 1:] = -1e9
1527
+ beam_scores = beam_scores.view((batch_size * num_beams,))
1528
+
1529
+ this_peer_finished = False # used by synced_gpus only
1530
+
1531
+ # init inference
1532
+ model_inputs, input_ids = self.prepare_inputs_for_generation(input_ids, batch_size, **model_kwargs)
1533
+ pred_start_index = input_ids.size(-1)
1534
+ outputs = self(
1535
+ **model_inputs,
1536
+ return_dict=True,
1537
+ output_attentions=output_attentions,
1538
+ output_hidden_states=output_hidden_states,
1539
+ )
1540
+
1541
+ # update model_kwargs
1542
+ model_kwargs["past_states"] = {
1543
+ "buffer_position": model_inputs["position"],
1544
+ "buffer_context": model_inputs["context"],
1545
+ "buffer_sample_ids": model_inputs["sample_ids"],
1546
+ "buffer_num_segments": model_inputs["num_segments"],
1547
+ "buffer_segments": model_inputs["segment"],
1548
+ "buffer": outputs.past_key_values,
1549
+ }
1550
+ model_kwargs["context"] = torch.ones(batch_beam_size, dtype=torch.bool, device=self.device).view(
1551
+ batch_beam_size, 1
1552
+ )
1553
+ model_kwargs["sample_ids"] = torch.zeros(batch_beam_size, dtype=torch.int32, device=self.device).view(
1554
+ batch_beam_size, 1
1555
+ )
1556
+ model_kwargs["num_segments"] = model_kwargs["num_segments"][:, -1:]
1557
+ model_kwargs["segment_rel_offset"] = model_kwargs["segment_rel_offset"][:, -1:]
1558
+ model_kwargs["past_key_values"] = outputs.past_key_values
1559
+
1560
+ ext_table_ids_cpu = model_inputs["ext_table_ids"].cpu()
1561
+ ext_table_sub_cpu = model_inputs["ext_table_sub"].cpu()
1562
+
1563
+ cur_len = 0
1564
+ while True:
1565
+ if synced_gpus:
1566
+ # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
1567
+ # The following logic allows an early break if all peers finished generating their sequence
1568
+ this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)
1569
+ # send 0.0 if we finished, 1.0 otherwise
1570
+ dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
1571
+ # did all peers finish? the reduced sum will be 0.0 then
1572
+ if this_peer_finished_flag.item() == 0.0:
1573
+ break
1574
+
1575
+ model_inputs, input_ids = self.prepare_inputs_for_generation(
1576
+ input_ids, batch_size, beam_scorer, **model_kwargs
1577
+ )
1578
+
1579
+ outputs = self(
1580
+ **model_inputs,
1581
+ return_dict=True,
1582
+ output_attentions=output_attentions,
1583
+ output_hidden_states=output_hidden_states,
1584
+ )
1585
+
1586
+ next_token_logits = outputs.logits[:, -1, :]
1587
+
1588
+ if all(beam_scorer._done):
1589
+ break
1590
+ # hack: adjust tokens for Marian. For Marian we have to make sure that the `pad_token_id`
1591
+ # cannot be generated both before and after the `nn.functional.log_softmax` operation.
1592
+ vocab_size = next_token_logits.shape[-1]
1593
+ next_token_logits = self.adjust_logits_during_generation(
1594
+ next_token_logits, batch_size, num_beams, vocab_size, ext_table_ids_cpu, **model_kwargs
1595
+ )
1596
+
1597
+ # repetition_penalty
1598
+ beam_scorer.apply_repetition_penalty(
1599
+ next_token_logits,
1600
+ batch_size,
1601
+ num_beams,
1602
+ model_inputs["input_ids"],
1603
+ repetition_penalty,
1604
+ pred_start_index,
1605
+ model_inputs["input_ids"].size(-1) - 1,
1606
+ None,
1607
+ )
1608
+
1609
+ _next_token_scores = nn.functional.log_softmax(
1610
+ next_token_logits, dim=-1
1611
+ ) # (batch_size * num_beams, vocab_size)
1612
+
1613
+ next_token_scores_processed = logits_processor(input_ids, _next_token_scores)
1614
+ # next_token_scores_processed = _next_token_scores
1615
+ next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(_next_token_scores)
1616
+
1617
+ # Store scores, attentions and hidden_states when required
1618
+ if return_dict_in_generate:
1619
+ if output_scores:
1620
+ scores += (next_token_scores_processed,)
1621
+ if output_attentions:
1622
+ decoder_attentions += (
1623
+ (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
1624
+ )
1625
+ if self.config.is_encoder_decoder:
1626
+ cross_attentions += (outputs.cross_attentions,)
1627
+
1628
+ if output_hidden_states:
1629
+ decoder_hidden_states += (
1630
+ (outputs.decoder_hidden_states,)
1631
+ if self.config.is_encoder_decoder
1632
+ else (outputs.hidden_states,)
1633
+ )
1634
+
1635
+ # reshape for beam search
1636
+ vocab_size = next_token_scores.shape[-1]
1637
+ next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size)
1638
+
1639
+ # Sample 2 next tokens for each beam (so we have some spare tokens and match output of beam search)
1640
+ next_token_scores, next_tokens = torch.topk(
1641
+ next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True
1642
+ )
1643
+
1644
+ beam_outputs = beam_scorer.process(
1645
+ batch_size,
1646
+ cur_len,
1647
+ _next_token_scores,
1648
+ next_token_scores,
1649
+ next_tokens,
1650
+ vocab_size=vocab_size,
1651
+ pad_token_id=pad_token_id,
1652
+ bos_token_id=bos_token_id,
1653
+ eos_token_id=eos_token_id,
1654
+ max_length=max_length,
1655
+ ext_table_ids_cpu=ext_table_ids_cpu,
1656
+ ext_table_sub_cpu=ext_table_sub_cpu,
1657
+ **model_kwargs,
1658
+ )
1659
+ if beam_outputs is None:
1660
+ break
1661
+ beam_idx = beam_outputs["next_beam_indices"]
1662
+ beam_scores = beam_outputs["next_beam_scores"]
1663
+
1664
+ input_ids = input_ids[beam_idx.tolist(), :]
1665
+ model_kwargs = self._update_model_kwargs_for_generation(outputs, model_inputs, **model_kwargs)
1666
+ if model_kwargs["past_states"] is not None:
1667
+ model_kwargs["past_states"] = self._reorder_cache(model_kwargs["past_states"], beam_idx)
1668
+
1669
+ if return_dict_in_generate and output_scores:
1670
+ beam_indices = tuple((beam_indices[beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices))))
1671
+
1672
+ cur_len += 1
1673
+
1674
+ if beam_scorer.is_done or cur_len == max_length + 1:
1675
+ if not synced_gpus:
1676
+ break
1677
+ else:
1678
+ this_peer_finished = True
1679
+
1680
+ sequence_outputs = beam_scorer.finalize()
1681
+
1682
+ return sequence_outputs
1683
+
1684
+ def _generate(
1685
+ self,
1686
+ inputs: Optional[torch.Tensor] = None,
1687
+ generation_config: Optional[GenerationConfig] = None,
1688
+ repetition_penalty: Optional[float] = 1.0,
1689
+ logits_processor: Optional[LogitsProcessorList] = None,
1690
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
1691
+ prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
1692
+ synced_gpus: Optional[bool] = None,
1693
+ streamer: Optional["BaseStreamer"] = None,
1694
+ **kwargs,
1695
+ ) -> List:
1696
+ r"""
1697
+ The generation of CPMBee.
1698
+ 1. It will use beam search as generation strategy.
1699
+ 2. It will use CpmBeeBeamSearchScorer as the beamsearch scorer.
1700
+ """
1701
+ if synced_gpus is None:
1702
+ if is_deepspeed_zero3_enabled() and dist.get_world_size() > 1:
1703
+ synced_gpus = True
1704
+ else:
1705
+ synced_gpus = False
1706
+
1707
+ # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
1708
+ self._validate_model_class()
1709
+
1710
+ # priority: `generation_config` argument > `model.generation_config` (the default generation config)
1711
+ if generation_config is None:
1712
+ # legacy: users may modify the model configuration to control generation -- update the generation config
1713
+ # model attribute accordingly, if it was created from the model config
1714
+ if self.generation_config._from_model_config:
1715
+ new_generation_config = GenerationConfig.from_model_config(self.config)
1716
+ if new_generation_config != self.generation_config:
1717
+ warnings.warn(
1718
+ "You have modified the pretrained model configuration to control generation. This is a"
1719
+ " deprecated strategy to control generation and will be removed soon, in a future version."
1720
+ " Please use a generation configuration file (see"
1721
+ " https://huggingface.co/docs/transformers/main_classes/text_generation)"
1722
+ )
1723
+ self.generation_config = new_generation_config
1724
+ generation_config = self.generation_config
1725
+
1726
+ generation_config = copy.deepcopy(generation_config)
1727
+ model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
1728
+ generation_config.validate()
1729
+ self._validate_model_kwargs(model_kwargs.copy())
1730
+
1731
+ # 2. Set generation parameters if not already defined
1732
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
1733
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
1734
+
1735
+ if generation_config.pad_token_id is None and generation_config.eos_token_id is not None:
1736
+ if model_kwargs.get("attention_mask", None) is None:
1737
+ logger.warning(
1738
+ "The attention mask and the pad token id were not set. As a consequence, you may observe "
1739
+ "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results."
1740
+ )
1741
+ eos_token_id = generation_config.eos_token_id
1742
+ if isinstance(eos_token_id, list):
1743
+ eos_token_id = eos_token_id[0]
1744
+ logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{eos_token_id} for open-end generation.")
1745
+ generation_config.pad_token_id = eos_token_id
1746
+
1747
+ # 3. Define model inputs
1748
+ # inputs_tensor has to be defined
1749
+ # model_input_name is defined if model-specific keyword input is passed
1750
+ # otherwise model_input_name is None
1751
+ # all model-specific keyword inputs are removed from `model_kwargs`
1752
+ inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
1753
+ inputs, generation_config.bos_token_id, model_kwargs
1754
+ )
1755
+ batch_size = inputs_tensor.shape[0]
1756
+
1757
+ # 4. Define other model kwargs
1758
+ model_kwargs["output_attentions"] = generation_config.output_attentions
1759
+ model_kwargs["output_hidden_states"] = generation_config.output_hidden_states
1760
+ model_kwargs["use_cache"] = generation_config.use_cache
1761
+
1762
+ accepts_attention_mask = "attention_mask" in set(inspect.signature(self.forward).parameters.keys())
1763
+ requires_attention_mask = "encoder_outputs" not in model_kwargs
1764
+
1765
+ if model_kwargs.get("attention_mask", None) is None and requires_attention_mask and accepts_attention_mask:
1766
+ model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
1767
+ inputs_tensor, generation_config.pad_token_id, generation_config.eos_token_id
1768
+ )
1769
+
1770
+ # decoder-only models should use left-padding for generation
1771
+ if not self.config.is_encoder_decoder:
1772
+ # If `input_ids` was given, check if the last id in any sequence is `pad_token_id`
1773
+ # Note: If using, `inputs_embeds` this check does not work, because we want to be more hands-off.
1774
+ if (
1775
+ generation_config.pad_token_id is not None
1776
+ and len(inputs_tensor.shape) == 2
1777
+ and torch.sum(inputs_tensor[:, -1] == generation_config.pad_token_id) > 0
1778
+ ):
1779
+ logger.warning(
1780
+ "A decoder-only architecture is being used, but right-padding was detected! For correct "
1781
+ "generation results, please set `padding_side='left'` when initializing the tokenizer."
1782
+ )
1783
+
1784
+ # 5. Prepare `input_ids` which will be used for auto-regressive generation
1785
+ input_ids = inputs_tensor if model_input_name == "input_ids" else model_kwargs.pop("input_ids")
1786
+
1787
+ if streamer is not None:
1788
+ streamer.put(input_ids.cpu())
1789
+
1790
+ # 6. Prepare `max_length` depending on other stopping criteria.
1791
+ input_ids_seq_length = input_ids.shape[-1]
1792
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
1793
+ if has_default_max_length and generation_config.max_new_tokens is None:
1794
+ warnings.warn(
1795
+ f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
1796
+ "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
1797
+ " recommend using `max_new_tokens` to control the maximum length of the generation.",
1798
+ UserWarning,
1799
+ )
1800
+ elif generation_config.max_new_tokens is not None:
1801
+ if not has_default_max_length:
1802
+ logger.warning(
1803
+ f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
1804
+ f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
1805
+ "Please refer to the documentation for more information. "
1806
+ "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
1807
+ )
1808
+ generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
1809
+
1810
+ if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length:
1811
+ raise ValueError(
1812
+ f"Unfeasible length constraints: the minimum length ({generation_config.min_length}) is larger than"
1813
+ f" the maximum length ({generation_config.max_length})"
1814
+ )
1815
+ if input_ids_seq_length >= generation_config.max_length:
1816
+ input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
1817
+ logger.warning(
1818
+ f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
1819
+ f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
1820
+ " increasing `max_new_tokens`."
1821
+ )
1822
+
1823
+ if streamer is not None and (generation_config.num_beams > 1):
1824
+ raise ValueError(
1825
+ "`streamer` cannot be used with beam search (yet!). Make sure that `num_beams` is set to 1."
1826
+ )
1827
+
1828
+ if self.device.type != input_ids.device.type:
1829
+ warnings.warn(
1830
+ "You are calling .generate() with the `input_ids` being on a device type different"
1831
+ f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
1832
+ f" is on {self.device.type}. You may experience unexpected behaviors or slower generation."
1833
+ " Please make sure that you have put `input_ids` to the"
1834
+ f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"
1835
+ " running `.generate()`.",
1836
+ UserWarning,
1837
+ )
1838
+
1839
+ # 7. prepare distribution pre_processing samplers
1840
+ logits_processor = self._get_logits_processor(
1841
+ generation_config=generation_config,
1842
+ input_ids_seq_length=input_ids_seq_length,
1843
+ encoder_input_ids=inputs_tensor,
1844
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
1845
+ logits_processor=logits_processor,
1846
+ )
1847
+
1848
+ # 8. prepare beam search scorer
1849
+ beam_scorer = CpmBeeBeamSearchScorer(
1850
+ batch_size=batch_size,
1851
+ num_beams=generation_config.num_beams,
1852
+ device=inputs_tensor.device,
1853
+ length_penalty=generation_config.length_penalty,
1854
+ do_early_stopping=generation_config.early_stopping,
1855
+ num_beam_hyps_to_keep=generation_config.num_return_sequences,
1856
+ max_length=generation_config.max_length,
1857
+ **kwargs,
1858
+ )
1859
+ # 9. interleave input_ids with `num_beams` additional sequences per batch
1860
+ input_ids, model_kwargs = self._expand_inputs_for_generation(
1861
+ input_ids=input_ids,
1862
+ expand_size=generation_config.num_beams,
1863
+ is_encoder_decoder=self.config.is_encoder_decoder,
1864
+ **model_kwargs,
1865
+ )
1866
+ # 10. run beam search
1867
+ return self.beam_search(
1868
+ input_ids,
1869
+ beam_scorer,
1870
+ repetition_penalty=repetition_penalty,
1871
+ logits_processor=logits_processor,
1872
+ pad_token_id=generation_config.pad_token_id,
1873
+ eos_token_id=generation_config.eos_token_id,
1874
+ output_scores=generation_config.output_scores,
1875
+ return_dict_in_generate=generation_config.return_dict_in_generate,
1876
+ synced_gpus=synced_gpus,
1877
+ **model_kwargs,
1878
+ )
1879
+
1880
+ @torch.no_grad()
1881
+ def generate(
1882
+ self,
1883
+ data_list: Union[Dict, List[Dict]],
1884
+ tokenizer: CpmBeeTokenizer,
1885
+ generation_config=None,
1886
+ **kwargs,
1887
+ ):
1888
+ """
1889
+ Override the generate for CPMBee. It will accept dict or list(dict) as input and returns dict or list(dict)
1890
+ with `<ans>` filled.
1891
+
1892
+ Parameters:
1893
+ data_list (`dict` or `list(dict)`):
1894
+ The sequence used as a prompt for the generation or as model inputs to the encoder. If dict, data_list
1895
+ will be wrapped as a list.
1896
+ tokenizer: (`CpmBeeTokenizer`):
1897
+ The tokenizer.
1898
+ generation_config (`~generation.GenerationConfig`, *optional*):
1899
+ The generation configuration to be used as base parametrization for the generation call. `**kwargs`
1900
+ passed to generate matching the attributes of `generation_config` will override them. If
1901
+ `generation_config` is not provided, the default will be used, which had the following loading
1902
+ priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
1903
+ configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
1904
+ default values, whose documentation should be checked to parameterize generation.
1905
+ """
1906
+ if isinstance(data_list, dict):
1907
+ data_list = [data_list]
1908
+ input_encoded = tokenizer(data_list, return_tensors="pt", padding=True, device=self.device)
1909
+ input_encoded.update(kwargs)
1910
+ input_encoded["generation_config"] = generation_config
1911
+
1912
+ decode_res = self._generate(**input_encoded)
1913
+
1914
+ for sent_id, result in enumerate(decode_res):
1915
+ ans_result_map: Dict[int, List[int]] = {}
1916
+ for raw_word_id, ans_id in result:
1917
+ if ans_id not in ans_result_map:
1918
+ ans_result_map[ans_id] = []
1919
+ ans_result_map[ans_id].append(raw_word_id)
1920
+
1921
+ answer_placeholders = input_encoded["other_info"][sent_id]["answer_placeholders"]
1922
+ ext_table = input_encoded["other_info"][sent_id]["ext_table"]
1923
+ data = data_list[sent_id]
1924
+ for ans_id, token_ids in ans_result_map.items():
1925
+ if token_ids[-1] == tokenizer.eos_token_id:
1926
+ token_ids = token_ids[:-1]
1927
+ text = tokenizer.decode(token_ids, ext_table)
1928
+ path = answer_placeholders[ans_id - 1]
1929
+
1930
+ if len(path) > 0:
1931
+ p = data["<ans>"]
1932
+ for part in path[:-1]:
1933
+ p = p[part]
1934
+ p[path[-1]] = text
1935
+ else:
1936
+ data["<ans>"] = text
1937
+ for ans_id in range(len(answer_placeholders)):
1938
+ if (ans_id + 1) not in ans_result_map:
1939
+ path = answer_placeholders[ans_id]
1940
+ p = data["<ans>"]
1941
+ for part in path[:-1]:
1942
+ p = p[part]
1943
+ p[path[-1]] = None
1944
+ return data_list
test_modeling_cpmbee.py ADDED
@@ -0,0 +1,183 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """ Testing suite for the PyTorch CpmBee model. """
16
+
17
+
18
+ import unittest
19
+
20
+ from transformers.testing_utils import is_torch_available, require_torch, tooslow
21
+
22
+ from ...generation.test_utils import torch_device
23
+ from ...test_configuration_common import ConfigTester
24
+ from ...test_modeling_common import ModelTesterMixin, ids_tensor
25
+ from ...test_pipeline_mixin import PipelineTesterMixin
26
+
27
+
28
+ if is_torch_available():
29
+ import torch
30
+
31
+ from transformers import (
32
+ CpmBeeConfig,
33
+ CpmBeeForCausalLM,
34
+ CpmBeeModel,
35
+ CpmBeeTokenizer,
36
+ )
37
+
38
+
39
+ @require_torch
40
+ class CpmBeeModelTester:
41
+ def __init__(
42
+ self,
43
+ parent,
44
+ batch_size=2,
45
+ seq_length=8,
46
+ is_training=True,
47
+ use_token_type_ids=False,
48
+ use_input_mask=False,
49
+ use_labels=False,
50
+ use_mc_token_ids=False,
51
+ vocab_size=99,
52
+ hidden_size=32,
53
+ num_hidden_layers=3,
54
+ num_attention_heads=4,
55
+ intermediate_size=37,
56
+ num_buckets=32,
57
+ max_distance=128,
58
+ position_bias_num_segment_buckets=32,
59
+ init_std=1.0,
60
+ return_dict=True,
61
+ ):
62
+ self.parent = parent
63
+ self.batch_size = batch_size
64
+ self.seq_length = seq_length
65
+ self.is_training = is_training
66
+ self.use_token_type_ids = use_token_type_ids
67
+ self.use_input_mask = use_input_mask
68
+ self.use_labels = use_labels
69
+ self.use_mc_token_ids = use_mc_token_ids
70
+ self.vocab_size = vocab_size
71
+ self.hidden_size = hidden_size
72
+ self.num_hidden_layers = num_hidden_layers
73
+ self.num_attention_heads = num_attention_heads
74
+ self.intermediate_size = intermediate_size
75
+ self.num_buckets = num_buckets
76
+ self.max_distance = max_distance
77
+ self.position_bias_num_segment_buckets = position_bias_num_segment_buckets
78
+ self.init_std = init_std
79
+ self.return_dict = return_dict
80
+
81
+ def prepare_config_and_inputs(self):
82
+ input_ids = {}
83
+ input_ids["input_ids"] = ids_tensor([self.batch_size, self.seq_length], self.vocab_size).type(torch.int32)
84
+ input_ids["use_cache"] = False
85
+
86
+ config = self.get_config()
87
+
88
+ return (config, input_ids)
89
+
90
+ def get_config(self):
91
+ return CpmBeeConfig(
92
+ vocab_size=self.vocab_size,
93
+ hidden_size=self.hidden_size,
94
+ num_hidden_layers=self.num_hidden_layers,
95
+ num_attention_heads=self.num_attention_heads,
96
+ dim_ff=self.intermediate_size,
97
+ position_bias_num_buckets=self.num_buckets,
98
+ position_bias_max_distance=self.max_distance,
99
+ position_bias_num_segment_buckets=self.position_bias_num_segment_buckets,
100
+ use_cache=True,
101
+ init_std=self.init_std,
102
+ return_dict=self.return_dict,
103
+ )
104
+
105
+ def create_and_check_cpmbee_model(self, config, input_ids, *args):
106
+ model = CpmBeeModel(config=config)
107
+ model.to(torch_device)
108
+ model.eval()
109
+
110
+ hidden_states = model(**input_ids).last_hidden_state
111
+
112
+ self.parent.assertEqual(hidden_states.shape, (self.batch_size, self.seq_length, config.hidden_size))
113
+
114
+ def create_and_check_lm_head_model(self, config, input_ids, *args):
115
+ model = CpmBeeForCausalLM(config)
116
+ model.to(torch_device)
117
+ input_ids["input_ids"] = input_ids["input_ids"].to(torch_device)
118
+ model.eval()
119
+
120
+ model_output = model(**input_ids)
121
+ self.parent.assertEqual(
122
+ model_output.logits.shape,
123
+ (self.batch_size, self.seq_length, config.vocab_size),
124
+ )
125
+
126
+ def prepare_config_and_inputs_for_common(self):
127
+ config, inputs_dict = self.prepare_config_and_inputs()
128
+ return config, inputs_dict
129
+
130
+
131
+ @require_torch
132
+ class CpmBeeModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
133
+ all_model_classes = (CpmBeeModel, CpmBeeForCausalLM) if is_torch_available() else ()
134
+ pipeline_model_mapping = (
135
+ {"feature-extraction": CpmBeeModel, "text-generation": CpmBeeForCausalLM} if is_torch_available() else {}
136
+ )
137
+
138
+ test_pruning = False
139
+ test_missing_keys = False
140
+ test_mismatched_shapes = False
141
+ test_head_masking = False
142
+ test_resize_embeddings = False
143
+
144
+ def setUp(self):
145
+ self.model_tester = CpmBeeModelTester(self)
146
+ self.config_tester = ConfigTester(self, config_class=CpmBeeConfig)
147
+
148
+ def test_config(self):
149
+ self.config_tester.create_and_test_config_common_properties()
150
+ self.config_tester.create_and_test_config_to_json_string()
151
+ self.config_tester.create_and_test_config_to_json_file()
152
+ self.config_tester.create_and_test_config_from_and_save_pretrained()
153
+ self.config_tester.check_config_can_be_init_without_params()
154
+ self.config_tester.check_config_arguments_init()
155
+
156
+ def test_inputs_embeds(self):
157
+ unittest.skip("CPMBee doesn't support input_embeds.")(self.test_inputs_embeds)
158
+
159
+ def test_retain_grad_hidden_states_attentions(self):
160
+ unittest.skip(
161
+ "CPMBee doesn't support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
162
+ So is attentions. We strongly recommand you use loss to tune model."
163
+ )(self.test_retain_grad_hidden_states_attentions)
164
+
165
+ def test_cpmbee_model(self):
166
+ config, inputs = self.model_tester.prepare_config_and_inputs()
167
+ self.model_tester.create_and_check_cpmbee_model(config, inputs)
168
+
169
+ def test_cpmbee_lm_head_model(self):
170
+ config, inputs = self.model_tester.prepare_config_and_inputs()
171
+ self.model_tester.create_and_check_lm_head_model(config, inputs)
172
+
173
+
174
+ @require_torch
175
+ class CpmBeeForCausalLMlIntegrationTest(unittest.TestCase):
176
+ @tooslow
177
+ def test_simple_generation(self):
178
+ texts = {"input": "今天天气不错,", "<ans>": ""}
179
+ model = CpmBeeForCausalLM.from_pretrained("openbmb/cpm-bee-10b")
180
+ tokenizer = CpmBeeTokenizer.from_pretrained("openbmb/cpm-bee-10b")
181
+ output_texts = model.generate(texts, tokenizer)
182
+ expected_output = {"input": "今天天气不错,", "<ans>": "适合睡觉。"}
183
+ self.assertEqual(expected_output["<ans>"], output_texts["<ans>"])
test_tokenization_cpmbee.py ADDED
@@ -0,0 +1,187 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """ Testing suite for the PyTorch CpmBee tokenizer. """
16
+
17
+ import os
18
+ import unittest
19
+
20
+ from transformers.models.cpmbee.tokenization_cpmbee import VOCAB_FILES_NAMES, CpmBeeTokenizer
21
+ from transformers.tokenization_utils import AddedToken
22
+
23
+ from ...test_tokenization_common import TokenizerTesterMixin
24
+
25
+
26
+ class CPMBeeTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
27
+ tokenizer_class = CpmBeeTokenizer
28
+ test_rust_tokenizer = False
29
+
30
+ def setUp(self):
31
+ super().setUp()
32
+
33
+ vocab_tokens = [
34
+ "<d>",
35
+ "</d>",
36
+ "<s>",
37
+ "</s>",
38
+ "</_>",
39
+ "<unk>",
40
+ "<pad>",
41
+ "<mask>",
42
+ "</n>",
43
+ "我",
44
+ "是",
45
+ "C",
46
+ "P",
47
+ "M",
48
+ "B",
49
+ "e",
50
+ "e",
51
+ ]
52
+ self.vocab_file = os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["vocab_file"])
53
+ vocab_tokens = list(set(vocab_tokens))
54
+ with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
55
+ vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
56
+
57
+ # override test_add_tokens_tokenizer because <...> is special token in CpmBeeTokenizer.
58
+ def test_add_tokens_tokenizer(self):
59
+ tokenizers = self.get_tokenizers(do_lower_case=False)
60
+ for tokenizer in tokenizers:
61
+ with self.subTest(f"{tokenizer.__class__.__name__}"):
62
+ vocab_size = tokenizer.vocab_size
63
+ all_size = len(tokenizer)
64
+
65
+ self.assertNotEqual(vocab_size, 0)
66
+
67
+ # We usually have added tokens from the start in tests because our vocab fixtures are
68
+ # smaller than the original vocabs - let's not assert this
69
+ # self.assertEqual(vocab_size, all_size)
70
+
71
+ new_toks = ["aaaaa bbbbbb", "cccccccccdddddddd"]
72
+ added_toks = tokenizer.add_tokens(new_toks)
73
+ vocab_size_2 = tokenizer.vocab_size
74
+ all_size_2 = len(tokenizer)
75
+
76
+ self.assertNotEqual(vocab_size_2, 0)
77
+ self.assertEqual(vocab_size, vocab_size_2)
78
+ self.assertEqual(added_toks, len(new_toks))
79
+ self.assertEqual(all_size_2, all_size + len(new_toks))
80
+
81
+ tokens = tokenizer.encode("aaaaa bbbbbb low cccccccccdddddddd l", add_special_tokens=False)
82
+
83
+ self.assertGreaterEqual(len(tokens), 4)
84
+ self.assertGreater(tokens[0], tokenizer.vocab_size - 1)
85
+ self.assertGreater(tokens[-2], tokenizer.vocab_size - 1)
86
+
87
+ new_toks_2 = {"eos_token": ">>>>|||<||<<|<<", "pad_token": "<<<<<|||;;;||;"}
88
+ added_toks_2 = tokenizer.add_special_tokens(new_toks_2)
89
+ vocab_size_3 = tokenizer.vocab_size
90
+ all_size_3 = len(tokenizer)
91
+
92
+ self.assertNotEqual(vocab_size_3, 0)
93
+ self.assertEqual(vocab_size, vocab_size_3)
94
+ self.assertEqual(added_toks_2, len(new_toks_2))
95
+ self.assertEqual(all_size_3, all_size_2 + len(new_toks_2))
96
+
97
+ tokens = tokenizer.encode(
98
+ ">>>>|||<||<<|<< aaaaabbbbbb low cccccccccdddddddd <<<<<|||;;;||; l", add_special_tokens=False
99
+ )
100
+
101
+ self.assertGreaterEqual(len(tokens), 6)
102
+ self.assertGreater(tokens[0], tokenizer.vocab_size - 1)
103
+ self.assertGreater(tokens[0], tokens[1])
104
+ self.assertGreater(tokens[-2], tokenizer.vocab_size - 1)
105
+ self.assertGreater(tokens[-2], tokens[-3])
106
+ self.assertEqual(tokens[0], tokenizer.eos_token_id)
107
+ self.assertEqual(tokens[-2], tokenizer.pad_token_id)
108
+
109
+ def test_added_tokens_do_lower_case(self):
110
+ tokenizers = self.get_tokenizers(do_lower_case=True)
111
+ for tokenizer in tokenizers:
112
+ with self.subTest(f"{tokenizer.__class__.__name__}"):
113
+ if not hasattr(tokenizer, "do_lower_case") or not tokenizer.do_lower_case:
114
+ continue
115
+
116
+ special_token = tokenizer.all_special_tokens[0]
117
+
118
+ text = special_token + " aaaaa bbbbbb low cccccccccdddddddd l " + special_token
119
+ text2 = special_token + " AAAAA BBBBBB low CCCCCCCCCDDDDDDDD l " + special_token
120
+
121
+ toks_before_adding = tokenizer.tokenize(text) # toks before adding new_toks
122
+
123
+ new_toks = ["aaaaa bbbbbb", "cccccccccdddddddd", "AAAAA BBBBBB", "CCCCCCCCCDDDDDDDD"]
124
+ added = tokenizer.add_tokens([AddedToken(tok, lstrip=True, rstrip=True) for tok in new_toks])
125
+
126
+ toks_after_adding = tokenizer.tokenize(text)
127
+ toks_after_adding2 = tokenizer.tokenize(text2)
128
+
129
+ # Rust tokenizers dont't lowercase added tokens at the time calling `tokenizer.add_tokens`,
130
+ # while python tokenizers do, so new_toks 0 and 2 would be treated as the same, so do new_toks 1 and 3.
131
+ self.assertIn(added, [2, 4])
132
+
133
+ self.assertListEqual(toks_after_adding, toks_after_adding2)
134
+ self.assertTrue(
135
+ len(toks_before_adding) > len(toks_after_adding), # toks_before_adding should be longer
136
+ )
137
+
138
+ # Check that none of the special tokens are lowercased
139
+ sequence_with_special_tokens = "A " + " yEs ".join(tokenizer.all_special_tokens) + " B"
140
+ # Convert the tokenized list to str as some special tokens are tokenized like normal tokens
141
+ # which have a prefix spacee e.g. the mask token of Albert, and cannot match the original
142
+ # special tokens exactly.
143
+ tokenized_sequence = "".join(tokenizer.tokenize(sequence_with_special_tokens))
144
+
145
+ for special_token in tokenizer.all_special_tokens:
146
+ self.assertTrue(special_token in tokenized_sequence)
147
+
148
+ tokenizers = self.get_tokenizers(do_lower_case=True)
149
+ for tokenizer in tokenizers:
150
+ with self.subTest(f"{tokenizer.__class__.__name__}"):
151
+ if hasattr(tokenizer, "do_lower_case") and tokenizer.do_lower_case:
152
+ continue
153
+
154
+ special_token = tokenizer.all_special_tokens[0]
155
+
156
+ text = special_token + " aaaaa bbbbbb low cccccccccdddddddd l " + special_token
157
+ text2 = special_token + " AAAAA BBBBBB low CCCCCCCCCDDDDDDDD l " + special_token
158
+
159
+ toks_before_adding = tokenizer.tokenize(text) # toks before adding new_toks
160
+
161
+ new_toks = ["aaaaa bbbbbb", "cccccccccdddddddd", "AAAAA BBBBBB", "CCCCCCCCCDDDDDDDD"]
162
+ added = tokenizer.add_tokens([AddedToken(tok, lstrip=True, rstrip=True) for tok in new_toks])
163
+ self.assertIn(added, [2, 4])
164
+
165
+ toks_after_adding = tokenizer.tokenize(text)
166
+ toks_after_adding2 = tokenizer.tokenize(text2)
167
+
168
+ self.assertEqual(len(toks_after_adding), len(toks_after_adding2)) # Length should still be the same
169
+ self.assertNotEqual(
170
+ toks_after_adding[1], toks_after_adding2[1]
171
+ ) # But at least the first non-special tokens should differ
172
+ self.assertTrue(
173
+ len(toks_before_adding) > len(toks_after_adding), # toks_before_adding should be longer
174
+ )
175
+
176
+ def test_pre_tokenization(self):
177
+ tokenizer = CpmBeeTokenizer.from_pretrained("openbmb/cpm-bee-10b")
178
+ texts = {"input": "你好,", "<ans>": ""}
179
+ tokens = tokenizer(texts)
180
+ tokens = tokens["input_ids"][0]
181
+
182
+ input_tokens = [6, 8, 7, 6, 65678, 7, 6, 10273, 246, 7, 6, 9, 7]
183
+ self.assertListEqual(tokens, input_tokens)
184
+
185
+ normalized_text = "<s><root></s><s>input</s><s>你好,</s><s><ans></s>"
186
+ reconstructed_text = tokenizer.decode(tokens)
187
+ self.assertEqual(reconstructed_text, normalized_text)
tokenization_cpmbee.py ADDED
@@ -0,0 +1,868 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """Tokenization classes for CpmBee."""
16
+ import json
17
+ import os
18
+ from typing import Any, Dict, List, Optional, Tuple, Union
19
+
20
+ import numpy as np
21
+ from typing_extensions import TypedDict
22
+
23
+ from ...tokenization_utils import PaddingStrategy, PreTrainedTokenizer, TensorType
24
+ from ...tokenization_utils_base import AddedToken, BatchEncoding, TextInput, TruncationStrategy
25
+ from ...utils import logging
26
+
27
+
28
+ logger = logging.get_logger(__name__)
29
+
30
+ VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
31
+
32
+ PRETRAINED_VOCAB_FILES_MAP = {
33
+ "vocab_file": {
34
+ "openbmb/cpm-bee-10b": "https://huggingface.co/openbmb/cpm-bee-10b/blob/main/vocab.txt",
35
+ "openbmb/cpm-bee-5b": "https://huggingface.co/openbmb/cpm-bee-5b/blob/main/vocab.txt",
36
+ "openbmb/cpm-bee-2b": "https://huggingface.co/openbmb/cpm-bee-2b/blob/main/vocab.txt",
37
+ "openbmb/cpm-bee-1b": "https://huggingface.co/openbmb/cpm-bee-1b/blob/main/vocab.txt",
38
+ },
39
+ }
40
+
41
+ PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
42
+ "openbmb/cpm-bee-10b": 4096,
43
+ "openbmb/cpm-bee-5b": 4096,
44
+ "openbmb/cpm-bee-2b": 4096,
45
+ "openbmb/cpm-bee-1b": 4096,
46
+ }
47
+
48
+
49
+ class _PrevExtTableStates(TypedDict):
50
+ ext_table: Dict[int, str]
51
+ token_id_table: Dict[str, Dict[int, int]]
52
+
53
+
54
+ CPMBeeInputType = Union[str, Dict[str, "CPMBeeInputType"]]
55
+
56
+
57
+ def rel_to_bucket(n_up: int, n_down: int, max_depth: int = 8):
58
+ ret = n_up * max_depth + n_down
59
+ if ret == 0:
60
+ return ret
61
+ else:
62
+ # bucket 1 is reserved for incontext samples
63
+ return ret + 1
64
+
65
+
66
+ class _DictTree(TypedDict):
67
+ value: str
68
+ children: List["_DictTree"]
69
+ depth: int
70
+ segment_id: int
71
+ need_predict: bool
72
+
73
+
74
+ class CpmBeeTokenizer(PreTrainedTokenizer):
75
+ """
76
+ Construct a CPMBee tokenizer.
77
+
78
+ Args:
79
+ vocab_file (`str`):
80
+ Path to the vocabulary file.
81
+ bos_token (`str`, *optional*, defaults to `"<s>"`):
82
+ The beginning of sequence token.
83
+ eos_token (`str`, *optional*, defaults to `"</s>"`):
84
+ The end of sequence token.
85
+ line_token (`str`, *optional*, defaults to `"\n"`):
86
+ The line token.
87
+ space_token (`str`, *optional*, defaults to `" "`):
88
+ The space token.
89
+ unk_token (`str`, *optional*, defaults to `"<unk>"`):
90
+ The unknown token.
91
+ mask_token (`str`, *optional*, defaults to `"<mask>"`):
92
+ The mask token.
93
+ pad_token (`str`, *optional*, defaults to `"<pad>"`):
94
+ The token used for padding.
95
+ padding_side (`str`, *optional*, defaults to `"left"`):
96
+ The padding side. CPM-Bee will use left padding by default.
97
+ """
98
+
99
+ vocab_files_names = VOCAB_FILES_NAMES
100
+ pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
101
+ max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
102
+ model_input_names: List[str] = [
103
+ "input_ids",
104
+ "attention_mask",
105
+ "input_id_sub",
106
+ "position",
107
+ "context",
108
+ "sample_ids",
109
+ "num_segments",
110
+ "segment",
111
+ "segment_rel_offset",
112
+ "segment_rel",
113
+ ]
114
+ add_prefix_space = False
115
+
116
+ def __init__(
117
+ self,
118
+ vocab_file,
119
+ bos_token="<s>",
120
+ eos_token="</s>",
121
+ line_token="\n",
122
+ space_token=" ",
123
+ unk_token="<unk>",
124
+ mask_token="<mask>",
125
+ pad_token="<pad>",
126
+ padding_side="left",
127
+ **kwargs,
128
+ ):
129
+ super().__init__(
130
+ bos_token=bos_token,
131
+ eos_token=eos_token,
132
+ line_token=line_token,
133
+ space_token=space_token,
134
+ unk_token=unk_token,
135
+ mask_token=mask_token,
136
+ pad_token=pad_token,
137
+ padding_side=padding_side,
138
+ **kwargs,
139
+ )
140
+
141
+ self.encoder: Dict[str, int] = {}
142
+
143
+ with open(vocab_file, "r", encoding="utf-8") as reader:
144
+ for token in reader.readlines():
145
+ token = token.rstrip("\n")
146
+ if len(token) == 0:
147
+ continue
148
+ self.encoder[token] = len(self.encoder)
149
+
150
+ self.encoder[" "] = self.encoder["</_>"]
151
+ self.encoder["\n"] = self.encoder["</n>"]
152
+ del self.encoder["</_>"]
153
+ del self.encoder["</n>"]
154
+
155
+ self.decoder = {v: k for k, v in self.encoder.items()}
156
+
157
+ self._max_word_len = max([len(x) for x in self.encoder.keys()])
158
+ self.cpmbee_special_tokens = {k: v for k, v in self.encoder.items() if k.startswith("<") and k.endswith(">")}
159
+
160
+ self.ext_table: Dict[int, str] = {}
161
+ self.ext_table_rev: Dict[str, int] = {}
162
+
163
+ self.token_id_table: Dict[str, Dict[int, int]] = {}
164
+ self.ext_special_tokens = []
165
+
166
+ self.ext_args_for_model = [
167
+ "input_id_subs",
168
+ "input_pos",
169
+ "context",
170
+ "segment_ids",
171
+ "segment_rel_offset",
172
+ "segment_rel",
173
+ "sample_ids",
174
+ "num_segments",
175
+ "predict_segments",
176
+ "answer_placeholders",
177
+ "ext_table",
178
+ "token_id_table",
179
+ ]
180
+
181
+ @property
182
+ def bod_token_id(self):
183
+ return self.encoder[self.bod_token]
184
+
185
+ @property
186
+ def eod_token_id(self):
187
+ return self.encoder[self.eod_token]
188
+
189
+ @property
190
+ def newline_id(self):
191
+ return self.encoder[self.line_token]
192
+
193
+ @property
194
+ def vocab_size(self) -> int:
195
+ return len(self.encoder)
196
+
197
+ def __len__(self):
198
+ """
199
+ Size of the full vocabulary with the added tokens.
200
+ """
201
+ return self.vocab_size + len(self.added_tokens_encoder)
202
+
203
+ def get_vocab(self):
204
+ return dict(self.encoder, **self.added_tokens_encoder)
205
+
206
+ def get_piece(self, text: str) -> str:
207
+ """
208
+ Match with maximum length.
209
+ """
210
+ len_text = len(text)
211
+ for i in range(len(text)):
212
+ sub = text[: len_text - i]
213
+ if (sub in self.encoder) or (sub in self.added_tokens_encoder):
214
+ return sub
215
+ return text[0]
216
+
217
+ def tokenize(self, text: TextInput, **kwargs) -> List[str]:
218
+ r"""
219
+ Override the `tokenize` to meet the needs of CPMBee:
220
+ 1. Mark the special token with `<` and `>`. The `<>` will be ignored.
221
+ 2. Split sentences by the marked special tokens.
222
+ 3. Record the marked special token by `ext_table` and `ext_table_rev`.
223
+ 4. Tokenize the sentence without special tokens.
224
+ """
225
+ for_cpmbee = kwargs.get("for_cpmbee", False)
226
+ all_special_tokens_extended = {
227
+ str(t): t for t in self.all_special_tokens_extended if isinstance(t, AddedToken)
228
+ }
229
+
230
+ sentence_split = [""]
231
+ is_special_token = False
232
+ for i, c in enumerate(text):
233
+ if is_special_token:
234
+ if c == "<":
235
+ tail = sentence_split.pop(-1)
236
+ sentence_split[-1] += tail
237
+ sentence_split.append(c)
238
+ elif c == ">":
239
+ # end of special token
240
+ sentence_split[-1] += c
241
+ if sentence_split[-1] == "<>":
242
+ continue
243
+ is_special_token = False
244
+ sentence_split.append("")
245
+ else:
246
+ sentence_split[-1] += c
247
+ else:
248
+ if c == "<":
249
+ is_special_token = True
250
+ sentence_split.append(c)
251
+ else:
252
+ sentence_split[-1] += c
253
+ if is_special_token:
254
+ tail = sentence_split.pop(-1)
255
+ sentence_split[-1] += tail
256
+
257
+ output_tokens = []
258
+ for i, part in enumerate(sentence_split):
259
+ if (i & 1) == 1:
260
+ # special token
261
+ output_tokens.append(part)
262
+ if for_cpmbee and (part not in self.encoder) and (part not in self.ext_table_rev):
263
+ self.ext_table_rev[part] = len(self.ext_table_rev) + self.vocab_size
264
+ self.ext_table[self.ext_table_rev[part]] = part
265
+ else:
266
+ output_tokens.extend(self._tokenize(part, for_cpmbee=for_cpmbee))
267
+
268
+ # drop spaces
269
+ for i, token in enumerate(output_tokens):
270
+ if token in self.added_tokens_encoder:
271
+ token = all_special_tokens_extended.get(token, None)
272
+ left = output_tokens[i - 1] if i > 0 else None
273
+ right = output_tokens[i + 1] if i < len(output_tokens) - 1 else None
274
+ if isinstance(token, AddedToken):
275
+ if token.rstrip and right:
276
+ # A bit counter-intuitive but we strip the left of the string
277
+ # since tok_extended.rstrip means the special token is eating all white spaces on its right
278
+ output_tokens[i + 1] = right.lstrip()
279
+ # Strip white spaces on the left
280
+ if token.lstrip and left:
281
+ output_tokens[i - 1] = left.rstrip() # Opposite here
282
+ else:
283
+ if right:
284
+ output_tokens[i + 1] = right.lstrip()
285
+ if left:
286
+ output_tokens[i - 1] = left.rstrip()
287
+
288
+ skipped_tokens = []
289
+ for token in output_tokens:
290
+ if not token:
291
+ continue
292
+ else:
293
+ skipped_tokens.append(token)
294
+
295
+ return skipped_tokens
296
+
297
+ def _tokenize(self, text, **kwargs):
298
+ """
299
+ Converts a string in a sequence of tokens (string), using the tokenizer. Split in words for word-based
300
+ vocabulary.
301
+
302
+ Do NOT take care of added tokens. Record the unk tokens and special tokens in `ext_table` and `ext_table_rev`.
303
+ """
304
+ for_cpmbee = kwargs.get("for_cpmbee", False)
305
+ output_tokens = []
306
+
307
+ part_st = 0
308
+ last_unk = None
309
+ while part_st < len(text):
310
+ piece = self.get_piece(text[part_st:])
311
+ if piece in self.encoder or self.added_tokens_encoder:
312
+ if last_unk is None:
313
+ output_tokens.append(piece)
314
+ else:
315
+ if for_cpmbee and (last_unk not in self.ext_table_rev):
316
+ self.ext_table_rev[last_unk] = len(self.ext_table_rev) + self.vocab_size
317
+ self.ext_table[self.ext_table_rev[last_unk]] = last_unk
318
+ output_tokens.append(last_unk)
319
+ output_tokens.append(piece)
320
+ last_unk = None
321
+ else:
322
+ if last_unk is None:
323
+ last_unk = piece
324
+ else:
325
+ last_unk += piece
326
+ part_st += len(piece)
327
+ if last_unk is not None:
328
+ # part end with UNK
329
+ if for_cpmbee and (last_unk not in self.ext_table_rev):
330
+ self.ext_table_rev[last_unk] = len(self.ext_table_rev) + self.vocab_size
331
+ self.ext_table[self.ext_table_rev[last_unk]] = last_unk
332
+ output_tokens.append(last_unk)
333
+
334
+ return output_tokens
335
+
336
+ def check(self, token):
337
+ return token in self.encoder
338
+
339
+ def convert_tokens_to_string(self, tokens: List[str]) -> str:
340
+ return "".join(tokens)
341
+
342
+ def _convert_token_to_id(self, token: str):
343
+ """Converts a token (str) in an id using the vocab and ext_table."""
344
+ if token in self.encoder:
345
+ return self.encoder.get(token)
346
+ elif token in self.ext_table_rev:
347
+ return self.ext_table_rev[token]
348
+ elif token in self.added_tokens_encoder:
349
+ return self.added_tokens_encoder[token]
350
+ else:
351
+ return self.unk_token_id
352
+
353
+ def _convert_id_to_token(self, index):
354
+ """Converts an index (integer) in a token (str) using the vocab and ext_table."""
355
+ if index in self.ext_table:
356
+ return self.ext_table[index]
357
+ elif index in self.added_tokens_decoder:
358
+ return self.added_tokens_decoder[index]
359
+ else:
360
+ if index >= 0:
361
+ return self.decoder[index]
362
+
363
+ def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
364
+ if os.path.isdir(save_directory):
365
+ vocab_file = os.path.join(
366
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
367
+ )
368
+ else:
369
+ vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
370
+ index = 0
371
+ self.encoder["</n>"] = self.encoder["\n"]
372
+ del self.encoder["\n"]
373
+ self.encoder["</_>"] = self.encoder[" "]
374
+ del self.encoder[" "]
375
+ with open(vocab_file, "w", encoding="utf-8") as writer:
376
+ for token, token_index in sorted(self.encoder.items(), key=lambda x: x[1]):
377
+ if index != token_index:
378
+ logger.warning(
379
+ f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
380
+ " Please check that the vocabulary is not corrupted!"
381
+ )
382
+ index = token_index
383
+ writer.write(token + "\n")
384
+ index += 1
385
+ return (vocab_file,)
386
+
387
+ def __call__(self, text, *args, **kwargs):
388
+ r"""
389
+ CPMBee `call` method will use `_tokenize_cpmbee` when the input type is dict.
390
+ """
391
+ if isinstance(text, dict):
392
+ return self._batch_tokenize_cpmbee([text], *args, **kwargs)
393
+ elif isinstance(text, (list, tuple)):
394
+ if isinstance(text[0], dict):
395
+ return self._batch_tokenize_cpmbee(text, *args, **kwargs)
396
+ else:
397
+ return super().__call__(text, *args, **kwargs)
398
+ else:
399
+ return super().__call__(text, *args, **kwargs)
400
+
401
+ # 分词
402
+ def _tokenize_cpmbee(self, data: TextInput, *args, **kwargs) -> List[str]:
403
+ """
404
+ A tokenize method to process dict data. Exclusive for CPMBee.
405
+ """
406
+ if isinstance(data, str):
407
+ data = json.loads(data)
408
+ if not isinstance(data, Dict):
409
+ raise TypeError(
410
+ "CpmBeeTokenizer input data should be dict or str in dict format, but got {}".format(type(data))
411
+ )
412
+
413
+ # 1. prepare answer placeholder
414
+ answer_placeholders = []
415
+
416
+ def _put_placeholder(data: Any, path: List[str] = []):
417
+ if isinstance(data, dict):
418
+ ret = {}
419
+ for k, v in data.items():
420
+ ret[k] = _put_placeholder(v, path + [k])
421
+ return ret
422
+ else:
423
+ answer_placeholders.append(path)
424
+ return "<ans_{}>".format(len(answer_placeholders))
425
+
426
+ data["<ans>"] = _put_placeholder(data["<ans>"])
427
+
428
+ (
429
+ input_ids,
430
+ input_id_subs,
431
+ context,
432
+ segment_ids,
433
+ segment_rel,
434
+ n_segments,
435
+ table_states,
436
+ ) = self.convert_data_to_id(data, shuffle_answer=False, max_depth=8)
437
+
438
+ # <ans> mapping from sub to id
439
+ sub_ans_map: Dict[int, int] = {}
440
+ for fake_id, token_sub in table_states["token_id_table"]["<ans>"].items():
441
+ token = table_states["ext_table"][fake_id]
442
+ if token.startswith("<ans_") and token.endswith(">"):
443
+ ans_id = int(token[5:-1])
444
+ sub_ans_map[token_sub] = ans_id
445
+
446
+ tmp_input_ids = []
447
+ tmp_input_sub = []
448
+ tmp_input_seg = []
449
+
450
+ # get predict segments
451
+ predict_segments: List[Tuple[int, int]] = []
452
+ for i in range(input_ids.shape[0]):
453
+ if context[i] == 0:
454
+ if input_ids[i] == self.encoder["<ans>"]:
455
+ # is ans
456
+ # (segment_id, ans_id)
457
+ predict_segments.append((segment_ids[i], sub_ans_map[input_id_subs[i]]))
458
+ else:
459
+ tmp_input_ids.append(input_ids[i])
460
+ tmp_input_sub.append(input_id_subs[i])
461
+ tmp_input_seg.append(segment_ids[i])
462
+
463
+ if len(predict_segments) == 0:
464
+ raise ValueError("No answer to predict")
465
+
466
+ input_ids = np.array(tmp_input_ids, dtype=np.int32) # all context
467
+ input_id_subs = np.array(tmp_input_sub, dtype=np.int32) # [0, 0, 0, 0, 1, 0, 0, 2, 0, ...]
468
+ context = np.full_like(tmp_input_ids, 1, dtype=np.int8) # [1, 1, 1, ...]
469
+ segment_ids = np.array(tmp_input_seg, dtype=np.int32) # [0, 0, 0, 1, 1, 1, 2, 2, 2, 2, ...]
470
+ sample_ids = np.zeros(input_ids.shape, dtype=np.int32) # [0, 0, 0, 0, ...]
471
+ segment_rel_offset = np.zeros(input_ids.shape, dtype=np.int32) # [0, 0, 0, ...]
472
+ num_segments = np.full(input_ids.shape, n_segments, dtype=np.int32) # [n_seg, n_seg, n_seg, ...]
473
+ input_pos = np.arange(input_ids.shape[0], dtype=np.int32) # [0, 1, 2, 3, 4, ...]
474
+
475
+ return (
476
+ self.prepare_for_model(
477
+ input_ids.tolist(),
478
+ input_id_subs=input_id_subs.tolist(),
479
+ input_pos=input_pos.tolist(),
480
+ context=context.tolist(),
481
+ segment_ids=segment_ids.tolist(),
482
+ segment_rel_offset=segment_rel_offset.tolist(),
483
+ segment_rel=segment_rel.tolist(),
484
+ sample_ids=sample_ids.tolist(),
485
+ num_segments=num_segments.tolist(),
486
+ **kwargs,
487
+ ),
488
+ predict_segments,
489
+ answer_placeholders,
490
+ table_states["ext_table"],
491
+ table_states["token_id_table"],
492
+ )
493
+
494
+ def _batch_tokenize_cpmbee(self, data_lst, *args, **kwargs):
495
+ """
496
+ Batched _token_cpmbee.
497
+ """
498
+ device = kwargs.get("device", "cpu")
499
+ return_tensors = kwargs.get("return_tensors", None)
500
+ batch_outputs = {}
501
+ segment_rel_pack = []
502
+ other_info = []
503
+
504
+ batch_ext_table_map: Dict[Tuple[int, int], int] = {}
505
+ batch_ext_table_ids: List[int] = []
506
+ batch_ext_table_sub: List[int] = []
507
+
508
+ for data in data_lst:
509
+ self.ext_table = {}
510
+ self.ext_table_rev = {}
511
+ self.token_id_table = {}
512
+ (outputs, predict_segments, answer_placeholders, ext_table, token_id_table) = self._tokenize_cpmbee(
513
+ data,
514
+ truncation=None,
515
+ padding=PaddingStrategy.DO_NOT_PAD.value,
516
+ max_length=None,
517
+ pad_to_multiple_of=None,
518
+ return_attention_mask=False,
519
+ return_tensors=None,
520
+ )
521
+ rev_ext_table = {}
522
+ for token, mp in token_id_table.items():
523
+ if token == "<ans>":
524
+ continue
525
+ token_id = self.encoder[token]
526
+ for fake_id, token_sub in mp.items():
527
+ if token_sub > 0:
528
+ if (token_id, token_sub) not in batch_ext_table_map:
529
+ batch_ext_table_map[(token_id, token_sub)] = len(batch_ext_table_ids) + self.vocab_size
530
+ batch_ext_table_ids.append(token_id)
531
+ batch_ext_table_sub.append(token_sub)
532
+ rev_ext_table[batch_ext_table_map[(token_id, token_sub)]] = ext_table[fake_id]
533
+ else:
534
+ rev_ext_table[token_id] = ext_table[fake_id]
535
+
536
+ segment_rel_pack.append(np.array(outputs.pop("segment_rel")))
537
+ other_info.append(
538
+ {
539
+ "predict_segments": predict_segments,
540
+ "answer_placeholders": answer_placeholders,
541
+ "ext_table": rev_ext_table,
542
+ }
543
+ )
544
+
545
+ for key, value in outputs.items():
546
+ if key not in batch_outputs:
547
+ batch_outputs[key] = []
548
+ batch_outputs[key].append(value)
549
+
550
+ max_length = max([len(item) for item in batch_outputs[self.model_input_names[0]]])
551
+ batch_size = len(batch_outputs[self.model_input_names[0]])
552
+ for i in range(batch_size):
553
+ inputs = {k: v[i] for k, v in batch_outputs.items()}
554
+
555
+ for k, v in inputs.items():
556
+ required_input = v
557
+
558
+ needs_to_be_padded = len(required_input) != max_length
559
+
560
+ if needs_to_be_padded:
561
+ difference = max_length - len(required_input)
562
+ batch_outputs[k][i] = [self.pad_token_id] * difference + required_input
563
+
564
+ max_num_rels = 0
565
+ for rel in segment_rel_pack:
566
+ max_num_rels = max(max_num_rels, rel.shape[0])
567
+ padded_rels = np.zeros((len(segment_rel_pack), max_num_rels), dtype=np.int32)
568
+ for i, rel in enumerate(segment_rel_pack):
569
+ padded_rels[i, : rel.shape[0]] = rel
570
+ batch_outputs["segment_rel"] = padded_rels
571
+ batch_outputs["batch_ext_table_ids"] = np.array(batch_ext_table_ids, dtype=np.int32)
572
+ batch_outputs["batch_ext_table_sub"] = np.array(batch_ext_table_sub, dtype=np.int32)
573
+ batch_outputs = BatchEncoding(batch_outputs, tensor_type=return_tensors)
574
+ if return_tensors == "pt":
575
+ batch_outputs = batch_outputs.to(device=device)
576
+ batch_outputs["other_info"] = other_info
577
+
578
+ return batch_outputs
579
+
580
+ def convert_data_to_id(
581
+ self,
582
+ data: Any,
583
+ prev_ext_states: Optional[_PrevExtTableStates] = None,
584
+ shuffle_answer: bool = True,
585
+ max_depth: int = 8,
586
+ ):
587
+ """
588
+ Parse a dict to data ids. Exclusive for CPMBee. It will
589
+ 1. parse the dict to segments and get segment_rel, which for calculating of position_bias.
590
+ 2. tokenize every segment.
591
+ """
592
+ root: _DictTree = {
593
+ "value": "<root>",
594
+ "children": [],
595
+ "depth": 0,
596
+ "segment_id": 0,
597
+ "need_predict": False,
598
+ }
599
+
600
+ segments = [root]
601
+
602
+ def _build_dict_tree(data: CPMBeeInputType, depth: int, need_predict: bool) -> List[_DictTree]:
603
+ if isinstance(data, dict):
604
+ ret_list: List[_DictTree] = []
605
+ curr_items = list(data.items())
606
+ if need_predict and shuffle_answer:
607
+ access_idx = np.arange(len(curr_items))
608
+ np.random.shuffle(access_idx)
609
+ curr_items = [curr_items[idx] for idx in access_idx]
610
+ for k, v in curr_items:
611
+ child_info: _DictTree = {
612
+ "value": k,
613
+ "children": [],
614
+ "depth": depth,
615
+ "segment_id": len(segments),
616
+ "need_predict": False, # only leaves are contexts
617
+ }
618
+ segments.append(child_info)
619
+ child_info["children"] = _build_dict_tree(
620
+ v, depth + 1, need_predict or (depth == 1 and k == "<ans>")
621
+ ) # elements in <root>.<ans>
622
+
623
+ ret_list.append(child_info)
624
+ return ret_list
625
+ else:
626
+ assert isinstance(data, str), "Invalid data {}".format(data)
627
+ ret: _DictTree = {
628
+ "value": data,
629
+ "children": [],
630
+ "depth": depth,
631
+ "segment_id": len(segments),
632
+ "need_predict": need_predict,
633
+ }
634
+ segments.append(ret)
635
+ return [ret]
636
+
637
+ root["children"] = _build_dict_tree(data, 1, False)
638
+
639
+ num_segments = len(segments)
640
+ segment_rel = np.zeros((num_segments * num_segments,), dtype=np.int32)
641
+
642
+ def _build_segment_rel(node: _DictTree) -> List[Tuple[int, int]]:
643
+ ret: List[Tuple[int, int]] = [(node["segment_id"], node["depth"])]
644
+ for child in node["children"]:
645
+ sub = _build_segment_rel(child)
646
+ for seg_id_1, depth_1 in sub:
647
+ for seg_id_2, depth_2 in ret:
648
+ n_up = min(depth_1 - node["depth"], max_depth - 1)
649
+ n_down = min(depth_2 - node["depth"], max_depth - 1)
650
+ segment_rel[seg_id_1 * num_segments + seg_id_2] = rel_to_bucket(
651
+ n_up, n_down, max_depth=max_depth
652
+ )
653
+ segment_rel[seg_id_2 * num_segments + seg_id_1] = rel_to_bucket(
654
+ n_down, n_up, max_depth=max_depth
655
+ )
656
+ ret.extend(sub)
657
+ return ret
658
+
659
+ _build_segment_rel(root)
660
+
661
+ input_ids: List[int] = []
662
+ input_id_subs: List[int] = []
663
+ segment_bound: List[Tuple[int, int]] = []
664
+
665
+ if prev_ext_states is not None:
666
+ self.ext_table = prev_ext_states["ext_table"]
667
+ self.token_id_table = prev_ext_states["token_id_table"]
668
+
669
+ for seg in segments:
670
+ # tokenize
671
+ tokens = self.convert_tokens_to_ids(self.tokenize(seg["value"], for_cpmbee=True))
672
+
673
+ token_id_subs = []
674
+ reid_token_ids = []
675
+ for idx in tokens:
676
+ if idx in self.ext_table:
677
+ # unk or special token
678
+ token = self.ext_table[idx]
679
+ if token.startswith("<") and token.endswith(">"):
680
+ # special token
681
+ if "_" in token:
682
+ token_name = token[1:-1].split("_", maxsplit=1)[0]
683
+ else:
684
+ token_name = token[1:-1]
685
+ token_name = "<{}>".format(token_name)
686
+ else:
687
+ token_name = "<unk>"
688
+
689
+ if token_name not in self.token_id_table:
690
+ self.token_id_table[token_name] = {}
691
+ if idx not in self.token_id_table[token_name]:
692
+ self.token_id_table[token_name][idx] = len(self.token_id_table[token_name])
693
+ if token_name not in self.encoder:
694
+ raise ValueError("Invalid token {}".format(token))
695
+ reid_token_ids.append(self.encoder[token_name])
696
+ token_id_subs.append(self.token_id_table[token_name][idx])
697
+ else:
698
+ reid_token_ids.append(idx)
699
+ token_id_subs.append(0)
700
+ tokens = [self.bos_token_id] + reid_token_ids
701
+ token_id_subs = [0] + token_id_subs
702
+ # eos_id 表示 no need_predict
703
+ if not seg["need_predict"]: # eos
704
+ tokens = tokens + [self.eos_token_id]
705
+ token_id_subs = token_id_subs + [0]
706
+ else:
707
+ # no eos
708
+ pass
709
+ begin = len(input_ids)
710
+ input_ids.extend(tokens)
711
+ input_id_subs.extend(token_id_subs)
712
+ end = len(input_ids)
713
+ segment_bound.append((begin, end))
714
+
715
+ ids = np.array(input_ids, dtype=np.int32)
716
+ id_subs = np.array(input_id_subs, dtype=np.int32)
717
+ segs = np.zeros((ids.shape[0],), dtype=np.int32) # 按segment_bound对seg编号
718
+ context = np.zeros((ids.shape[0],), dtype=np.int8)
719
+ for i, (begin, end) in enumerate(segment_bound):
720
+ if not segments[i]["need_predict"]:
721
+ context[begin:end] = 1
722
+ segs[begin:end] = i
723
+
724
+ curr_ext_table_states: _PrevExtTableStates = {
725
+ "ext_table": self.ext_table,
726
+ "token_id_table": self.token_id_table,
727
+ }
728
+ return ids, id_subs, context, segs, segment_rel, num_segments, curr_ext_table_states
729
+
730
+ def prepare_for_model(
731
+ self,
732
+ ids: List[int],
733
+ pair_ids: Optional[List[int]] = None,
734
+ add_special_tokens: bool = True,
735
+ padding: Union[bool, str, PaddingStrategy] = False,
736
+ truncation: Union[bool, str, TruncationStrategy] = None,
737
+ max_length: Optional[int] = None,
738
+ stride: int = 0,
739
+ pad_to_multiple_of: Optional[int] = None,
740
+ return_tensors: Optional[Union[str, TensorType]] = None,
741
+ return_token_type_ids: Optional[bool] = None,
742
+ return_attention_mask: Optional[bool] = None,
743
+ return_overflowing_tokens: bool = False,
744
+ return_special_tokens_mask: bool = False,
745
+ return_length: bool = False,
746
+ verbose: bool = True,
747
+ prepend_batch_axis: bool = False,
748
+ **kwargs,
749
+ ) -> BatchEncoding:
750
+ """
751
+ Prepares a sequence of input id, or a pair of sequences of inputs ids so that it can be used by the model. It
752
+ adds special tokens, truncates sequences if overflowing while taking into account the special tokens and
753
+ manages a moving window (with user defined stride) for overflowing tokens. Please Note, for *pair_ids*
754
+ different than `None` and *truncation_strategy = longest_first* or `True`, it is not possible to return
755
+ overflowing tokens. Such a combination of arguments will raise an error.
756
+
757
+ Args:
758
+ ids (`List[int]`):
759
+ Tokenized input ids of the first sequence. Can be obtained from a string by chaining the `tokenize` and
760
+ `convert_tokens_to_ids` methods.
761
+ pair_ids (`List[int]`, *optional*):
762
+ Tokenized input ids of the second sequence. Can be obtained from a string by chaining the `tokenize`
763
+ and `convert_tokens_to_ids` methods.
764
+ """
765
+
766
+ # Backward compatibility for 'truncation_strategy', 'pad_to_max_length'
767
+ padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies(
768
+ padding=padding,
769
+ truncation=truncation,
770
+ max_length=max_length,
771
+ pad_to_multiple_of=pad_to_multiple_of,
772
+ verbose=verbose,
773
+ **kwargs,
774
+ )
775
+
776
+ pair = bool(pair_ids is not None)
777
+ len_ids = len(ids)
778
+ len_pair_ids = len(pair_ids) if pair else 0
779
+
780
+ if return_token_type_ids and not add_special_tokens:
781
+ raise ValueError(
782
+ "Asking to return token_type_ids while setting add_special_tokens to False "
783
+ "results in an undefined behavior. Please set add_special_tokens to True or "
784
+ "set return_token_type_ids to None."
785
+ )
786
+
787
+ if (
788
+ return_overflowing_tokens
789
+ and truncation_strategy == TruncationStrategy.LONGEST_FIRST
790
+ and pair_ids is not None
791
+ ):
792
+ raise ValueError(
793
+ "Not possible to return overflowing tokens for pair of sequences with the "
794
+ "`longest_first`. Please select another truncation strategy than `longest_first`, "
795
+ "for instance `only_second` or `only_first`."
796
+ )
797
+
798
+ # Load from model defaults
799
+ if return_token_type_ids is None:
800
+ return_token_type_ids = "token_type_ids" in self.model_input_names
801
+ if return_attention_mask is None:
802
+ return_attention_mask = "attention_mask" in self.model_input_names
803
+
804
+ encoded_inputs = {}
805
+
806
+ # Compute the total size of the returned encodings
807
+ total_len = len_ids + len_pair_ids + (self.num_special_tokens_to_add(pair=pair) if add_special_tokens else 0)
808
+
809
+ # Truncation: Handle max sequence length
810
+ overflowing_tokens = []
811
+ if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE and max_length and total_len > max_length:
812
+ ids, pair_ids, overflowing_tokens = self.truncate_sequences(
813
+ ids,
814
+ pair_ids=pair_ids,
815
+ num_tokens_to_remove=total_len - max_length,
816
+ truncation_strategy=truncation_strategy,
817
+ stride=stride,
818
+ )
819
+
820
+ if return_overflowing_tokens:
821
+ encoded_inputs["overflowing_tokens"] = overflowing_tokens
822
+ encoded_inputs["num_truncated_tokens"] = total_len - max_length
823
+
824
+ # Add special tokens
825
+ if add_special_tokens:
826
+ sequence = self.build_inputs_with_special_tokens(ids, pair_ids)
827
+ token_type_ids = self.create_token_type_ids_from_sequences(ids, pair_ids)
828
+ else:
829
+ sequence = ids + pair_ids if pair else ids
830
+ token_type_ids = [0] * len(ids) + ([0] * len(pair_ids) if pair else [])
831
+
832
+ # Build output dictionary
833
+ encoded_inputs["input_ids"] = sequence
834
+ if return_token_type_ids:
835
+ encoded_inputs["token_type_ids"] = token_type_ids
836
+ if return_special_tokens_mask:
837
+ if add_special_tokens:
838
+ encoded_inputs["special_tokens_mask"] = self.get_special_tokens_mask(ids, pair_ids)
839
+ else:
840
+ encoded_inputs["special_tokens_mask"] = [0] * len(sequence)
841
+
842
+ # Check lengths
843
+ self._eventual_warn_about_too_long_sequence(encoded_inputs["input_ids"], max_length, verbose)
844
+
845
+ # Padding
846
+ if padding_strategy != PaddingStrategy.DO_NOT_PAD or return_attention_mask:
847
+ encoded_inputs = self.pad(
848
+ encoded_inputs,
849
+ max_length=max_length,
850
+ padding=padding_strategy.value,
851
+ pad_to_multiple_of=pad_to_multiple_of,
852
+ return_attention_mask=return_attention_mask,
853
+ )
854
+
855
+ if return_length:
856
+ encoded_inputs["length"] = len(encoded_inputs["input_ids"])
857
+
858
+ # for CPMBee, encode all the model arguments
859
+ for arg in self.ext_args_for_model:
860
+ v = kwargs.get(arg, None)
861
+ if v is not None:
862
+ encoded_inputs[arg] = v
863
+
864
+ batch_outputs = BatchEncoding(
865
+ encoded_inputs, tensor_type=return_tensors, prepend_batch_axis=prepend_batch_axis
866
+ )
867
+
868
+ return batch_outputs