aigencydev commited on
Commit
e14df4c
·
verified ·
1 Parent(s): e4cd16d

Standart Qwen3 mimarisine geciliyor: model her arac tarafindan yuklenebilir

Browse files

Olculdu: Erk-14B'nin 443 tensorunun 443'u de Qwen3-14B ile ayni isimde ve
mimari alanlarin hepsi esit (hidden_size 5120, 40 katman, 40/8 bas, head_dim 128,
intermediate 17408, vocab 151936, rope_theta 1e6, max_pos 40960). Yani
ErkForCausalLM, Qwen3 kodunun yeniden adlandirilmisiydi; islevsel fark yoktu.

Bedeli buyuktu:
- llama.cpp convert_hf_to_gguf.py model_type'a gore dagitim yapar; "erk" tanimsiz
oldugu icin GGUF donusumu bastan basarisiz oluyordu.
- vLLM, SGLang, TGI, Ollama, LM Studio ya reddediyor ya trust_remote_code istiyor.
- Ozel kod yalnizca transformers 5.x API'lerini kullandigi icin 4.x surumlerde
model hic acilmiyordu.
- use_cache false idi: KV onbellegi varsayilan kapali, uretim cok yavas.

Bu islem yalnizca config ve ozel kod dosyalarini degistirir; agirliklara
dokunulmamistir.

Files changed (3) hide show
  1. config.json +9 -58
  2. configuration_erk.py +0 -100
  3. modeling_erk.py +0 -538
config.json CHANGED
@@ -1,79 +1,30 @@
1
  {
2
  "architectures": [
3
- "ErkForCausalLM"
4
  ],
5
  "attention_bias": false,
6
  "attention_dropout": 0.0,
7
  "bos_token_id": 151643,
8
- "dtype": "float32",
9
  "eos_token_id": 151645,
10
  "head_dim": 128,
11
  "hidden_act": "silu",
12
  "hidden_size": 5120,
13
  "initializer_range": 0.02,
14
  "intermediate_size": 17408,
15
- "layer_types": [
16
- "full_attention",
17
- "full_attention",
18
- "full_attention",
19
- "full_attention",
20
- "full_attention",
21
- "full_attention",
22
- "full_attention",
23
- "full_attention",
24
- "full_attention",
25
- "full_attention",
26
- "full_attention",
27
- "full_attention",
28
- "full_attention",
29
- "full_attention",
30
- "full_attention",
31
- "full_attention",
32
- "full_attention",
33
- "full_attention",
34
- "full_attention",
35
- "full_attention",
36
- "full_attention",
37
- "full_attention",
38
- "full_attention",
39
- "full_attention",
40
- "full_attention",
41
- "full_attention",
42
- "full_attention",
43
- "full_attention",
44
- "full_attention",
45
- "full_attention",
46
- "full_attention",
47
- "full_attention",
48
- "full_attention",
49
- "full_attention",
50
- "full_attention",
51
- "full_attention",
52
- "full_attention",
53
- "full_attention",
54
- "full_attention",
55
- "full_attention"
56
- ],
57
  "max_position_embeddings": 40960,
58
  "max_window_layers": 40,
59
- "model_type": "erk",
60
  "num_attention_heads": 40,
61
  "num_hidden_layers": 40,
62
  "num_key_value_heads": 8,
63
- "pad_token_id": null,
64
  "rms_norm_eps": 1e-06,
65
- "rope_parameters": {
66
- "rope_theta": 1000000,
67
- "rope_type": "default"
68
- },
69
  "sliding_window": null,
70
  "tie_word_embeddings": false,
71
- "transformers_version": "5.13.0",
72
- "use_cache": false,
 
73
  "use_sliding_window": false,
74
- "vocab_size": 151936,
75
- "auto_map": {
76
- "AutoConfig": "configuration_erk.ErkConfig",
77
- "AutoModelForCausalLM": "modeling_erk.ErkForCausalLM"
78
- }
79
- }
 
