Upload modeling_xlstm.py with huggingface_hub

modeling_xlstm.py

@@ -0,0 +1,214 @@
+from typing import Optional, Sequence, Tuple, Union
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutputWithPast
+from xlstm.components.init import small_init_init_
+from xlstm.utils import WeightDecayOptimGroupMixin
+from xlstm.xlstm_block_stack import xLSTMBlockStack as _xLSTMBlockStack
+
+from .configuration_xlstm import xLSTMConfig
+
+
+class xLSTMPreTrainedModel(PreTrainedModel):
+    """Base class for all models."""
+
+    config_class = xLSTMConfig
+
+
+class xLSTMBlockStack(_xLSTMBlockStack):
+    """Small wrapper to expose hidden states"""
+
+    def forward(
+        self, x: torch.Tensor, **kwargs
+    ) -> Tuple[torch.Tensor, Sequence[torch.Tensor]]:
+        hidden_states = ()
+        for block in self.blocks:
+            x = block(x, **kwargs)
+            hidden_states += (x,)
+
+        x = self.post_blocks_norm(x)
+
+        return x, hidden_states
+
+
+class xLSTMModel(xLSTMPreTrainedModel):
+    def __init__(self, config: xLSTMConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.token_embedding = nn.Embedding(
+            num_embeddings=config.vocab_size, embedding_dim=config.embedding_dim
+        )
+        _config = config.to_xlstm_config()
+
+        self.emb_dropout = (
+            nn.Dropout(_config.dropout)
+            if _config.add_embedding_dropout
+            else nn.Identity()
+        )
+
+        self.xlstm_block_stack = xLSTMBlockStack(config=_config)
+
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        output_hidden_states: Optional[bool] = None,
+        return_dict=Optional[bool],
+    ) -> Union[Tuple, BaseModelOutput]:
+        token_embedding = self.token_embedding(input_ids)
+        x = self.emb_dropout(token_embedding)
+        x, hidden_states = self.xlstm_block_stack(x)
+
+        if output_hidden_states:
+            hidden_states = (token_embedding,) + hidden_states
+
+        if not return_dict:
+            return x, hidden_states
+
+        return BaseModelOutput(
+            last_hidden_state=x,
+            hidden_states=hidden_states if output_hidden_states else None,
+        )
+
+
+class xLSTMForCausalLM(xLSTMPreTrainedModel, WeightDecayOptimGroupMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: xLSTMConfig, **kwargs):
+        super().__init__(config)
+        self.config = config
+        self.vocab_size = config.vocab_size
+
+        self.model = xLSTMModel(config)
+
+        self.lm_head = nn.Linear(
+            in_features=config.embedding_dim,
+            out_features=config.vocab_size,
+            bias=False,
+        )
+
+        self.post_init()
+        # TODO: Add option for up-projection
+
+    def get_input_embeddings(self):
+        return self.model.token_embedding
+
+    def set_input_embeddings(self, value: nn.Module):
+        self.model.token_embedding = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def reset_parameters(self):
+        self.model.xlstm_block_stack.reset_parameters()
+
+        small_init_init_(
+            self.get_input_embeddings().weight, dim=self.config.embedding_dim
+        )
+
+        if not self.config.tie_word_embeddings:
+            small_init_init_(
+                self.get_output_embeddings().weight, dim=self.config.embedding_dim
+            )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output = self.model(
+            input_ids,
+            output_hidden_states=output_hidden_states,
+        )
+
+        hidden_state = output[0]
+
+        logits = self.lm_head(hidden_state)
+        logits = logits.float()
+
+        loss = None
+
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            loss_fct = nn.CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + output[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            hidden_states=output.hidden_states,
+        )
+
+    def step(
+        self,
+        idx: torch.Tensor,
+        state: dict[str, dict[str, tuple[torch.Tensor, ...]]] = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, dict[str, dict[str, tuple[torch.Tensor, ...]]]]:
+        x = self.token_embedding(idx)
+        x = self.emb_dropout(x)
+        x, state = self.xlstm_block_stack.step(x, state=state, **kwargs)
+        logits = self.lm_head(x)
+        return logits, state
+
+    def _create_weight_decay_optim_groups(
+        self, **kwargs
+    ) -> tuple[Sequence[nn.Parameter], Sequence[nn.Parameter]]:
+        weight_decay, no_weight_decay = super()._create_weight_decay_optim_groups(
+            **kwargs
+        )
+        # remove token embedding and add it to the correct group, accrording to the config
+        weight_decay = list(weight_decay)
+        removed = 0
+        for idx in range(len(weight_decay)):
+            if weight_decay[idx - removed] is self.get_input_embeddings().weight:
+                weight_decay.pop(idx - removed)
+                removed += 1
+        weight_decay = tuple(weight_decay)
+
+        # TODO: Fix this
+        # if self.config.weight_decay_on_embedding:
+        if True:
+            weight_decay += (self.get_input_embeddings().weight,)
+        else:
+            no_weight_decay += (self.get_input_embeddings().weight,)
+
+        return weight_decay, no_weight_decay
+
+    def resize_token_embeddings(self, new_num_tokens: int) -> nn.Embedding:
+        new_embeddings = nn.Embedding(
+            new_num_tokens, self.token_embedding.embedding_dim
+        )
+        self.token_embedding = new_embeddings.to(self.device)
+        return new_embeddings
+
+    def tie_weights(self):
+        self.get_output_embeddings().weight = self.get_input_embeddings().weight
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        **kwargs,
+    ):
+        model_inputs = {
+            "input_ids": input_ids.to(self.device),
+        }
+        return model_inputs