File size: 2,455 Bytes
e8ac551 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | from __future__ import annotations
import torch
import torch.nn.functional as F
from torch import nn
from transformers import GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput
from .configuration_hanse import HanseConfig
from .modeling_hanse_layers import HanseBlock, HanseRMSNorm, initialize_weights
class HanseForCausalLM(PreTrainedModel, GenerationMixin):
config_class = HanseConfig
base_model_prefix = "hanse"
_tied_weights_keys = ["lm_head.weight"]
_supports_assign_param_buffer = False
def __init__(self, config: HanseConfig) -> None:
super().__init__(config)
self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
self.blocks = nn.ModuleList(
HanseBlock(config, kind) for kind in config.layer_pattern
)
self.final_norm = HanseRMSNorm(config.hidden_size, config.norm_eps)
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
self.apply(lambda module: initialize_weights(module, config.num_layers))
self.tie_weights()
def get_input_embeddings(self) -> nn.Embedding:
return self.token_embedding
def set_input_embeddings(self, value: nn.Embedding) -> None:
self.token_embedding = value
def get_output_embeddings(self) -> nn.Linear:
return self.lm_head
def set_output_embeddings(self, value: nn.Linear) -> None:
self.lm_head = value
def forward(
self,
input_ids: torch.Tensor,
labels: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
use_cache: bool = False,
**_: object,
) -> CausalLMOutput:
del attention_mask, use_cache
if input_ids.ndim != 2:
raise ValueError("input_ids must have shape [batch, sequence]")
if input_ids.size(1) > self.config.max_seq_len:
raise ValueError("input exceeds max_seq_len")
hidden = self.token_embedding(input_ids)
for block in self.blocks:
hidden = block(hidden)
logits = self.lm_head(self.final_norm(hidden))
loss = None
if labels is not None:
loss = F.cross_entropy(
logits[:, :-1].float().reshape(-1, self.config.vocab_size),
labels[:, 1:].reshape(-1),
ignore_index=-100,
)
return CausalLMOutput(loss=loss, logits=logits)
|