1
  {
2
  "architectures": [
3
+ "Qwen3ForCausalLM"
4
  ],
5
  "attention_bias": false,
6
  "attention_dropout": 0.0,
7
  "bos_token_id": 151643,
 
8
  "eos_token_id": 151645,
9
  "head_dim": 128,
10
  "hidden_act": "silu",
11
  "hidden_size": 5120,
12
  "initializer_range": 0.02,
13
  "intermediate_size": 17408,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  "max_position_embeddings": 40960,
15
  "max_window_layers": 40,
16
+ "model_type": "qwen3",
17
  "num_attention_heads": 40,
18
  "num_hidden_layers": 40,
19
  "num_key_value_heads": 8,
 
20
  "rms_norm_eps": 1e-06,
21
+ "rope_scaling": null,
22
+ "rope_theta": 1000000,
 
 
23
  "sliding_window": null,
24
  "tie_word_embeddings": false,
25
+ "torch_dtype": "bfloat16",
26
+ "transformers_version": "4.51.0",
27
+ "use_cache": true,
28
  "use_sliding_window": false,
29
+ "vocab_size": 151936
30
+ }
 
 
 
 
configuration_erk.py DELETED
@@ -1,100 +0,0 @@
1
- # Copyright 2026 eCloud Yazılım Teknolojileri. Based on Apache-2.0 licensed transformer architecture.
2
- #
3
- # you may not use this file except in compliance with the License.
4
- # You may obtain a copy of the License at
5
- #
6
- # http://www.apache.org/licenses/LICENSE-2.0
7
- #
8
- # Unless required by applicable law or agreed to in writing, software
9
- # distributed under the License is distributed on an "AS IS" BASIS,
10
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11
- # See the License for the specific language governing permissions and
12
- # limitations under the License.
13
- """Erk model configuration"""
14
-
15
- from huggingface_hub.dataclasses import strict
16
-
17
- from transformers.configuration_utils import PreTrainedConfig
18
- from transformers.modeling_rope_utils import RopeParameters
19
- from transformers.utils import auto_docstring
20
-
21
-
22
- @auto_docstring(checkpoint="ecloudtech/erk")
23
- @strict
24
- class ErkConfig(PreTrainedConfig):
25
- r"""
26
- ```python
27
- >>> from transformers import ErkModel, ErkConfig
28
-
29
- >>> # Initializing a Erk style configuration
30
- >>> configuration = ErkConfig()
31
-
32
- >>> # Initializing a model from the Erk-8B style configuration
33
- >>> model = ErkModel(configuration)
34
-
35
- >>> # Accessing the model configuration
36
- >>> configuration = model.config
37
- ```
38
- """
39
-
40
- model_type = "erk"
41
- keys_to_ignore_at_inference = ["past_key_values"]
42
-
43
- # Default tensor parallel plan for base model `Erk`
44
- base_model_tp_plan = {
45
- "layers.*.self_attn.q_proj": "colwise",
46
- "layers.*.self_attn.k_proj": "colwise",
47
- "layers.*.self_attn.v_proj": "colwise",
48
- "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
49
- "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
50
- "layers.*.self_attn.o_proj": "rowwise",
51
- "layers.*.mlp.gate_proj": "colwise",
52
- "layers.*.mlp.up_proj": "colwise",
53
- "layers.*.mlp.down_proj": "rowwise",
54
- }
55
- base_model_pp_plan = {
56
- "embed_tokens": (["input_ids"], ["inputs_embeds"]),
57
- "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
58
- "norm": (["hidden_states"], ["hidden_states"]),
59
- }
60
-
61
- vocab_size: int = 151936
62
- hidden_size: int = 4096
63
- intermediate_size: int = 22016
64
- num_hidden_layers: int = 32
65
- num_attention_heads: int = 32
66
- num_key_value_heads: int | None = 32
67
- head_dim: int = 128
68
- hidden_act: str = "silu"
69
- max_position_embeddings: int = 32768
70
- initializer_range: float = 0.02
71
- rms_norm_eps: float = 1e-6
72
- use_cache: bool = True
73
- tie_word_embeddings: bool = False
74
- rope_parameters: RopeParameters | dict | None = None
75
- attention_bias: bool = False
76
- use_sliding_window: bool = False
77
- sliding_window: int | None = 4096
78
- max_window_layers: int = 28
79
- layer_types: list[str] | None = None
80
- attention_dropout: float | int = 0.0
81
- pad_token_id: int | None = None
82
- bos_token_id: int | None = None
83
- eos_token_id: int | list[int] | None = None
84
-
85
- def __post_init__(self, **kwargs):
86
- self.sliding_window = self.sliding_window if self.use_sliding_window else None
87
- if self.num_key_value_heads is None:
88
- self.num_key_value_heads = self.num_attention_heads
89
-
90
- if self.layer_types is None:
91
- self.layer_types = [
92
- "sliding_attention"
93
- if self.sliding_window is not None and i >= self.max_window_layers
94
- else "full_attention"
95
- for i in range(self.num_hidden_layers)
96
- ]
97
- super().__post_init__(**kwargs)
98
-
99
-
100
- __all__ = ["ErkConfig"]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
modeling_erk.py DELETED
@@ -1,538 +0,0 @@
1
- # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
- # This file was automatically generated from src/transformers/models/erk/modular_erk.py.
3
- # Do NOT edit this file manually as any edits will be overwritten by the generation of
4
- # the file from the modular. If any change should be done, please apply the change to the
5
- # modular_erk.py file directly. One of our CI enforces this.
6
- # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
- # Copyright 2026 eCloud Yazılım Teknolojileri. Based on Apache-2.0 licensed transformer architecture.
8
- #
9
- # you may not use this file except in compliance with the License.
10
- # You may obtain a copy of the License at
11
- #
12
- # http://www.apache.org/licenses/LICENSE-2.0
13
- #
14
- # Unless required by applicable law or agreed to in writing, software
15
- # distributed under the License is distributed on an "AS IS" BASIS,
16
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
- # See the License for the specific language governing permissions and
18
- # limitations under the License.
19
-
20
- from collections.abc import Callable
21
- from typing import Optional
22
-
23
- import torch
24
- from torch import nn
25
-
26
- from transformers.activations import ACT2FN
27
- from transformers.cache_utils import Cache, DynamicCache
28
- from transformers.generation import GenerationMixin
29
- from transformers.integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
30
- from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
31
- from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
32
- from transformers.modeling_layers import (
33
- GenericForQuestionAnswering,
34
- GenericForSequenceClassification,
35
- GenericForTokenClassification,
36
- GradientCheckpointingLayer,
37
- )
38
- from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
39
- from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
40
- from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
41
- from transformers.processing_utils import Unpack
42
- from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
43
- from transformers.utils.generic import maybe_autocast, merge_with_config_defaults
44
- from transformers.utils.output_capturing import capture_outputs
45
- from .configuration_erk import ErkConfig
46
-
47
-
48
- @use_kernel_forward_from_hub("RMSNorm")
49
- class ErkRMSNorm(nn.Module):
50
- def __init__(self, hidden_size, eps: float = 1e-6) -> None:
51
- """
52
- ErkRMSNorm is equivalent to T5LayerNorm
53
- """
54
- super().__init__()
55
- self.weight = nn.Parameter(torch.ones(hidden_size))
56
- self.variance_epsilon = eps
57
-
58
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
59
- input_dtype = hidden_states.dtype
60
- hidden_states = hidden_states.to(torch.float32)
61
- variance = hidden_states.pow(2).mean(-1, keepdim=True)
62
- hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
63
- return self.weight * hidden_states.to(input_dtype)
64
-
65
- def extra_repr(self):
66
- return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
67
-
68
-
69
- class ErkMLP(nn.Module):
70
- def __init__(self, config):
71
- super().__init__()
72
- self.config = config
73
- self.hidden_size = config.hidden_size
74
- self.intermediate_size = config.intermediate_size
75
- self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
76
- self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
77
- self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
78
- self.act_fn = ACT2FN[config.hidden_act]
79
-
80
- def forward(self, x):
81
- down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
82
- return down_proj
83
-
84
-
85
- class ErkRotaryEmbedding(nn.Module):
86
- inv_freq: torch.Tensor # fix linting for `register_buffer`
87
-
88
- def __init__(self, config: ErkConfig, device=None):
89
- super().__init__()
90
- self.max_seq_len_cached = config.max_position_embeddings
91
- self.original_max_seq_len = config.max_position_embeddings
92
-
93
- self.config = config
94
-
95
- self.rope_type = self.config.rope_parameters["rope_type"]
96
- rope_init_fn: Callable = self.compute_default_rope_parameters
97
- if self.rope_type != "default":
98
- rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
99
- inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
100
-
101
- self.register_buffer("inv_freq", inv_freq, persistent=False)
102
- self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
103
-
104
- @staticmethod
105
- def compute_default_rope_parameters(
106
- config: ErkConfig | None = None,
107
- device: Optional["torch.device"] = None,
108
- seq_len: int | None = None,
109
- ) -> tuple["torch.Tensor", float]:
110
- """
111
- Computes the inverse frequencies according to the original RoPE implementation
112
- Args:
113
- config ([`~transformers.PreTrainedConfig`]):
114
- The model configuration.
115
- device (`torch.device`):
116
- The device to use for initialization of the inverse frequencies.
117
- seq_len (`int`, *optional*):
118
- The current sequence length. Unused for this type of RoPE.
119
- Returns:
120
- Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
121
- post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
122
- """
123
- base = config.rope_parameters["rope_theta"]
124
- dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
125
-
126
- attention_factor = 1.0 # Unused in this type of RoPE
127
-
128
- # Compute the inverse frequencies
129
- inv_freq = 1.0 / (
130
- base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
131
- )
132
- return inv_freq, attention_factor
133
-
134
- @torch.no_grad()
135
- @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
136
- def forward(self, x, position_ids):
137
- inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
138
- position_ids_expanded = position_ids[:, None, :].float()
139
-
140
- device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
141
- with maybe_autocast(device_type=device_type, enabled=False): # Force float32
142
- freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
143
- emb = torch.cat((freqs, freqs), dim=-1)
144
- cos = emb.cos() * self.attention_scaling
145
- sin = emb.sin() * self.attention_scaling
146
-
147
- return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
148
-
149
-
150
- def rotate_half(x):
151
- """Rotates half the hidden dims of the input."""
152
- x1 = x[..., : x.shape[-1] // 2]
153
- x2 = x[..., x.shape[-1] // 2 :]
154
- return torch.cat((-x2, x1), dim=-1)
155
-
156
-
157
- @use_kernel_func_from_hub("rotary_pos_emb")
158
- def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
159
- """Applies Rotary Position Embedding to the query and key tensors.
160
-
161
- Args:
162
- q (`torch.Tensor`): The query tensor.
163
- k (`torch.Tensor`): The key tensor.
164
- cos (`torch.Tensor`): The cosine part of the rotary embedding.
165
- sin (`torch.Tensor`): The sine part of the rotary embedding.
166
- unsqueeze_dim (`int`, *optional*, defaults to 1):
167
- The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
168
- sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
169
- that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
170
- k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
171
- cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
172
- the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
173
- Returns:
174
- `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
175
- """
176
- cos = cos.unsqueeze(unsqueeze_dim)
177
- sin = sin.unsqueeze(unsqueeze_dim)
178
- q_embed = (q * cos) + (rotate_half(q) * sin)
179
- k_embed = (k * cos) + (rotate_half(k) * sin)
180
- return q_embed, k_embed
181
-
182
-
183
- def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
184
- """
185
- This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
186
- num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
187
- """
188
- batch, num_key_value_heads, slen, head_dim = hidden_states.shape
189
- if n_rep == 1:
190
- return hidden_states
191
- hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
192
- return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
193
-
194
-
195
- def eager_attention_forward(
196
- module: nn.Module,
197
- query: torch.Tensor,
198
- key: torch.Tensor,
199
- value: torch.Tensor,
200
- attention_mask: torch.Tensor | None,
201
- scaling: float,
202
- dropout: float = 0.0,
203
- **kwargs: Unpack[TransformersKwargs],
204
- ):
205
- key_states = repeat_kv(key, module.num_key_value_groups)
206
- value_states = repeat_kv(value, module.num_key_value_groups)
207
-
208
- attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
209
- if attention_mask is not None:
210
- attn_weights = attn_weights + attention_mask
211
-
212
- attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
213
- attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
214
- attn_output = torch.matmul(attn_weights, value_states)
215
- attn_output = attn_output.transpose(1, 2).contiguous()
216
-
217
- return attn_output, attn_weights
218
-
219
-
220
- @use_kernelized_func(apply_rotary_pos_emb)
221
- class ErkAttention(nn.Module):
222
- """Multi-headed attention from 'Attention Is All You Need' paper"""
223
-
224
- def __init__(self, config: ErkConfig, layer_idx: int):
225
- super().__init__()
226
- self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None
227
- self.config = config
228
- self.layer_idx = layer_idx
229
- self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
230
- self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
231
- self.scaling = self.head_dim**-0.5
232
- self.attention_dropout = config.attention_dropout
233
- self.is_causal = True
234
-
235
- self.q_proj = nn.Linear(
236
- config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
237
- )
238
- self.k_proj = nn.Linear(
239
- config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
240
- )
241
- self.v_proj = nn.Linear(
242
- config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
243
- )
244
- self.o_proj = nn.Linear(
245
- config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
246
- )
247
- self.q_norm = ErkRMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim!
248
- self.k_norm = ErkRMSNorm(self.head_dim, eps=config.rms_norm_eps) # thus post q_norm does not need reshape
249
- self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None
250
-
251
- def forward(
252
- self,
253
- hidden_states: torch.Tensor,
254
- position_embeddings: tuple[torch.Tensor, torch.Tensor],
255
- attention_mask: torch.Tensor | None,
256
- past_key_values: Cache | None = None,
257
- **kwargs: Unpack[FlashAttentionKwargs],
258
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
259
- input_shape = hidden_states.shape[:-1]
260
- hidden_shape = (*input_shape, -1, self.head_dim)
261
-
262
- query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
263
- key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
264
- value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
265
-
266
- cos, sin = position_embeddings
267
- query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
268
-
269
- if past_key_values is not None:
270
- key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
271
-
272
- attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
273
- self.config._attn_implementation, eager_attention_forward
274
- )
275
-
276
- attn_output, attn_weights = attention_interface(
277
- self,
278
- query_states,
279
- key_states,
280
- value_states,
281
- attention_mask,
282
- dropout=0.0 if not self.training else self.attention_dropout,
283
- scaling=self.scaling,
284
- sliding_window=self.sliding_window, # diff with Llama
285
- **kwargs,
286
- )
287
-
288
- attn_output = attn_output.reshape(*input_shape, -1).contiguous()
289
- attn_output = self.o_proj(attn_output)
290
- return attn_output, attn_weights
291
-
292
-
293
- class ErkDecoderLayer(GradientCheckpointingLayer):
294
- def __init__(self, config: ErkConfig, layer_idx: int):
295
- super().__init__()
296
- self.hidden_size = config.hidden_size
297
-
298
- self.self_attn = ErkAttention(config=config, layer_idx=layer_idx)
299
-
300
- self.mlp = ErkMLP(config)
301
- self.input_layernorm = ErkRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
302
- self.post_attention_layernorm = ErkRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
303
-
304
- def forward(
305
- self,
306
- hidden_states: torch.Tensor,
307
- attention_mask: torch.Tensor | None = None,
308
- position_ids: torch.LongTensor | None = None,
309
- past_key_values: Cache | None = None,
310
- use_cache: bool | None = False,
311
- position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
312
- **kwargs: Unpack[TransformersKwargs],
313
- ) -> torch.Tensor:
314
- residual = hidden_states
315
- hidden_states = self.input_layernorm(hidden_states)
316
- # Self Attention
317
- hidden_states, _ = self.self_attn(
318
- hidden_states=hidden_states,
319
- attention_mask=attention_mask,
320
- position_ids=position_ids,
321
- past_key_values=past_key_values,
322
- use_cache=use_cache,
323
- position_embeddings=position_embeddings,
324
- **kwargs,
325
- )
326
- hidden_states = residual + hidden_states
327
-
328
- # Fully Connected
329
- residual = hidden_states
330
- hidden_states = self.post_attention_layernorm(hidden_states)
331
- hidden_states = self.mlp(hidden_states)
332
- hidden_states = residual + hidden_states
333
- return hidden_states
334
-
335
-
336
- @auto_docstring
337
- class ErkPreTrainedModel(PreTrainedModel):
338
- config: ErkConfig
339
- base_model_prefix = "model"
340
- supports_gradient_checkpointing = True
341
- _no_split_modules = ["ErkDecoderLayer"]
342
- _skip_keys_device_placement = ["past_key_values"]
343
- _supports_flash_attn = True
344
- _supports_sdpa = True
345
- _supports_flex_attn = True
346
-
347
- _can_compile_fullgraph = True
348
- _supports_attention_backend = True
349
- _can_record_outputs = {
350
- "hidden_states": ErkDecoderLayer,
351
- "attentions": ErkAttention,
352
- }
353
-
354
-
355
- @auto_docstring
356
- class ErkModel(ErkPreTrainedModel):
357
- def __init__(self, config: ErkConfig):
358
- super().__init__(config)
359
- self.padding_idx = config.pad_token_id
360
- self.vocab_size = config.vocab_size
361
-
362
- self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
363
- self.layers = nn.ModuleList(
364
- [ErkDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
365
- )
366
- self.norm = ErkRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
367
- self.rotary_emb = ErkRotaryEmbedding(config=config)
368
- self.gradient_checkpointing = False
369
- self.has_sliding_layers = "sliding_attention" in self.config.layer_types
370
-
371
- # Initialize weights and apply final processing
372
- self.post_init()
373
-
374
- @merge_with_config_defaults
375
- @capture_outputs
376
- @auto_docstring
377
- def forward(
378
- self,
379
- input_ids: torch.LongTensor | None = None,
380
- attention_mask: torch.Tensor | None = None,
381
- position_ids: torch.LongTensor | None = None,
382
- past_key_values: Cache | None = None,
383
- inputs_embeds: torch.FloatTensor | None = None,
384
- use_cache: bool | None = None,
385
- **kwargs: Unpack[TransformersKwargs],
386
- ) -> BaseModelOutputWithPast:
387
- if (input_ids is None) ^ (inputs_embeds is not None):
388
- raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
389
-
390
- if inputs_embeds is None:
391
- inputs_embeds = self.embed_tokens(input_ids)
392
-
393
- if use_cache and past_key_values is None:
394
- past_key_values = DynamicCache(config=self.config)
395
-
396
- if position_ids is None:
397
- past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
398
- position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
399
- position_ids = position_ids.unsqueeze(0)
400
-
401
- # It may already have been prepared by e.g. `generate`
402
- if not isinstance(causal_mask_mapping := attention_mask, dict):
403
- # Prepare mask arguments
404
- mask_kwargs = {
405
- "config": self.config,
406
- "inputs_embeds": inputs_embeds,
407
- "attention_mask": attention_mask,
408
- "past_key_values": past_key_values,
409
- "position_ids": position_ids,
410
- }
411
- # Create the masks
412
- causal_mask_mapping = {
413
- "full_attention": create_causal_mask(**mask_kwargs),
414
- }
415
- # The sliding window alternating layers are not always activated depending on the config
416
- if self.has_sliding_layers:
417
- causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
418
-
419
- hidden_states = inputs_embeds
420
- position_embeddings = self.rotary_emb(hidden_states, position_ids)
421
-
422
- for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
423
- hidden_states = decoder_layer(
424
- hidden_states,
425
- attention_mask=causal_mask_mapping[self.config.layer_types[i]],
426
- position_embeddings=position_embeddings,
427
- position_ids=position_ids,
428
- past_key_values=past_key_values,
429
- use_cache=use_cache,
430
- **kwargs,
431
- )
432
-
433
- hidden_states = self.norm(hidden_states)
434
- return BaseModelOutputWithPast(
435
- last_hidden_state=hidden_states,
436
- past_key_values=past_key_values if use_cache else None,
437
- )
438
-
439
-
440
- @auto_docstring
441
- class ErkForCausalLM(ErkPreTrainedModel, GenerationMixin):
442
- _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
443
- _tp_plan = {"lm_head": "colwise_gather_output"}
444
- _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
445
-
446
- def __init__(self, config):
447
- super().__init__(config)
448
- self.model = ErkModel(config)
449
- self.vocab_size = config.vocab_size
450
- self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
451
-
452
- # Initialize weights and apply final processing
453
- self.post_init()
454
-
455
- @can_return_tuple
456
- @auto_docstring
457
- def forward(
458
- self,
459
- input_ids: torch.LongTensor | None = None,
460
- attention_mask: torch.Tensor | None = None,
461
- position_ids: torch.LongTensor | None = None,
462
- past_key_values: Cache | None = None,
463
- inputs_embeds: torch.FloatTensor | None = None,
464
- labels: torch.LongTensor | None = None,
465
- use_cache: bool | None = None,
466
- logits_to_keep: int | torch.Tensor = 0,
467
- **kwargs: Unpack[TransformersKwargs],
468
- ) -> CausalLMOutputWithPast:
469
- r"""
470
- labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
471
- Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
472
- config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
473
- (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
474
-
475
- Example:
476
-
477
- ```python
478
- >>> from transformers import AutoTokenizer, ErkForCausalLM
479
-
480
- >>> model = ErkForCausalLM.from_pretrained("ecloudtech/erk")
481
- >>> tokenizer = AutoTokenizer.from_pretrained("ecloudtech/erk")
482
-
483
- >>> prompt = "Hey, are you conscious? Can you talk to me?"
484
- >>> inputs = tokenizer(prompt, return_tensors="pt")
485
-
486
- >>> # Generate
487
- >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
488
- >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
489
- "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
490
- ```"""
491
- outputs: BaseModelOutputWithPast = self.model(
492
- input_ids=input_ids,
493
- attention_mask=attention_mask,
494
- position_ids=position_ids,
495
- past_key_values=past_key_values,
496
- inputs_embeds=inputs_embeds,
497
- use_cache=use_cache,
498
- **kwargs,
499
- )
500
-
501
- hidden_states = outputs.last_hidden_state
502
- # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
503
- slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
504
- logits = self.lm_head(hidden_states[:, slice_indices, :])
505
-
506
- loss = None
507
- if labels is not None:
508
- loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
509
-
510
- return CausalLMOutputWithPast(
511
- loss=loss,
512
- logits=logits,
513
- past_key_values=outputs.past_key_values,
514
- hidden_states=outputs.hidden_states,
515
- attentions=outputs.attentions,
516
- )
517
-
518
-
519
- class ErkForSequenceClassification(GenericForSequenceClassification, ErkPreTrainedModel):
520
- pass
521
-
522
-
523
- class ErkForTokenClassification(GenericForTokenClassification, ErkPreTrainedModel):
524
- pass
525
-
526
-
527
- class ErkForQuestionAnswering(GenericForQuestionAnswering, ErkPreTrainedModel):
528
- base_model_prefix = "transformer" # For BC, where `transformer` was used instead of `model`
529
-
530
-
531
- __all__ = [
532
- "ErkForCausalLM",
533
- "ErkForQuestionAnswering",
534
- "ErkPreTrainedModel",
535
- "ErkModel",
536
- "ErkForSequenceClassification",
537
- "ErkForTokenClassification",
538
- ]