LoneStriker commited on
Commit
b06b086
1 Parent(s): 9170220

Upload folder using huggingface_hub

Browse files
README.md ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ datasets:
4
+ - teknium/OpenHermes-2.5
5
+ language:
6
+ - en
7
+ ---
8
+ ## Training
9
+ - 8x A6000s
10
+ - [Forked version of unsloth](https://github.com/serp-ai/unsloth) for efficient training
11
+ - Sequence Length: 4096
12
+ - Effective batch size: 128
13
+ - Learning Rate: 2e-5 with linear decay
14
+ - Epochs: 1
15
+ - [Base model](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2) trained with QLoRA (rank 64, alpha 16) and MoE adapters/routers trained in bf16
16
+ - Num Experts: 16
17
+ - Top K: 4
18
+ - Adapter Dim: 512
19
+
20
+ ## Prompt Format
21
+ ```
22
+ <|im_start|>system\n{message}<|im_end|>\n<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n
23
+ ```
24
+
25
+ ## Usage
26
+ ```python
27
+ from transformers import AutoModelForCausalLM, AutoTokenizer
28
+
29
+ tokenizer = AutoTokenizer.from_pretrained("serpdotai/sparsetral-16x7B-v2", trust_remote_code=True)
30
+ model = AutoModelForCausalLM.from_pretrained("serpdotai/sparsetral-16x7B-v2", device_map="auto", trust_remote_code=True).eval()
31
+
32
+ system_str = "<|im_start|>system\n{message}<|im_end|>\n"
33
+ user_str = "<|im_start|>user\n{message}<|im_end|>\n"
34
+ assistant_str = "<|im_start|>assistant\n{message}<|im_end|>\n"
35
+
36
+ def construct_prompt(messages):
37
+ prompt = ""
38
+ for message in messages:
39
+ if message["from"] in ["human", "user"]:
40
+ prompt += user_str.format(
41
+ message=message["value"]
42
+ )
43
+ elif message["from"] in ["gpt", "assistant"]:
44
+ prompt += assistant_str.format(
45
+ message=message["value"]
46
+ )
47
+ elif message["from"] in ["system", "instruction"]:
48
+ prompt += system_str.format(
49
+ message=message["value"]
50
+ )
51
+ else:
52
+ raise ValueError(
53
+ f"Unknown message type: {message['from']}"
54
+ )
55
+ return prompt + "<|im_start|>assistant\n"
56
+
57
+ system = "You are a helpful assistant who will help the user to the best of their ability. If you don't know something, say \"I don't know\""
58
+ user = "Are you sentient?"
59
+
60
+ messages = [
61
+ {"from": "system", "value": system},
62
+ {"from": "user", "value": user},
63
+ ]
64
+
65
+ prompt = construct_prompt(messages)
66
+ inputs = tokenizer(prompt, return_tensors="pt")
67
+ inputs = inputs.to(model.device)
68
+ pred = model.generate(**inputs, max_length=4096, do_sample=True, top_k=50, top_p=0.99, temperature=0.9, num_return_sequences=1)
69
+ print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))
70
+ ```
71
+
72
+ ## Other Information
73
+ Paper reference: [Parameter-Efficient Sparsity Crafting from Dense to Mixture-of-Experts for Instruction Tuning on General Tasks](https://arxiv.org/abs/2401.02731)
74
+
75
+ [Original Paper repo](https://github.com/wuhy68/Parameter-Efficient-MoE)
76
+
77
+ [Forked repo with mistral support (sparsetral)](https://github.com/serp-ai/Parameter-Efficient-MoE)
78
+
79
+ If you are interested in faster inferencing, check out our [fork of vLLM](https://github.com/serp-ai/vllm) that adds sparsetral support
config.json ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "adapter_dim": 512,
3
+ "adapter_dropout": 0.0,
4
+ "architectures": [
5
+ "modeling_sparsetral.MistralForCausalLM"
6
+ ],
7
+ "attention_dropout": 0.0,
8
+ "bos_token_id": 1,
9
+ "eos_token_id": 2,
10
+ "hidden_act": "silu",
11
+ "hidden_size": 4096,
12
+ "initializer_range": 0.02,
13
+ "intermediate_size": 14336,
14
+ "max_position_embeddings": 32768,
15
+ "model_type": "sparsetral",
16
+ "moe_dtype": "bfloat16",
17
+ "moe_scaling": 1,
18
+ "num_attention_heads": 32,
19
+ "num_experts": 16,
20
+ "num_hidden_layers": 32,
21
+ "num_key_value_heads": 8,
22
+ "output_router_logits": false,
23
+ "pretraining_tp": 1,
24
+ "rms_norm_eps": 1e-05,
25
+ "rope_theta": 1000000.0,
26
+ "router_aux_loss_coef": 0.01,
27
+ "sliding_window": null,
28
+ "tie_word_embeddings": false,
29
+ "topk": 4,
30
+ "torch_dtype": "bfloat16",
31
+ "transformers_version": "4.37.2",
32
+ "use_cache": true,
33
+ "vocab_size": 32000,
34
+ "auto_map": {
35
+ "AutoConfig": "configuration_sparsetral.SparsetralConfig",
36
+ "AutoModel": "modeling_sparsetral.MistralModel",
37
+ "AutoModelForCausalLM": "modeling_sparsetral.MistralForCausalLM"
38
+ }
39
+ }
configuration_sparsetral.py ADDED
@@ -0,0 +1,165 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Mistral AI 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
+ """ Sparsetral model configuration"""
16
+
17
+ from transformers.configuration_utils import PretrainedConfig
18
+ from transformers.utils import logging
19
+
20
+
21
+ logger = logging.get_logger(__name__)
22
+
23
+
24
+ class SparsetralConfig(PretrainedConfig):
25
+ r"""
26
+ This is the configuration class to store the configuration of a [`MistralModel`]. It is used to instantiate an
27
+ Mistral model according to the specified arguments, defining the model architecture. Instantiating a configuration
28
+ with the defaults will yield a similar configuration to that of the Mistral-7B-v0.1 or Mistral-7B-Instruct-v0.1.
29
+
30
+ [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1)
31
+ [mistralai/Mistral-7B-Instruct-v0.1](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1)
32
+
33
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
34
+ documentation from [`PretrainedConfig`] for more information.
35
+
36
+
37
+ Args:
38
+ vocab_size (`int`, *optional*, defaults to 32000):
39
+ Vocabulary size of the Mistral model. Defines the number of different tokens that can be represented by the
40
+ `inputs_ids` passed when calling [`MistralModel`]
41
+ hidden_size (`int`, *optional*, defaults to 4096):
42
+ Dimension of the hidden representations.
43
+ intermediate_size (`int`, *optional*, defaults to 14336):
44
+ Dimension of the MLP representations.
45
+ num_hidden_layers (`int`, *optional*, defaults to 32):
46
+ Number of hidden layers in the Transformer encoder.
47
+ num_attention_heads (`int`, *optional*, defaults to 32):
48
+ Number of attention heads for each attention layer in the Transformer encoder.
49
+ num_key_value_heads (`int`, *optional*, defaults to 8):
50
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
51
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
52
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
53
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
54
+ by meanpooling all the original heads within that group. For more details checkout [this
55
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8`.
56
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
57
+ The non-linear activation function (function or string) in the decoder.
58
+ max_position_embeddings (`int`, *optional*, defaults to `4096*32`):
59
+ The maximum sequence length that this model might ever be used with. Mistral's sliding window attention
60
+ allows sequence of up to 4096*32 tokens.
61
+ initializer_range (`float`, *optional*, defaults to 0.02):
62
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
63
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
64
+ The epsilon used by the rms normalization layers.
65
+ use_cache (`bool`, *optional*, defaults to `True`):
66
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
67
+ relevant if `config.is_decoder=True`.
68
+ pad_token_id (`int`, *optional*):
69
+ The id of the padding token.
70
+ bos_token_id (`int`, *optional*, defaults to 1):
71
+ The id of the "beginning-of-sequence" token.
72
+ eos_token_id (`int`, *optional*, defaults to 2):
73
+ The id of the "end-of-sequence" token.
74
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
75
+ Whether the model's input and output word embeddings should be tied.
76
+ rope_theta (`float`, *optional*, defaults to 10000.0):
77
+ The base period of the RoPE embeddings.
78
+ sliding_window (`int`, *optional*, defaults to 4096):
79
+ Sliding window attention window size. If not specified, will default to `4096`.
80
+ attention_dropout (`float`, *optional*, defaults to 0.0):
81
+ The dropout ratio for the attention probabilities.
82
+
83
+ ```python
84
+ >>> from transformers import MistralModel, MistralConfig
85
+
86
+ >>> # Initializing a Mistral 7B style configuration
87
+ >>> configuration = MistralConfig()
88
+
89
+ >>> # Initializing a model from the Mistral 7B style configuration
90
+ >>> model = MistralModel(configuration)
91
+
92
+ >>> # Accessing the model configuration
93
+ >>> configuration = model.config
94
+ ```"""
95
+
96
+ model_type = "mistral"
97
+ keys_to_ignore_at_inference = ["past_key_values"]
98
+
99
+ def __init__(
100
+ self,
101
+ vocab_size=32000,
102
+ hidden_size=4096,
103
+ intermediate_size=14336,
104
+ num_hidden_layers=32,
105
+ num_attention_heads=32,
106
+ num_key_value_heads=8,
107
+ hidden_act="silu",
108
+ max_position_embeddings=32768,
109
+ initializer_range=0.02,
110
+ rms_norm_eps=1e-6,
111
+ use_cache=True,
112
+ pad_token_id=None,
113
+ bos_token_id=1,
114
+ eos_token_id=2,
115
+ tie_word_embeddings=False,
116
+ rope_theta=10000.0,
117
+ sliding_window=4096,
118
+ attention_dropout=0.0,
119
+ moe_dtype="bfloat16",
120
+ moe_scaling=1.0,
121
+ num_experts=16,
122
+ topk=4,
123
+ output_router_logits=False,
124
+ adapter_dim=512,
125
+ adapter_dropout=0.0,
126
+ router_aux_loss_coef=0.01,
127
+ **kwargs,
128
+ ):
129
+ self.vocab_size = vocab_size
130
+ self.max_position_embeddings = max_position_embeddings
131
+ self.hidden_size = hidden_size
132
+ self.intermediate_size = intermediate_size
133
+ self.num_hidden_layers = num_hidden_layers
134
+ self.num_attention_heads = num_attention_heads
135
+ self.sliding_window = sliding_window
136
+
137
+ # for backward compatibility
138
+ if num_key_value_heads is None:
139
+ num_key_value_heads = num_attention_heads
140
+
141
+ self.num_key_value_heads = num_key_value_heads
142
+ self.hidden_act = hidden_act
143
+ self.initializer_range = initializer_range
144
+ self.rms_norm_eps = rms_norm_eps
145
+ self.use_cache = use_cache
146
+ self.rope_theta = rope_theta
147
+ self.attention_dropout = attention_dropout
148
+
149
+ self.moe_dtype = moe_dtype
150
+ self.moe_scaling = moe_scaling
151
+ self.num_experts = num_experts
152
+ self.topk = topk
153
+ self.output_router_logits = output_router_logits
154
+
155
+ self.adapter_dim = adapter_dim
156
+ self.adapter_dropout = adapter_dropout
157
+ self.router_aux_loss_coef = router_aux_loss_coef
158
+
159
+ super().__init__(
160
+ pad_token_id=pad_token_id,
161
+ bos_token_id=bos_token_id,
162
+ eos_token_id=eos_token_id,
163
+ tie_word_embeddings=tie_word_embeddings,
164
+ **kwargs,
165
+ )
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 1,
4
+ "eos_token_id": 2,
5
+ "transformers_version": "4.37.2"
6
+ }
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
modeling_sparsetral.py ADDED
@@ -0,0 +1,1648 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 Mistral model."""
21
+ import inspect
22
+ import math
23
+ import warnings
24
+ from dataclasses import dataclass
25
+ from typing import List, Optional, Tuple, Union
26
+
27
+ import torch
28
+ import torch.nn.functional as F
29
+ import torch.utils.checkpoint
30
+ from torch import nn
31
+ from torch.nn import CrossEntropyLoss
32
+
33
+ from transformers.activations import ACT2FN
34
+ from transformers.cache_utils import Cache, DynamicCache
35
+ from transformers.modeling_attn_mask_utils import (
36
+ _prepare_4d_causal_attention_mask,
37
+ _prepare_4d_causal_attention_mask_for_sdpa,
38
+ )
39
+ from transformers.modeling_outputs import (
40
+ BaseModelOutputWithPast,
41
+ CausalLMOutputWithPast,
42
+ SequenceClassifierOutputWithPast,
43
+ )
44
+ from transformers.modeling_utils import PreTrainedModel
45
+ from transformers.utils import (
46
+ add_start_docstrings,
47
+ add_start_docstrings_to_model_forward,
48
+ is_flash_attn_2_available,
49
+ is_flash_attn_greater_or_equal_2_10,
50
+ logging,
51
+ replace_return_docstrings,
52
+ ModelOutput,
53
+ )
54
+ from .configuration_sparsetral import SparsetralConfig
55
+
56
+
57
+ if is_flash_attn_2_available():
58
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
59
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
60
+
61
+ _flash_supports_window_size = "window_size" in list(
62
+ inspect.signature(flash_attn_func).parameters
63
+ )
64
+
65
+
66
+ logger = logging.get_logger(__name__)
67
+
68
+ _CONFIG_FOR_DOC = "SparsetralConfig"
69
+
70
+
71
+ @dataclass
72
+ class MoEModelOutputWithPast(ModelOutput):
73
+ last_hidden_state: torch.FloatTensor = None
74
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
75
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
76
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
77
+ router_logits: Optional[Tuple[torch.FloatTensor]] = None
78
+
79
+
80
+ @dataclass
81
+ class MoECausalLMOutputWithPast(ModelOutput):
82
+ loss: Optional[torch.FloatTensor] = None
83
+ aux_loss: Optional[torch.FloatTensor] = None
84
+ logits: torch.FloatTensor = None
85
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
86
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
87
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
88
+ router_logits: Optional[Tuple[torch.FloatTensor]] = None
89
+
90
+
91
+ # Copied from transformers.models.llama.modeling_llama._get_unpad_data
92
+ def _get_unpad_data(attention_mask):
93
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
94
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
95
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
96
+ cu_seqlens = F.pad(
97
+ torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)
98
+ )
99
+ return (
100
+ indices,
101
+ cu_seqlens,
102
+ max_seqlen_in_batch,
103
+ )
104
+
105
+
106
+ # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Mistral
107
+ class MistralRMSNorm(nn.Module):
108
+ def __init__(self, hidden_size, eps=1e-6):
109
+ """
110
+ MistralRMSNorm is equivalent to T5LayerNorm
111
+ """
112
+ super().__init__()
113
+ self.weight = nn.Parameter(torch.ones(hidden_size))
114
+ self.variance_epsilon = eps
115
+
116
+ def forward(self, hidden_states):
117
+ input_dtype = hidden_states.dtype
118
+ hidden_states = hidden_states.to(torch.float32)
119
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
120
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
121
+ return self.weight * hidden_states.to(input_dtype)
122
+
123
+
124
+ # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Mistral
125
+ class MistralRotaryEmbedding(nn.Module):
126
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
127
+ super().__init__()
128
+
129
+ self.dim = dim
130
+ self.max_position_embeddings = max_position_embeddings
131
+ self.base = base
132
+ inv_freq = 1.0 / (
133
+ self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
134
+ )
135
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
136
+
137
+ # Build here to make `torch.jit.trace` work.
138
+ self._set_cos_sin_cache(
139
+ seq_len=max_position_embeddings,
140
+ device=self.inv_freq.device,
141
+ dtype=torch.get_default_dtype(),
142
+ )
143
+
144
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
145
+ self.max_seq_len_cached = seq_len
146
+ t = torch.arange(
147
+ self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
148
+ )
149
+
150
+ freqs = torch.outer(t, self.inv_freq)
151
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
152
+ emb = torch.cat((freqs, freqs), dim=-1)
153
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
154
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
155
+
156
+ def forward(self, x, seq_len=None):
157
+ # x: [bs, num_attention_heads, seq_len, head_size]
158
+ if seq_len > self.max_seq_len_cached:
159
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
160
+
161
+ return (
162
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
163
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
164
+ )
165
+
166
+
167
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
168
+ def rotate_half(x):
169
+ """Rotates half the hidden dims of the input."""
170
+ x1 = x[..., : x.shape[-1] // 2]
171
+ x2 = x[..., x.shape[-1] // 2 :]
172
+ return torch.cat((-x2, x1), dim=-1)
173
+
174
+
175
+ # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
176
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
177
+ """Applies Rotary Position Embedding to the query and key tensors.
178
+
179
+ Args:
180
+ q (`torch.Tensor`): The query tensor.
181
+ k (`torch.Tensor`): The key tensor.
182
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
183
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
184
+ position_ids (`torch.Tensor`):
185
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
186
+ used to pass offsetted position ids when working with a KV-cache.
187
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
188
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
189
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
190
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
191
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
192
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
193
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
194
+ Returns:
195
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
196
+ """
197
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
198
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
199
+ q_embed = (q * cos) + (rotate_half(q) * sin)
200
+ k_embed = (k * cos) + (rotate_half(k) * sin)
201
+ return q_embed, k_embed
202
+
203
+
204
+ # Mistral MoE
205
+ def load_balancing_loss_func(
206
+ gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2
207
+ ) -> float:
208
+ r"""
209
+ Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
210
+
211
+ See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
212
+ function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
213
+ experts is too unbalanced.
214
+
215
+ Args:
216
+ gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):
217
+ Logits from the `gate`, should be a tuple of tensors. Shape: [batch_size, seqeunce_length, num_experts].
218
+ num_experts (`int`, *optional*):
219
+ Number of experts
220
+
221
+ Returns:
222
+ The auxiliary loss.
223
+ """
224
+ if gate_logits is None:
225
+ return 0
226
+
227
+ if isinstance(gate_logits, tuple):
228
+ # cat along the layers?
229
+ compute_device = gate_logits[0].device
230
+ gate_logits = torch.cat(
231
+ [gate.to(compute_device) for gate in gate_logits], dim=0
232
+ )
233
+
234
+ routing_weights, selected_experts = torch.topk(gate_logits, top_k, dim=-1)
235
+ routing_weights = routing_weights.softmax(dim=-1)
236
+
237
+ # cast the expert indices to int64, otherwise one-hot encoding will fail
238
+ if selected_experts.dtype != torch.int64:
239
+ selected_experts = selected_experts.to(torch.int64)
240
+
241
+ if len(selected_experts.shape) == 2:
242
+ selected_experts = selected_experts.unsqueeze(2)
243
+
244
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
245
+
246
+ # For a given token, determine if it was routed to a given expert.
247
+ expert_mask = torch.max(expert_mask, axis=-2).values
248
+
249
+ # cast to float32 otherwise mean will fail
250
+ expert_mask = expert_mask.to(torch.float32)
251
+ tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)
252
+
253
+ router_prob_per_group_and_expert = torch.mean(routing_weights, axis=-1)
254
+ return torch.mean(
255
+ tokens_per_group_and_expert * router_prob_per_group_and_expert.unsqueeze(-1)
256
+ ) * (num_experts**2)
257
+
258
+
259
+ class ParallelAdapterMLP(nn.Module):
260
+ def __init__(self, config, adapter_dim, adapter_scaling):
261
+ super().__init__()
262
+ self.config = config
263
+ self.intermediate_size = config.intermediate_size
264
+ self.hidden_size = config.hidden_size
265
+ self.adapter_down = nn.Linear(self.hidden_size, adapter_dim, bias=False)
266
+ self.adapter_up = nn.Linear(adapter_dim, self.hidden_size, bias=False)
267
+ self.adapter_act = nn.GELU()
268
+
269
+ self.adapter_dropout = nn.Dropout(p=config.adapter_dropout)
270
+ self.adapter_scaling = adapter_scaling
271
+
272
+ def forward(self, x):
273
+ x = self.adapter_dropout(x)
274
+ x = self.adapter_scaling * self.adapter_up(
275
+ self.adapter_act(self.adapter_down(x))
276
+ )
277
+ return x
278
+
279
+
280
+ class SparsetralGateAdapter(nn.Module):
281
+ def __init__(self, config: SparsetralConfig):
282
+ super().__init__()
283
+
284
+ self.intermediate_size = config.intermediate_size
285
+ self.hidden_size = config.hidden_size
286
+
287
+ # Step 1: Router
288
+ self.num_experts = config.num_experts
289
+ self.topk = config.topk
290
+ self.router = nn.Linear(config.hidden_size, self.num_experts, bias=False)
291
+ self.dtype = getattr(torch, config.moe_dtype)
292
+
293
+ # Step 2: Get the experts
294
+ self.expert_indicies = list(range(self.num_experts))
295
+ self.experts = nn.ModuleList(
296
+ [
297
+ ParallelAdapterMLP(config, config.adapter_dim, config.moe_scaling)
298
+ for _ in self.expert_indicies
299
+ ]
300
+ )
301
+
302
+ def forward(self, input_hidden_states, output_hidden_states, router_hidden_states):
303
+ orig_shape = output_hidden_states.shape
304
+ input_hidden_states = input_hidden_states.view(
305
+ -1, input_hidden_states.shape[-1]
306
+ )
307
+ output_hidden_states = output_hidden_states.view(
308
+ -1, output_hidden_states.shape[-1]
309
+ )
310
+ router_hidden_states = router_hidden_states.view(
311
+ -1, router_hidden_states.shape[-1]
312
+ )
313
+
314
+ router_logits = self.router(router_hidden_states)
315
+
316
+ routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
317
+ routing_weights, selected_experts = torch.topk(
318
+ routing_weights, self.topk, dim=-1
319
+ )
320
+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
321
+
322
+ final_hidden_states = None
323
+ for expert_idx in self.expert_indicies:
324
+ expert_layer = self.experts[expert_idx]
325
+ expert_mask = selected_experts == expert_idx
326
+ expert_weights = (routing_weights * expert_mask).sum(dim=-1, keepdim=True)
327
+
328
+ current_hidden_states = (
329
+ expert_layer(input_hidden_states)
330
+ .add_(output_hidden_states)
331
+ .mul_(expert_weights)
332
+ )
333
+ if final_hidden_states is None:
334
+ final_hidden_states = current_hidden_states
335
+ else:
336
+ final_hidden_states.add_(current_hidden_states)
337
+
338
+ return final_hidden_states.view(*orig_shape), router_logits
339
+
340
+
341
+ class MistralMLP(nn.Module):
342
+ def __init__(self, config):
343
+ super().__init__()
344
+ self.config = config
345
+ self.hidden_size = config.hidden_size
346
+ self.intermediate_size = config.intermediate_size
347
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
348
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
349
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
350
+ self.act_fn = ACT2FN[config.hidden_act]
351
+
352
+ self.moe_adapter = SparsetralGateAdapter(config)
353
+
354
+ def forward(self, x):
355
+ router_hidden_states = x
356
+ up_proj = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
357
+ down_proj = self.down_proj(up_proj)
358
+ down_proj, router_logits = self.moe_adapter(
359
+ down_proj, down_proj, router_hidden_states
360
+ )
361
+
362
+ return down_proj, router_logits
363
+
364
+
365
+ # Copied from transformers.models.llama.modeling_llama.repeat_kv
366
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
367
+ """
368
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
369
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
370
+ """
371
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
372
+ if n_rep == 1:
373
+ return hidden_states
374
+ hidden_states = hidden_states[:, :, None, :, :].expand(
375
+ batch, num_key_value_heads, n_rep, slen, head_dim
376
+ )
377
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
378
+
379
+
380
+ class MistralAttention(nn.Module):
381
+ """
382
+ Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
383
+ and "Generating Long Sequences with Sparse Transformers".
384
+ """
385
+
386
+ def __init__(self, config: SparsetralConfig, layer_idx: Optional[int] = None):
387
+ super().__init__()
388
+ self.config = config
389
+ self.layer_idx = layer_idx
390
+ if layer_idx is None:
391
+ logger.warning_once(
392
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
393
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
394
+ "when creating this class."
395
+ )
396
+
397
+ self.hidden_size = config.hidden_size
398
+ self.num_heads = config.num_attention_heads
399
+ self.head_dim = self.hidden_size // self.num_heads
400
+ self.num_key_value_heads = config.num_key_value_heads
401
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
402
+ self.max_position_embeddings = config.max_position_embeddings
403
+ self.rope_theta = config.rope_theta
404
+ self.is_causal = True
405
+ self.attention_dropout = config.attention_dropout
406
+
407
+ if (self.head_dim * self.num_heads) != self.hidden_size:
408
+ raise ValueError(
409
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
410
+ f" and `num_heads`: {self.num_heads})."
411
+ )
412
+ self.q_proj = nn.Linear(
413
+ self.hidden_size, self.num_heads * self.head_dim, bias=False
414
+ )
415
+ self.k_proj = nn.Linear(
416
+ self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
417
+ )
418
+ self.v_proj = nn.Linear(
419
+ self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
420
+ )
421
+ self.o_proj = nn.Linear(
422
+ self.num_heads * self.head_dim, self.hidden_size, bias=False
423
+ )
424
+
425
+ self.rotary_emb = MistralRotaryEmbedding(
426
+ self.head_dim,
427
+ max_position_embeddings=self.max_position_embeddings,
428
+ base=self.rope_theta,
429
+ )
430
+
431
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
432
+ return (
433
+ tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
434
+ .transpose(1, 2)
435
+ .contiguous()
436
+ )
437
+
438
+ def forward(
439
+ self,
440
+ hidden_states: torch.Tensor,
441
+ attention_mask: Optional[torch.Tensor] = None,
442
+ position_ids: Optional[torch.LongTensor] = None,
443
+ past_key_value: Optional[Cache] = None,
444
+ output_attentions: bool = False,
445
+ use_cache: bool = False,
446
+ **kwargs,
447
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
448
+ if "padding_mask" in kwargs:
449
+ warnings.warn(
450
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
451
+ )
452
+ bsz, q_len, _ = hidden_states.size()
453
+
454
+ query_states = self.q_proj(hidden_states)
455
+ key_states = self.k_proj(hidden_states)
456
+ value_states = self.v_proj(hidden_states)
457
+
458
+ query_states = query_states.view(
459
+ bsz, q_len, self.num_heads, self.head_dim
460
+ ).transpose(1, 2)
461
+ key_states = key_states.view(
462
+ bsz, q_len, self.num_key_value_heads, self.head_dim
463
+ ).transpose(1, 2)
464
+ value_states = value_states.view(
465
+ bsz, q_len, self.num_key_value_heads, self.head_dim
466
+ ).transpose(1, 2)
467
+
468
+ kv_seq_len = key_states.shape[-2]
469
+ if past_key_value is not None:
470
+ if self.layer_idx is None:
471
+ raise ValueError(
472
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
473
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
474
+ "with a layer index."
475
+ )
476
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
477
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
478
+ query_states, key_states = apply_rotary_pos_emb(
479
+ query_states, key_states, cos, sin, position_ids
480
+ )
481
+
482
+ if past_key_value is not None:
483
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
484
+ key_states, value_states = past_key_value.update(
485
+ key_states, value_states, self.layer_idx, cache_kwargs
486
+ )
487
+
488
+ # repeat k/v heads if n_kv_heads < n_heads
489
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
490
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
491
+
492
+ attn_weights = torch.matmul(
493
+ query_states, key_states.transpose(2, 3)
494
+ ) / math.sqrt(self.head_dim)
495
+
496
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
497
+ raise ValueError(
498
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
499
+ f" {attn_weights.size()}"
500
+ )
501
+
502
+ if attention_mask is not None:
503
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
504
+ raise ValueError(
505
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
506
+ )
507
+
508
+ attn_weights = attn_weights + attention_mask
509
+
510
+ # upcast attention to fp32
511
+ attn_weights = nn.functional.softmax(
512
+ attn_weights, dim=-1, dtype=torch.float32
513
+ ).to(query_states.dtype)
514
+ attn_weights = nn.functional.dropout(
515
+ attn_weights, p=self.attention_dropout, training=self.training
516
+ )
517
+ attn_output = torch.matmul(attn_weights, value_states)
518
+
519
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
520
+ raise ValueError(
521
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
522
+ f" {attn_output.size()}"
523
+ )
524
+
525
+ attn_output = attn_output.transpose(1, 2).contiguous()
526
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
527
+
528
+ attn_output = self.o_proj(attn_output)
529
+
530
+ if not output_attentions:
531
+ attn_weights = None
532
+
533
+ return attn_output, attn_weights, past_key_value
534
+
535
+
536
+ class MistralFlashAttention2(MistralAttention):
537
+ """
538
+ Mistral flash attention module. This module inherits from `MistralAttention` as the weights of the module stays
539
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
540
+ flash attention and deal with padding tokens in case the input contains any of them.
541
+ """
542
+
543
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
544
+ def __init__(self, *args, **kwargs):
545
+ super().__init__(*args, **kwargs)
546
+
547
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
548
+ # 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.
549
+ # 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).
550
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
551
+
552
+ def forward(
553
+ self,
554
+ hidden_states: torch.Tensor,
555
+ attention_mask: Optional[torch.Tensor] = None,
556
+ position_ids: Optional[torch.LongTensor] = None,
557
+ past_key_value: Optional[Cache] = None,
558
+ output_attentions: bool = False,
559
+ use_cache: bool = False,
560
+ **kwargs,
561
+ ):
562
+ if "padding_mask" in kwargs:
563
+ warnings.warn(
564
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
565
+ )
566
+
567
+ # overwrite attention_mask with padding_mask
568
+ attention_mask = kwargs.pop("padding_mask")
569
+ bsz, q_len, _ = hidden_states.size()
570
+
571
+ query_states = self.q_proj(hidden_states)
572
+ key_states = self.k_proj(hidden_states)
573
+ value_states = self.v_proj(hidden_states)
574
+
575
+ query_states = query_states.view(
576
+ bsz, q_len, self.num_heads, self.head_dim
577
+ ).transpose(1, 2)
578
+ key_states = key_states.view(
579
+ bsz, q_len, self.num_key_value_heads, self.head_dim
580
+ ).transpose(1, 2)
581
+ value_states = value_states.view(
582
+ bsz, q_len, self.num_key_value_heads, self.head_dim
583
+ ).transpose(1, 2)
584
+
585
+ kv_seq_len = key_states.shape[-2]
586
+ if past_key_value is not None:
587
+ if self.layer_idx is None:
588
+ raise ValueError(
589
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
590
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
591
+ "with a layer index."
592
+ )
593
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
594
+
595
+ # Because the input can be padded, the absolute sequence length depends on the max position id.
596
+ rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
597
+ cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
598
+
599
+ query_states, key_states = apply_rotary_pos_emb(
600
+ query_states, key_states, cos, sin, position_ids
601
+ )
602
+
603
+ use_sliding_windows = (
604
+ _flash_supports_window_size
605
+ and getattr(self.config, "sliding_window", None) is not None
606
+ and kv_seq_len > self.config.sliding_window
607
+ )
608
+
609
+ if not _flash_supports_window_size:
610
+ logger.warning_once(
611
+ "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"
612
+ " make sure to upgrade flash-attn library."
613
+ )
614
+
615
+ if past_key_value is not None:
616
+ # Activate slicing cache only if the config has a value `sliding_windows` attribute
617
+ cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0
618
+ if (
619
+ getattr(self.config, "sliding_window", None) is not None
620
+ and kv_seq_len > self.config.sliding_window
621
+ and cache_has_contents
622
+ ):
623
+ slicing_tokens = 1 - self.config.sliding_window
624
+
625
+ past_key = past_key_value[self.layer_idx][0]
626
+ past_value = past_key_value[self.layer_idx][1]
627
+
628
+ past_key = past_key[:, :, slicing_tokens:, :].contiguous()
629
+ past_value = past_value[:, :, slicing_tokens:, :].contiguous()
630
+
631
+ if past_key.shape[-2] != self.config.sliding_window - 1:
632
+ raise ValueError(
633
+ f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
634
+ f" {past_key.shape}"
635
+ )
636
+
637
+ if attention_mask is not None:
638
+ attention_mask = attention_mask[:, slicing_tokens:]
639
+ attention_mask = torch.cat(
640
+ [attention_mask, torch.ones_like(attention_mask[:, -1:])],
641
+ dim=-1,
642
+ )
643
+
644
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
645
+ key_states, value_states = past_key_value.update(
646
+ key_states, value_states, self.layer_idx, cache_kwargs
647
+ )
648
+
649
+ # repeat k/v heads if n_kv_heads < n_heads
650
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
651
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
652
+ dropout_rate = 0.0 if not self.training else self.attention_dropout
653
+
654
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
655
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
656
+ # cast them back in float16 just to be sure everything works as expected.
657
+ input_dtype = query_states.dtype
658
+ if input_dtype == torch.float32:
659
+ if torch.is_autocast_enabled():
660
+ target_dtype = torch.get_autocast_gpu_dtype()
661
+ # Handle the case where the model is quantized
662
+ elif hasattr(self.config, "_pre_quantization_dtype"):
663
+ target_dtype = self.config._pre_quantization_dtype
664
+ else:
665
+ target_dtype = self.q_proj.weight.dtype
666
+
667
+ logger.warning_once(
668
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
669
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
670
+ f" {target_dtype}."
671
+ )
672
+
673
+ query_states = query_states.to(target_dtype)
674
+ key_states = key_states.to(target_dtype)
675
+ value_states = value_states.to(target_dtype)
676
+
677
+ # Reashape to the expected shape for Flash Attention
678
+ query_states = query_states.transpose(1, 2)
679
+ key_states = key_states.transpose(1, 2)
680
+ value_states = value_states.transpose(1, 2)
681
+
682
+ attn_output = self._flash_attention_forward(
683
+ query_states,
684
+ key_states,
685
+ value_states,
686
+ attention_mask,
687
+ q_len,
688
+ dropout=dropout_rate,
689
+ use_sliding_windows=use_sliding_windows,
690
+ )
691
+
692
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
693
+ attn_output = self.o_proj(attn_output)
694
+
695
+ if not output_attentions:
696
+ attn_weights = None
697
+
698
+ return attn_output, attn_weights, past_key_value
699
+
700
+ def _flash_attention_forward(
701
+ self,
702
+ query_states,
703
+ key_states,
704
+ value_states,
705
+ attention_mask,
706
+ query_length,
707
+ dropout=0.0,
708
+ softmax_scale=None,
709
+ use_sliding_windows=False,
710
+ ):
711
+ """
712
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
713
+ first unpad the input, then computes the attention scores and pad the final attention scores.
714
+
715
+ Args:
716
+ query_states (`torch.Tensor`):
717
+ Input query states to be passed to Flash Attention API
718
+ key_states (`torch.Tensor`):
719
+ Input key states to be passed to Flash Attention API
720
+ value_states (`torch.Tensor`):
721
+ Input value states to be passed to Flash Attention API
722
+ attention_mask (`torch.Tensor`):
723
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
724
+ position of padding tokens and 1 for the position of non-padding tokens.
725
+ dropout (`int`, *optional*):
726
+ Attention dropout
727
+ softmax_scale (`float`, *optional*):
728
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
729
+ use_sliding_windows (`bool`, *optional*):
730
+ Whether to activate sliding window attention.
731
+ """
732
+ if not self._flash_attn_uses_top_left_mask:
733
+ causal = self.is_causal
734
+ else:
735
+ # 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__.
736
+ causal = self.is_causal and query_length != 1
737
+
738
+ # Contains at least one padding token in the sequence
739
+ if attention_mask is not None:
740
+ batch_size = query_states.shape[0]
741
+ (
742
+ query_states,
743
+ key_states,
744
+ value_states,
745
+ indices_q,
746
+ cu_seq_lens,
747
+ max_seq_lens,
748
+ ) = self._upad_input(
749
+ query_states, key_states, value_states, attention_mask, query_length
750
+ )
751
+
752
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
753
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
754
+
755
+ if not use_sliding_windows:
756
+ attn_output_unpad = flash_attn_varlen_func(
757
+ query_states,
758
+ key_states,
759
+ value_states,
760
+ cu_seqlens_q=cu_seqlens_q,
761
+ cu_seqlens_k=cu_seqlens_k,
762
+ max_seqlen_q=max_seqlen_in_batch_q,
763
+ max_seqlen_k=max_seqlen_in_batch_k,
764
+ dropout_p=dropout,
765
+ softmax_scale=softmax_scale,
766
+ causal=causal,
767
+ )
768
+ else:
769
+ attn_output_unpad = flash_attn_varlen_func(
770
+ query_states,
771
+ key_states,
772
+ value_states,
773
+ cu_seqlens_q=cu_seqlens_q,
774
+ cu_seqlens_k=cu_seqlens_k,
775
+ max_seqlen_q=max_seqlen_in_batch_q,
776
+ max_seqlen_k=max_seqlen_in_batch_k,
777
+ dropout_p=dropout,
778
+ softmax_scale=softmax_scale,
779
+ causal=causal,
780
+ window_size=(
781
+ self.config.sliding_window,
782
+ self.config.sliding_window,
783
+ ),
784
+ )
785
+
786
+ attn_output = pad_input(
787
+ attn_output_unpad, indices_q, batch_size, query_length
788
+ )
789
+ else:
790
+ if not use_sliding_windows:
791
+ attn_output = flash_attn_func(
792
+ query_states,
793
+ key_states,
794
+ value_states,
795
+ dropout,
796
+ softmax_scale=softmax_scale,
797
+ causal=causal,
798
+ )
799
+ else:
800
+ attn_output = flash_attn_func(
801
+ query_states,
802
+ key_states,
803
+ value_states,
804
+ dropout,
805
+ softmax_scale=softmax_scale,
806
+ causal=causal,
807
+ window_size=(
808
+ self.config.sliding_window,
809
+ self.config.sliding_window,
810
+ ),
811
+ )
812
+
813
+ return attn_output
814
+
815
+ def _upad_input(
816
+ self, query_layer, key_layer, value_layer, attention_mask, query_length
817
+ ):
818
+ batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
819
+
820
+ # On the first iteration we need to properly re-create the padding mask
821
+ # by slicing it on the proper place
822
+ if kv_seq_len != attention_mask.shape[-1]:
823
+ attention_mask_num_tokens = attention_mask.shape[-1]
824
+ attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
825
+
826
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
827
+
828
+ key_layer = index_first_axis(
829
+ key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
830
+ )
831
+ value_layer = index_first_axis(
832
+ value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
833
+ )
834
+
835
+ if query_length == kv_seq_len:
836
+ query_layer = index_first_axis(
837
+ query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim),
838
+ indices_k,
839
+ )
840
+ cu_seqlens_q = cu_seqlens_k
841
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
842
+ indices_q = indices_k
843
+ elif query_length == 1:
844
+ max_seqlen_in_batch_q = 1
845
+ cu_seqlens_q = torch.arange(
846
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
847
+ ) # There is a memcpy here, that is very bad.
848
+ indices_q = cu_seqlens_q[:-1]
849
+ query_layer = query_layer.squeeze(1)
850
+ else:
851
+ # The -q_len: slice assumes left padding.
852
+ attention_mask = attention_mask[:, -query_length:]
853
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
854
+ query_layer, attention_mask
855
+ )
856
+
857
+ return (
858
+ query_layer,
859
+ key_layer,
860
+ value_layer,
861
+ indices_q,
862
+ (cu_seqlens_q, cu_seqlens_k),
863
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
864
+ )
865
+
866
+
867
+ # Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Mistral
868
+ class MistralSdpaAttention(MistralAttention):
869
+ """
870
+ Mistral attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
871
+ `MistralAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
872
+ SDPA API.
873
+ """
874
+
875
+ # Adapted from MistralAttention.forward
876
+ def forward(
877
+ self,
878
+ hidden_states: torch.Tensor,
879
+ attention_mask: Optional[torch.Tensor] = None,
880
+ position_ids: Optional[torch.LongTensor] = None,
881
+ past_key_value: Optional[Cache] = None,
882
+ output_attentions: bool = False,
883
+ use_cache: bool = False,
884
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
885
+ if output_attentions:
886
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
887
+ logger.warning_once(
888
+ "MistralModel is using MistralSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
889
+ 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
890
+ )
891
+ return super().forward(
892
+ hidden_states=hidden_states,
893
+ attention_mask=attention_mask,
894
+ position_ids=position_ids,
895
+ past_key_value=past_key_value,
896
+ output_attentions=output_attentions,
897
+ use_cache=use_cache,
898
+ )
899
+
900
+ bsz, q_len, _ = hidden_states.size()
901
+
902
+ query_states = self.q_proj(hidden_states)
903
+ key_states = self.k_proj(hidden_states)
904
+ value_states = self.v_proj(hidden_states)
905
+
906
+ query_states = query_states.view(
907
+ bsz, q_len, self.num_heads, self.head_dim
908
+ ).transpose(1, 2)
909
+ key_states = key_states.view(
910
+ bsz, q_len, self.num_key_value_heads, self.head_dim
911
+ ).transpose(1, 2)
912
+ value_states = value_states.view(
913
+ bsz, q_len, self.num_key_value_heads, self.head_dim
914
+ ).transpose(1, 2)
915
+
916
+ kv_seq_len = key_states.shape[-2]
917
+ if past_key_value is not None:
918
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
919
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
920
+
921
+ query_states, key_states = apply_rotary_pos_emb(
922
+ query_states, key_states, cos, sin, position_ids
923
+ )
924
+
925
+ if past_key_value is not None:
926
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
927
+ key_states, value_states = past_key_value.update(
928
+ key_states, value_states, self.layer_idx, cache_kwargs
929
+ )
930
+
931
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
932
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
933
+
934
+ if attention_mask is not None:
935
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
936
+ raise ValueError(
937
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
938
+ )
939
+
940
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
941
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
942
+ if query_states.device.type == "cuda" and attention_mask is not None:
943
+ query_states = query_states.contiguous()
944
+ key_states = key_states.contiguous()
945
+ value_states = value_states.contiguous()
946
+
947
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
948
+ query_states,
949
+ key_states,
950
+ value_states,
951
+ attn_mask=attention_mask,
952
+ dropout_p=self.attention_dropout if self.training else 0.0,
953
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
954
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
955
+ )
956
+
957
+ attn_output = attn_output.transpose(1, 2).contiguous()
958
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
959
+
960
+ attn_output = self.o_proj(attn_output)
961
+
962
+ return attn_output, None, past_key_value
963
+
964
+
965
+ MISTRAL_ATTENTION_CLASSES = {
966
+ "eager": MistralAttention,
967
+ "flash_attention_2": MistralFlashAttention2,
968
+ "sdpa": MistralSdpaAttention,
969
+ }
970
+
971
+
972
+ class MistralDecoderLayer(nn.Module):
973
+ def __init__(self, config: SparsetralConfig, layer_idx: int):
974
+ super().__init__()
975
+ self.config = config
976
+ self.hidden_size = config.hidden_size
977
+
978
+ self.self_attn = MISTRAL_ATTENTION_CLASSES[config._attn_implementation](
979
+ config, layer_idx
980
+ )
981
+
982
+ self.mlp = MistralMLP(config)
983
+ self.input_layernorm = MistralRMSNorm(
984
+ config.hidden_size, eps=config.rms_norm_eps
985
+ )
986
+ self.post_attention_layernorm = MistralRMSNorm(
987
+ config.hidden_size, eps=config.rms_norm_eps
988
+ )
989
+
990
+ def forward(
991
+ self,
992
+ hidden_states: torch.Tensor,
993
+ attention_mask: Optional[torch.Tensor] = None,
994
+ position_ids: Optional[torch.LongTensor] = None,
995
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
996
+ output_attentions: Optional[bool] = False,
997
+ output_router_logits: Optional[bool] = False,
998
+ use_cache: Optional[bool] = False,
999
+ **kwargs,
1000
+ ) -> Tuple[
1001
+ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
1002
+ ]:
1003
+ if "padding_mask" in kwargs:
1004
+ warnings.warn(
1005
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
1006
+ )
1007
+ """
1008
+ Args:
1009
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
1010
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
1011
+ `(batch, sequence_length)` where padding elements are indicated by 0.
1012
+ output_attentions (`bool`, *optional*):
1013
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
1014
+ returned tensors for more detail.
1015
+ use_cache (`bool`, *optional*):
1016
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
1017
+ (see `past_key_values`).
1018
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
1019
+ """
1020
+
1021
+ residual = hidden_states
1022
+
1023
+ hidden_states = self.input_layernorm(hidden_states)
1024
+
1025
+ # Self Attention
1026
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
1027
+ hidden_states=hidden_states,
1028
+ attention_mask=attention_mask,
1029
+ position_ids=position_ids,
1030
+ past_key_value=past_key_value,
1031
+ output_attentions=output_attentions,
1032
+ use_cache=use_cache,
1033
+ )
1034
+ hidden_states = residual + hidden_states
1035
+
1036
+ # Fully Connected
1037
+ residual = hidden_states
1038
+ hidden_states = self.post_attention_layernorm(hidden_states)
1039
+ hidden_states, router_logits = self.mlp(hidden_states)
1040
+ hidden_states = residual + hidden_states
1041
+
1042
+ outputs = (hidden_states,)
1043
+
1044
+ if output_attentions:
1045
+ outputs += (self_attn_weights,)
1046
+
1047
+ if use_cache:
1048
+ outputs += (present_key_value,)
1049
+
1050
+ if output_router_logits:
1051
+ outputs += (router_logits,)
1052
+
1053
+ return outputs
1054
+
1055
+
1056
+ MISTRAL_START_DOCSTRING = r"""
1057
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
1058
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
1059
+ etc.)
1060
+
1061
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
1062
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
1063
+ and behavior.
1064
+
1065
+ Parameters:
1066
+ config ([`SparsetralConfig`]):
1067
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
1068
+ load the weights associated with the model, only the configuration. Check out the
1069
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
1070
+ """
1071
+
1072
+
1073
+ @add_start_docstrings(
1074
+ "The bare Mistral Model outputting raw hidden-states without any specific head on top.",
1075
+ MISTRAL_START_DOCSTRING,
1076
+ )
1077
+ class MistralPreTrainedModel(PreTrainedModel):
1078
+ config_class = SparsetralConfig
1079
+ base_model_prefix = "model"
1080
+ supports_gradient_checkpointing = True
1081
+ _no_split_modules = ["MistralDecoderLayer"]
1082
+ _skip_keys_device_placement = "past_key_values"
1083
+ _supports_flash_attn_2 = True
1084
+ _supports_sdpa = True
1085
+ _supports_cache_class = True
1086
+
1087
+ def _init_weights(self, module):
1088
+ std = self.config.initializer_range
1089
+ if isinstance(module, nn.Linear):
1090
+ module.weight.data.normal_(mean=0.0, std=std)
1091
+ if module.bias is not None:
1092
+ module.bias.data.zero_()
1093
+ elif isinstance(module, nn.Embedding):
1094
+ module.weight.data.normal_(mean=0.0, std=std)
1095
+ if module.padding_idx is not None:
1096
+ module.weight.data[module.padding_idx].zero_()
1097
+
1098
+ def _set_gradient_checkpointing(self, module, value=False):
1099
+ if isinstance(module, MistralModel):
1100
+ module.gradient_checkpointing = value
1101
+
1102
+
1103
+ MISTRAL_INPUTS_DOCSTRING = r"""
1104
+ Args:
1105
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
1106
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
1107
+ it.
1108
+
1109
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1110
+ [`PreTrainedTokenizer.__call__`] for details.
1111
+
1112
+ [What are input IDs?](../glossary#input-ids)
1113
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
1114
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
1115
+
1116
+ - 1 for tokens that are **not masked**,
1117
+ - 0 for tokens that are **masked**.
1118
+
1119
+ [What are attention masks?](../glossary#attention-mask)
1120
+
1121
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1122
+ [`PreTrainedTokenizer.__call__`] for details.
1123
+
1124
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
1125
+ `past_key_values`).
1126
+
1127
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
1128
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
1129
+ information on the default strategy.
1130
+
1131
+ - 1 indicates the head is **not masked**,
1132
+ - 0 indicates the head is **masked**.
1133
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1134
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
1135
+ config.n_positions - 1]`.
1136
+
1137
+ [What are position IDs?](../glossary#position-ids)
1138
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
1139
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
1140
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
1141
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
1142
+
1143
+ Two formats are allowed:
1144
+ - a [`~cache_utils.Cache`] instance;
1145
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
1146
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
1147
+ cache format.
1148
+
1149
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
1150
+ legacy cache format will be returned.
1151
+
1152
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
1153
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
1154
+ of shape `(batch_size, sequence_length)`.
1155
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
1156
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
1157
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
1158
+ model's internal embedding lookup matrix.
1159
+ use_cache (`bool`, *optional*):
1160
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
1161
+ `past_key_values`).
1162
+ output_attentions (`bool`, *optional*):
1163
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
1164
+ tensors for more detail.
1165
+ output_hidden_states (`bool`, *optional*):
1166
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
1167
+ more detail.
1168
+ return_dict (`bool`, *optional*):
1169
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
1170
+ """
1171
+
1172
+
1173
+ @add_start_docstrings(
1174
+ "The bare Mistral Model outputting raw hidden-states without any specific head on top.",
1175
+ MISTRAL_START_DOCSTRING,
1176
+ )
1177
+ class MistralModel(MistralPreTrainedModel):
1178
+ """
1179
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MistralDecoderLayer`]
1180
+
1181
+ Args:
1182
+ config: MistralConfig
1183
+ """
1184
+
1185
+ def __init__(self, config: SparsetralConfig):
1186
+ super().__init__(config)
1187
+ self.padding_idx = config.pad_token_id
1188
+ self.vocab_size = config.vocab_size
1189
+
1190
+ self.embed_tokens = nn.Embedding(
1191
+ config.vocab_size, config.hidden_size, self.padding_idx
1192
+ )
1193
+ self.layers = nn.ModuleList(
1194
+ [
1195
+ MistralDecoderLayer(config, layer_idx)
1196
+ for layer_idx in range(config.num_hidden_layers)
1197
+ ]
1198
+ )
1199
+ self._attn_implementation = config._attn_implementation
1200
+ self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1201
+
1202
+ self.gradient_checkpointing = False
1203
+ # Initialize weights and apply final processing
1204
+ self.post_init()
1205
+
1206
+ def get_input_embeddings(self):
1207
+ return self.embed_tokens
1208
+
1209
+ def set_input_embeddings(self, value):
1210
+ self.embed_tokens = value
1211
+
1212
+ @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
1213
+ def forward(
1214
+ self,
1215
+ input_ids: torch.LongTensor = None,
1216
+ attention_mask: Optional[torch.Tensor] = None,
1217
+ position_ids: Optional[torch.LongTensor] = None,
1218
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1219
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1220
+ use_cache: Optional[bool] = None,
1221
+ output_attentions: Optional[bool] = None,
1222
+ output_hidden_states: Optional[bool] = None,
1223
+ output_router_logits: Optional[bool] = None,
1224
+ return_dict: Optional[bool] = None,
1225
+ ) -> Union[Tuple, MoEModelOutputWithPast]:
1226
+ output_attentions = (
1227
+ output_attentions
1228
+ if output_attentions is not None
1229
+ else self.config.output_attentions
1230
+ )
1231
+ output_hidden_states = (
1232
+ output_hidden_states
1233
+ if output_hidden_states is not None
1234
+ else self.config.output_hidden_states
1235
+ )
1236
+ output_router_logits = (
1237
+ output_router_logits
1238
+ if output_router_logits is not None
1239
+ else self.config.output_router_logits
1240
+ )
1241
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1242
+
1243
+ return_dict = (
1244
+ return_dict if return_dict is not None else self.config.use_return_dict
1245
+ )
1246
+
1247
+ # retrieve input_ids and inputs_embeds
1248
+ if input_ids is not None and inputs_embeds is not None:
1249
+ raise ValueError(
1250
+ "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
1251
+ )
1252
+ elif input_ids is not None:
1253
+ batch_size, seq_length = input_ids.shape
1254
+ elif inputs_embeds is not None:
1255
+ batch_size, seq_length, _ = inputs_embeds.shape
1256
+ else:
1257
+ raise ValueError(
1258
+ "You have to specify either decoder_input_ids or decoder_inputs_embeds"
1259
+ )
1260
+
1261
+ if self.gradient_checkpointing and self.training:
1262
+ if use_cache:
1263
+ logger.warning_once(
1264
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1265
+ )
1266
+ use_cache = False
1267
+
1268
+ seq_length_with_past = seq_length
1269
+ past_key_values_length = 0
1270
+
1271
+ if past_key_values is not None:
1272
+ past_key_values_length = past_key_values[0][0].shape[2]
1273
+ seq_length_with_past = seq_length_with_past + past_key_values_length
1274
+
1275
+ if use_cache:
1276
+ use_legacy_cache = not isinstance(past_key_values, Cache)
1277
+ if use_legacy_cache:
1278
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1279
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
1280
+
1281
+ if position_ids is None:
1282
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1283
+ position_ids = torch.arange(
1284
+ past_key_values_length,
1285
+ seq_length + past_key_values_length,
1286
+ dtype=torch.long,
1287
+ device=device,
1288
+ )
1289
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
1290
+ else:
1291
+ position_ids = position_ids.view(-1, seq_length).long()
1292
+
1293
+ if inputs_embeds is None:
1294
+ inputs_embeds = self.embed_tokens(input_ids)
1295
+
1296
+ if (
1297
+ attention_mask is not None
1298
+ and self._attn_implementation == "flash_attention_2"
1299
+ and use_cache
1300
+ ):
1301
+ is_padding_right = attention_mask[:, -1].sum().item() != batch_size
1302
+ if is_padding_right:
1303
+ raise ValueError(
1304
+ "You are attempting to perform batched generation with padding_side='right'"
1305
+ " this may lead to unexpected behaviour for Flash Attention version of Mistral. Make sure to "
1306
+ " call `tokenizer.padding_side = 'left'` before tokenizing the input. "
1307
+ )
1308
+
1309
+ if self._attn_implementation == "flash_attention_2":
1310
+ # 2d mask is passed through the layers
1311
+ attention_mask = (
1312
+ attention_mask
1313
+ if (attention_mask is not None and 0 in attention_mask)
1314
+ else None
1315
+ )
1316
+ elif self._attn_implementation == "sdpa" and not output_attentions:
1317
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
1318
+ # the manual implementation that requires a 4D causal mask in all cases.
1319
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
1320
+ attention_mask,
1321
+ (batch_size, seq_length),
1322
+ inputs_embeds,
1323
+ past_key_values_length,
1324
+ )
1325
+ else:
1326
+ # 4d mask is passed through the layers
1327
+ attention_mask = _prepare_4d_causal_attention_mask(
1328
+ attention_mask,
1329
+ (batch_size, seq_length),
1330
+ inputs_embeds,
1331
+ past_key_values_length,
1332
+ sliding_window=self.config.sliding_window,
1333
+ )
1334
+
1335
+ hidden_states = inputs_embeds
1336
+
1337
+ # decoder layers
1338
+ all_hidden_states = () if output_hidden_states else None
1339
+ all_self_attns = () if output_attentions else None
1340
+ all_router_logits = () if output_router_logits else None
1341
+ next_decoder_cache = None
1342
+
1343
+ for decoder_layer in self.layers:
1344
+ if output_hidden_states:
1345
+ all_hidden_states += (hidden_states,)
1346
+
1347
+ if self.gradient_checkpointing and self.training:
1348
+
1349
+ def create_custom_forward(module):
1350
+ def custom_forward(*inputs):
1351
+ # None for past_key_value
1352
+ return module(
1353
+ *inputs, output_attentions, output_router_logits, None
1354
+ )
1355
+
1356
+ return custom_forward
1357
+
1358
+ layer_outputs = torch.utils.checkpoint.checkpoint(
1359
+ create_custom_forward(decoder_layer),
1360
+ hidden_states,
1361
+ attention_mask,
1362
+ position_ids,
1363
+ None,
1364
+ )
1365
+ else:
1366
+ layer_outputs = decoder_layer(
1367
+ hidden_states,
1368
+ attention_mask=attention_mask,
1369
+ position_ids=position_ids,
1370
+ past_key_value=past_key_values,
1371
+ output_attentions=output_attentions,
1372
+ output_router_logits=output_router_logits,
1373
+ use_cache=use_cache,
1374
+ )
1375
+
1376
+ hidden_states = layer_outputs[0]
1377
+
1378
+ if use_cache:
1379
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1380
+
1381
+ if output_attentions:
1382
+ all_self_attns += (layer_outputs[1],)
1383
+
1384
+ if output_router_logits:
1385
+ all_router_logits += (layer_outputs[-1],)
1386
+
1387
+ hidden_states = self.norm(hidden_states)
1388
+
1389
+ # add hidden states from the last decoder layer
1390
+ if output_hidden_states:
1391
+ all_hidden_states += (hidden_states,)
1392
+
1393
+ next_cache = None
1394
+ if use_cache:
1395
+ next_cache = (
1396
+ next_decoder_cache.to_legacy_cache()
1397
+ if use_legacy_cache
1398
+ else next_decoder_cache
1399
+ )
1400
+
1401
+ if not return_dict:
1402
+ return tuple(
1403
+ v
1404
+ for v in [
1405
+ hidden_states,
1406
+ next_cache,
1407
+ all_hidden_states,
1408
+ all_self_attns,
1409
+ all_router_logits,
1410
+ ]
1411
+ if v is not None
1412
+ )
1413
+ return MoEModelOutputWithPast(
1414
+ last_hidden_state=hidden_states,
1415
+ past_key_values=next_cache,
1416
+ hidden_states=all_hidden_states,
1417
+ attentions=all_self_attns,
1418
+ router_logits=all_router_logits,
1419
+ )
1420
+
1421
+
1422
+ class MistralForCausalLM(MistralPreTrainedModel):
1423
+ _tied_weights_keys = ["lm_head.weight"]
1424
+
1425
+ def __init__(self, config):
1426
+ super().__init__(config)
1427
+ self.config = config
1428
+ self.model = MistralModel(config)
1429
+ self.vocab_size = config.vocab_size
1430
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1431
+
1432
+ # Initialize weights and apply final processing
1433
+ self.post_init()
1434
+
1435
+ def get_input_embeddings(self):
1436
+ return self.model.embed_tokens
1437
+
1438
+ def set_input_embeddings(self, value):
1439
+ self.model.embed_tokens = value
1440
+
1441
+ def get_output_embeddings(self):
1442
+ return self.lm_head
1443
+
1444
+ def set_output_embeddings(self, new_embeddings):
1445
+ self.lm_head = new_embeddings
1446
+
1447
+ def set_decoder(self, decoder):
1448
+ self.model = decoder
1449
+
1450
+ def get_decoder(self):
1451
+ return self.model
1452
+
1453
+ @add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
1454
+ @replace_return_docstrings(
1455
+ output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
1456
+ )
1457
+ def forward(
1458
+ self,
1459
+ input_ids: torch.LongTensor = None,
1460
+ attention_mask: Optional[torch.Tensor] = None,
1461
+ position_ids: Optional[torch.LongTensor] = None,
1462
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1463
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1464
+ labels: Optional[torch.LongTensor] = None,
1465
+ use_cache: Optional[bool] = None,
1466
+ output_attentions: Optional[bool] = None,
1467
+ output_hidden_states: Optional[bool] = None,
1468
+ output_router_logits: Optional[bool] = None,
1469
+ return_dict: Optional[bool] = None,
1470
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1471
+ r"""
1472
+ Args:
1473
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1474
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1475
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1476
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1477
+
1478
+ Returns:
1479
+
1480
+ Example:
1481
+
1482
+ ```python
1483
+ >>> from transformers import AutoTokenizer, MistralForCausalLM
1484
+
1485
+ >>> model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
1486
+ >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
1487
+
1488
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1489
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1490
+
1491
+ >>> # Generate
1492
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1493
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1494
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1495
+ ```"""
1496
+
1497
+ output_attentions = (
1498
+ output_attentions
1499
+ if output_attentions is not None
1500
+ else self.config.output_attentions
1501
+ )
1502
+ output_hidden_states = (
1503
+ output_hidden_states
1504
+ if output_hidden_states is not None
1505
+ else self.config.output_hidden_states
1506
+ )
1507
+ output_router_logits = (
1508
+ output_router_logits
1509
+ if output_router_logits is not None
1510
+ else self.config.output_router_logits
1511
+ )
1512
+ return_dict = (
1513
+ return_dict if return_dict is not None else self.config.use_return_dict
1514
+ )
1515
+
1516
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1517
+ outputs = self.model(
1518
+ input_ids=input_ids,
1519
+ attention_mask=attention_mask,
1520
+ position_ids=position_ids,
1521
+ past_key_values=past_key_values,
1522
+ inputs_embeds=inputs_embeds,
1523
+ use_cache=use_cache,
1524
+ output_attentions=output_attentions,
1525
+ output_hidden_states=output_hidden_states,
1526
+ output_router_logits=output_router_logits,
1527
+ return_dict=return_dict,
1528
+ )
1529
+
1530
+ hidden_states = outputs[0]
1531
+ logits = self.lm_head(hidden_states)
1532
+ logits = logits.float()
1533
+
1534
+ loss = None
1535
+ if labels is not None:
1536
+ # Shift so that tokens < n predict n
1537
+ shift_logits = logits[..., :-1, :].contiguous()
1538
+ shift_labels = labels[..., 1:].contiguous()
1539
+ # Flatten the tokens
1540
+ loss_fct = CrossEntropyLoss()
1541
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1542
+ shift_labels = shift_labels.view(-1)
1543
+ # Enable model parallelism
1544
+ shift_labels = shift_labels.to(shift_logits.device)
1545
+ loss = loss_fct(shift_logits, shift_labels)
1546
+
1547
+ aux_loss = None
1548
+ if output_router_logits:
1549
+ aux_loss = load_balancing_loss_func(
1550
+ outputs.router_logits if return_dict else outputs[-1],
1551
+ self.config.num_experts,
1552
+ self.config.topk,
1553
+ )
1554
+ if labels is not None:
1555
+ loss += self.config.router_aux_loss_coef * aux_loss
1556
+
1557
+ if not return_dict:
1558
+ output = (logits,) + outputs[1:]
1559
+ if output_router_logits:
1560
+ output = (aux_loss,) + output
1561
+ return (loss,) + output if loss is not None else output
1562
+
1563
+ return MoECausalLMOutputWithPast(
1564
+ loss=loss,
1565
+ aux_loss=aux_loss,
1566
+ logits=logits,
1567
+ past_key_values=outputs.past_key_values,
1568
+ hidden_states=outputs.hidden_states,
1569
+ attentions=outputs.attentions,
1570
+ router_logits=outputs.router_logits,
1571
+ )
1572
+
1573
+ def prepare_inputs_for_generation(
1574
+ self,
1575
+ input_ids,
1576
+ past_key_values=None,
1577
+ attention_mask=None,
1578
+ inputs_embeds=None,
1579
+ **kwargs,
1580
+ ):
1581
+ # Omit tokens covered by past_key_values
1582
+ if past_key_values is not None:
1583
+ if isinstance(past_key_values, Cache):
1584
+ cache_length = past_key_values.get_seq_length()
1585
+ past_length = past_key_values.seen_tokens
1586
+ max_cache_length = past_key_values.get_max_length()
1587
+ else:
1588
+ cache_length = past_length = past_key_values[0][0].shape[2]
1589
+ max_cache_length = None
1590
+
1591
+ # Keep only the unprocessed tokens:
1592
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1593
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
1594
+ # input)
1595
+ if (
1596
+ attention_mask is not None
1597
+ and attention_mask.shape[1] > input_ids.shape[1]
1598
+ ):
1599
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1600
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1601
+ # input_ids based on the past_length.
1602
+ elif past_length < input_ids.shape[1]:
1603
+ input_ids = input_ids[:, past_length:]
1604
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1605
+
1606
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1607
+ if (
1608
+ max_cache_length is not None
1609
+ and attention_mask is not None
1610
+ and cache_length + input_ids.shape[1] > max_cache_length
1611
+ ):
1612
+ attention_mask = attention_mask[:, -max_cache_length:]
1613
+
1614
+ position_ids = kwargs.get("position_ids", None)
1615
+ if attention_mask is not None and position_ids is None:
1616
+ # create position_ids on the fly for batch generation
1617
+ position_ids = attention_mask.long().cumsum(-1) - 1
1618
+ position_ids.masked_fill_(attention_mask == 0, 1)
1619
+ if past_key_values:
1620
+ position_ids = position_ids[:, -1].unsqueeze(-1)
1621
+
1622
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1623
+ if inputs_embeds is not None and past_key_values is None:
1624
+ model_inputs = {"inputs_embeds": inputs_embeds}
1625
+ else:
1626
+ model_inputs = {"input_ids": input_ids}
1627
+
1628
+ model_inputs.update(
1629
+ {
1630
+ "position_ids": position_ids,
1631
+ "past_key_values": past_key_values,
1632
+ "use_cache": kwargs.get("use_cache"),
1633
+ "attention_mask": attention_mask,
1634
+ }
1635
+ )
1636
+ return model_inputs
1637
+
1638
+ @staticmethod
1639
+ def _reorder_cache(past_key_values, beam_idx):
1640
+ reordered_past = ()
1641
+ for layer_past in past_key_values:
1642
+ reordered_past += (
1643
+ tuple(
1644
+ past_state.index_select(0, beam_idx.to(past_state.device))
1645
+ for past_state in layer_past
1646
+ ),
1647
+ )
1648
+ return reordered_past
output.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:97806fc876aeab07336c5cb986f7bec567b8974704c892233f81b85000eb0286
3
+ size 2982860388
special_tokens_map.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "</s>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "unk_token": {
17
+ "content": "<unk>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ }
23
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:dadfd56d766715c61d2ef780a525ab43b8e6da4de6865bda3d95fdef5e134055
3
+ size 493443
tokenizer_config.json ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": true,
3
+ "add_eos_token": false,
4
+ "added_tokens_decoder": {
5
+ "0": {
6
+ "content": "<unk>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "1": {
14
+ "content": "<s>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "2": {
22
+ "content": "</s>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ }
29
+ },
30
+ "additional_special_tokens": [],
31
+ "bos_token": "<s>",
32
+ "chat_template": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}",
33
+ "clean_up_tokenization_spaces": false,
34
+ "eos_token": "</s>",
35
+ "legacy": true,
36
+ "model_max_length": 1000000000000000019884624838656,
37
+ "pad_token": null,
38
+ "sp_model_kwargs": {},
39
+ "spaces_between_special_tokens": false,
40
+ "tokenizer_class": "LlamaTokenizer",
41
+ "unk_token": "<unk>",
42
+ "use_default_system_prompt": false
43
+ }