Upload folder using huggingface_hub

generation_config.json

@@ -1,6 +1,6 @@
 {
-  "_from_model_config": true,
   "bos_token_id": 151643,
   "eos_token_id": 151643,
+  "max_new_tokens": 2048,
   "transformers_version": "4.57.3"
 }

lance_ai_model.py

@@ -1,11 +1,19 @@
 import math
+from functools import partial
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
 from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
-from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING
+from transformers.cache_utils import Cache, DynamicCache
+
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, logging
+from transformers.modeling_layers import GradientCheckpointingLayer
 
 class LanceAIConfig(PretrainedConfig):
     model_type = "lance_ai"
@@ -52,6 +60,7 @@ class LanceAIConfig(PretrainedConfig):
         self.bos_token_id = bos_token_id
         self.eos_token_id = eos_token_id
 
+
 class LanceAIRMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-6):
         super().__init__()
@@ -65,6 +74,7 @@ class LanceAIRMSNorm(nn.Module):
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         return self.weight * hidden_states.to(input_dtype)
 
+
 class LanceAIRotaryEmbedding(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -84,11 +94,13 @@ class LanceAIRotaryEmbedding(nn.Module):
         sin = emb.sin()
         return cos.to(dtype=x.dtype), sin.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):
     cos = cos.unsqueeze(1)
     sin = sin.unsqueeze(1)
@@ -96,6 +108,7 @@ def apply_rotary_pos_emb(q, k, cos, sin):
     k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed
 
+
 def repeat_kv(hidden_states, n_rep):
     batch, num_key_value_heads, slen, head_dim = hidden_states.shape
     if n_rep == 1:
@@ -103,6 +116,33 @@ def repeat_kv(hidden_states, n_rep):
     hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
     return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = F.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
 class LanceAIAttention(nn.Module):
     def __init__(self, config, layer_idx):
         super().__init__()
@@ -121,39 +161,50 @@ class LanceAIAttention(nn.Module):
         self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True)
         self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
 
-    def forward(self, hidden_states, attention_mask=None, position_embeddings=None, past_key_values=None, use_cache=False):
-        batch_size, seq_len, _ = hidden_states.shape
-        query_states = self.q_proj(hidden_states).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = self.k_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
-        value_states = self.v_proj(hidden_states).view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ):
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
 
         cos, sin = position_embeddings
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
 
         if past_key_values is not None:
-            if past_key_values[0] is not None:
-                key_states = torch.cat([past_key_values[0], key_states], dim=2)
-            if past_key_values[1] is not None:
-                value_states = torch.cat([past_key_values[1], value_states], dim=2)
-
-        past = (key_states, value_states) if use_cache else None
-
-        key_states = repeat_kv(key_states, self.num_key_value_groups)
-        value_states = repeat_kv(value_states, self.num_key_value_groups)
-
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling
-
-        if attention_mask is not None:
-            attn_weights = attn_weights + attention_mask
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            if self.config._attn_implementation in ALL_ATTENTION_FUNCTIONS:
+                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=getattr(self, 'sliding_window', None),
+            **kwargs,
+        )
 
-        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
-        attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
-        attn_output = torch.matmul(attn_weights, value_states)
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.reshape(batch_size, seq_len, -1)
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
         attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
 
-        return attn_output, past
 
 class LanceAIMLP(nn.Module):
     def __init__(self, config):
@@ -167,7 +218,8 @@ class LanceAIMLP(nn.Module):
     def forward(self, x):
         return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
 
-class LanceAIDecoderLayer(nn.Module):
+
+class LanceAIDecoderLayer(GradientCheckpointingLayer):
     def __init__(self, config, layer_idx):
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -176,15 +228,24 @@ class LanceAIDecoderLayer(nn.Module):
         self.input_layernorm = LanceAIRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = LanceAIRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-    def forward(self, hidden_states, attention_mask=None, position_embeddings=None, past_key_values=None, use_cache=False):
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        position_embeddings=None,
+        past_key_values=None,
+        cache_position=None,
+        **kwargs,
+    ):
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
-        hidden_states, past_kv = self.self_attn(
-            hidden_states,
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
             attention_mask=attention_mask,
             position_embeddings=position_embeddings,
             past_key_values=past_key_values,
-            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
         )
         hidden_states = residual + hidden_states
 
@@ -193,7 +254,8 @@ class LanceAIDecoderLayer(nn.Module):
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
 
-        return hidden_states, past_kv
+        return hidden_states
+
 
 class LanceAIPreTrainedModel(PreTrainedModel):
     config_class = LanceAIConfig
@@ -203,6 +265,9 @@ class LanceAIPreTrainedModel(PreTrainedModel):
     _skip_keys_device_placement = ["past_key_values"]
     _supports_flash_attn = True
     _supports_sdpa = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = True
+
 
 class LanceAIModel(LanceAIPreTrainedModel):
     def __init__(self, config):
@@ -220,49 +285,30 @@ class LanceAIModel(LanceAIPreTrainedModel):
         self.gradient_checkpointing = False
         self.post_init()
 
