adjusts

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2dc6c386af412163c51f18f97152117040b6464f9e64159ef464d50471ceda1c
-size 1101168184
+oid sha256:a469bc2dde18f9b248e83bf2d86d1d23fa32b8ad9646317d58e9d513b37b3120
+size 832480488

src/pre-training.py

@@ -35,7 +35,7 @@ if __name__ == "__main__":
     dataloader = create_train_dataloader(
         folder_path,
         tokenizer,
-        batch_size=20,
+        batch_size=5,
         max_length=1024,
         drop_last=True,
         num_workers=10

src/tynerox/modeling.py

@@ -9,108 +9,37 @@ from typing import Optional, Literal, Union, Tuple
 
 
 class PositionalEncoding(nn.Module):
-    """
-    Implements positional encoding (sinusoidal or rotary).
-    """
-    def __init__(
-        self,
-        embed_dim: int,
-        context_length: int,
-        dropout: float = 0.1,
-        encoding_type: Literal['sinusoidal', 'rotary'] = 'rotary',
-    ):
+    def __init__(self, embed_dim, context_length):
         super().__init__()
-        if embed_dim <= 0 or context_length <= 0:
-            raise ValueError("embed_dim and context_length must be positive integers")
-        if not 0 <= dropout < 1:
-            raise ValueError("dropout must be between 0 and 1")
-
-        self.dropout = nn.Dropout(dropout)
-        self.encoding_type = encoding_type.lower()
-        self.max_seq_len = context_length
+        if embed_dim % 2 != 0:
+            raise ValueError("embed_dim must be even for rotary")
         self.embed_dim = embed_dim
+        self._build_table(context_length)
 
-        if self.encoding_type == 'sinusoidal':
-            pe = self._create_sinusoidal_embeddings(context_length, embed_dim)
-            self.register_buffer('pe', pe.unsqueeze(0), persistent=True)
-        elif self.encoding_type == 'rotary':
-            if embed_dim % 2 != 0:
-                raise ValueError("embed_dim must be even for rotary encoding")
-            # inv_freq of size D/2
-            inv_freq = 1.0 / (10000 ** (torch.arange(0, embed_dim, 2).float() / embed_dim))
-            self.register_buffer('inv_freq', inv_freq, persistent=True)
-        else:
-            raise ValueError("Unsupported encoding_type: 'sinusoidal' or 'rotary'")
+    def _build_table(self, length):
+        # inv_freq[j] = 1 / 10000^(2j/embed_dim)
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, self.embed_dim, 2).float() / self.embed_dim))
+        positions = torch.arange(length).float().unsqueeze(1)  # [L,1]
+        freqs = positions * inv_freq.unsqueeze(0)             # [L, D/2]
+        self.register_buffer('sin_table', freqs.sin(), persistent=True)
+        self.register_buffer('cos_table', freqs.cos(), persistent=True)
 
