Upload modeling_alinlight.py with huggingface_hub

modeling_alinlight.py

@@ -0,0 +1,245 @@
+# -*- coding: utf-8 -*-
+# Copyright 2026 EngineerGL Research.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List, Union
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
+from transformers import GenerationMixin
+from configuration_alinlight import AlinlightConfig # Импортируем конфиг из соседнего файла
+
+class AlinlightRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+    def forward(self, x):
+        input_dtype = x.dtype
+        x = x.to(torch.float32)
+        variance = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return self.weight * x.to(input_dtype)
+
+class AlinlightRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        super().__init__()
+        self.dim = dim
+        self.base = base
+        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):
+        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)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        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)
+        return self.cos_cached[:seq_len].to(dtype=x.dtype), self.sin_cached[:seq_len].to(dtype=x.dtype)
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+class AlinlightMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+class AlinlightAttention(nn.Module):
+    def __init__(self, config, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.sliding_window = config.sliding_window
+        self.attention_dropout = config.attention_dropout
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, output_attentions=False, use_cache=False, cos_sin=None):
+        bsz, q_len, _ = hidden_states.size()
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if cos_sin is not None:
+            cos, sin = cos_sin
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        if past_key_value is not None:
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        # Truncation logic for sliding window
+        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
+        
+        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).reshape(bsz, self.num_heads, key_states.shape[-2], self.head_dim)
+            value_states = value_states[:, :, None, :, :].expand(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)
+
+        # Use Scaled Dot Product Attention
+        attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=None, dropout_p=0.0, is_causal=True)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
+        return self.o_proj(attn_output), None, past_key_value
+
+class AlinlightDecoderLayer(nn.Module):
+    def __init__(self, config, layer_idx: int):
+        super().__init__()
+        self.self_attn = AlinlightAttention(config, layer_idx=layer_idx)
+        self.mlp = AlinlightMLP(config)
+        self.input_layernorm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, output_attentions=False, use_cache=False, cos_sin=None):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _, present_key_value = self.self_attn(hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cos_sin)
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states, None, present_key_value
+
+class AlinlightModel(PreTrainedModel):
+    config_class = AlinlightConfig
+    def __init__(self, config: AlinlightConfig):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        self.layers = nn.ModuleList([AlinlightDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+        self.norm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        
+        scaling_factor = 1.0
+        if config.rope_scaling and config.rope_scaling.get("type") == "linear":
+             scaling_factor = config.rope_scaling.get("factor", 1.0)
+             
+        self.rotary_emb = AlinlightRotaryEmbedding(config.hidden_size // config.num_attention_heads, max_position_embeddings=config.max_position_embeddings, base=config.rope_theta, scaling_factor=scaling_factor)
+
+    def forward(self, input_ids=None, past_key_values=None, use_cache=None, **kwargs):
+        if input_ids is not None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        else:
+             inputs_embeds = kwargs.get("inputs_embeds")
+
+        seq_len = inputs_embeds.shape[1]
+        if past_key_values is not None:
+             seq_len += past_key_values[0][0].shape[2]
+             
+        cos, sin = self.rotary_emb(inputs_embeds, seq_len=seq_len)
+        
+        position_ids = kwargs.get("position_ids")
+        if position_ids is None:
+             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)
+
+        hidden_states = inputs_embeds
+        next_decoder_cache = () if use_cache else None
+
+        for idx, layer in enumerate(self.layers):
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+            layer_outputs = layer(hidden_states, position_ids=position_ids, past_key_value=past_key_value, use_cache=use_cache, cos_sin=(cos, sin))
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2],)
+
+        hidden_states = self.norm(hidden_states)
+        
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache
+        )
+
+class AlinlightForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = AlinlightConfig
+    _keys_to_ignore_on_load_missing = ["model.rotary_emb.inv_freq"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = AlinlightModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.model.embed_tokens.weight
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        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:
+                 position_ids = torch.arange(input_ids.shape[1], dtype=torch.long, device=input_ids.device).unsqueeze(0)
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": True,
+            "position_ids": position_ids
+        }
+
+    def forward(self, input_ids=None, past_key_values=None, labels=None, **kwargs):
+        outputs = self.model(input_ids=input_ids, past_key_values=past_key_values, **kwargs)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
+
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values)