use fast tokenizer, fix transformers v5 inference issues

modeling_deepseek.py

@@ -44,7 +44,11 @@ from transformers.utils import (add_start_docstrings,
                                 is_flash_attn_2_available,
                                 is_flash_attn_greater_or_equal_2_10, logging,
                                 replace_return_docstrings)
-from transformers.utils.import_utils import is_torch_fx_available
+try:
+    from transformers.utils.import_utils import is_torch_fx_available
+except ImportError:
+    def is_torch_fx_available() -> bool:
+        return hasattr(torch, "fx")
 
 from .configuration_deepseek import DeepseekV3Config
 

modeling_kimi_k25.py

@@ -64,6 +64,7 @@ from transformers.models.llava.modeling_llava import \
 from transformers.utils import is_flash_attn_2_available
 
 from .configuration_kimi_k25 import KimiK25Config
+from .configuration_deepseek import DeepseekV3Config
 from .modeling_deepseek import DeepseekV3ForCausalLM
 
 # Flash attention imports
@@ -245,6 +246,39 @@ def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False):
                                    axis=0)
     return pos_embed
 
+def _first_layer_key_first_token_vector(past_key_values):
+    """``past_key_values[0][0][..., 0]`` for LLaVA-style cache masking (shape ``[batch, heads, seq]``).
+    Legacy caches are ``list`` of ``(key, value)`` per layer. Transformers v4.36+ / v5 use ``Cache`` (e.g.
+    ``DynamicCache``) with per-layer ``.keys`` tensors instead of subscripting ``[0][0]``.
+    """
+    if isinstance(past_key_values, Cache):
+        layers = getattr(past_key_values, "layers", None) or []
+        if not layers:
+            return None
+        layer0 = layers[0]
+        keys = getattr(layer0, "keys", None)
+        if keys is None or keys.numel() == 0 or keys.ndim < 4:
+            return None
+        return keys[:, :, :, 0]
+    return past_key_values[0][0][:, :, :, 0]
+
+
+def _first_layer_past_seq_length(past_key_values):
+    """Layer-0 KV cache sequence length (BHSD keys: ``shape[2] == seq_len``).
+    """
+    if isinstance(past_key_values, Cache):
+        try:
+            return int(past_key_values.get_seq_length(0))
+        except Exception:
+            return None
+    try:
+        k0 = past_key_values[0][0]
+        if k0 is None or k0.ndim < 3:
+            return None
+        return int(k0.shape[2])
+    except Exception:
+        return None
+
 
 class Learnable2DInterpPosEmbDivided_fixed(nn.Module):
 
@@ -636,6 +670,7 @@ class MoonViT3dPretrainedModel(PreTrainedModel):
     model_type = 'moonvit3d'
     _no_split_modules = ['PackingTransformer']
     _supports_flash_attn_2 = True
+    _supports_flash_attn = True
     _supports_sdpa = True
 
     def __init__(self, config, *inputs, **kwargs):
@@ -772,6 +807,7 @@ class KimiK25PreTrainedModel(PreTrainedModel):
     ]
     _skip_keys_device_placement = "past_key_values"
     _supports_flash_attn_2 = True
+    _supports_flash_attn = True
     _supports_sdpa = False
 
     def _init_weights(self, module):
@@ -872,9 +908,10 @@ class KimiK25ForConditionalGeneration(KimiK25PreTrainedModel):
 
     def get_decoder(self):
         return self.language_model.get_decoder()