-    def _create_sinusoidal_embeddings(self, max_seq_len: int, dim: int) -> torch.Tensor:
-        position = torch.arange(max_seq_len).unsqueeze(1).float()
-        div_term = torch.exp(torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim))
-        pe = torch.zeros(max_seq_len, dim)
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        return pe
+    def apply_rotary(self, x, sin, cos):
+        # x: [B, T, D], sin/cos: [1, T, D/2]
+        x_pairs = x.view(*x.shape[:-1], -1, 2)   # [..., D/2, 2]
+        x1, x2 = x_pairs[..., 0], x_pairs[..., 1]
+        y1 = x1 * cos - x2 * sin
+        y2 = x1 * sin + x2 * cos
+        x_rot = torch.stack([y1, y2], dim=-1)   # [..., D/2, 2]
+        return x_rot.flatten(-2)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x shape: [B, T, D]
-        if self.encoding_type == 'sinusoidal':
-            seq_len = x.size(1)
-            x = x + self.pe[:, :seq_len, :]
-        else:
-            # rotary: split even/odd dims and apply rotary
-            seq_len = x.size(1)
-            positions = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
-            # freqs of shape [T, D/2]
-            freqs = torch.einsum('i , j -> i j', positions, self.inv_freq)
-            x = self.apply_rotary(x, freqs)
-        return self.dropout(x)
-
-    @staticmethod
-    def apply_rotary(x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
-        # x: [B, T, D], emb: [T, D/2]
-        x1, x2 = x.chunk(2, dim=-1)              # each [B, T, D/2]
-        emb_sin = emb.sin()[None, :, :]          # [1, T, D/2]
-        emb_cos = emb.cos()[None, :, :]          # [1, T, D/2]
-        # apply rotary
-        rotated1 = x1 * emb_cos + x2 * emb_sin
-        rotated2 = x2 * emb_cos - x1 * emb_sin
-        return torch.cat([rotated1, rotated2], dim=-1)  # [B, T, D]
-
-
-class PositionalEmbedding(nn.Module):
-    """
-    Combines token embedding with positional encoding.
-    """
-    def __init__(
-        self,
-        vocab_size: int,
-        embed_dim: int,
-        context_length: int,
-        dropout: float = 0.05,
-        encoding_type: Literal['sinusoidal', 'rotary'] = 'rotary'
-    ):
-        super().__init__()
-        if vocab_size <= 0 or embed_dim <= 0 or context_length <= 0:
-            raise ValueError("vocab_size, embed_dim, context_length must be > 0")
-
-        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
-        self.scale = math.sqrt(embed_dim)
-        self.pos_encoding = PositionalEncoding(
-            embed_dim=embed_dim,
-            context_length=context_length,
-            dropout=dropout,
-            encoding_type=encoding_type
-        )
-
-    def forward(self, input_ids: torch.LongTensor) -> torch.Tensor:
-        # input_ids: [B, T]
-        x = self.token_embedding(input_ids) * self.scale  # [B, T, D]
-        return self.pos_encoding(x)
 
 
 def get_alibi_slopes(n_heads: int) -> torch.Tensor:
     def _get_slopes(n):
         base = 2 ** (-8.0 / n)
         return torch.tensor([base ** (i + 1) for i in range(n)])
+
     if math.log2(n_heads).is_integer():
         return _get_slopes(n_heads)
     m = 2 ** math.floor(math.log2(n_heads))
@@ -118,6 +47,7 @@ def get_alibi_slopes(n_heads: int) -> torch.Tensor:
     extra = _get_slopes(2 * m)[::2][: n_heads - m]
     return torch.cat([slopes, extra], dim=0)
 
+
 # -----------------------------------------------------------------------------
 # Feed-Forward
 # -----------------------------------------------------------------------------
@@ -135,6 +65,7 @@ class FeedForward(nn.Module):
         x_up, x_gate = x_fc1.chunk(2, dim=-1)
         return self.fc2(x_up * self.activation(x_gate))
 
+
 # -----------------------------------------------------------------------------
 # Attention-Free Transformer (AFT) Simple
 # -----------------------------------------------------------------------------
@@ -143,6 +74,7 @@ class AFTSimple(nn.Module):
     def __init__(
         self,
         embed_dim: int,
+        max_position_embeddings: int,
         activation=torch.sigmoid,
         causal: bool = True,
     ):
@@ -151,6 +83,9 @@ class AFTSimple(nn.Module):
         self.causal = causal
         self.activation = activation
 
+        # Rotary PE (dropout=0 para não afetar Q/K)
+        self.rotary = PositionalEncoding(embed_dim, max_position_embeddings)
+
         self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=False)
         self.project = nn.Linear(embed_dim, embed_dim)
 
@@ -159,53 +94,75 @@ class AFTSimple(nn.Module):
         x: torch.Tensor,
         past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        # x: [B, T_new, D]
+        """
+        x: [B, T_new, D]
+        past_key_values: (K_past, V_past), cada um [B, T_past, D]
+        """
         B, T_new, D = x.shape