-    def _past_seq_length(self, past_key_values):
-        """Get sequence length from past cache (tuple or DynamicCache)"""
-        if past_key_values is None:
-            return 0
-        if hasattr(past_key_values, 'get_seq_length'):
-            return past_key_values.get_seq_length()
-        if len(past_key_values) > 0 and past_key_values[0] is not None:
-            if past_key_values[0][0] is not None:
-                return past_key_values[0][0].shape[2]
-        return 0
-
-    def _convert_past(self, past_key_values, use_cache):
-        """Convert DynamicCache to our tuple format if needed"""
-        if past_key_values is None or not use_cache:
-            return past_key_values, use_cache
-        if isinstance(past_key_values, list):
-            return past_key_values, use_cache
-        # DynamicCache -> list of (k, v) tuples
-        if hasattr(past_key_values, 'get_seq_length'):
-            converted = []
-            for i in range(len(self.layers)):
-                if i < len(past_key_values):
-                    kv = past_key_values[i]
-                    converted.append((kv[0], kv[1]))
-                else:
-                    converted.append((None, None))
-            return converted, use_cache
-        return past_key_values, use_cache
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=None):
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        use_cache=None,
+        **kwargs,
+    ):
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("Specify exactly one of input_ids or inputs_embeds")
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
-        batch_size, seq_len = inputs_embeds.shape[:2]
 
-        past_key_values, use_cache = self._convert_past(past_key_values, use_cache)
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
 
         if position_ids is None:
-            past_seen_tokens = self._past_seq_length(past_key_values)
-            position_ids = torch.arange(seq_len, device=inputs_embeds.device) + past_seen_tokens
-            position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
 
+        cache_position = position_ids[0]
         position_embeddings = self.rotary_emb(inputs_embeds, position_ids)
 
         if attention_mask is not None:
@@ -274,31 +320,31 @@ class LanceAIModel(LanceAIPreTrainedModel):
             attention_mask = self._make_causal_mask(inputs_embeds, past_key_values)
 
         hidden_states = inputs_embeds
-        new_past = [] if use_cache else None
-
         for i, layer in enumerate(self.layers):
-            layer_past = past_key_values[i] if past_key_values is not None else None
-            hidden_states, layer_past = layer(
+            hidden_states = layer(
                 hidden_states,
                 attention_mask=attention_mask,
                 position_embeddings=position_embeddings,
-                past_key_values=layer_past,
-                use_cache=use_cache,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                **kwargs,
             )
-            if use_cache:
-                new_past.append(layer_past)
 
         hidden_states = self.norm(hidden_states)
-        return hidden_states, new_past
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
 
-    def _make_causal_mask(self, inputs_embeds, past_key_values=None):
-        batch_size, seq_len, _ = inputs_embeds.shape
-        past_len = self._past_seq_length(past_key_values)
+    def _make_causal_mask(self, input_embeds, past_key_values=None):
+        batch_size, seq_len, _ = input_embeds.shape
+        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
         total_len = past_len + seq_len
-        mask = torch.full((seq_len, total_len), float('-inf'), device=inputs_embeds.device)
+        mask = torch.full((seq_len, total_len), float('-inf'), device=input_embeds.device)
         mask = torch.triu(mask, diagonal=1 + past_len)
         return mask[None, None, :, :]
 
+
 class LanceAI(LanceAIPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
 
@@ -310,45 +356,47 @@ class LanceAI(LanceAIPreTrainedModel, GenerationMixin):
 
         self.post_init()
 
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, **kwargs):
-        hidden_states, past = self.model(
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        logits_to_keep=0,
+        **kwargs,
+    ):
+        outputs = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
             position_ids=position_ids,
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
             use_cache=use_cache,
+            **kwargs,
         )
 
-        logits = self.lm_head(hidden_states)
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
 
         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, shift_logits.size(-1)),
-                shift_labels.view(-1),
-                ignore_index=-100,
-            )
-
-        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=past)
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
 
-    def _past_seq_len(self, past_key_values):
-        if past_key_values is None:
-            return 0
-        if hasattr(past_key_values, 'get_seq_length'):
-            return past_key_values.get_seq_length()
-        if isinstance(past_key_values, list) and len(past_key_values) > 0:
-            if past_key_values[0] is not None:
-                k, v = past_key_values[0]
-                if k is not None:
-                    return k.shape[2]
-        return 0
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+        )
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
-        if self._past_seq_len(past_key_values) > 0:
-            input_ids = input_ids[:, -1:]
+        if past_key_values is not None:
+            past_length = past_key_values.get_seq_length() if hasattr(past_key_values, 'get_seq_length') else 0
+            if past_length > 0:
+                input_ids = input_ids[:, -1:]
         return {
             "input_ids": input_ids,
             "attention_mask": attention_mask,
@@ -357,13 +405,17 @@ class LanceAI(LanceAIPreTrainedModel, GenerationMixin):
         }
 
     def _reorder_cache(self, past_key_values, beam_idx):
+        if hasattr(past_key_values, 'reorder_cache'):
+            past_key_values.reorder_cache(beam_idx)
+            return past_key_values
         reordered = []
         for layer_past in past_key_values:
             layer_k, layer_v = layer_past
             reordered.append((layer_k.index_select(0, beam_idx), layer_v.index_select(0, beam_idx)))
         return reordered
 
+
 CONFIG_MAPPING.register("lance_ai", LanceAIConfig)
 MODEL_FOR_CAUSAL_LM_MAPPING.register(LanceAIConfig, LanceAI)
 LanceAIConfig.register_for_auto_class("AutoConfig")
-LanceAI.register_for_auto_class("AutoModelForCausalLM")
+LanceAI.register_for_auto_class("AutoModelForCausalLM")