-
-    def tie_weights(self):
-        return self.language_model.tie_weights()
+    
+    def tie_weights(self, *args, **kwargs):
+        # Transformers >=5 passes ``missing_keys`` / ``recompute_mapping``; forward for the text backbone only.
+        return self.language_model.tie_weights(*args, **kwargs)
 
     def resize_token_embeddings(self,
                                 new_num_tokens: int | None = None,
@@ -1100,42 +1137,43 @@ class KimiK25ForConditionalGeneration(KimiK25PreTrainedModel):
             # generation with cache
             elif (past_key_values is not None and pixel_values is not None
                   and input_ids.shape[1] == 1):
-                # Retrieve the first layer to inspect the logits and mask out the hidden states
-                # that are set to 0
-                first_layer_past_key_value = past_key_values[0][0][:, :, :, 0]
-
-                # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941
-                batch_index, non_attended_tokens = torch.where(
-                    first_layer_past_key_value.float().sum(-2) == 0)
-
-                # Get the target length
-                target_length = input_ids.shape[1]
-                past_length = first_layer_past_key_value.shape[-1]
-
-                extended_attention_mask = torch.ones(
-                    (attention_mask.shape[0], past_length),
-                    dtype=attention_mask.dtype,
-                    device=attention_mask.device,
-                )
-
-                # Filter out only the tokens that can be un-attended, this can happen
-                # if one uses Llava + Fused modules where the cache on the
-                # first iteration is already big enough, or if one passes custom cache
-                valid_indices = non_attended_tokens < extended_attention_mask.size(
-                    -1)
-                new_batch_index = batch_index[valid_indices]
-                new_non_attended_tokens = non_attended_tokens[valid_indices]
-
-                # Zero-out the places where we don't need to attend
-                extended_attention_mask[new_batch_index,
-                                        new_non_attended_tokens] = 0
-
-                attention_mask = torch.cat(
-                    (extended_attention_mask, attention_mask[:,
-                                                             -target_length:]),
-                    dim=1)
-                position_ids = torch.sum(attention_mask,
-                                         dim=1).unsqueeze(-1) - 1
+                first_layer_past_key_value = _first_layer_key_first_token_vector(
+                    past_key_values)
+                if first_layer_past_key_value is not None:
+                    # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941
+                    batch_index, non_attended_tokens = torch.where(
+                        first_layer_past_key_value.float().sum(-2) == 0)
+
+                    # Get the target length
+                    target_length = input_ids.shape[1]
+                    past_length = _first_layer_past_seq_length(past_key_values)
+                    if past_length is None:
+                        past_length = int(first_layer_past_key_value.shape[-1])
+
+                    extended_attention_mask = torch.ones(
+                        (attention_mask.shape[0], past_length),
+                        dtype=attention_mask.dtype,
+                        device=attention_mask.device,
+                    )
+
+                    # Filter out only the tokens that can be un-attended, this can happen
+                    # if one uses Llava + Fused modules where the cache on the
+                    # first iteration is already big enough, or if one passes custom cache
+                    valid_indices = non_attended_tokens < extended_attention_mask.size(
+                        -1)
+                    new_batch_index = batch_index[valid_indices]
+                    new_non_attended_tokens = non_attended_tokens[valid_indices]
+
+                    # Zero-out the places where we don't need to attend
+                    extended_attention_mask[new_batch_index,
+                                            new_non_attended_tokens] = 0
+
+                    attention_mask = torch.cat(
+                        (extended_attention_mask, attention_mask[:,
+                                                                 -target_length:]),
+                        dim=1)
+                    position_ids = torch.sum(attention_mask,
+                                             dim=1).unsqueeze(-1) - 1
 
         outputs = self.language_model(
             attention_mask=attention_mask,
@@ -1228,6 +1266,13 @@ class KimiK25ForConditionalGeneration(KimiK25PreTrainedModel):
             if past_key_values:
                 position_ids = position_ids[:, -input_ids.shape[1]:]
 
+        # Generation (especially transformers v5) may supply ``position_ids`` for the full sequence while
+        # ``input_ids`` here is only the new suffix (e.g. length 1). RoPE must index with the current step length.
+        if past_key_values is not None and position_ids is not None:
+            cur_len = input_ids.shape[1]
+            if position_ids.shape[-1] > cur_len:
+                position_ids = position_ids[..., -cur_len:]
+                
         # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
         if inputs_embeds is not None and past_key_values is None:
             model_inputs = {"inputs_embeds": inputs_embeds}

tokenization_kimi_fast.py

@@ -0,0 +1,124 @@
+import os
+from typing import Optional
+
+from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
+
+from .tool_declaration_ts import encode_tools_to_typescript_style
+
+
+class TikTokenTokenizerFast(PreTrainedTokenizerFast):
+    vocab_files_names = {
+        "tokenizer_file": "tokenizer.json",
+        "vocab_file": "tiktoken.model",
+    }
+    model_input_names = ["input_ids", "attention_mask"]
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
+        # we need to find tokenizer.json from original path for our custom tokenizer.
+        kwargs["model_root"] = str(pretrained_model_name_or_path)
+        return super().from_pretrained(pretrained_model_name_or_path, *inputs,
+                                       **kwargs)
+
+    def __init__(
+        self,
+        tokenizer_file=None,
+        vocab_file=None,
+        model_root=None,
+        bos_token="[BOS]",
+        eos_token="[EOS]",
+        unk_token="[UNK]",
+        pad_token="[PAD]",
+        **kwargs,
+    ):
+        if model_root is None:
+            raise ValueError("model_root is required")
+        tokenizer_file = os.path.join(model_root, "tokenizer.json")
+        vocab_file = os.path.join(model_root, "tiktoken.model")
+        if not (os.path.isfile(tokenizer_file) and os.path.isfile(vocab_file)):
+            raise ValueError(f"Missing tokenizer files under: {model_root}")
+        self._tokenizer_dir = model_root
+        super().__init__(
+            tokenizer_file=tokenizer_file,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            **kwargs,
+        )
+        self.vocab_file = vocab_file
+
+    @property
+    def vocab_size(self) -> int:
+        """Return the vocabulary size."""
+        return self.backend_tokenizer.get_vocab_size()
+
+    def _sort_tools(self, tools):
+        """Deep sort tools for deterministic output."""
+        if isinstance(tools, dict):
+            return {k: self._sort_tools(v) for k, v in sorted(tools.items())}
+        if isinstance(tools, list):
+            return [self._sort_tools(item) for item in tools]
+        return tools
+
+    def save_vocabulary(self,
+                        save_directory: str,
+                        filename_prefix: Optional[str] = None) -> tuple:
+        """Save the tokenizer vocabulary."""
+        if not os.path.isdir(save_directory):
+            raise ValueError(
+                f"Vocabulary path ({save_directory}) should be a directory")
+
+        # Save tokenizer.json
+        tokenizer_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") +
+            "tokenizer.json")
+        self.backend_tokenizer.save(tokenizer_file)
+
+        # Also copy tiktoken.model if available
+        vocab_files = []
+        if self.vocab_file and os.path.isfile(self.vocab_file):
+            vocab_file = os.path.join(
+                save_directory,
+                (filename_prefix + "-" if filename_prefix else "") +
+                "tiktoken.model")
+            if os.path.abspath(self.vocab_file) != os.path.abspath(vocab_file):
+                import shutil
+                shutil.copy(self.vocab_file, vocab_file)
+            vocab_files.append(vocab_file)
+
+        return (tokenizer_file, ) + tuple(vocab_files)
+
+    def apply_chat_template(self,
+                            conversation,
+                            tools=None,
+                            tokenize=False,
+                            add_generation_prompt=True,
+                            thinking: bool = True,
+                            preserve_thinking: bool = False,
+                            **kwargs):
+        """Apply chat template with TypeScript tools support."""
+        tools = self._sort_tools(tools)
+
+        # Convert tools to TypeScript style string if tools are provided
+        tools_ts_str = None
+        if tools:
+            try:
+                tools_ts_str = encode_tools_to_typescript_style(tools)
+
+            except Exception as e:
+                print(f"Failed to convert tools to TypeScript style: {e}")
+                tools_ts_str = None
+
+        # Store the TypeScript string in kwargs so it can be accessed by the template
+        if tools_ts_str is not None:
+            kwargs['tools_ts_str'] = tools_ts_str
+        return super().apply_chat_template(
+            conversation,
+            tools=tools,
+            tokenize=tokenize,
+            add_generation_prompt=add_generation_prompt,
+            thinking=thinking,
+            preserve_thinking=preserve_thinking,
+            **kwargs)

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57ec7040095cadc25269b917f95ba026e1b2b7b2e5c0540ce0a9afe8afb06d2e
+size 19591764

tokenizer_config.json

@@ -205,12 +205,12 @@
   "extra_special_tokens": {},
   "model_max_length": 1000000000000000019884624838656,
   "pad_token": "[PAD]",
-  "tokenizer_class": "TikTokenTokenizer",
   "unk_token": "[UNK]",
+  "tokenizer_class": "TikTokenTokenizerFast",
   "auto_map": {
     "AutoTokenizer": [
-      "tokenization_kimi.TikTokenTokenizer",
-      null
+      null,
+      "tokenization_kimi_fast.TikTokenTokenizerFast"
     ]
   }
 }