-        if D != self.embed_dim:
-            raise ValueError(f"Input dim ({D}) != embed_dim ({self.embed_dim})")
+        assert D == self.embed_dim, f"Embedding dimension mismatch: {D} != {self.embed_dim}"
 
-        qkv = self.qkv(x)                 # [B, T_new, 3*D]
-        Q, K_new, V_new = qkv.chunk(3, dim=-1)  # each [B, T_new, D]
+        # 1) Projeções lineares
+        qkv = self.qkv(x)               # [B, T_new, 3*D]
+        Q, K_new, V_new = qkv.chunk(3, dim=-1)
 
-        # concatenate past if provided
+        # 2) Calcular sin/cos para as posições de Q/K_new
+        #    Se houver cache, shift nas posições; senão, 0..T_new-1
         if past_key_values is not None:
             K_past, V_past = past_key_values
-            K = torch.cat([K_past, K_new], dim=1)  # [B, T_all, D]
-            V = torch.cat([V_past, V_new], dim=1)
+            T_past = K_past.size(1)
         else:
-            K, V = K_new, V_new
+            T_past = 0
 
-        # compute attention-free aggregate
-        softmax_k = F.softmax(K, dim=1)          # [B, T_all, D]
-        weighted_v = softmax_k * V               # [B, T_all, D]
+        # obtem sin/cos para posições [T_past .. T_past+T_new-1]
+        device, dtype = Q.device, Q.dtype
+        pos = torch.arange(T_past, T_past + T_new, device=device)
+        sin = self.rotary.sin_table[pos].unsqueeze(0).to(dtype=dtype)  # [1, T_new, D/2]
+        cos = self.rotary.cos_table[pos].unsqueeze(0).to(dtype=dtype)
+
+        # 3) Aplicar RoPE em Q e K_new
+        Q_rot     = self.rotary.apply_rotary(Q,     sin, cos)  # [B, T_new, D]
+        K_new_rot = self.rotary.apply_rotary(K_new, sin, cos)  # [B, T_new, D]
+
+        # 4) Concatena cache (já rotacionado) com K_new_rot
+        if past_key_values is not None:
+            K = torch.cat([K_past, K_new_rot], dim=1)  # [B, T_all, D]
+            V = torch.cat([V_past, V_new],   dim=1)
+        else:
+            K, V = K_new_rot, V_new
+
+        # 5) Atenção “free” sobre V
+        softmax_k = F.softmax(K, dim=1)               # [B, T_all, D]
+        weighted_v = softmax_k * V                    # [B, T_all, D]
 
         if self.causal:
-            context = torch.cumsum(weighted_v, dim=1)  # [B, T_all, D]
+            context = torch.cumsum(weighted_v, dim=1) # [B, T_all, D]
         else:
-            total = weighted_v.sum(dim=1, keepdim=True)  # [B, 1, D]
-            context = total.expand(-1, K.size(1), -1)    # [B, T_all, D]
+            total = weighted_v.sum(dim=1, keepdim=True)  # [B,1,D]
+            context = total.expand(-1, K.size(1), -1)    # [B,T_all,D]
+
+        # 6) Seleciona apenas as posições novas
+        context_new = context[:, -T_new:, :]           # [B, T_new, D]
 
-        # slice only the new positions
-        context_new = context[:, -T_new:, :]     # [B, T_new, D]
-        gate = self.activation(Q)                # [B, T_new, D]
-        Y = gate * context_new                   # [B, T_new, D]
-        Y = self.project(Y)                      # [B, T_new, D]
+        # 7) Gating e projeção final
+        gate = self.activation(Q_rot)                  # [B, T_new, D]
+        Y = gate * context_new                         # [B, T_new, D]
+        Y = self.project(Y)                            # [B, T_new, D]
 
-        # return output and updated cache
         return Y, (K, V)
 
