Upload modeling_limon.py

modeling_limon.py

@@ -3,7 +3,13 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import PreTrainedModel
-from .configuration_limon import LimonConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+# Умный импорт: работает и локально, и на Hugging Face
+try:
+    from .configuration_limon import LimonConfig
+except ImportError:
+    from configuration_limon import LimonConfig
 
 class TimeConditionedAttention(nn.Module):
     def __init__(self, config):
@@ -30,7 +36,7 @@ class TimeConditionedAttention(nn.Module):
 class VectorFieldV2(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.anchor_strength = config.anchor_strength
+        self.anchor_strength = getattr(config, "anchor_strength", 0.1)
         self.ln1 = nn.LayerNorm(config.hidden_size, elementwise_affine=False)
         self.attn = TimeConditionedAttention(config)
         self.ln2 = nn.LayerNorm(config.hidden_size, elementwise_affine=False)
@@ -72,29 +78,91 @@ class ODESolverV2(nn.Module):
 
 class LimonFlowV1Model(PreTrainedModel):
     config_class = LimonConfig
+    
+    # Жесткий запрет на попытки HF создать DynamicCache
+    _supports_cache_class = False
+
     def __init__(self, config):
         super().__init__(config)
         self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.pos_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
-        self.ode_solver = ODESolverV2(VectorFieldV2(config), config.integration_steps)
+        
+        max_pos = getattr(config, "max_position_embeddings", getattr(config, "max_seq_len", 256))
+        self.pos_embeddings = nn.Embedding(max_pos, config.hidden_size)
+        
+        steps = getattr(config, "integration_steps", 6)
+        self.ode_solver = ODESolverV2(VectorFieldV2(config), steps)
         self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        
+        # ХАК для обхода внутренних проверок HF
+        self.config.num_hidden_layers = 1
+        
         self.post_init()
 
-    def forward(self, input_ids, labels=None, attention_mask=None, **kwargs):
-        batch_size, seq_len = input_ids.shape
-        pos = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
-        x = self.embeddings(input_ids) + self.pos_embeddings(pos)
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings = value
+
+    def forward(
+        self, 
+        input_ids=None, 
+        attention_mask=None, 
+        inputs_embeds=None, 
+        labels=None, 
+        past_key_values=None, 
+        use_cache=None, 
+        output_attentions=None, 
+        output_hidden_states=None, 
+        return_dict=None, 
+        **kwargs
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        if input_ids is not None:
+            batch_size, seq_len = input_ids.shape
+            device = input_ids.device
+            x = self.embeddings(input_ids)
+        elif inputs_embeds is not None:
+            batch_size, seq_len, _ = inputs_embeds.shape
+            device = inputs_embeds.device
+            x = inputs_embeds
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        pos = torch.arange(seq_len, device=device).unsqueeze(0)
+        x = x + self.pos_embeddings(pos)
+        
         x = self.ode_solver(x)
         logits = self.head(x)
         
         loss = None
         if labels is not None:
-            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1))
+            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))
+        
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        # ИСПОЛЬЗУЕМ ПРАВИЛЬНЫЙ КЛАСС (WithPast)
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
 
-        return {
-            "logits": logits,
-            "loss": loss
-        }
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs):
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
 
-    def prepare_inputs_for_generation(self, input_ids, **kwargs):
-        return {"input_ids": input_ids}
+        model_inputs.update({
+            "attention_mask": attention_mask,
+            "use_cache": False,
+        })
+        return model_inputs