Update modeling_alinlight.py

modeling_alinlight.py

@@ -13,11 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# -*- coding: utf-8 -*-
-# Copyright 2026 EngineerGL Research.
-
-# -*- coding: utf-8 -*-
-# Copyright 2026 EngineerGL Research.
 
 import math
 import torch
@@ -49,15 +44,15 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.max_position_embeddings = max_position_embeddings
         self.scaling_factor = scaling_factor
         
-        # Вычисляем частоты сразу
+
         inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32).to(device) / self.dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         
-        # Инициализируем кэш
+ 
         self._set_cos_sin_cache(seq_len=max_position_embeddings, device=device, dtype=torch.get_default_dtype())
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
-        # RoPE кэш считаем во float32 для точности, потом кастим
+ 
         t = torch.arange(seq_len, device=device, dtype=torch.int64).type_as(self.inv_freq)
         t = t / self.scaling_factor
         freqs = torch.outer(t, self.inv_freq)
@@ -66,7 +61,7 @@ class AlinlightRotaryEmbedding(nn.Module):
         self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
 
     def forward(self, x, seq_len=None):
-        # Если вдруг длина последовательности больше кэша (редкий случай)
+
         if seq_len > self.cos_cached.shape[0]:
              self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
         
@@ -145,15 +140,14 @@ class AlinlightAttention(nn.Module):
             key_states = torch.cat([past_key_value[0], key_states], dim=2)
             value_states = torch.cat([past_key_value[1], value_states], dim=2)
 
-        # === TPU/Efficiency Optimization: Physical Truncation ===
-        # Вместо сложной маски просто обрезаем старые токены, вышедшие за пределы окна.
+
         if self.sliding_window is not None and key_states.shape[2] > self.sliding_window:
             key_states = key_states[:, :, -self.sliding_window:, :]
             value_states = value_states[:, :, -self.sliding_window:, :]
 
         past_key_value = (key_states, value_states) if use_cache else None
         
-        # GQA / MQA Expansion
+ 
         if self.num_key_value_groups > 1:
             key_states = key_states[:, :, None, :, :].expand(
                 bsz, self.num_key_value_heads, self.num_key_value_groups, key_states.shape[-2], self.head_dim
@@ -163,10 +157,7 @@ class AlinlightAttention(nn.Module):
                 bsz, self.num_key_value_heads, self.num_key_value_groups, value_states.shape[-2], self.head_dim
             ).reshape(bsz, self.num_heads, value_states.shape[-2], self.head_dim)
 
-        # === Fix for "Blind" Generation ===
-        # is_causal=True нужно ТОЛЬКО при обучении (или prefill), когда q_len > 1.
-        # Когда мы генерируем по 1 токену (q_len == 1), мы смотрим на весь кэш (прошлое).
-        # Если поставить True при q_len=1, SDPA может замаскировать всё.
+
         is_causal = q_len > 1
 
         attn_output = F.scaled_dot_product_attention(
@@ -232,17 +223,17 @@ class AlinlightModel(PreTrainedModel):
         else:
              inputs_embeds = kwargs.get("inputs_embeds")
 
-        # Определяем длину последовательности для RoPE
+        
         seq_len = inputs_embeds.shape[1]
         if past_key_values is not None:
              seq_len += past_key_values[0][0].shape[2]
              
-        # Получаем cos/sin
+        
         cos, sin = self.rotary_emb(inputs_embeds, seq_len=seq_len)
         
         position_ids = kwargs.get("position_ids")
         if position_ids is None:
-             # Генерируем position_ids сами
+             
              position_ids = torch.arange(seq_len - inputs_embeds.shape[1], seq_len, dtype=torch.long, device=inputs_embeds.device)
              position_ids = position_ids.unsqueeze(0).expand(inputs_embeds.shape[0], -1)
 
@@ -278,10 +269,6 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
         self.model = AlinlightModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         
-        # === ВАЖНЕЙШЕЕ ИСПРАВЛЕНИЕ ===
-        # Связываем веса ТОЛЬКО если это разрешено в конфиге.
-        # Раньше это было принудительно, что ломало генерацию для моделей типа Mistral/Llama3,
-        # где tie_word_embeddings=False.
         if config.tie_word_embeddings:
             self.lm_head.weight = self.model.embed_tokens.weight
             
@@ -293,18 +280,18 @@ class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
     def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
-        # Если есть кэш, нам нужен только последний токен из input_ids
+
         if past_key_values:
             input_ids = input_ids[:, -1:]
         
         position_ids = kwargs.get("position_ids", None)
         if position_ids is None:
             if past_key_values:
-                 # Если есть кэш, позиция = длина кэша
+
                  past_length = past_key_values[0][0].shape[2]
                  position_ids = torch.tensor([[past_length]], dtype=torch.long, device=input_ids.device)
             else:
-                 # Если кэша нет (первый токен), позиции = [0, 1, 2...]
+
                  position_ids = torch.arange(input_ids.shape[1], dtype=torch.long, device=input_ids.device).unsqueeze(0)
 
         return {