+
 # -----------------------------------------------------------------------------
 # Flash Attention with ALiBi and KV-cache
 # -----------------------------------------------------------------------------
 
 class FlashAttention(nn.Module):
     def __init__(
-        self,
-        embed_dim: int,
-        num_heads: int,
-        window_size: int,
-        causal: bool = True,
-        qkv_bias: bool = False,
+            self,
+            embed_dim: int,
+            num_heads: int,
+            window_size: int,
+            causal: bool = True,
+            qkv_bias: bool = False,
     ):
         super().__init__()
         assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
@@ -221,9 +178,9 @@ class FlashAttention(nn.Module):
         self.register_buffer('alibi', get_alibi_slopes(num_heads))
 
     def forward(
-        self,
-        x: torch.Tensor,
-        past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+            self,
+            x: torch.Tensor,
+            past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
     ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         # x: [B, T_new, D]
         B, T_new, _ = x.size()
@@ -246,11 +203,12 @@ class FlashAttention(nn.Module):
             return_attn_probs=False,
         )
         # attn_out: [B, T_new, H, Dh]
-        out = attn_out.contiguous().view(B, T_new, -1)   # [B, T_new, D]
-        y = self.out_proj(out)                           # [B, T_new, D]
+        out = attn_out.contiguous().view(B, T_new, -1)  # [B, T_new, D]
+        y = self.out_proj(out)  # [B, T_new, D]
 
         return y, (k, v)
 
+
 # -----------------------------------------------------------------------------
 # Transformer Blocks and Model
 # -----------------------------------------------------------------------------
@@ -259,7 +217,7 @@ class TransformerBlock(nn.Module):
     def __init__(self, config, att_global: bool = True):
         super().__init__()
         if att_global:
-            self.attn = AFTSimple(embed_dim=config.d_model, causal=config.causal)
+            self.attn = AFTSimple(embed_dim=config.d_model, causal=config.causal, max_position_embeddings=config.max_position_embeddings)
         else:
             self.attn = FlashAttention(
                 embed_dim=config.d_model,
@@ -268,18 +226,18 @@ class TransformerBlock(nn.Module):
                 causal=config.causal,
                 qkv_bias=True,
             )
