Update modeling_easyformer.py
Browse files- modeling_easyformer.py +6 -17
modeling_easyformer.py
CHANGED
|
@@ -52,6 +52,7 @@ class EasyFormerBlock(nn.Module):
|
|
| 52 |
x = x + self.ffn(self.ln2(x))
|
| 53 |
return x
|
| 54 |
|
|
|
|
| 55 |
class EasyFormerPreTrainedModel(PreTrainedModel):
|
| 56 |
config_class = EasyFormerConfig
|
| 57 |
base_model_prefix = "easyformer"
|
|
@@ -63,6 +64,10 @@ class EasyFormerLMHeadModel(EasyFormerPreTrainedModel, GenerationMixin):
|
|
| 63 |
base_model_prefix = "easyformer"
|
| 64 |
_tied_weights_keys = ["lm_head.weight"]
|
| 65 |
all_tied_weights_keys = {"lm_head.weight": "tok_emb.weight"}
|
|
|
|
|
|
|
|
|
|
|
|
|
| 66 |
|
| 67 |
def __init__(self, config):
|
| 68 |
super().__init__(config)
|
|
@@ -75,7 +80,6 @@ class EasyFormerLMHeadModel(EasyFormerPreTrainedModel, GenerationMixin):
|
|
| 75 |
)
|
| 76 |
self.ln_f = RMSNorm(config.d_model)
|
| 77 |
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
|
| 78 |
-
# weight tying
|
| 79 |
self.lm_head.weight = self.tok_emb.weight
|
| 80 |
self.post_init()
|
| 81 |
|
|
@@ -92,24 +96,9 @@ class EasyFormerLMHeadModel(EasyFormerPreTrainedModel, GenerationMixin):
|
|
| 92 |
def set_output_embeddings(self, new_embeddings):
|
| 93 |
self.lm_head = new_embeddings
|
| 94 |
|
| 95 |
-
def tie_weights(self):
|
| 96 |
self.lm_head.weight = self.tok_emb.weight
|
| 97 |
|
| 98 |
-
def resize_token_embeddings(self, new_num_tokens=None):
|
| 99 |
-
# при необходимости — пересоздать эмбеддинги
|
| 100 |
-
if new_num_tokens is None or new_num_tokens == self.cfg.vocab_size:
|
| 101 |
-
return self.tok_emb
|
| 102 |
-
old = self.tok_emb
|
| 103 |
-
new = nn.Embedding(new_num_tokens, self.cfg.d_model).to(old.weight.device)
|
| 104 |
-
n = min(old.num_embeddings, new_num_tokens)
|
| 105 |
-
with torch.no_grad():
|
| 106 |
-
new.weight[:n] = old.weight[:n]
|
| 107 |
-
self.tok_emb = new
|
| 108 |
-
self.lm_head = nn.Linear(self.cfg.d_model, new_num_tokens, bias=False).to(old.weight.device)
|
| 109 |
-
self.lm_head.weight = self.tok_emb.weight
|
| 110 |
-
self.cfg.vocab_size = new_num_tokens
|
| 111 |
-
return self.tok_emb
|
| 112 |
-
|
| 113 |
# --- forward ---
|
| 114 |
def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
|
| 115 |
B, T = input_ids.shape
|
|
|
|
| 52 |
x = x + self.ffn(self.ln2(x))
|
| 53 |
return x
|
| 54 |
|
| 55 |
+
# ------------------------- HF-ОБЁРТКА -----------------------------
|
| 56 |
class EasyFormerPreTrainedModel(PreTrainedModel):
|
| 57 |
config_class = EasyFormerConfig
|
| 58 |
base_model_prefix = "easyformer"
|
|
|
|
| 64 |
base_model_prefix = "easyformer"
|
| 65 |
_tied_weights_keys = ["lm_head.weight"]
|
| 66 |
all_tied_weights_keys = {"lm_head.weight": "tok_emb.weight"}
|
| 67 |
+
_supports_cache_class = False
|
| 68 |
+
_supports_flash_attn_2 = False
|
| 69 |
+
_supports_sdpa = False
|
| 70 |
+
main_input_name = "input_ids"
|
| 71 |
|
| 72 |
def __init__(self, config):
|
| 73 |
super().__init__(config)
|
|
|
|
| 80 |
)
|
| 81 |
self.ln_f = RMSNorm(config.d_model)
|
| 82 |
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
|
|
|
|
| 83 |
self.lm_head.weight = self.tok_emb.weight
|
| 84 |
self.post_init()
|
| 85 |
|
|
|
|
| 96 |
def set_output_embeddings(self, new_embeddings):
|
| 97 |
self.lm_head = new_embeddings
|
| 98 |
|
| 99 |
+
def tie_weights(self, recompute_mapping=False, **kwargs):
|
| 100 |
self.lm_head.weight = self.tok_emb.weight
|
| 101 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 102 |
# --- forward ---
|
| 103 |
def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
|
| 104 |
B, T = input_ids.shape
|