|
|
|
|
|
|
|
|
from copy import deepcopy |
|
|
from dataclasses import dataclass, field |
|
|
from typing import Literal, Optional, Union |
|
|
|
|
|
import torch |
|
|
from torch import nn |
|
|
|
|
|
from .blocks.mlstm.block import mLSTMBlock, mLSTMBlockConfig |
|
|
from .blocks.slstm.block import sLSTMBlock, sLSTMBlockConfig |
|
|
from .components.ln import LayerNorm |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class xLSTMBlockStackConfig: |
|
|
mlstm_block: Optional[mLSTMBlockConfig] = None |
|
|
slstm_block: Optional[sLSTMBlockConfig] = None |
|
|
|
|
|
context_length: int = -1 |
|
|
num_blocks: int = 1 |
|
|
embedding_dim: int = 128 |
|
|
add_post_blocks_norm: bool = True |
|
|
bias: bool = False |
|
|
dropout: float = 0.0 |
|
|
|
|
|
|
|
|
|
|
|
slstm_at: Union[list[int], Literal["all"]] = field(default_factory=list) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
_block_map: str = None |
|
|
|
|
|
@property |
|
|
def block_map(self) -> list[int]: |
|
|
return list(map(int, self._block_map.split(","))) |
|
|
|
|
|
def _create_block_map(self) -> str: |
|
|
"""Creates the block map, that specifies which block is used at which position.""" |
|
|
block_map = [0] * self.num_blocks |
|
|
|
|
|
for slstm_position_idx in self.slstm_at: |
|
|
assert slstm_position_idx < self.num_blocks, f"Invalid slstm position {slstm_position_idx}" |
|
|
block_map[slstm_position_idx] = 1 |
|
|
|
|
|
block_map_str = ",".join(map(str, block_map)) |
|
|
|
|
|
return block_map_str |
|
|
|
|
|
def __post_init__(self): |
|
|
if self.mlstm_block is None: |
|
|
self.slstm_at = "all" |
|
|
if self.slstm_at == "all": |
|
|
self.slstm_at = list(range(self.num_blocks)) |
|
|
|
|
|
if self.mlstm_block is not None: |
|
|
self.mlstm_block.mlstm.embedding_dim = self.embedding_dim |
|
|
self.mlstm_block.mlstm.bias = self.bias |
|
|
self.mlstm_block.mlstm.dropout = self.dropout |
|
|
self.mlstm_block.mlstm.context_length = self.context_length |
|
|
|
|
|
self.mlstm_block._num_blocks = self.num_blocks |
|
|
|
|
|
self.mlstm_block.__post_init__() |
|
|
|
|
|
if self.slstm_block is not None: |
|
|
self.slstm_block.slstm.dropout = self.dropout |
|
|
self.slstm_block.slstm.embedding_dim = self.embedding_dim |
|
|
self.slstm_block._num_blocks = self.num_blocks |
|
|
self.slstm_block.__post_init__() |
|
|
|
|
|
self._block_map = self._create_block_map() |
|
|
|
|
|
|
|
|
class xLSTMBlockStack(nn.Module): |
|
|
config_class = xLSTMBlockStackConfig |
|
|
|
|
|
def __init__(self, config: xLSTMBlockStackConfig): |
|
|
super().__init__() |
|
|
self.config = config |
|
|
|
|
|
self.blocks = self._create_blocks(config=config) |
|
|
if config.add_post_blocks_norm: |
|
|
self.post_blocks_norm = LayerNorm(ndim=config.embedding_dim) |
|
|
else: |
|
|
self.post_blocks_norm = nn.Identity() |
|
|
|
|
|
def _create_blocks(self, config: xLSTMBlockStackConfig): |
|
|
|
|
|
blocks = [] |
|
|
for block_idx, block_type_int in enumerate(config.block_map): |
|
|
if block_type_int == 0: |
|
|
config = deepcopy(self.config.mlstm_block) |
|
|
if hasattr(config, "_block_idx"): |
|
|
config._block_idx = block_idx |
|
|
config.__post_init__() |
|
|
blocks.append(mLSTMBlock(config=config)) |
|
|
elif block_type_int == 1: |
|
|
config = deepcopy(self.config.slstm_block) |
|
|
if hasattr(config, "_block_idx"): |
|
|
config._block_idx = block_idx |
|
|
config.__post_init__() |
|
|
blocks.append(sLSTMBlock(config=config)) |
|
|
else: |
|
|
raise ValueError(f"Invalid block type {block_type_int}") |
|
|
|
|
|
return nn.ModuleList(blocks) |
|
|
|
|
|
def reset_parameters(self) -> None: |
|
|
for block in self.blocks: |
|
|
block.reset_parameters() |
|
|
if not isinstance(self.post_blocks_norm, nn.Identity): |
|
|
self.post_blocks_norm.reset_parameters() |
|
|
|
|
|
def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: |
|
|
|
|
|
for block in self.blocks: |
|
|
x = block(x, **kwargs) |
|
|
|
|
|
x = self.post_blocks_norm(x) |
|
|
|
|
|
return x |
|
|
|
|
|
def step( |
|
|
self, x: torch.Tensor, state: dict[str, dict[str, tuple[torch.Tensor, ...]]] = None |
|
|
) -> tuple[torch.Tensor, dict[str, dict[str, tuple[torch.Tensor, ...]]]]: |
|
|
if state is None: |
|
|
state = {} |
|
|
|
|
|
for block_idx, block in enumerate(self.blocks): |
|
|
x, state[f"block_{block_idx}"] = block.step(x, **state.get(f"block_{block_idx}", {})) |
|
|
|
|
|
x = self.post_blocks_norm(x) |
|
|
|
|
|
return x, state |
|
|
|