+
         self.ff = nn.Sequential(
             FeedForward(config.d_model),
-            FeedForward(config.d_model),
         )
         self.ln1 = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
         self.ln2 = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
         self.drop = nn.Dropout(config.dropout)
 
     def forward(
-        self,
-        x: torch.Tensor,
-        past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+            self,
+            x: torch.Tensor,
+            past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
     ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
         # Attention + residual
         residual = x
@@ -295,6 +253,7 @@ class TransformerBlock(nn.Module):
 
         return x, present
 
+
 class ResidualBlocks(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -307,10 +266,10 @@ class ResidualBlocks(nn.Module):
         self.final_ln = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
 
     def forward(
-        self,
-        x: torch.Tensor,
-        past_key_values: Optional[Tuple[Tuple[torch.Tensor,torch.Tensor], ...]] = None
-    ) -> Tuple[torch.Tensor, Tuple[Tuple[torch.Tensor,torch.Tensor], ...]]:
+            self,
+            x: torch.Tensor,
+            past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None
+    ) -> Tuple[torch.Tensor, Tuple[Tuple[torch.Tensor, torch.Tensor], ...]]:
         new_past = []
         for i, layer in enumerate(self.layers):
             pkv = None if past_key_values is None else past_key_values[i]
@@ -319,6 +278,7 @@ class ResidualBlocks(nn.Module):
         x = self.final_ln(x)
         return x, tuple(new_past)
 
+
 # -----------------------------------------------------------------------------
 # Configuration and Model
 # -----------------------------------------------------------------------------
@@ -327,19 +287,19 @@ class TyneRoxConfig(PretrainedConfig):
     model_type = "tynerox"
 
     def __init__(
-        self,
-        vocab_size: int = 30522,
-        context_length: int = 2048,
-        d_model: int = 1024,
-        num_heads: int = 16,
-        window_size: int = 512,
-        num_hidden_layers: int = 12,
-        causal: bool = True,
-        dropout: float = 0.1,
-        layer_norm_eps: float = 1e-5,
-        tie_word_embeddings: bool = False,
-        pad_token_id:int = 0,
-        **kwargs
+            self,
+            vocab_size: int = 30522,
+            context_length: int = 2048,
+            d_model: int = 1024,
+            num_heads: int = 16,
+            window_size: int = 512,
+            num_hidden_layers: int = 12,
+            causal: bool = True,
+            dropout: float = 0.1,
+            layer_norm_eps: float = 1e-5,
+            tie_word_embeddings: bool = False,
+            pad_token_id: int = 0,
+            **kwargs
     ):
         super().__init__(**kwargs)
         self.vocab_size = vocab_size
@@ -354,27 +314,29 @@ class TyneRoxConfig(PretrainedConfig):
         self.tie_word_embeddings = tie_word_embeddings
         self.pad_token_id = pad_token_id
 
+
 class TyneRoxModel(PreTrainedModel, GenerationMixin):
     config_class = TyneRoxConfig
 
     def __init__(self, config: TyneRoxConfig):
         super().__init__(config)
-        self.embed = PositionalEmbedding(
-            config.vocab_size,
-            config.d_model,
-            config.max_position_embeddings,
-            dropout=config.dropout,
-            encoding_type='rotary'
-        )
+        self.scale = math.sqrt(config.d_model)
+        self.embed = nn.Embedding(config.vocab_size, config.d_model)
         self.transformer = ResidualBlocks(config)
         self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        # Correction #2: tie embeddings if requested
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.embed.weight
+
         self.post_init()
 
+    # Correction #3: fix get/set input embeddings to use .weight
     def get_input_embeddings(self):
-        return self.embed.token_embedding
+        return self.embed
 
     def set_input_embeddings(self, value):
-        self.embed.token_embedding = value
+        self.embed = value
 
     def get_output_embeddings(self):
         return self.lm_head
@@ -383,24 +345,18 @@ class TyneRoxModel(PreTrainedModel, GenerationMixin):
         self.lm_head = value
 
     def forward(
-        self,
-        input_ids: torch.LongTensor,
-        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: bool = True,
-        return_dict: bool = True,
-        **kwargs
+            self,
+            input_ids: torch.LongTensor,
+            past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: bool = True,
+            return_dict: bool = True,
+            **kwargs
     ) -> Union[Tuple, CausalLMOutputWithPast]:
-        # 1) Embeddings
-        x = self.embed(input_ids)  # [B, T, D]
-
-        # 2) Transformer blocks with KV-cache
+        x = self.embed(input_ids) * self.scale
         x, new_past = self.transformer(x, past_key_values=past_key_values)
+        logits = self.lm_head(x)
 
-        # 3) Project to vocabulary logits
-        logits = self.lm_head(x)  # [B, T, V]
-
-        # 4) Compute loss if labels provided
         loss = None
         if labels is not None:
             shift_logits = logits[:, :-1, :].contiguous()
@@ -411,7 +367,6 @@ class TyneRoxModel(PreTrainedModel, GenerationMixin):
                 ignore_index=-100,
             )
 
-        # 5) Return standardized output
         if not return_dict:
             output = (logits, new_past) if use_cache else (logits,)
             return ((loss,) + output) if loss is not None else output
@@ -429,21 +384,19 @@ class TyneRoxModel(PreTrainedModel, GenerationMixin):
     ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
         reordered = []
         for k, v in past_key_values:
-            # ambos têm batch dim = dim 0
             reordered.append((k.index_select(0, beam_idx),
                               v.index_select(0, beam_idx)))
         return tuple(reordered)
 
     def prepare_inputs_for_generation(
-        self,
-        input_ids: torch.LongTensor,
-        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
-        **kwargs
+            self,
+            input_ids: torch.LongTensor,
+            past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+            **kwargs
     ) -> dict:
-        # at generation time, only feed in the last token
         if past_key_values is not None:
             input_ids = input_ids[:, -1:].contiguous()
         return {
             "input_ids": input_ids,
             "past_key_values": past_key_values,
-        }
+        }