qwerwxy commited on
Commit
b28e08b
1 Parent(s): cfa8206

add config

Browse files
config.json ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "AquilaMoeForCausalLM"
4
+ ],
5
+ "attention_dropout": 0.0,
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_aquilamoe.AquilaMoeConfig",
8
+ "AutoModelForCausalLM": "modeling_aquilamoe.AquilaMoeForCausalLM"
9
+ },
10
+ "bos_token_id": 151849,
11
+ "eos_token_id": 151850,
12
+ "hidden_act": "silu",
13
+ "hidden_size": 5120,
14
+ "initializer_range": 0.02,
15
+ "intermediate_size": 20480,
16
+ "max_position_embeddings": 4096,
17
+ "model_type": "aquila3",
18
+ "num_attention_heads": 40,
19
+ "num_experts_per_tok": 2,
20
+ "num_hidden_layers": 40,
21
+ "num_key_value_heads": 8,
22
+ "num_local_experts": 8,
23
+ "output_router_logits": false,
24
+ "pad_token_id": 151643,
25
+ "rms_norm_eps": 1e-05,
26
+ "rope_theta": 1000000.0,
27
+ "router_aux_loss_coef": 0.001,
28
+ "sliding_window": null,
29
+ "tie_word_embeddings": false,
30
+ "torch_dtype": "bfloat16",
31
+ "transformers_version": "4.40.1",
32
+ "use_cache": true,
33
+ "vocab_size": 151851
34
+ }
configuration_aquilamoe.py ADDED
@@ -0,0 +1,142 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """ AquilaMoE model configuration"""
15
+
16
+ from transformers import PretrainedConfig
17
+ from transformers.utils import logging
18
+
19
+
20
+ logger = logging.get_logger(__name__)
21
+
22
+
23
+ class AquilaMoeConfig(PretrainedConfig):
24
+ r"""
25
+ Args:
26
+ vocab_size (`int`, *optional*, defaults to 32000):
27
+ Vocabulary size of the AquilaMoE model. Defines the number of different tokens that can be represented by the
28
+ `inputs_ids` passed when calling [`AquilaMoE`]
29
+ hidden_size (`int`, *optional*, defaults to 4096):
30
+ Dimension of the hidden representations.
31
+ intermediate_size (`int`, *optional*, defaults to 14336):
32
+ Dimension of the MLP representations.
33
+ num_hidden_layers (`int`, *optional*, defaults to 32):
34
+ Number of hidden layers in the Transformer encoder.
35
+ num_attention_heads (`int`, *optional*, defaults to 32):
36
+ Number of attention heads for each attention layer in the Transformer encoder.
37
+ num_key_value_heads (`int`, *optional*, defaults to 8):
38
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
39
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
40
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
41
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
42
+ by meanpooling all the original heads within that group. For more details checkout [this
43
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8`.
44
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
45
+ The non-linear activation function (function or string) in the decoder.
46
+ max_position_embeddings (`int`, *optional*, defaults to `4096*32`):
47
+ The maximum sequence length that this model might ever be used with. AquilaMoE's sliding window attention
48
+ allows sequence of up to 4096*32 tokens.
49
+ initializer_range (`float`, *optional*, defaults to 0.02):
50
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
51
+ rms_norm_eps (`float`, *optional*, defaults to 1e-05):
52
+ The epsilon used by the rms normalization layers.
53
+ use_cache (`bool`, *optional*, defaults to `True`):
54
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
55
+ relevant if `config.is_decoder=True`.
56
+ pad_token_id (`int`, *optional*):
57
+ The id of the padding token.
58
+ bos_token_id (`int`, *optional*, defaults to 1):
59
+ The id of the "beginning-of-sequence" token.
60
+ eos_token_id (`int`, *optional*, defaults to 2):
61
+ The id of the "end-of-sequence" token.
62
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
63
+ Whether the model's input and output word embeddings should be tied.
64
+ rope_theta (`float`, *optional*, defaults to 1000000.0):
65
+ The base period of the RoPE embeddings.
66
+ sliding_window (`int`, *optional*, defaults to 4096):
67
+ Sliding window attention window size. If not specified, will default to `4096`.
68
+ attention_dropout (`float`, *optional*, defaults to 0.0):
69
+ The dropout ratio for the attention probabilities.
70
+ num_experts_per_tok (`int`, *optional*, defaults to 2):
71
+ The number of experts to root per-token, can be also interpreted as the `top-p` routing
72
+ parameter
73
+ num_local_experts (`int`, *optional*, defaults to 8):
74
+ Number of experts per Sparse MLP layer.
75
+ output_router_logits (`bool`, *optional*, defaults to `False`):
76
+ Whether or not the router logits should be returned by the model. Enabeling this will also
77
+ allow the model to output the auxiliary loss. See [here]() for more details
78
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
79
+ The aux loss factor for the total loss.
80
+
81
+ """
82
+
83
+ model_type = "aquila3"
84
+ keys_to_ignore_at_inference = ["past_key_values"]
85
+
86
+ def __init__(
87
+ self,
88
+ vocab_size=32000,
89
+ hidden_size=4096,
90
+ intermediate_size=14336,
91
+ num_hidden_layers=32,
92
+ num_attention_heads=32,
93
+ num_key_value_heads=8,
94
+ hidden_act="silu",
95
+ max_position_embeddings=4096 * 32,
96
+ initializer_range=0.02,
97
+ rms_norm_eps=1e-5,
98
+ use_cache=True,
99
+ pad_token_id=None,
100
+ bos_token_id=1,
101
+ eos_token_id=2,
102
+ tie_word_embeddings=False,
103
+ rope_theta=1e6,
104
+ sliding_window=4096,
105
+ attention_dropout=0.0,
106
+ num_experts_per_tok=2,
107
+ num_local_experts=8,
108
+ output_router_logits=False,
109
+ router_aux_loss_coef=0.001,
110
+ **kwargs,
111
+ ):
112
+ self.vocab_size = vocab_size
113
+ self.max_position_embeddings = max_position_embeddings
114
+ self.hidden_size = hidden_size
115
+ self.intermediate_size = intermediate_size
116
+ self.num_hidden_layers = num_hidden_layers
117
+ self.num_attention_heads = num_attention_heads
118
+ self.sliding_window = sliding_window
119
+
120
+ # for backward compatibility
121
+ if num_key_value_heads is None:
122
+ num_key_value_heads = num_attention_heads
123
+
124
+ self.num_key_value_heads = num_key_value_heads
125
+ self.hidden_act = hidden_act
126
+ self.initializer_range = initializer_range
127
+ self.rms_norm_eps = rms_norm_eps
128
+ self.use_cache = use_cache
129
+ self.rope_theta = rope_theta
130
+ self.attention_dropout = attention_dropout
131
+
132
+ self.num_experts_per_tok = num_experts_per_tok
133
+ self.num_local_experts = num_local_experts
134
+ self.output_router_logits = output_router_logits
135
+ self.router_aux_loss_coef = router_aux_loss_coef
136
+ super().__init__(
137
+ pad_token_id=pad_token_id,
138
+ bos_token_id=bos_token_id,
139
+ eos_token_id=eos_token_id,
140
+ tie_word_embeddings=tie_word_embeddings,
141
+ **kwargs,
142
+ )
modeling_aquilamoe.py ADDED
@@ -0,0 +1,1444 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """ PyTorch AquilaMoE model."""
21
+ import inspect
22
+ import math
23
+ import warnings
24
+ from typing import List, Optional, Tuple, Union
25
+
26
+ import torch
27
+ import torch.nn.functional as F
28
+ import torch.utils.checkpoint
29
+ from torch import nn
30
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
31
+
32
+ from transformers.activations import ACT2FN
33
+ from transformers.cache_utils import Cache, DynamicCache
34
+ from transformers.modeling_attn_mask_utils import (
35
+ _prepare_4d_causal_attention_mask,
36
+ )
37
+ from transformers.modeling_outputs import (
38
+ MoeCausalLMOutputWithPast,
39
+ MoeModelOutputWithPast,
40
+ SequenceClassifierOutputWithPast,
41
+ )
42
+ from transformers.modeling_utils import PreTrainedModel
43
+ from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_13
44
+ from transformers.utils import (
45
+ add_start_docstrings,
46
+ add_start_docstrings_to_model_forward,
47
+ is_flash_attn_2_available,
48
+ is_flash_attn_greater_or_equal_2_10,
49
+ logging,
50
+ replace_return_docstrings,
51
+ )
52
+ from transformers.utils.import_utils import is_torch_fx_available
53
+ from .configuration_aquilamoe import AquilaMoeConfig
54
+
55
+
56
+ if is_flash_attn_2_available():
57
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
58
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
59
+
60
+ _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
61
+
62
+ # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
63
+ # It means that the function will not be traced through and simply appear as a node in the graph.
64
+ if is_torch_fx_available():
65
+ if not is_torch_greater_or_equal_than_1_13:
66
+ import torch.fx
67
+
68
+ _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
69
+
70
+
71
+ logger = logging.get_logger(__name__)
72
+
73
+ _CONFIG_FOR_DOC = "AquilaMoEConfig"
74
+
75
+
76
+ def load_balancing_loss_func(gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2) -> float:
77
+ r"""
78
+ Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
79
+
80
+ See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
81
+ function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
82
+ experts is too unbalanced.
83
+
84
+ Args:
85
+ gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):
86
+ Logits from the `gate`, should be a tuple of tensors. Shape: [batch_size, seqeunce_length, num_experts].
87
+ num_experts (`int`, *optional*):
88
+ Number of experts
89
+
90
+ Returns:
91
+ The auxiliary loss.
92
+ """
93
+ if gate_logits is None:
94
+ return 0
95
+
96
+ if isinstance(gate_logits, tuple):
97
+ # cat along the layers?
98
+ compute_device = gate_logits[0].device
99
+ gate_logits = torch.cat([gate.to(compute_device) for gate in gate_logits], dim=0)
100
+
101
+ routing_weights, selected_experts = torch.topk(gate_logits, top_k, dim=-1)
102
+ routing_weights = routing_weights.softmax(dim=-1)
103
+
104
+ # cast the expert indices to int64, otherwise one-hot encoding will fail
105
+ if selected_experts.dtype != torch.int64:
106
+ selected_experts = selected_experts.to(torch.int64)
107
+
108
+ if len(selected_experts.shape) == 2:
109
+ selected_experts = selected_experts.unsqueeze(2)
110
+
111
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
112
+
113
+ # For a given token, determine if it was routed to a given expert.
114
+ expert_mask = torch.max(expert_mask, axis=-2).values
115
+
116
+ # cast to float32 otherwise mean will fail
117
+ expert_mask = expert_mask.to(torch.float32)
118
+ tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)
119
+
120
+ router_prob_per_group_and_expert = torch.mean(routing_weights, axis=-1)
121
+ return torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert.unsqueeze(-1)) * (num_experts**2)
122
+
123
+
124
+ # Copied from transformers.models.llama.modeling_llama._get_unpad_data
125
+ def _get_unpad_data(attention_mask):
126
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
127
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
128
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
129
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
130
+ return (
131
+ indices,
132
+ cu_seqlens,
133
+ max_seqlen_in_batch,
134
+ )
135
+
136
+
137
+ # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->AquilaMoE
138
+ class AquilaMoeRMSNorm(nn.Module):
139
+ def __init__(self, hidden_size, eps=1e-6):
140
+ """
141
+ AquilaMoeRMSNorm is equivalent to T5LayerNorm
142
+ """
143
+ super().__init__()
144
+ self.weight = nn.Parameter(torch.ones(hidden_size))
145
+ self.variance_epsilon = eps
146
+
147
+ def forward(self, hidden_states):
148
+ input_dtype = hidden_states.dtype
149
+ hidden_states = hidden_states.to(torch.float32)
150
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
151
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
152
+ return self.weight * hidden_states.to(input_dtype)
153
+
154
+
155
+ # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->AquilaMoE
156
+ class AquilaMoeRotaryEmbedding(nn.Module):
157
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
158
+ super().__init__()
159
+
160
+ self.dim = dim
161
+ self.max_position_embeddings = max_position_embeddings
162
+ self.base = base
163
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
164
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
165
+
166
+ # Build here to make `torch.jit.trace` work.
167
+ self._set_cos_sin_cache(
168
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
169
+ )
170
+
171
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
172
+ self.max_seq_len_cached = seq_len
173
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
174
+
175
+ freqs = torch.outer(t, self.inv_freq)
176
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
177
+ emb = torch.cat((freqs, freqs), dim=-1)
178
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
179
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
180
+
181
+ def forward(self, x, seq_len=None):
182
+ # x: [bs, num_attention_heads, seq_len, head_size]
183
+ if seq_len > self.max_seq_len_cached:
184
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
185
+
186
+ return (
187
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
188
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
189
+ )
190
+
191
+
192
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
193
+ def rotate_half(x):
194
+ """Rotates half the hidden dims of the input."""
195
+ x1 = x[..., : x.shape[-1] // 2]
196
+ x2 = x[..., x.shape[-1] // 2 :]
197
+ return torch.cat((-x2, x1), dim=-1)
198
+
199
+
200
+ # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
201
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
202
+ """Applies Rotary Position Embedding to the query and key tensors.
203
+
204
+ Args:
205
+ q (`torch.Tensor`): The query tensor.
206
+ k (`torch.Tensor`): The key tensor.
207
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
208
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
209
+ position_ids (`torch.Tensor`):
210
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
211
+ used to pass offsetted position ids when working with a KV-cache.
212
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
213
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
214
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
215
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
216
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
217
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
218
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
219
+ Returns:
220
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
221
+ """
222
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
223
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
224
+ q_embed = (q * cos) + (rotate_half(q) * sin)
225
+ k_embed = (k * cos) + (rotate_half(k) * sin)
226
+ return q_embed, k_embed
227
+
228
+
229
+ # Copied from transformers.models.llama.modeling_llama.repeat_kv
230
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
231
+ """
232
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
233
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
234
+ """
235
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
236
+ if n_rep == 1:
237
+ return hidden_states
238
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
239
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
240
+
241
+
242
+ class AquilaMoeAttention(nn.Module):
243
+ """
244
+ Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
245
+ and "Generating Long Sequences with Sparse Transformers".
246
+ """
247
+
248
+ def __init__(self, config: AquilaMoeConfig, layer_idx: Optional[int] = None):
249
+ super().__init__()
250
+ self.config = config
251
+ self.layer_idx = layer_idx
252
+ if layer_idx is None:
253
+ logger.warning_once(
254
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
255
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
256
+ "when creating this class."
257
+ )
258
+
259
+ self.hidden_size = config.hidden_size
260
+ self.num_heads = config.num_attention_heads
261
+ self.head_dim = self.hidden_size // self.num_heads
262
+ self.num_key_value_heads = config.num_key_value_heads
263
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
264
+ self.max_position_embeddings = config.max_position_embeddings
265
+ self.rope_theta = config.rope_theta
266
+ self.is_causal = True
267
+ self.attention_dropout = config.attention_dropout
268
+
269
+ if (self.head_dim * self.num_heads) != self.hidden_size:
270
+ raise ValueError(
271
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
272
+ f" and `num_heads`: {self.num_heads})."
273
+ )
274
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
275
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
276
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
277
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
278
+
279
+ self.rotary_emb = AquilaMoeRotaryEmbedding(
280
+ self.head_dim,
281
+ max_position_embeddings=self.max_position_embeddings,
282
+ base=self.rope_theta,
283
+ )
284
+
285
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
286
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
287
+
288
+ def forward(
289
+ self,
290
+ hidden_states: torch.Tensor,
291
+ attention_mask: Optional[torch.Tensor] = None,
292
+ position_ids: Optional[torch.LongTensor] = None,
293
+ past_key_value: Optional[Cache] = None,
294
+ output_attentions: bool = False,
295
+ use_cache: bool = False,
296
+ **kwargs,
297
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
298
+ if "padding_mask" in kwargs:
299
+ warnings.warn(
300
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
301
+ )
302
+ bsz, q_len, _ = hidden_states.size()
303
+
304
+ query_states = self.q_proj(hidden_states)
305
+ key_states = self.k_proj(hidden_states)
306
+ value_states = self.v_proj(hidden_states)
307
+
308
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
309
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
310
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
311
+
312
+ kv_seq_len = key_states.shape[-2]
313
+ if past_key_value is not None:
314
+ if self.layer_idx is None:
315
+ raise ValueError(
316
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
317
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
318
+ "with a layer index."
319
+ )
320
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
321
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
322
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
323
+
324
+ if past_key_value is not None:
325
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
326
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
327
+
328
+ # repeat k/v heads if n_kv_heads < n_heads
329
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
330
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
331
+
332
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
333
+
334
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
335
+ raise ValueError(
336
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
337
+ f" {attn_weights.size()}"
338
+ )
339
+
340
+ if attention_mask is not None:
341
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
342
+ raise ValueError(
343
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
344
+ )
345
+
346
+ attn_weights = attn_weights + attention_mask
347
+
348
+ # upcast attention to fp32
349
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
350
+ attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
351
+ attn_output = torch.matmul(attn_weights, value_states)
352
+
353
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
354
+ raise ValueError(
355
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
356
+ f" {attn_output.size()}"
357
+ )
358
+
359
+ attn_output = attn_output.transpose(1, 2).contiguous()
360
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
361
+
362
+ attn_output = self.o_proj(attn_output)
363
+
364
+ if not output_attentions:
365
+ attn_weights = None
366
+
367
+ return attn_output, attn_weights, past_key_value
368
+
369
+
370
+ # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2 with Mistral->AquilaMoe
371
+ class AquilaMoeFlashAttention2(AquilaMoeAttention):
372
+ """
373
+ AquilaMoE flash attention module. This module inherits from `AquilaMoeAttention` as the weights of the module stays
374
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
375
+ flash attention and deal with padding tokens in case the input contains any of them.
376
+ """
377
+
378
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
379
+ def __init__(self, *args, **kwargs):
380
+ super().__init__(*args, **kwargs)
381
+
382
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
383
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
384
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
385
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
386
+
387
+ def forward(
388
+ self,
389
+ hidden_states: torch.Tensor,
390
+ attention_mask: Optional[torch.Tensor] = None,
391
+ position_ids: Optional[torch.LongTensor] = None,
392
+ past_key_value: Optional[Cache] = None,
393
+ output_attentions: bool = False,
394
+ use_cache: bool = False,
395
+ **kwargs,
396
+ ):
397
+ if "padding_mask" in kwargs:
398
+ warnings.warn(
399
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
400
+ )
401
+
402
+ # overwrite attention_mask with padding_mask
403
+ attention_mask = kwargs.pop("padding_mask")
404
+ bsz, q_len, _ = hidden_states.size()
405
+
406
+ query_states = self.q_proj(hidden_states)
407
+ key_states = self.k_proj(hidden_states)
408
+ value_states = self.v_proj(hidden_states)
409
+
410
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
411
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
412
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
413
+
414
+ kv_seq_len = key_states.shape[-2]
415
+ if past_key_value is not None:
416
+ if self.layer_idx is None:
417
+ raise ValueError(
418
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
419
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
420
+ "with a layer index."
421
+ )
422
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
423
+
424
+ # Because the input can be padded, the absolute sequence length depends on the max position id.
425
+ rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
426
+ cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
427
+
428
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
429
+
430
+ use_sliding_windows = (
431
+ _flash_supports_window_size
432
+ and getattr(self.config, "sliding_window", None) is not None
433
+ and kv_seq_len > self.config.sliding_window
434
+ )
435
+
436
+ if not _flash_supports_window_size:
437
+ logger.warning_once(
438
+ "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"
439
+ " make sure to upgrade flash-attn library."
440
+ )
441
+
442
+ if past_key_value is not None:
443
+ # Activate slicing cache only if the config has a value `sliding_windows` attribute
444
+ cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0
445
+ if (
446
+ getattr(self.config, "sliding_window", None) is not None
447
+ and kv_seq_len > self.config.sliding_window
448
+ and cache_has_contents
449
+ ):
450
+ slicing_tokens = 1 - self.config.sliding_window
451
+
452
+ past_key = past_key_value[self.layer_idx][0]
453
+ past_value = past_key_value[self.layer_idx][1]
454
+
455
+ past_key = past_key[:, :, slicing_tokens:, :].contiguous()
456
+ past_value = past_value[:, :, slicing_tokens:, :].contiguous()
457
+
458
+ if past_key.shape[-2] != self.config.sliding_window - 1:
459
+ raise ValueError(
460
+ f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
461
+ f" {past_key.shape}"
462
+ )
463
+
464
+ if attention_mask is not None:
465
+ attention_mask = attention_mask[:, slicing_tokens:]
466
+ attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)
467
+
468
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
469
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
470
+
471
+ # repeat k/v heads if n_kv_heads < n_heads
472
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
473
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
474
+ dropout_rate = 0.0 if not self.training else self.attention_dropout
475
+
476
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
477
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
478
+ # cast them back in float16 just to be sure everything works as expected.
479
+ input_dtype = query_states.dtype
480
+ if input_dtype == torch.float32:
481
+ # Handle the case where the model is quantized
482
+ if hasattr(self.config, "_pre_quantization_dtype"):
483
+ target_dtype = self.config._pre_quantization_dtype
484
+ else:
485
+ target_dtype = self.q_proj.weight.dtype
486
+
487
+ logger.warning_once(
488
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
489
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
490
+ f" {target_dtype}."
491
+ )
492
+
493
+ query_states = query_states.to(target_dtype)
494
+ key_states = key_states.to(target_dtype)
495
+ value_states = value_states.to(target_dtype)
496
+
497
+ # Reashape to the expected shape for Flash Attention
498
+ query_states = query_states.transpose(1, 2)
499
+ key_states = key_states.transpose(1, 2)
500
+ value_states = value_states.transpose(1, 2)
501
+
502
+ attn_output = self._flash_attention_forward(
503
+ query_states,
504
+ key_states,
505
+ value_states,
506
+ attention_mask,
507
+ q_len,
508
+ dropout=dropout_rate,
509
+ use_sliding_windows=use_sliding_windows,
510
+ )
511
+
512
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
513
+ attn_output = self.o_proj(attn_output)
514
+
515
+ if not output_attentions:
516
+ attn_weights = None
517
+
518
+ return attn_output, attn_weights, past_key_value
519
+
520
+ def _flash_attention_forward(
521
+ self,
522
+ query_states,
523
+ key_states,
524
+ value_states,
525
+ attention_mask,
526
+ query_length,
527
+ dropout=0.0,
528
+ softmax_scale=None,
529
+ use_sliding_windows=False,
530
+ ):
531
+ """
532
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
533
+ first unpad the input, then computes the attention scores and pad the final attention scores.
534
+
535
+ Args:
536
+ query_states (`torch.Tensor`):
537
+ Input query states to be passed to Flash Attention API
538
+ key_states (`torch.Tensor`):
539
+ Input key states to be passed to Flash Attention API
540
+ value_states (`torch.Tensor`):
541
+ Input value states to be passed to Flash Attention API
542
+ attention_mask (`torch.Tensor`):
543
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
544
+ position of padding tokens and 1 for the position of non-padding tokens.
545
+ dropout (`int`, *optional*):
546
+ Attention dropout
547
+ softmax_scale (`float`, *optional*):
548
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
549
+ use_sliding_windows (`bool`, *optional*):
550
+ Whether to activate sliding window attention.
551
+ """
552
+ if not self._flash_attn_uses_top_left_mask:
553
+ causal = self.is_causal
554
+ else:
555
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
556
+ causal = self.is_causal and query_length != 1
557
+
558
+ # Contains at least one padding token in the sequence
559
+ if attention_mask is not None:
560
+ batch_size = query_states.shape[0]
561
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
562
+ query_states, key_states, value_states, attention_mask, query_length
563
+ )
564
+
565
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
566
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
567
+
568
+ if not use_sliding_windows:
569
+ attn_output_unpad = flash_attn_varlen_func(
570
+ query_states,
571
+ key_states,
572
+ value_states,
573
+ cu_seqlens_q=cu_seqlens_q,
574
+ cu_seqlens_k=cu_seqlens_k,
575
+ max_seqlen_q=max_seqlen_in_batch_q,
576
+ max_seqlen_k=max_seqlen_in_batch_k,
577
+ dropout_p=dropout,
578
+ softmax_scale=softmax_scale,
579
+ causal=causal,
580
+ )
581
+ else:
582
+ attn_output_unpad = flash_attn_varlen_func(
583
+ query_states,
584
+ key_states,
585
+ value_states,
586
+ cu_seqlens_q=cu_seqlens_q,
587
+ cu_seqlens_k=cu_seqlens_k,
588
+ max_seqlen_q=max_seqlen_in_batch_q,
589
+ max_seqlen_k=max_seqlen_in_batch_k,
590
+ dropout_p=dropout,
591
+ softmax_scale=softmax_scale,
592
+ causal=causal,
593
+ window_size=(self.config.sliding_window, self.config.sliding_window),
594
+ )
595
+
596
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
597
+ else:
598
+ if not use_sliding_windows:
599
+ attn_output = flash_attn_func(
600
+ query_states,
601
+ key_states,
602
+ value_states,
603
+ dropout,
604
+ softmax_scale=softmax_scale,
605
+ causal=causal,
606
+ )
607
+ else:
608
+ attn_output = flash_attn_func(
609
+ query_states,
610
+ key_states,
611
+ value_states,
612
+ dropout,
613
+ softmax_scale=softmax_scale,
614
+ causal=causal,
615
+ window_size=(self.config.sliding_window, self.config.sliding_window),
616
+ )
617
+
618
+ return attn_output
619
+
620
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
621
+ batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
622
+
623
+ # On the first iteration we need to properly re-create the padding mask
624
+ # by slicing it on the proper place
625
+ if kv_seq_len != attention_mask.shape[-1]:
626
+ attention_mask_num_tokens = attention_mask.shape[-1]
627
+ attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
628
+
629
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
630
+
631
+ key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
632
+ value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
633
+
634
+ if query_length == kv_seq_len:
635
+ query_layer = index_first_axis(
636
+ query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
637
+ )
638
+ cu_seqlens_q = cu_seqlens_k
639
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
640
+ indices_q = indices_k
641
+ elif query_length == 1:
642
+ max_seqlen_in_batch_q = 1
643
+ cu_seqlens_q = torch.arange(
644
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
645
+ ) # There is a memcpy here, that is very bad.
646
+ indices_q = cu_seqlens_q[:-1]
647
+ query_layer = query_layer.squeeze(1)
648
+ else:
649
+ # The -q_len: slice assumes left padding.
650
+ attention_mask = attention_mask[:, -query_length:]
651
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
652
+
653
+ return (
654
+ query_layer,
655
+ key_layer,
656
+ value_layer,
657
+ indices_q,
658
+ (cu_seqlens_q, cu_seqlens_k),
659
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
660
+ )
661
+
662
+
663
+ class AquilaMoeBLockSparseTop2MLP(nn.Module):
664
+ def __init__(self, config: AquilaMoeConfig):
665
+ super().__init__()
666
+ self.ffn_dim = config.intermediate_size
667
+ self.hidden_dim = config.hidden_size
668
+
669
+ self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
670
+ self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
671
+ self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
672
+
673
+ self.act_fn = ACT2FN[config.hidden_act]
674
+
675
+ def forward(self, hidden_states, routing_weights):
676
+ current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states)
677
+ current_hidden_states = self.w2(current_hidden_states)
678
+ return routing_weights * current_hidden_states
679
+
680
+
681
+ MISTRAL_ATTENTION_CLASSES = {
682
+ "eager": AquilaMoeAttention,
683
+ "flash_attention_2": AquilaMoeFlashAttention2,
684
+ }
685
+
686
+
687
+ class AquilaMoeSparseMoeBlock(nn.Module):
688
+ """
689
+ This implementation is
690
+ strictly equivalent to standard MoE with full capacity (no
691
+ dropped tokens). It's faster since it formulates MoE operations
692
+ in terms of block-sparse operations to accomodate imbalanced
693
+ assignments of tokens to experts, whereas standard MoE either
694
+ (1) drop tokens at the cost of reduced performance or (2) set
695
+ capacity factor to number of experts and thus waste computation
696
+ and memory on padding.
697
+ """
698
+
699
+ def __init__(self, config):
700
+ super().__init__()
701
+ self.hidden_dim = config.hidden_size
702
+ self.ffn_dim = config.intermediate_size
703
+ self.num_experts = config.num_local_experts
704
+ self.top_k = config.num_experts_per_tok
705
+
706
+ # gating
707
+ self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
708
+
709
+ self.experts = nn.ModuleList([AquilaMoeBLockSparseTop2MLP(config) for _ in range(self.num_experts)])
710
+
711
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
712
+ """ """
713
+ batch_size, sequence_length, hidden_dim = hidden_states.shape
714
+ hidden_states = hidden_states.view(-1, hidden_dim)
715
+ # router_logits: (batch * sequence_length, n_experts)
716
+ router_logits = self.gate(hidden_states)
717
+
718
+ routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
719
+ routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
720
+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
721
+ # we cast back to the input dtype
722
+ routing_weights = routing_weights.to(hidden_states.dtype)
723
+
724
+ final_hidden_states = torch.zeros(
725
+ (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
726
+ )
727
+
728
+ # One hot encode the selected experts to create an expert mask
729
+ # this will be used to easily index which expert is going to be sollicitated
730
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
731
+
732
+ # Loop over all available experts in the model and perform the computation on each expert
733
+ for expert_idx in range(self.num_experts):
734
+ expert_layer = self.experts[expert_idx]
735
+ idx, top_x = torch.where(expert_mask[expert_idx])
736
+
737
+ if top_x.shape[0] == 0:
738
+ continue
739
+
740
+ # in torch it is faster to index using lists than torch tensors
741
+ top_x_list = top_x.tolist()
742
+ idx_list = idx.tolist()
743
+
744
+ # Index the correct hidden states and compute the expert hidden state for
745
+ # the current expert. We need to make sure to multiply the output hidden
746
+ # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
747
+ current_state = hidden_states[None, top_x_list].reshape(-1, hidden_dim)
748
+ current_hidden_states = expert_layer(current_state, routing_weights[top_x_list, idx_list, None])
749
+
750
+ # However `index_add_` only support torch tensors for indexing so we'll use
751
+ # the `top_x` tensor here.
752
+ final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
753
+ final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
754
+ return final_hidden_states, router_logits
755
+
756
+
757
+ class AquilaMoeDecoderLayer(nn.Module):
758
+ def __init__(self, config: AquilaMoeConfig, layer_idx: int):
759
+ super().__init__()
760
+ self.hidden_size = config.hidden_size
761
+
762
+ self.self_attn = MISTRAL_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
763
+
764
+ self.block_sparse_moe = AquilaMoeSparseMoeBlock(config)
765
+ self.input_layernorm = AquilaMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
766
+ self.post_attention_layernorm = AquilaMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
767
+
768
+ def forward(
769
+ self,
770
+ hidden_states: torch.Tensor,
771
+ attention_mask: Optional[torch.Tensor] = None,
772
+ position_ids: Optional[torch.LongTensor] = None,
773
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
774
+ output_attentions: Optional[bool] = False,
775
+ output_router_logits: Optional[bool] = False,
776
+ use_cache: Optional[bool] = False,
777
+ **kwargs,
778
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
779
+ if "padding_mask" in kwargs:
780
+ warnings.warn(
781
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
782
+ )
783
+ """
784
+ Args:
785
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
786
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
787
+ `(batch, sequence_length)` where padding elements are indicated by 0.
788
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
789
+ output_attentions (`bool`, *optional*):
790
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
791
+ returned tensors for more detail.
792
+ output_router_logits (`bool`, *optional*):
793
+ Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
794
+ should not be returned during inference.
795
+ use_cache (`bool`, *optional*):
796
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
797
+ (see `past_key_values`).
798
+ """
799
+
800
+ residual = hidden_states
801
+
802
+ hidden_states = self.input_layernorm(hidden_states)
803
+
804
+ # Self Attention
805
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
806
+ hidden_states=hidden_states,
807
+ attention_mask=attention_mask,
808
+ position_ids=position_ids,
809
+ past_key_value=past_key_value,
810
+ output_attentions=output_attentions,
811
+ use_cache=use_cache,
812
+ )
813
+ hidden_states = residual + hidden_states
814
+
815
+ # Fully Connected
816
+ residual = hidden_states
817
+ hidden_states = self.post_attention_layernorm(hidden_states)
818
+ hidden_states, router_logits = self.block_sparse_moe(hidden_states)
819
+ hidden_states = residual + hidden_states
820
+
821
+ outputs = (hidden_states,)
822
+
823
+ if output_attentions:
824
+ outputs += (self_attn_weights,)
825
+
826
+ if use_cache:
827
+ outputs += (present_key_value,)
828
+
829
+ if output_router_logits:
830
+ outputs += (router_logits,)
831
+
832
+ return outputs
833
+
834
+
835
+ AquilaMoE_START_DOCSTRING = r"""
836
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
837
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
838
+ etc.)
839
+
840
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
841
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
842
+ and behavior.
843
+
844
+ Parameters:
845
+ config ([`AquilaMoeConfig`]):
846
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
847
+ load the weights associated with the model, only the configuration. Check out the
848
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
849
+ """
850
+
851
+
852
+ @add_start_docstrings(
853
+ "The bare AquilaMoE Model outputting raw hidden-states without any specific head on top.",
854
+ AquilaMoE_START_DOCSTRING,
855
+ )
856
+ # Copied from transformers.models.mistral.modeling_mistral.MistralPreTrainedModel with Mistral->AquilaMoE
857
+ class AquilaMoePreTrainedModel(PreTrainedModel):
858
+ config_class = AquilaMoeConfig
859
+ base_model_prefix = "model"
860
+ supports_gradient_checkpointing = True
861
+ _no_split_modules = ["AquilaMoeDecoderLayer"]
862
+ _skip_keys_device_placement = "past_key_values"
863
+ _supports_flash_attn_2 = True
864
+ _supports_cache_class = True
865
+
866
+ def _init_weights(self, module):
867
+ std = self.config.initializer_range
868
+ if isinstance(module, nn.Linear):
869
+ module.weight.data.normal_(mean=0.0, std=std)
870
+ if module.bias is not None:
871
+ module.bias.data.zero_()
872
+ elif isinstance(module, nn.Embedding):
873
+ module.weight.data.normal_(mean=0.0, std=std)
874
+ if module.padding_idx is not None:
875
+ module.weight.data[module.padding_idx].zero_()
876
+
877
+
878
+ AquilaMoE_INPUTS_DOCSTRING = r"""
879
+ Args:
880
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
881
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
882
+ it.
883
+
884
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
885
+ [`PreTrainedTokenizer.__call__`] for details.
886
+
887
+ [What are input IDs?](../glossary#input-ids)
888
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
889
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
890
+
891
+ - 1 for tokens that are **not masked**,
892
+ - 0 for tokens that are **masked**.
893
+
894
+ [What are attention masks?](../glossary#attention-mask)
895
+
896
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
897
+ [`PreTrainedTokenizer.__call__`] for details.
898
+
899
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
900
+ `past_key_values`).
901
+
902
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
903
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
904
+ information on the default strategy.
905
+
906
+ - 1 indicates the head is **not masked**,
907
+ - 0 indicates the head is **masked**.
908
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
909
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
910
+ config.n_positions - 1]`.
911
+
912
+ [What are position IDs?](../glossary#position-ids)
913
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
914
+ Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
915
+ `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
916
+ `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
917
+
918
+ Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
919
+ blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
920
+
921
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
922
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
923
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
924
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
925
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
926
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
927
+ model's internal embedding lookup matrix.
928
+ use_cache (`bool`, *optional*):
929
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
930
+ `past_key_values`).
931
+ output_attentions (`bool`, *optional*):
932
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
933
+ tensors for more detail.
934
+ output_hidden_states (`bool`, *optional*):
935
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
936
+ more detail.
937
+ output_router_logits (`bool`, *optional*):
938
+ Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
939
+ should not be returned during inference.
940
+ return_dict (`bool`, *optional*):
941
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
942
+ """
943
+
944
+
945
+ @add_start_docstrings(
946
+ "The bare AquilaMoe Model outputting raw hidden-states without any specific head on top.",
947
+ AquilaMoE_START_DOCSTRING,
948
+ )
949
+ # Copied from transformers.models.mistral.modeling_mistral.MistralModel with MISTRAL->AquilaMoe,Mistral->AquilaMoe
950
+ class AquilaMoeModel(AquilaMoePreTrainedModel):
951
+ """
952
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`AquilaMoeDecoderLayer`]
953
+
954
+ Args:
955
+ config: AquilaMoeConfig
956
+ """
957
+
958
+ def __init__(self, config: AquilaMoeConfig):
959
+ super().__init__(config)
960
+ self.padding_idx = config.pad_token_id
961
+ self.vocab_size = config.vocab_size
962
+
963
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
964
+ self.layers = nn.ModuleList(
965
+ [AquilaMoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
966
+ )
967
+ self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
968
+ self.norm = AquilaMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
969
+
970
+ self.gradient_checkpointing = False
971
+ # Initialize weights and apply final processing
972
+ self.post_init()
973
+
974
+ def get_input_embeddings(self):
975
+ return self.embed_tokens
976
+
977
+ def set_input_embeddings(self, value):
978
+ self.embed_tokens = value
979
+
980
+ # Ignore copy
981
+ @add_start_docstrings_to_model_forward(AquilaMoE_INPUTS_DOCSTRING)
982
+ def forward(
983
+ self,
984
+ input_ids: torch.LongTensor = None,
985
+ attention_mask: Optional[torch.Tensor] = None,
986
+ position_ids: Optional[torch.LongTensor] = None,
987
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
988
+ inputs_embeds: Optional[torch.FloatTensor] = None,
989
+ use_cache: Optional[bool] = None,
990
+ output_attentions: Optional[bool] = None,
991
+ output_hidden_states: Optional[bool] = None,
992
+ output_router_logits: Optional[bool] = None,
993
+ return_dict: Optional[bool] = None,
994
+ ) -> Union[Tuple, MoeModelOutputWithPast]:
995
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
996
+ output_router_logits = (
997
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
998
+ )
999
+ output_hidden_states = (
1000
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1001
+ )
1002
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1003
+
1004
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1005
+
1006
+ # retrieve input_ids and inputs_embeds
1007
+ if input_ids is not None and inputs_embeds is not None:
1008
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
1009
+ elif input_ids is not None:
1010
+ batch_size, seq_length = input_ids.shape
1011
+ elif inputs_embeds is not None:
1012
+ batch_size, seq_length, _ = inputs_embeds.shape
1013
+ else:
1014
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
1015
+
1016
+ past_key_values_length = 0
1017
+
1018
+ if self.gradient_checkpointing and self.training:
1019
+ if use_cache:
1020
+ logger.warning_once(
1021
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1022
+ )
1023
+ use_cache = False
1024
+
1025
+ if use_cache:
1026
+ use_legacy_cache = not isinstance(past_key_values, Cache)
1027
+ if use_legacy_cache:
1028
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1029
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
1030
+
1031
+ if position_ids is None:
1032
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1033
+ position_ids = torch.arange(
1034
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
1035
+ )
1036
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
1037
+ else:
1038
+ position_ids = position_ids.view(-1, seq_length).long()
1039
+
1040
+ if inputs_embeds is None:
1041
+ inputs_embeds = self.embed_tokens(input_ids)
1042
+
1043
+ if attention_mask is not None and self._use_flash_attention_2 and use_cache:
1044
+ is_padding_right = attention_mask[:, -1].sum().item() != batch_size
1045
+ if is_padding_right:
1046
+ raise ValueError(
1047
+ "You are attempting to perform batched generation with padding_side='right'"
1048
+ " this may lead to unexpected behaviour for Flash Attention version of AquilaMoE. Make sure to "
1049
+ " call `tokenizer.padding_side = 'left'` before tokenizing the input. "
1050
+ )
1051
+
1052
+ if self._use_flash_attention_2:
1053
+ # 2d mask is passed through the layers
1054
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
1055
+ else:
1056
+ # 4d mask is passed through the layers
1057
+ attention_mask = _prepare_4d_causal_attention_mask(
1058
+ attention_mask,
1059
+ (batch_size, seq_length),
1060
+ inputs_embeds,
1061
+ past_key_values_length,
1062
+ sliding_window=self.config.sliding_window,
1063
+ )
1064
+
1065
+ hidden_states = inputs_embeds
1066
+
1067
+ # decoder layers
1068
+ all_hidden_states = () if output_hidden_states else None
1069
+ all_self_attns = () if output_attentions else None
1070
+ all_router_logits = () if output_router_logits else None
1071
+ next_decoder_cache = None
1072
+
1073
+ for decoder_layer in self.layers:
1074
+ if output_hidden_states:
1075
+ all_hidden_states += (hidden_states,)
1076
+
1077
+ if self.gradient_checkpointing and self.training:
1078
+ layer_outputs = self._gradient_checkpointing_func(
1079
+ decoder_layer.__call__,
1080
+ hidden_states,
1081
+ attention_mask,
1082
+ position_ids,
1083
+ past_key_values,
1084
+ output_attentions,
1085
+ output_router_logits,
1086
+ use_cache,
1087
+ )
1088
+ else:
1089
+ layer_outputs = decoder_layer(
1090
+ hidden_states,
1091
+ attention_mask=attention_mask,
1092
+ position_ids=position_ids,
1093
+ past_key_value=past_key_values,
1094
+ output_attentions=output_attentions,
1095
+ output_router_logits=output_router_logits,
1096
+ use_cache=use_cache,
1097
+ )
1098
+
1099
+ hidden_states = layer_outputs[0]
1100
+
1101
+ if use_cache:
1102
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1103
+
1104
+ if output_attentions:
1105
+ all_self_attns += (layer_outputs[1],)
1106
+
1107
+ if output_router_logits:
1108
+ all_router_logits += (layer_outputs[-1],)
1109
+
1110
+ hidden_states = self.norm(hidden_states)
1111
+
1112
+ # add hidden states from the last decoder layer
1113
+ if output_hidden_states:
1114
+ all_hidden_states += (hidden_states,)
1115
+
1116
+ next_cache = None
1117
+ if use_cache:
1118
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
1119
+
1120
+ if not return_dict:
1121
+ return tuple(
1122
+ v
1123
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
1124
+ if v is not None
1125
+ )
1126
+ return MoeModelOutputWithPast(
1127
+ last_hidden_state=hidden_states,
1128
+ past_key_values=next_cache,
1129
+ hidden_states=all_hidden_states,
1130
+ attentions=all_self_attns,
1131
+ router_logits=all_router_logits,
1132
+ )
1133
+
1134
+
1135
+ class AquilaMoeForCausalLM(AquilaMoePreTrainedModel):
1136
+ _tied_weights_keys = ["lm_head.weight"]
1137
+
1138
+ def __init__(self, config):
1139
+ super().__init__(config)
1140
+ self.model = AquilaMoeModel(config)
1141
+ self.vocab_size = config.vocab_size
1142
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1143
+ self.router_aux_loss_coef = config.router_aux_loss_coef
1144
+ self.num_experts = config.num_local_experts
1145
+ self.num_experts_per_tok = config.num_experts_per_tok
1146
+ # Initialize weights and apply final processing
1147
+ self.post_init()
1148
+
1149
+ def get_input_embeddings(self):
1150
+ return self.model.embed_tokens
1151
+
1152
+ def set_input_embeddings(self, value):
1153
+ self.model.embed_tokens = value
1154
+
1155
+ def get_output_embeddings(self):
1156
+ return self.lm_head
1157
+
1158
+ def set_output_embeddings(self, new_embeddings):
1159
+ self.lm_head = new_embeddings
1160
+
1161
+ def set_decoder(self, decoder):
1162
+ self.model = decoder
1163
+
1164
+ def get_decoder(self):
1165
+ return self.model
1166
+
1167
+ @add_start_docstrings_to_model_forward(AquilaMoE_INPUTS_DOCSTRING)
1168
+ @replace_return_docstrings(output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1169
+ # Ignore copy
1170
+ def forward(
1171
+ self,
1172
+ input_ids: torch.LongTensor = None,
1173
+ attention_mask: Optional[torch.Tensor] = None,
1174
+ position_ids: Optional[torch.LongTensor] = None,
1175
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1176
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1177
+ labels: Optional[torch.LongTensor] = None,
1178
+ use_cache: Optional[bool] = None,
1179
+ output_attentions: Optional[bool] = None,
1180
+ output_hidden_states: Optional[bool] = None,
1181
+ output_router_logits: Optional[bool] = None,
1182
+ return_dict: Optional[bool] = None,
1183
+ ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
1184
+ r"""
1185
+ Args:
1186
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1187
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1188
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1189
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.voc
1190
+ ```"""
1191
+
1192
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1193
+ output_router_logits = (
1194
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
1195
+ )
1196
+
1197
+ output_hidden_states = (
1198
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1199
+ )
1200
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1201
+
1202
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1203
+ outputs = self.model(
1204
+ input_ids=input_ids,
1205
+ attention_mask=attention_mask,
1206
+ position_ids=position_ids,
1207
+ past_key_values=past_key_values,
1208
+ inputs_embeds=inputs_embeds,
1209
+ use_cache=use_cache,
1210
+ output_attentions=output_attentions,
1211
+ output_hidden_states=output_hidden_states,
1212
+ output_router_logits=output_router_logits,
1213
+ return_dict=return_dict,
1214
+ )
1215
+
1216
+ hidden_states = outputs[0]
1217
+ logits = self.lm_head(hidden_states)
1218
+ logits = logits.float()
1219
+
1220
+ loss = None
1221
+ if labels is not None:
1222
+ # Shift so that tokens < n predict n
1223
+ shift_logits = logits[..., :-1, :].contiguous()
1224
+ shift_labels = labels[..., 1:].contiguous()
1225
+ # Flatten the tokens
1226
+ loss_fct = CrossEntropyLoss()
1227
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1228
+ shift_labels = shift_labels.view(-1)
1229
+ # Enable model parallelism
1230
+ shift_labels = shift_labels.to(shift_logits.device)
1231
+ loss = loss_fct(shift_logits, shift_labels)
1232
+
1233
+ aux_loss = None
1234
+ if output_router_logits:
1235
+ aux_loss = load_balancing_loss_func(
1236
+ outputs.router_logits if return_dict else outputs[-1], self.num_experts, self.num_experts_per_tok
1237
+ )
1238
+ if labels is not None:
1239
+ loss += self.router_aux_loss_coef * aux_loss
1240
+
1241
+ if not return_dict:
1242
+ output = (logits,) + outputs[1:]
1243
+ if output_router_logits:
1244
+ output = (aux_loss,) + output
1245
+ return (loss,) + output if loss is not None else output
1246
+
1247
+ return MoeCausalLMOutputWithPast(
1248
+ loss=loss,
1249
+ aux_loss=aux_loss,
1250
+ logits=logits,
1251
+ past_key_values=outputs.past_key_values,
1252
+ hidden_states=outputs.hidden_states,
1253
+ attentions=outputs.attentions,
1254
+ router_logits=outputs.router_logits,
1255
+ )
1256
+
1257
+ def prepare_inputs_for_generation(
1258
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
1259
+ ):
1260
+ # Omit tokens covered by past_key_values
1261
+ if past_key_values is not None:
1262
+ if isinstance(past_key_values, Cache):
1263
+ cache_length = past_key_values.get_seq_length()
1264
+ past_length = past_key_values.seen_tokens
1265
+ max_cache_length = past_key_values.get_max_length()
1266
+ else:
1267
+ cache_length = past_length = past_key_values[0][0].shape[2]
1268
+ max_cache_length = None
1269
+
1270
+ # Keep only the unprocessed tokens:
1271
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1272
+ # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as
1273
+ # input)
1274
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
1275
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1276
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1277
+ # input_ids based on the past_length.
1278
+ elif past_length < input_ids.shape[1]:
1279
+ input_ids = input_ids[:, past_length:]
1280
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1281
+
1282
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1283
+ if (
1284
+ max_cache_length is not None
1285
+ and attention_mask is not None
1286
+ and cache_length + input_ids.shape[1] > max_cache_length
1287
+ ):
1288
+ attention_mask = attention_mask[:, -max_cache_length:]
1289
+
1290
+ position_ids = kwargs.get("position_ids", None)
1291
+ if attention_mask is not None and position_ids is None:
1292
+ # create position_ids on the fly for batch generation
1293
+ position_ids = attention_mask.long().cumsum(-1) - 1
1294
+ position_ids.masked_fill_(attention_mask == 0, 1)
1295
+ if past_key_values:
1296
+ position_ids = position_ids[:, -input_ids.shape[1] :]
1297
+
1298
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1299
+ if inputs_embeds is not None and past_key_values is None:
1300
+ model_inputs = {"inputs_embeds": inputs_embeds}
1301
+ else:
1302
+ model_inputs = {"input_ids": input_ids}
1303
+
1304
+ model_inputs.update(
1305
+ {
1306
+ "position_ids": position_ids,
1307
+ "past_key_values": past_key_values,
1308
+ "use_cache": kwargs.get("use_cache"),
1309
+ "attention_mask": attention_mask,
1310
+ }
1311
+ )
1312
+ return model_inputs
1313
+
1314
+ @staticmethod
1315
+ def _reorder_cache(past_key_values, beam_idx):
1316
+ reordered_past = ()
1317
+ for layer_past in past_key_values:
1318
+ reordered_past += (
1319
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
1320
+ )
1321
+ return reordered_past
1322
+
1323
+
1324
+ @add_start_docstrings(
1325
+ """
1326
+ The AquilaMoe Model transformer with a sequence classification head on top (linear layer).
1327
+
1328
+ [`AquilaMoeForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1329
+ (e.g. GPT-2) do.
1330
+
1331
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1332
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1333
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1334
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1335
+ each row of the batch).
1336
+ """,
1337
+ AquilaMoE_START_DOCSTRING,
1338
+ )
1339
+ # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->AquilaMoe, LLAMA->AquilaMoe
1340
+ class AquilaMoeForSequenceClassification(AquilaMoePreTrainedModel):
1341
+ def __init__(self, config):
1342
+ super().__init__(config)
1343
+ self.num_labels = config.num_labels
1344
+ self.model = AquilaMoeModel(config)
1345
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1346
+
1347
+ # Initialize weights and apply final processing
1348
+ self.post_init()
1349
+
1350
+ def get_input_embeddings(self):
1351
+ return self.model.embed_tokens
1352
+
1353
+ def set_input_embeddings(self, value):
1354
+ self.model.embed_tokens = value
1355
+
1356
+ @add_start_docstrings_to_model_forward(AquilaMoE_INPUTS_DOCSTRING)
1357
+ def forward(
1358
+ self,
1359
+ input_ids: torch.LongTensor = None,
1360
+ attention_mask: Optional[torch.Tensor] = None,
1361
+ position_ids: Optional[torch.LongTensor] = None,
1362
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1363
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1364
+ labels: Optional[torch.LongTensor] = None,
1365
+ use_cache: Optional[bool] = None,
1366
+ output_attentions: Optional[bool] = None,
1367
+ output_hidden_states: Optional[bool] = None,
1368
+ return_dict: Optional[bool] = None,
1369
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1370
+ r"""
1371
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1372
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1373
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1374
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1375
+ """
1376
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1377
+
1378
+ transformer_outputs = self.model(
1379
+ input_ids,
1380
+ attention_mask=attention_mask,
1381
+ position_ids=position_ids,
1382
+ past_key_values=past_key_values,
1383
+ inputs_embeds=inputs_embeds,
1384
+ use_cache=use_cache,
1385
+ output_attentions=output_attentions,
1386
+ output_hidden_states=output_hidden_states,
1387
+ return_dict=return_dict,
1388
+ )
1389
+ hidden_states = transformer_outputs[0]
1390
+ logits = self.score(hidden_states)
1391
+
1392
+ if input_ids is not None:
1393
+ batch_size = input_ids.shape[0]
1394
+ else:
1395
+ batch_size = inputs_embeds.shape[0]
1396
+
1397
+ if self.config.pad_token_id is None and batch_size != 1:
1398
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1399
+ if self.config.pad_token_id is None:
1400
+ sequence_lengths = -1
1401
+ else:
1402
+ if input_ids is not None:
1403
+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
1404
+ logits.device
1405
+ )
1406
+ else:
1407
+ sequence_lengths = -1
1408
+
1409
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1410
+
1411
+ loss = None
1412
+ if labels is not None:
1413
+ labels = labels.to(logits.device)
1414
+ if self.config.problem_type is None:
1415
+ if self.num_labels == 1:
1416
+ self.config.problem_type = "regression"
1417
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1418
+ self.config.problem_type = "single_label_classification"
1419
+ else:
1420
+ self.config.problem_type = "multi_label_classification"
1421
+
1422
+ if self.config.problem_type == "regression":
1423
+ loss_fct = MSELoss()
1424
+ if self.num_labels == 1:
1425
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1426
+ else:
1427
+ loss = loss_fct(pooled_logits, labels)
1428
+ elif self.config.problem_type == "single_label_classification":
1429
+ loss_fct = CrossEntropyLoss()
1430
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1431
+ elif self.config.problem_type == "multi_label_classification":
1432
+ loss_fct = BCEWithLogitsLoss()
1433
+ loss = loss_fct(pooled_logits, labels)
1434
+ if not return_dict:
1435
+ output = (pooled_logits,) + transformer_outputs[1:]
1436
+ return ((loss,) + output) if loss is not None else output
1437
+
1438
+ return SequenceClassifierOutputWithPast(
1439
+ loss=loss,
1440
+ logits=pooled_logits,
1441
+ past_key_values=transformer_outputs.past_key_values,
1442
+ hidden_states=transformer_outputs.hidden_states,
1443
+ attentions=transformer_outputs.attentions,
1444
+ )
tokenization.py ADDED
@@ -0,0 +1,276 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Alibaba Cloud.
2
+ #
3
+ # This source code is licensed under the license found in the
4
+ # LICENSE file in the root directory of this source tree.
5
+
6
+ """Tokenization classes for QWen."""
7
+
8
+ import base64
9
+ import logging
10
+ import os
11
+ import unicodedata
12
+ from typing import Collection, Dict, List, Set, Tuple, Union
13
+
14
+ import tiktoken
15
+ from transformers import PreTrainedTokenizer, AddedToken
16
+
17
+ logger = logging.getLogger(__name__)
18
+
19
+
20
+ VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.tiktoken"}
21
+
22
+ PAT_STR = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
23
+ ENDOFTEXT = "<|endoftext|>"
24
+ IMSTART = "<|im_start|>"
25
+ IMEND = "<|im_end|>"
26
+ # as the default behavior is changed to allow special tokens in
27
+ # regular texts, the surface forms of special tokens need to be
28
+ # as different as possible to minimize the impact
29
+ EXTRAS = tuple((f"<|extra_{i}|>" for i in range(205)))
30
+ # changed to use actual index to avoid misconfiguration with vocabulary expansion
31
+ SPECIAL_START_ID = 151643
32
+ SPECIAL_TOKENS = tuple(
33
+ enumerate(
34
+ (
35
+ (
36
+ ENDOFTEXT,
37
+ IMSTART,
38
+ IMEND,
39
+ )
40
+ + EXTRAS
41
+ ),
42
+ start=SPECIAL_START_ID,
43
+ )
44
+ )
45
+ SPECIAL_TOKENS_SET = set(t for i, t in SPECIAL_TOKENS)
46
+
47
+
48
+ def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
49
+ with open(tiktoken_bpe_file, "rb") as f:
50
+ contents = f.read()
51
+ return {
52
+ base64.b64decode(token): int(rank)
53
+ for token, rank in (line.split() for line in contents.splitlines() if line)
54
+ }
55
+
56
+
57
+ class QWenTokenizer(PreTrainedTokenizer):
58
+ """QWen tokenizer."""
59
+
60
+ vocab_files_names = VOCAB_FILES_NAMES
61
+
62
+ def __init__(
63
+ self,
64
+ vocab_file,
65
+ errors="replace",
66
+ extra_vocab_file=None,
67
+ **kwargs,
68
+ ):
69
+ super().__init__(**kwargs)
70
+
71
+ # how to handle errors in decoding UTF-8 byte sequences
72
+ # use ignore if you are in streaming inference
73
+ self.errors = errors
74
+
75
+ self.mergeable_ranks = _load_tiktoken_bpe(vocab_file) # type: Dict[bytes, int]
76
+ self.special_tokens = {
77
+ token: index
78
+ for index, token in SPECIAL_TOKENS
79
+ }
80
+
81
+ # try load extra vocab from file
82
+ if extra_vocab_file is not None:
83
+ used_ids = set(self.mergeable_ranks.values()) | set(self.special_tokens.values())
84
+ extra_mergeable_ranks = _load_tiktoken_bpe(extra_vocab_file)
85
+ for token, index in extra_mergeable_ranks.items():
86
+ if token in self.mergeable_ranks:
87
+ logger.info(f"extra token {token} exists, skipping")
88
+ continue
89
+ if index in used_ids:
90
+ logger.info(f'the index {index} for extra token {token} exists, skipping')
91
+ continue
92
+ self.mergeable_ranks[token] = index
93
+ # the index may be sparse after this, but don't worry tiktoken.Encoding will handle this
94
+
95
+ enc = tiktoken.Encoding(
96
+ "Qwen",
97
+ pat_str=PAT_STR,
98
+ mergeable_ranks=self.mergeable_ranks,
99
+ special_tokens=self.special_tokens,
100
+ )
101
+ assert (
102
+ len(self.mergeable_ranks) + len(self.special_tokens) == enc.n_vocab
103
+ ), f"{len(self.mergeable_ranks) + len(self.special_tokens)} != {enc.n_vocab} in encoding"
104
+
105
+ self.decoder = {
106
+ v: k for k, v in self.mergeable_ranks.items()
107
+ } # type: dict[int, bytes|str]
108
+ self.decoder.update({v: k for k, v in self.special_tokens.items()})
109
+
110
+ self.tokenizer = enc # type: tiktoken.Encoding
111
+
112
+ self.eod_id = self.tokenizer.eot_token
113
+ self.im_start_id = self.special_tokens[IMSTART]
114
+ self.im_end_id = self.special_tokens[IMEND]
115
+
116
+ def __getstate__(self):
117
+ # for pickle lovers
118
+ state = self.__dict__.copy()
119
+ del state["tokenizer"]
120
+ return state
121
+
122
+ def __setstate__(self, state):
123
+ # tokenizer is not python native; don't pass it; rebuild it
124
+ self.__dict__.update(state)
125
+ enc = tiktoken.Encoding(
126
+ "Qwen",
127
+ pat_str=PAT_STR,
128
+ mergeable_ranks=self.mergeable_ranks,
129
+ special_tokens=self.special_tokens,
130
+ )
131
+ self.tokenizer = enc
132
+
133
+ def __len__(self) -> int:
134
+ return self.tokenizer.n_vocab
135
+
136
+ def get_vocab(self) -> Dict[bytes, int]:
137
+ return self.mergeable_ranks
138
+
139
+ def convert_tokens_to_ids(
140
+ self, tokens: Union[bytes, str, List[Union[bytes, str]]]
141
+ ) -> List[int]:
142
+ ids = []
143
+ if isinstance(tokens, (str, bytes)):
144
+ if tokens in self.special_tokens:
145
+ return self.special_tokens[tokens]
146
+ else:
147
+ return self.mergeable_ranks.get(tokens)
148
+ for token in tokens:
149
+ if token in self.special_tokens:
150
+ ids.append(self.special_tokens[token])
151
+ else:
152
+ ids.append(self.mergeable_ranks.get(token))
153
+ return ids
154
+
155
+ def _add_tokens(
156
+ self,
157
+ new_tokens: Union[List[str], List[AddedToken]],
158
+ special_tokens: bool = False,
159
+ ) -> int:
160
+ if not special_tokens and new_tokens:
161
+ raise ValueError("Adding regular tokens is not supported")
162
+ for token in new_tokens:
163
+ surface_form = token.content if isinstance(token, AddedToken) else token
164
+ if surface_form not in SPECIAL_TOKENS_SET:
165
+ raise ValueError("Adding unknown special tokens is not supported")
166
+ return 0
167
+
168
+ def save_vocabulary(self, save_directory: str, **kwargs) -> Tuple[str]:
169
+ """
170
+ Save only the vocabulary of the tokenizer (vocabulary).
171
+
172
+ Returns:
173
+ `Tuple(str)`: Paths to the files saved.
174
+ """
175
+ file_path = os.path.join(save_directory, "tokenizer.tiktoken")
176
+ with open(file_path, "w", encoding="utf8") as w:
177
+ for k, v in self.mergeable_ranks.items():
178
+ line = base64.b64encode(k).decode("utf8") + " " + str(v) + "\n"
179
+ w.write(line)
180
+ return (file_path,)
181
+
182
+ def tokenize(
183
+ self,
184
+ text: str,
185
+ allowed_special: Union[Set, str] = "all",
186
+ disallowed_special: Union[Collection, str] = (),
187
+ **kwargs,
188
+ ) -> List[Union[bytes, str]]:
189
+ """
190
+ Converts a string in a sequence of tokens.
191
+
192
+ Args:
193
+ text (`str`):
194
+ The sequence to be encoded.
195
+ allowed_special (`Literal["all"]` or `set`):
196
+ The surface forms of the tokens to be encoded as special tokens in regular texts.
197
+ Default to "all".
198
+ disallowed_special (`Literal["all"]` or `Collection`):
199
+ The surface forms of the tokens that should not be in regular texts and trigger errors.
200
+ Default to an empty tuple.
201
+
202
+ kwargs (additional keyword arguments, *optional*):
203
+ Will be passed to the underlying model specific encode method.
204
+
205
+ Returns:
206
+ `List[bytes|str]`: The list of tokens.
207
+ """
208
+ tokens = []
209
+ text = unicodedata.normalize("NFC", text)
210
+
211
+ # this implementation takes a detour: text -> token id -> token surface forms
212
+ for t in self.tokenizer.encode(
213
+ text, allowed_special=allowed_special, disallowed_special=disallowed_special
214
+ ):
215
+ tokens.append(self.decoder[t])
216
+ return tokens
217
+
218
+ def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str:
219
+ """
220
+ Converts a sequence of tokens in a single string.
221
+ """
222
+ text = ""
223
+ temp = b""
224
+ for t in tokens:
225
+ if isinstance(t, str):
226
+ if temp:
227
+ text += temp.decode("utf-8", errors=self.errors)
228
+ temp = b""
229
+ text += t
230
+ elif isinstance(t, bytes):
231
+ temp += t
232
+ else:
233
+ raise TypeError("token should only be of type types or str")
234
+ if temp:
235
+ text += temp.decode("utf-8", errors=self.errors)
236
+ return text
237
+
238
+ @property
239
+ def vocab_size(self):
240
+ return self.tokenizer.n_vocab
241
+
242
+ def _convert_id_to_token(self, index: int) -> Union[bytes, str]:
243
+ """Converts an id to a token, special tokens included"""
244
+ if index in self.decoder:
245
+ return self.decoder[index]
246
+ raise ValueError("unknown ids")
247
+
248
+ def _convert_token_to_id(self, token: Union[bytes, str]) -> int:
249
+ """Converts a token to an id using the vocab, special tokens included"""
250
+ if token in self.special_tokens:
251
+ return self.special_tokens[token]
252
+ if token in self.mergeable_ranks:
253
+ return self.mergeable_ranks[token]
254
+ raise ValueError("unknown token")
255
+
256
+ def _tokenize(self, text: str, **kwargs):
257
+ """
258
+ Converts a string in a sequence of tokens (string), using the tokenizer. Split in words for word-based
259
+ vocabulary or sub-words for sub-word-based vocabularies (BPE/SentencePieces/WordPieces).
260
+
261
+ Do NOT take care of added tokens.
262
+ """
263
+ raise NotImplementedError
264
+
265
+ def _decode(
266
+ self,
267
+ token_ids: Union[int, List[int]],
268
+ skip_special_tokens: bool = False,
269
+ errors: str = None,
270
+ **kwargs,
271
+ ) -> str:
272
+ if isinstance(token_ids, int):
273
+ token_ids = [token_ids]
274
+ if skip_special_tokens:
275
+ token_ids = [i for i in token_ids if i < self.eod_id]
276
+ return self.tokenizer.decode(token_ids, errors=errors or self.errors)
tokenizer_config.json ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model_max_length": 8192,
3
+ "tokenizer_class": "QWenTokenizer",
4
+ "auto_map": {
5
+ "AutoTokenizer": [
6
+ "tokenization.QWenTokenizer",
7
+ null
8
+ ]
9
+ }
10
+ }