rxnn 0.1.0__py3-none-any.whl

src/transformers/layers.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn as nn
+from attention import MultiHeadAttention
+from ff import FeedForward, GatedFeedForward
+from moe import MoeFeedForward, GatedMoeFeedForward
+
+
+class ReactiveTransformerLayer(nn.Module):
+    """Reactive Transformer layer - extending the classic Transformer layer with Memory Cross-Attention"""
+
+    def __init__(
+            self,
+            embed_dim: int,
+            ff_dim: int,
+            self_attention: MultiHeadAttention,
+            memory_cross_attention: MultiHeadAttention,
+            use_rms_norm: bool = False,
+            use_post_norm: bool = False,
+            ff_activation: nn.Module = nn.GELU(),
+            ff_dropout: float = 0.1,
+            use_gated: bool = False,
+            use_moe: bool = False,
+            num_experts: int = 1,
+            moe_top_k: int = 1,
+            *args,
+            **kwargs,
+    ):
+        super(ReactiveTransformerLayer, self).__init__(*args, **kwargs)
+
+        self.attention = self_attention
+
+        self.memory_cross_attention = memory_cross_attention
+
+        if use_gated:
+            if use_moe:
+                self.ff = GatedMoeFeedForward(embed_dim, ff_dim, num_experts, ff_activation, top_k=moe_top_k,
+                                              dropout=ff_dropout)
+            else:
+                self.ff = GatedFeedForward(embed_dim, ff_dim, ff_activation, dropout=ff_dropout)
+        else:
+            if use_moe:
+                self.ff = MoeFeedForward(embed_dim, ff_dim, num_experts, ff_activation, top_k=moe_top_k,
+                                         dropout=ff_dropout)
+            else:
+                self.ff = FeedForward(embed_dim, ff_dim, ff_activation, dropout=ff_dropout)
+
+        if use_rms_norm:
+            self.norm1 = nn.RMSNorm(embed_dim)
+            self.norm2 = nn.RMSNorm(embed_dim)
+            self.norm3 = nn.RMSNorm(embed_dim)
+        else:
+            self.norm1 = nn.LayerNorm(embed_dim)
+            self.norm2 = nn.LayerNorm(embed_dim)
+            self.norm3 = nn.LayerNorm(embed_dim)
+        self.use_post_norm = use_post_norm
+
+    def trainable_cross_attention_(self, is_trainable: bool):
+        for param in self.memory_cross_attention.parameters():
+            param.requires_grad_(is_trainable)
+
+    def forward(self, x: torch.Tensor, stm: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        # First step, self-attention
+        residual = x
+        if not self.use_post_norm:
+            x = self.norm1(x)
+        x = self.attention(x, x, x, mask=mask)
+        x = residual + x
+        if self.use_post_norm:
+            x = self.norm1(x)
+        # Second step, Memory cross-attention
+        residual = x
+        if not self.use_post_norm:
+            x = self.norm2(x)
+        x = self.memory_cross_attention(x, stm, stm)
+        x = residual + x
+        if self.use_post_norm:
+            x = self.norm2(x)
+
+        # Third step, Feed Forward network
+        residual = x
+        if not self.use_post_norm:
+            x = self.norm3(x)
+        x = self.ff(x)
+        x = residual + x
+        if self.use_post_norm:
+            x = self.norm3(x)
+        return x
+
+
+class ClassicTransformerLayer(nn.Module):
+    """Classic Transformer layer - classic decoder-only/encoder-only Transformer layer with self-attention and Feed-Forward network."""
+
+    def __init__(
+            self,
+            embed_dim: int,
+            ff_dim: int,
+            self_attention: MultiHeadAttention,
+            use_rms_norm: bool = False,
+            use_post_norm: bool = False,
+            ff_activation: nn.Module = nn.GELU(),
+            ff_dropout: float = 0.1,
+            use_gated: bool = False,
+            use_moe: bool = False,
+            num_experts: int = 1,
+            moe_top_k: int = 1,
+            *args,
+            **kwargs,
+    ):
+        super(ClassicTransformerLayer, self).__init__(*args, **kwargs)
+
+        self.attention = self_attention
+
+        if use_gated:
+            if use_moe:
+                self.ff = GatedMoeFeedForward(embed_dim, ff_dim, num_experts, top_k=moe_top_k, dropout=ff_dropout)
+            else:
+                self.ff = GatedFeedForward(embed_dim, ff_dim, dropout=ff_dropout)
+        else:
+            if use_moe:
+                self.ff = MoeFeedForward(embed_dim, ff_dim, num_experts, ff_activation, top_k=moe_top_k,
+                                         dropout=ff_dropout)
+            else:
+                self.ff = FeedForward(embed_dim, ff_dim, ff_activation, dropout=ff_dropout)
+
+        if use_rms_norm:
+            self.norm1 = nn.RMSNorm(embed_dim)
+            self.norm2 = nn.RMSNorm(embed_dim)
+        else:
+            self.norm1 = nn.LayerNorm(embed_dim)
+            self.norm2 = nn.LayerNorm(embed_dim)
+        self.use_post_norm = use_post_norm
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        # First step, self-attention
+        residual = x
+        if not self.use_post_norm:
+            x = self.norm1(x)
+        x = self.attention(x, x, x, mask=mask)
+        x = residual + x
+        if self.use_post_norm:
+            x = self.norm1(x)
+        # Second step, Feed Forward network
+        residual = x
+        if not self.use_post_norm:
+            x = self.norm2(x)
+        x = self.ff(x)
+        x = residual + x
+        if self.use_post_norm:
+            x = self.norm2(x)
+        return x

src/transformers/mask.py

@@ -0,0 +1,10 @@
+import torch
+
+
+def create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
+    """Create a causal (lower triangular) attention mask for a given sequence length."""
+    # Create a lower triangular matrix of ones
+    mask = torch.tril(torch.ones((seq_len, seq_len), device=device))
+    # Expand the mask to have the shape (1, 1, seq_len, seq_len)
+    mask = mask.unsqueeze(0).unsqueeze(0)
+    return mask

src/transformers/models.py

@@ -0,0 +1,168 @@
+import torch
+import torch.nn as nn
+from positional import AbsolutePositionalEmbedding
+from mask import create_causal_mask
+from src.memory.stm import ShortTermMemory
+
+
+class ReactiveTransformerBase(nn.Module):
+    """Base class for Reactive Transformer models - common logic for both decoders and encoders."""
+
+    def __init__(
+            self,
+            stm: ShortTermMemory,
+            embedding: nn.Embedding,
+            own_layers: nn.ModuleList,
+            shared_layers: nn.ModuleList = None,
+            absolute_embedding: AbsolutePositionalEmbedding = None,
+            use_flash_attention: bool = False,
+            *args,
+            **kwargs,
+    ):
+        super(ReactiveTransformerBase, self).__init__(*args, **kwargs)
+
+        self.embedding = embedding
+        self.stm = stm
+        self.pos_embedding = absolute_embedding
+        self.use_flash_attention = use_flash_attention
+
+        self.shared_layers = shared_layers
+        self.layers = own_layers
+        self.num_shared_layers = len(shared_layers) if shared_layers else 0
+        self.num_own_layers = len(own_layers) if own_layers else 0
+
+    def trainable_cross_attention_(self, is_trainable: bool):
+        for i in range(self.num_shared_layers):
+            self.shared_layers[i].trainable_cross_attention_(is_trainable)
+        for i in range(self.num_own_layers):
+            self.layers[i].trainable_cross_attention_(is_trainable)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Shared logic for encoders and decoders - apply embeddings and positional encoding
+        x = self.embedding(x)
+        if self.pos_embedding is not None:
+            x = self.pos_embedding(x)
+        return x
+
+
+class ReactiveTransformerDecoder(ReactiveTransformerBase):
+    """Reactive Transformer decoder - extending the classic Transformer decoder with Memory Cross-Attention"""
+
+    def __init__(self, embed_dim: int, vocab_size: int, *args, **kwargs):
+        super(ReactiveTransformerDecoder, self).__init__(*args, **kwargs)
+        self.head = nn.Linear(embed_dim, vocab_size)
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
+        x = super().forward(x)  # apply embeddings
+        seq_len = x.size(1)
+        if not self.use_flash_attention:
+            mask = create_causal_mask(seq_len, device=x.device)
+            if attention_mask is not None:
+                mask &= attention_mask.unsqueeze(1).unsqueeze(1).bool()
+        elif attention_mask is not None:
+            mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
+        else:
+            mask = None
+        # Process shared layers
+        if self.shared_layers is not None:
+            for i in range(self.num_shared_layers):
+                layer_stm = self.stm(i)
+                x = self.shared_layers[i](x, layer_stm, mask=mask)
+        # Process own layers
+        for i in range(self.num_own_layers):
+            layer_stm = self.stm(i)
+            x = self.layers[i](x, layer_stm, mask=mask)
+        return self.head(x)
+
+
+class ReactiveTransformerEncoder(ReactiveTransformerBase):
+    """Reactive Transformer encoder - extending the classic Transformer encoder with Memory Cross-Attention"""
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+        x = super().forward(x)  # apply embeddings
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
+
+        hidden_states = []
+        # Process shared layers
+        if self.shared_layers is not None:
+            for i in range(self.num_shared_layers):
+                layer_stm = self.stm(i)
+                x = self.shared_layers[i](x, layer_stm, mask=attention_mask)
+                hidden_states.append(x)
+        # Process own layers
+        for i in range(self.num_own_layers):
+            layer_stm = self.stm(i)
+            x = self.layers[i](x, layer_stm, mask=attention_mask)
+            hidden_states.append(x)
+        return x, torch.stack(hidden_states)
+
+
+class ClassicTransformerBase(nn.Module):
+    """Base class for Classic Transformer models - common logic for both decoders and encoders."""
+
+    def __init__(
+            self,
+            embedding: nn.Embedding,
+            layers: nn.ModuleList,
+            absolute_embedding: AbsolutePositionalEmbedding = None,
+            use_flash_attention: bool = False,
+            *args,
+            **kwargs,
+    ):
+        super(ClassicTransformerBase, self).__init__(*args, **kwargs)
+
+        self.embedding = embedding
+        self.pos_embedding = absolute_embedding
+        self.use_flash_attention = use_flash_attention
+
+        self.layers = layers
+        self.num_layers = len(layers) if layers else 0
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Shared logic for encoders and decoders - apply embeddings and positional encoding
+        x = self.embedding(x)
+        if self.pos_embedding is not None:
+            x = self.pos_embedding(x)
+        return x
+
+
+class ClassicTransformerDecoder(ClassicTransformerBase):
+    """Classic Transformer decoder - for decoder-only Transformer models"""
+
+    def __init__(self, embed_dim: int, vocab_size: int, *args, **kwargs):
+        super(ClassicTransformerDecoder, self).__init__(*args, **kwargs)
+        self.head = nn.Linear(embed_dim, vocab_size)
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
+        x = super().forward(x)  # apply embeddings
+        seq_len = x.size(1)
+        if not self.use_flash_attention:
+            mask = create_causal_mask(seq_len, device=x.device)
+            if attention_mask is not None:
+                mask &= attention_mask.unsqueeze(1).unsqueeze(1).bool()
+        elif attention_mask is not None:
+            mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
+        else:
+            mask = None
+
+        # Process layers
+        for i in range(self.num_layers):
+            x = self.layers[i](x, mask=mask)
+        return self.head(x)
+
+
+class ClassicTransformerEncoder(ClassicTransformerBase):
+    """Classic Transformer encoder - for encoder-only Transformer models"""
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+        x = super().forward(x)  # apply embeddings
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
+
+        hidden_states = []
+        # Process own layers
+        for i in range(self.num_own_layers):
+            x = self.layers[i](x, mask=attention_mask)
+            hidden_states.append(x)
+        return x, torch.stack(hidden_states)

src/transformers/moe.py

@@ -0,0 +1,139 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MoeRouter(nn.Module):
+    """Mixture-of-Experts Router layer - computes routing weights for each expert."""
+
+    def __init__(self, embed_dim: int, num_experts: int, top_k: int = 1, *args, **kwargs):
+        super(MoeRouter, self).__init__(*args, **kwargs)
+        self.top_k = top_k
+        self.num_experts = num_experts
+        self.gate = nn.Linear(embed_dim, num_experts, bias=False)
+        self.aux_loss = 0.0  # For expert load balancing
+
+    def forward(self, x: torch.Tensor):
+        # x shape: [batch_size*seq_len, embed_dim]
+        logits = self.gate(x)
+        probs = F.softmax(logits, dim=-1)
+
+        # Expert load balancing loss
+        if self.training:
+            probs_for_bal = F.softmax(logits, dim=0)
+            self.aux_loss = (probs_for_bal.mean(dim=0) *
+                             torch.log(probs_for_bal.mean(dim=0) + 1e-9)).sum()
+
+        top_k_weights, top_k_indices = probs.topk(self.top_k, dim=-1)
+        top_k_weights = top_k_weights / (top_k_weights.sum(dim=-1, keepdim=True) + 1e-9)
+
+        return top_k_weights, top_k_indices
+
+
+class MoeFeedForward(nn.Module):
+    """Mixture-of-Experts Feed-Forward layer - combines multiple experts into a single model."""
+
+    def __init__(
+            self,
+            embed_dim: int,
+            hidden_dim: int,
+            num_experts: int,
+            activation: nn.Module,
+            top_k: int = 1,
+            dropout: float = 0.0,
+            *args,
+            **kwargs
+    ):
+        super(MoeFeedForward, self).__init__(*args, **kwargs)
+        self.embed_dim = embed_dim
+        self.num_experts = num_experts
+        self.top_k = top_k
+
+        self.router = MoeRouter(embed_dim, num_experts, top_k)
+
+        # Batch all expert parameters together
+        self.w1 = nn.Parameter(torch.empty(num_experts, embed_dim, self._w1_dim_factor(hidden_dim)))
+        self.b1 = nn.Parameter(torch.zeros(num_experts, self._w1_dim_factor(hidden_dim)))
+        self.w2 = nn.Parameter(torch.empty(num_experts, hidden_dim, embed_dim))
+        self.b2 = nn.Parameter(torch.zeros(num_experts, embed_dim))
+        self.activation = activation
+        self.dropout = nn.Dropout(dropout)
+
+        # Initialize parameters
+        self._init_linear_parameters()
+        nn.init.zeros_(self.b1)
+        nn.init.zeros_(self.b2)
+
+    def _init_linear_parameters(self):
+        nn.init.kaiming_normal_(self.w1, nonlinearity='relu')
+        nn.init.kaiming_normal_(self.w2, nonlinearity='relu')
+
+    def _w1_dim_factor(self, hidden_dim: int) -> int:
+        return hidden_dim
+
+    def _activate(self, h: torch.Tensor):
+        return self.activation(h)
+
+    def forward(self, x: torch.Tensor):
+        orig_shape = x.shape
+        x = x.view(-1, self.embed_dim)  # [batch*seq_len, embed_dim]
+
+        # Get routing weights and indices
+        weights, indices = self.router(x)  # [batch*seq_len, top_k]
+
+        # Create expert masks and combine it with masks
+        mask = F.one_hot(indices, self.num_experts).float()  # [batch*seq_len, top_k, num_experts]
+        weights = (weights.unsqueeze(-1) * mask).sum(dim=1)  # [batch*seq_len, num_experts]
+
+        # Expert computation
+        x = x.unsqueeze(1).expand(-1, self.num_experts, -1)  # [batch*seq_len, num_experts, embed_dim]
+
+        # First linear layer
+        h = torch.einsum('bie,ieh->bih', x, self.w1) + self.b1  # [batch*seq_len, num_experts, hidden_dim]
+        h = self._activate(h)
+        h = self.dropout(h)
+
+        # Second linear layer (projection back to embed_dim)
+        out = torch.einsum('bih,ihe->bie', h, self.w2) + self.b2  # [batch*seq_len, num_experts, embed_dim]
+
+        # Weighted sum of expert outputs
+        out = (out * weights.unsqueeze(-1)).sum(dim=1)  # [batch*seq_len, embed_dim]
+
+        return out.view(*orig_shape)
+
+
+class GatedMoeFeedForward(MoeFeedForward):
+    """Gated Mixture-of-Experts Feed-Forward layer - enable GLU-based activations for MoE"""
+
+    def __init__(
+            self,
+            embed_dim: int,
+            hidden_dim: int,
+            num_experts: int,
+            activation: nn.Module = nn.SiLU(),
+            top_k: int = 1,
+            dropout: float = 0.1,
+            *args,
+            **kwargs
+    ):
+        super(GatedMoeFeedForward, self).__init__(
+            embed_dim=embed_dim,
+            hidden_dim=hidden_dim,
+            num_experts=num_experts,
+            activation=activation,
+            top_k=top_k,
+            dropout=dropout,
+            *args,
+            **kwargs
+        )
+
+    def _init_linear_parameters(self):
+        nn.init.kaiming_normal_(self.w1, nonlinearity='relu')
+        nn.init.kaiming_normal_(self.w2, nonlinearity='linear')
+
+    def _w1_dim_factor(self, hidden_dim: int) -> int:
+        return 2 * hidden_dim
+
+    def _activate(self, h: torch.Tensor):
+        a, b = h.chunk(2, dim=-1)
+        return a * self.activation(b)

src/transformers/positional.py

@@ -0,0 +1,105 @@
+import torch
+from torch import nn
+import math
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    """Rotary Positional Embedding (RoPE) layer - recommended for positional encoding"""
+
+    def __init__(self, dim: int, max_seq_len: int = 4096, base: int = 10000, *args, **kwargs):
+        super(RotaryPositionalEmbedding, self).__init__(*args, **kwargs)
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.base = base
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+        self.register_buffer('cache', None, persistent=False)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        device = q.device
+        seq_len = q.size(-2)
+        # Prepare RoPE Frequencies
+        freqs = self._prepare_freqs(seq_len, device)
+
+        # Apply the rotation to the queries
+        q_embed = self._rotate(q, freqs)
+        # Apply the rotation to the keys
+        k_embed = self._rotate(k, freqs)
+
+        return q_embed, k_embed
+
+    def forward_one(self, q: torch.Tensor) -> torch.Tensor:
+        device = q.device
+        seq_len = q.size(-2)
+        # Prepare RoPE Frequencies
+        freqs = self._prepare_freqs(seq_len, device)
+
+        # Apply the rotation to the queries
+        q_embed = self._rotate(q, freqs)
+
+        return q_embed
+
+    def _prepare_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        if self.cache is None or self.cache.size(1) < seq_len:
+            t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
+            freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+            self.cache = freqs
+            return freqs[None, None, :, :]
+        else:
+            return self.cache[None, None, :, :]
+
+    def _rotate(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+        x1 = x[..., 0::2]
+        x2 = x[..., 1::2]
+        # Apply the rotation
+        x_rotated1 = x1 * torch.cos(freqs) - x2 * torch.sin(freqs)
+        x_rotated2 = x1 * torch.sin(freqs) + x2 * torch.cos(freqs)
+        # Concatenate the rotated parts back together
+        x_rotated = torch.cat((x_rotated1, x_rotated2), dim=-1)
+        return x_rotated
+
+
+class AbsolutePositionalEmbedding(nn.Module):
+    """Absolute Positional Embedding layer (legacy) - not recommended for memory-augmented Reactive Transformers"""
+
+    def __init__(self, max_seq_len: int, embed_dim: int, *args, **kwargs):
+        super(AbsolutePositionalEmbedding, self).__init__(*args, **kwargs)
+        self.max_seq_len = max_seq_len
+        self.embed_dim = embed_dim
+        self.position_embeddings = nn.Embedding(max_seq_len, embed_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Create position indices
+        seq_len = x.size(1)
+        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(x.size(0), -1)
+        # Get position embeddings
+        pos_embeddings = self.position_embeddings(positions)
+        # Add position embeddings to the input embeddings
+        return x + pos_embeddings
+
+
+class RelativePositionalEmbedding(nn.Module):
+    """Relative Positional Embedding layer (legacy) - not compatible with Flash Attention and not recommended for positional encoding"""
+
+    def __init__(self, max_seq_len: int, embed_dim: int, *args, **kwargs):
+        super(RelativePositionalEmbedding, self).__init__(*args, **kwargs)
+        self.max_seq_len = max_seq_len
+        self.embed_dim = embed_dim
+        self.position_embeddings = nn.Embedding(2 * max_seq_len - 1, embed_dim)
+        self.embed_dim_sqrt = math.sqrt(embed_dim)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
+        q_len = q.size(2)
+        k_len = k.size(2)
+
+        # Create relative position indices
+        indices = torch.arange(q_len, device=q.device)[:, None] - torch.arange(k_len, device=k.device)[None, :]
+        indices += self.max_seq_len - 1  # Shift to non-negative
+        indices = torch.clamp(indices, 0, 2 * self.max_seq_len - 2)
+
+        # Get embeddings
+        rel_emb = self.position_embeddings(indices)
+
+        rel_emb = rel_emb.permute(2, 0, 1)
+        rel_pos_bias = torch.einsum('bhqd, dqk -> bhqk', q, rel_emb)
+        return rel_pos_bias / self.embed_dim_sqrt

src/transformers/sampler.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Iterator
+
+def sample(
+    logits: torch.Tensor,
+    temperature: float = 1.0,
+    top_k: int = None,
+    top_p: float = None,
+) -> torch.Tensor:
+    if temperature <= 0:
+        raise ValueError("Temperature must be > 0")
+
+    # Apply temperature
+    logits = logits / temperature
+
+    # Apply top-k filtering
+    if top_k is not None and top_k > 0:
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = float('-inf')
+
+    # Apply top-p (nucleus) sampling
+    if top_p is not None and 0 < top_p <= 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above threshold
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # Shift right to keep first token above threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # Scatter sorted indices back to original positions
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            dim=-1,
+            index=sorted_indices,
+            src=sorted_indices_to_remove
+        )
+        logits[indices_to_remove] = float('-inf')
+
+    # Convert to probabilities
+    probs = F.softmax(logits, dim=-1)
+
+    # Sample from distribution
+    return torch.multinomial(probs, num_samples=1)
+
+class Sampler:
+    def __init__(self, model: nn.Module, device: torch.device, end_token_id: int):
+        self.model = model.to(device)
+        self.device = device
+        self.end_token_id = end_token_id
+
+    def _generate_token(
+        self,
+        input_ids: torch.Tensor,
+        temperature: float,
+        top_k: int,
+        top_p: float ,
+        attention_mask: torch.Tensor,
+    ) -> tuple[int, torch.Tensor, torch.Tensor]:
+        # Forward pass to get next token logits
+        outputs = self.model(input_ids, attention_mask=attention_mask)
+        next_token_logits = outputs[:, -1, :]  # Get logits for next token
+        # Apply sampling
+        next_token = sample(
+            next_token_logits,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+        )
+        next_token = next_token.item()  # Extract scalar token
+        next_token_ten = torch.tensor([[next_token]], device=self.device)
+        next_input_ids = torch.cat([input_ids, next_token_ten], dim=1)
+        new_one = torch.ones(1, 1, dtype=torch.bool, device=self.device)
+        next_mask = torch.cat([attention_mask, new_one], dim=1) if attention_mask is not None else None
+        # Yield the generated token
+        return (
+            next_token,
+            next_input_ids,
+            next_mask
+        )
+
+    def __call__(
+        self,
+        initial_tokens: torch.Tensor,
+        temperature: float = 1.0,
+        top_k: int = None,
+        top_p: float = None,
+        max_seq_len: int = 50,
+        attention_mask: torch.Tensor = None,
+        no_grad: bool = True,
+    ) -> Iterator[int]:
+        # Convert initial tokens to tensor and move to device
+        input_ids = initial_tokens
+
+        if no_grad:
+            with torch.no_grad():
+                for _ in range(max_seq_len):
+                    next_token, input_ids, attention_mask = self._generate_token(input_ids, temperature, top_k, top_p, attention_mask)
+                    yield next_token
+                    if next_token == self.end_token_id:
+                        break
+        else:
+            for _ in range(max_seq_len):
+                next_token, input_ids, attention_mask = self._generate_token(input_ids, temperature, top_k, top_p, attention_mask)
+                yield next_token
+                if next_token == self.end_token_id:
+                    break

src/utils.py

@@ -0,0 +1,14 @@
+import torch
+
+def human_format(num: int):
+    num = float('{:.3g}'.format(num))
+    magnitude = 0
+    while abs(num) >= 1000:
+        magnitude += 1
+        num /= 1000.0
+    return '{}{}'.format('{:f}'.format(num).rstrip('0').rstrip('.'), ['', 'K', 'M', 'B', 'T'][magnitude])
+
+
+def get_model_size(model: torch.nn.Module):
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return f'Model params {human_format(trainable_params)}'