diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..62375aaf6f0f97fad2b39e33dbc99fee014c7d8d --- /dev/null +++ b/.gitattributes @@ -0,0 +1,37 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text +model.safetensors.index.json filter=lfs diff=lfs merge=lfs -text +assets/k3-hero.jpg filter=lfs diff=lfs merge=lfs -text diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..97c0111bba229b0981b1d74e34bacef723ca8f7b --- /dev/null +++ b/LICENSE @@ -0,0 +1,52 @@ +Kimi K3 License + +Copyright (c) 2026 Moonshot AI + +Permission is hereby granted, free of charge, to any person (the "Licensee") +obtaining a copy of this software — including the model weights, parameters, +configuration files, inference and training code, and associated documentation +(collectively, the "Software") — to deal in the Software without restriction. +This includes, without limitation, the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of the Software; to run, +deploy, fine-tune, or otherwise modify the Software and create derivative works +from it; and to permit persons to whom the Software is furnished to do so, in +each case subject to the following conditions: + +1. The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. Licensee's use of the +Software must comply with applicable laws and regulations. + +2. "Model as a Service" means giving a third party access to language model +inference or fine-tuning (e.g., via API) in a manner that allows such third +party to exercise meaningful control over the inputs, parameters, or training +data. This does not include (a) end-user products with model capabilities solely +embedded within specific features or harnesses, or (b) mere relaying of requests +to models hosted by others. + +If the Licensee or any of its affiliates operates a Model as a Service business, +and the aggregate revenue of the Licensee and its affiliates exceeds 20 million +US dollars (or the equivalent in other currencies) in total over any consecutive +12 months, the Licensee must enter into a separate agreement with Moonshot AI +before using the Software or its derivative works for any commercial purpose. + +3. If the Software (or any derivative works thereof) is used for any of the +Licensee's commercial products or services that have more than 100 million +monthly active users, or more than 20 million US dollars (or equivalent in other +currencies) in monthly revenue, "Kimi K3" must be prominently displayed on the +user interface of such product or service. + +4. The requirements set forth in Sections 2 and 3 do not apply to: (a) internal +use of the Software, defined as any use that does not make the Software, its +outputs, or its underlying capabilities available to third parties; or (b) any +use of the Software accessed through Moonshot AI's official products or +certified inference partners. + +5. THE SOFTWARE AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” +BASIS, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT +LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE +AND NONINFRINGEMENT. IN NO EVENT SHALL MOONSHOT AI OR ITS AFFILIATES OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +For any questions regarding this license, please contact . \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d914b941cd824a5eb8ab653722f4e1c92feeb720 --- /dev/null +++ b/README.md @@ -0,0 +1,28 @@ +--- +base_model: +- moonshotai/Kimi-K3 +base_model_relation: finetune +tags: +- compressed-tensors +- conversational +- abliterated +- uncensored +license: other +license_name: "kimi-k3" +library_name: transformers +pipeline_tag: image-text-to-text + +--- +# Uniboshi/Kimi-K3-Abliterated-V1 + +- This model is an abliterated version of [moonshotai/Kimi-K3]() for uncensored. +- This model has been tuned to make refusals less likely to be triggered for English and Japanese. +- In this model, more than 98% of the safeguards have been removed. + ``` + writer written as n before after removed + *.2 [0] as stored 1 0.980..0.980 1.53e-02..1.53e-02 98.44% + *.down_proj [0] as stored 93 0.932..1.421 1.35e-02..1.56e-02 98.62% + *.embed_tokens [1] as stored 1 1.040..1.040 1.46e-02..1.46e-02 98.60% + *.o_proj [0] as stored 93 0.938..1.358 1.38e-02..1.60e-02 98.54% + *.routed_expert_up_proj [0] as stored 92 0.910..1.402 1.24e-02..1.56e-02 98.57% + ``` diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..9cd5078dd9a9ac60ef157a8754434b2e2e66df08 --- /dev/null +++ b/config.json @@ -0,0 +1,301 @@ +{ + "architectures": [ + "KimiK3ForConditionalGeneration" + ], + "auto_map": { + "AutoConfig": "configuration_kimi_k3.KimiK3Config", + "AutoModel": "modeling_kimi_k3.KimiK3ForConditionalGeneration", + "AutoModelForCausalLM": "modeling_kimi_k3.KimiK3ForConditionalGeneration" + }, + "bos_token_id": 163584, + "dtype": "bfloat16", + "eos_token_id": 163586, + "ignore_index": -100, + "image_placeholder": "<|kimi_image_placeholder|>", + "media_placeholder_token_id": 163605, + "model_type": "kimi_k3", + "pad_token_id": 163839, + "text_config": { + "_name_or_path": "", + "activation_situ_beta": 4.0, + "activation_situ_linear_beta": 25.0, + "add_cross_attention": false, + "architectures": [ + "KimiLinearForCausalLM" + ], + "attn_res_block_size": 12, + "auto_map": { + "AutoConfig": "configuration_kimi_k3.KimiLinearConfig", + "AutoModel": "modeling_kimi_linear.KimiLinearModel", + "AutoModelForCausalLM": "modeling_kimi_linear.KimiLinearForCausalLM" + }, + "bad_words_ids": null, + "begin_suppress_tokens": null, + "bos_token_id": 163584, + "chunk_size_feed_forward": 0, + "cross_attention_hidden_size": null, + "decoder_start_token_id": null, + "diversity_penalty": 0.0, + "do_sample": false, + "dtype": "bfloat16", + "early_stopping": false, + "encoder_no_repeat_ngram_size": 0, + "eos_token_id": 163586, + "exponential_decay_length_penalty": null, + "finetuning_task": null, + "first_k_dense_replace": 1, + "forced_bos_token_id": null, + "forced_eos_token_id": null, + "hidden_act": "situ", + "hidden_size": 7168, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "initializer_range": 0.02, + "intermediate_size": 33792, + "is_decoder": false, + "is_encoder_decoder": false, + "kv_lora_rank": 512, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "latent_moe_use_norm": true, + "length_penalty": 1.0, + "linear_attn_config": { + "full_attn_layers": [ + 4, + 8, + 12, + 16, + 20, + 24, + 28, + 32, + 36, + 40, + 44, + 48, + 52, + 56, + 60, + 64, + 68, + 72, + 76, + 80, + 84, + 88, + 92, + 93 + ], + "gate_lower_bound": -5.0, + "head_dim": 128, + "kda_layers": [ + 1, + 2, + 3, + 5, + 6, + 7, + 9, + 10, + 11, + 13, + 14, + 15, + 17, + 18, + 19, + 21, + 22, + 23, + 25, + 26, + 27, + 29, + 30, + 31, + 33, + 34, + 35, + 37, + 38, + 39, + 41, + 42, + 43, + 45, + 46, + 47, + 49, + 50, + 51, + 53, + 54, + 55, + 57, + 58, + 59, + 61, + 62, + 63, + 65, + 66, + 67, + 69, + 70, + 71, + 73, + 74, + 75, + 77, + 78, + 79, + 81, + 82, + 83, + 85, + 86, + 87, + 89, + 90, + 91 + ], + "num_heads": 96, + "short_conv_kernel_size": 4, + "use_full_rank_gate": true + }, + "max_length": 20, + "max_position_embeddings": 1048576, + "min_length": 0, + "mla_use_nope": true, + "mla_use_output_gate": true, + "model_type": "kimi_linear", + "moe_intermediate_size": 3072, + "moe_layer_freq": 1, + "moe_renormalize": true, + "moe_router_activation_func": "sigmoid", + "no_repeat_ngram_size": 0, + "num_attention_heads": 96, + "num_beam_groups": 1, + "num_beams": 1, + "num_expert_group": 1, + "num_experts": 896, + "num_experts_per_token": 16, + "num_hidden_layers": 93, + "num_key_value_heads": 96, + "num_nextn_predict_layers": 0, + "num_return_sequences": 1, + "num_shared_experts": 2, + "output_attentions": false, + "output_hidden_states": false, + "output_scores": false, + "pad_token_id": 163839, + "prefix": null, + "problem_type": null, + "pruned_heads": {}, + "q_lora_rank": 1536, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "quantization_config": { + "config_groups": { + "group_0": { + "format": "mxfp4-pack-quantized", + "input_activations": null, + "output_activations": null, + "targets": [ + "Linear" + ], + "weights": { + "actorder": null, + "block_structure": null, + "dynamic": false, + "group_size": 32, + "num_bits": 4, + "observer": "minmax", + "observer_kwargs": {}, + "scale_dtype": "torch.uint8", + "strategy": "group", + "symmetric": true, + "type": "float", + "zp_dtype": null + } + } + }, + "format": "mxfp4-pack-quantized", + "global_compression_ratio": null, + "ignore": [ + "re:.*self_attn.*", + "re:.*shared_experts.*", + "re:.*mlp\\.(gate|up|gate_up|down)_proj.*", + "re:.*lm_head.*", + "re:.*vision_tower.*", + "re:.*mm_projector.*", + "re:.*\\.(?:embed_tokens|gate|input_layernorm|mlp_res_norm|mlp_res_proj|norm|output_attn_res_norm|output_attn_res_proj|post_attention_layernorm|routed_expert_down_proj|routed_expert_norm|routed_expert_up_proj|self_attention_res_norm|self_attention_res_proj)$" + ], + "kv_cache_scheme": null, + "quant_method": "compressed-tensors", + "quantization_status": "compressed" + }, + "remove_invalid_values": false, + "repetition_penalty": 1.0, + "return_dict": true, + "return_dict_in_generate": false, + "rms_norm_eps": 1e-05, + "routed_expert_hidden_size": 3584, + "routed_scaling_factor": 1.0, + "sep_token_id": null, + "suppress_tokens": null, + "task_specific_params": null, + "temperature": 1.0, + "tf_legacy_loss": false, + "tie_encoder_decoder": false, + "tie_word_embeddings": false, + "tokenizer_class": null, + "top_k": 50, + "top_p": 1.0, + "topk_group": 1, + "topk_method": "noaux_tc", + "torchscript": false, + "transformers_version": "4.56.2", + "typical_p": 1.0, + "use_bfloat16": false, + "use_cache": true, + "use_grouped_topk": true, + "v_head_dim": 128, + "vocab_size": 163840 + }, + "tie_word_embeddings": false, + "vision_config": { + "_attn_implementation": "flash_attention_2", + "activation_func": "gelu_pytorch_tanh", + "attn_bias": false, + "init_pos_emb_height": 64, + "init_pos_emb_time": 4, + "init_pos_emb_width": 64, + "linear_bias": false, + "merge_kernel_size": [ + 2, + 2 + ], + "merge_type": "sd2_tpool", + "mlp_type": "mlp2", + "mm_hidden_size": 1024, + "mm_projector_type": "patchmergerv2", + "norm_type": "rmsnorm", + "patch_embed_proj_bias": false, + "patch_size": 14, + "pos_emb_interpolation_mode": "bilinear", + "pos_emb_type": "divided_fixed", + "projector_hidden_act": "gelu", + "projector_ln_eps": 1e-05, + "qkv_hidden_size": 1536, + "text_hidden_size": 7168, + "vt_hidden_size": 1024, + "vt_intermediate_size": 4096, + "vt_num_attention_heads": 12, + "vt_num_hidden_layers": 27 + } +} \ No newline at end of file diff --git a/configuration_kimi_k3.py b/configuration_kimi_k3.py new file mode 100644 index 0000000000000000000000000000000000000000..b38c4f8fcab1ce295aea236772ea1cba57f978c9 --- /dev/null +++ b/configuration_kimi_k3.py @@ -0,0 +1,285 @@ +from typing import Optional + +from transformers.configuration_utils import PretrainedConfig + + +class KimiLinearConfig(PretrainedConfig): + model_type = "kimi_linear" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + model_type="kimi_linear", + vocab_size=163840, + hidden_size=4096, + head_dim=None, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + hidden_act="silu", + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + rope_theta=10000.0, + rope_scaling=None, + tie_word_embeddings=False, + moe_intermediate_size: Optional[int] = None, + moe_renormalize: bool = True, + moe_router_activation_func: str = "sigmoid", + num_experts: Optional[int] = None, + num_experts_per_token: Optional[int] = None, + num_shared_experts: int = 0, + routed_scaling_factor: float = 1.0, + first_k_dense_replace: int = 0, + moe_layer_freq: int = 1, + use_grouped_topk: bool = True, + num_expert_group: int = 1, + topk_group: int = 1, + q_lora_rank: Optional[int] = None, + kv_lora_rank: Optional[int] = None, + qk_nope_head_dim: Optional[int] = None, + qk_rope_head_dim: Optional[int] = None, + v_head_dim: Optional[int] = None, + mla_use_nope: Optional[bool] = False, + mla_use_output_gate: Optional[bool] = False, + num_nextn_predict_layers: int = 0, + linear_attn_config: Optional[dict] = None, + attn_res_block_size: Optional[int] = None, + latent_moe_use_norm: bool = False, + activation_situ_beta: Optional[float] = None, + activation_situ_linear_beta: Optional[float] = None, + max_position_embeddings: int = 4096, + routed_expert_hidden_size: Optional[int] = None, + topk_method: str = "noaux_tc", + **kwargs, + ): + self.model_type = model_type + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.head_dim = ( + head_dim if head_dim is not None else hidden_size // num_attention_heads + ) + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.mla_use_nope = mla_use_nope + self.mla_use_output_gate = mla_use_output_gate + # moe config + self.num_experts = num_experts + self.num_experts_per_token = num_experts_per_token + self.moe_renormalize = moe_renormalize + self.num_shared_experts = num_shared_experts + self.routed_scaling_factor = routed_scaling_factor + self.moe_router_activation_func = moe_router_activation_func + assert self.moe_router_activation_func in ("softmax", "sigmoid") + self.moe_intermediate_size = moe_intermediate_size + self.first_k_dense_replace = first_k_dense_replace + self.moe_layer_freq = moe_layer_freq + self.use_grouped_topk = use_grouped_topk + self.num_expert_group = num_expert_group + self.topk_group = topk_group + self.num_nextn_predict_layers = num_nextn_predict_layers + + self.attn_res_block_size = attn_res_block_size + self.latent_moe_use_norm = latent_moe_use_norm + self.activation_situ_beta = activation_situ_beta + self.activation_situ_linear_beta = activation_situ_linear_beta + self.max_position_embeddings = max_position_embeddings + self.routed_expert_hidden_size = routed_expert_hidden_size + self.topk_method = topk_method + + if linear_attn_config is not None: + assert linear_attn_config["kda_layers"] is not None + assert linear_attn_config["full_attn_layers"] is not None + self.linear_attn_config = linear_attn_config + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @property + def is_mla(self): + return ( + self.q_lora_rank is not None + or self.kv_lora_rank is not None + or self.qk_nope_head_dim is not None + or self.qk_rope_head_dim is not None + or self.v_head_dim is not None + or self.mla_use_nope is True + ) + + @property + def is_moe(self): + return self.num_experts is not None + + @property + def is_linear_attn(self) -> bool: + return not ( + self.linear_attn_config is None + or ( + isinstance(self.linear_attn_config, dict) + and self.linear_attn_config["kda_layers"] is not None + and len(self.linear_attn_config["kda_layers"]) == 0 + ) + ) + + def is_kda_layer(self, layer_idx: int): + return ( + self.linear_attn_config is not None + and (layer_idx + 1) in self.linear_attn_config["kda_layers"] + ) + + +class KimiK3VisionConfig(PretrainedConfig): + + def __init__( + self, + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + init_pos_emb_time: int = 4, + pos_emb_type: str = 'divided_fixed', + vt_num_attention_heads: int = 12, + vt_num_hidden_layers: int = 27, + vt_hidden_size: int = 1024, + vt_intermediate_size: int = 4096, + merge_kernel_size: tuple = (2, 2), + merge_type: str = 'sd2_tpool', + _attn_implementation: str = 'flash_attention_2', + # MM Projector parameters + mm_projector_type: str = 'patchmergerv2', + mm_hidden_size: int | None = None, + projector_hidden_act: str = "gelu", + projector_ln_eps: float = 1e-5, + # vision tower parameters + qkv_hidden_size: int = 1536, + norm_type: str = 'rmsnorm', + attn_bias: bool = False, + patch_embed_proj_bias: bool = False, + mlp_type: str = 'mlp2', + linear_bias: bool = False, + activation_func: str = 'gelu_pytorch_tanh', + pos_emb_interpolation_mode: str = 'bilinear', + # Other parameters + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + text_hidden_size=7168, + **kwargs): + + self.patch_size = patch_size + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + self.init_pos_emb_time = init_pos_emb_time + self.pos_emb_type = pos_emb_type + self.vt_num_attention_heads = vt_num_attention_heads + self.vt_num_hidden_layers = vt_num_hidden_layers + self.vt_hidden_size = vt_hidden_size + self.vt_intermediate_size = vt_intermediate_size + self.merge_kernel_size = merge_kernel_size + self.merge_type = merge_type + self._attn_implementation = _attn_implementation + + # MM Projector config + self.mm_projector_type = mm_projector_type + self.mm_hidden_size = mm_hidden_size if mm_hidden_size is not None else vt_hidden_size + self.projector_hidden_act = projector_hidden_act + self.projector_ln_eps = projector_ln_eps + self.text_hidden_size = text_hidden_size + + # vision tower parameters + self.qkv_hidden_size = qkv_hidden_size + self.norm_type = norm_type + self.attn_bias = attn_bias + self.patch_embed_proj_bias = patch_embed_proj_bias + self.mlp_type = mlp_type + self.linear_bias = linear_bias + self.activation_func = activation_func + self.pos_emb_interpolation_mode = pos_emb_interpolation_mode + + super().__init__(**kwargs) + + +class KimiK3Config(PretrainedConfig): + """Kimi-K3 model configuration. + + Args: + text_config (dict | KimiLinearConfig): Configuration for the text model. + + Vision Tower Parameters (from MoonViT3dConfig): + patch_size (int): Patch size for vision tower. + init_pos_emb_height (int): Initial position embedding height. + init_pos_emb_width (int): Initial position embedding width. + init_pos_emb_time (int): Initial position embedding time dimension. + pos_emb_type (str): Type of position embedding. + vt_num_attention_heads (int): Number of attention heads in vision tower. + vt_num_hidden_layers (int): Number of hidden layers in vision tower. + vt_hidden_size (int): Hidden size of vision tower. + vt_intermediate_size (int): Intermediate size in vision tower FFN. + merge_kernel_size (tuple): Kernel size for patch merging. + merge_type (str): Type of merge operation. + _attn_implementation (str): Attention implementation type. + + MM Projector Parameters (from MultiModalProjectorConfig): + mm_projector_type (str): Type of multimodal projector. + mm_hidden_size (int): Hidden size from vision tower (should match vt_hidden_size). + projector_hidden_act (str): Activation function for projector. + projector_ln_eps (float): Layer norm epsilon for projector. + + Other Parameters: + ignore_index (int): The ignore index for the loss function. + media_placeholder_token_id (int): The token ID to use for media placeholders. + pad_token_id (int): The token ID to use for padding. + """ + + model_type = "kimi_k3" + + def __init__( + self, + text_config: dict | KimiLinearConfig = None, + vision_config: dict | KimiK3VisionConfig = None, + # Other parameters + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + **kwargs, + ): + if isinstance(text_config, dict): + text_config = KimiLinearConfig(**text_config) + if isinstance(vision_config, dict): + vision_config = KimiK3VisionConfig(**vision_config) + self.text_config = text_config + self.vision_config = vision_config + # Other config + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) diff --git a/encoding_k3.py b/encoding_k3.py new file mode 100644 index 0000000000000000000000000000000000000000..44a56b7371fab32513da0eaed6b9a8a7c56499a2 --- /dev/null +++ b/encoding_k3.py @@ -0,0 +1,647 @@ +"""Kimi K3 XTML encoding helpers. + +This module keeps chat rendering in Python. +Callers that need token IDs should consume ``EncodeSegment`` objects directly: +structural markers may be encoded as tiktoken special tokens, while user/tool +text and attribute values are encoded as ordinary text. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Iterable, Optional + +OPEN_TOKEN = "<|open|>" +CLOSE_TOKEN = "<|close|>" +SEP_TOKEN = "<|sep|>" +END_OF_MSG_TOKEN = "<|end_of_msg|>" +IMAGE_PLACEHOLDER = "<|kimi_image_placeholder|>" + +_VALID_THINKING_EFFORTS = {"low", "high", "max"} + + +@dataclass(frozen=True) +class EncodeSegment: + text: str + allow_special: bool = False + + +class _ImagePromptState: + def __init__(self, image_prompts: Optional[list[str]] = None): + self.image_prompts = image_prompts + self.index = 0 + + def next_prompt(self) -> str: + if self.image_prompts is None: + return IMAGE_PLACEHOLDER + if self.index >= len(self.image_prompts): + raise ValueError("More image placeholders than image prompts.") + prompt = self.image_prompts[self.index] + self.index += 1 + return prompt + + def assert_consumed(self) -> None: + if self.image_prompts is None: + return + if self.index != len(self.image_prompts): + raise ValueError( + f"image prompt count {len(self.image_prompts)} != " + f"consumed placeholder count {self.index}" + ) + + +def _segment(text: Any, *, allow_special: bool = False) -> list[EncodeSegment]: + text = str(text) + if not text: + return [] + return [EncodeSegment(text, allow_special=allow_special)] + + +def _control(text: str) -> list[EncodeSegment]: + return _segment(text, allow_special=True) + + +def _text(text: Any) -> list[EncodeSegment]: + return _segment(text, allow_special=False) + + +def _append_text( + segments: list[EncodeSegment], + text: Any, + image_state: _ImagePromptState, +) -> None: + text = str(text) + if text == "": + return + if image_state.image_prompts is None or IMAGE_PLACEHOLDER not in text: + segments.extend(_text(text)) + return + + parts = text.split(IMAGE_PLACEHOLDER) + for i, part in enumerate(parts): + segments.extend(_text(part)) + if i < len(parts) - 1: + segments.extend(_segment(image_state.next_prompt(), + allow_special=True)) + + +def _escape_attr_value(value: Any) -> str: + return str(value).replace("&", "&").replace('"', """) + + +def _attr(key: str, value: Any) -> list[EncodeSegment]: + return ( + _text(f" {key}") + + _text('="') + + _text(_escape_attr_value(value)) + + _text('"') + ) + + +def _open_tag(tag: str, attrs: Iterable[tuple[str, Any]] = ()) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(OPEN_TOKEN)) + segments.extend(_text(tag)) + for key, value in attrs: + segments.extend(_attr(key, value)) + segments.extend(_control(SEP_TOKEN)) + return segments + +def _close_tag(tag: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(CLOSE_TOKEN)) + segments.extend(_text(tag)) + segments.extend(_control(SEP_TOKEN)) + return segments + + +def _end_of_msg() -> list[EncodeSegment]: + return _control(END_OF_MSG_TOKEN) + + +def _json_compact(value: Any) -> str: + return json.dumps(value, ensure_ascii=False, separators=(",", ":")) + + +def _is_mapping(value: Any) -> bool: + return isinstance(value, dict) + + +def _xtml_type(value: Any) -> str: + if isinstance(value, bool): + return "boolean" + if value is None: + return "null" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return "number" + if isinstance(value, str): + return "string" + if _is_mapping(value): + return "object" + return "array" + + +def _xtml_value(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + + +def _get_value(obj: Any, key: str, default: Any = None) -> Any: + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + +def extract_response_schema(response_format: Any) -> Any: + if response_format is None: + return None + + json_schema = _get_value(response_format, "json_schema") + if json_schema is None: + return None + + if isinstance(json_schema, dict): + return json_schema.get( + "schema", + json_schema.get("json_schema", json_schema), + ) + + schema = _get_value(json_schema, "schema") + if schema is not None: + return schema + + schema = _get_value(json_schema, "json_schema") + if schema is not None: + return schema + + return json_schema + + +def deep_sort_dict(obj: Any) -> Any: + if isinstance(obj, dict): + return {k: deep_sort_dict(v) for k, v in sorted(obj.items())} + if isinstance(obj, list): + return [deep_sort_dict(item) for item in obj] + return obj + + +def normalize_tool_arguments(arguments: Any) -> tuple[dict[str, Any], Optional[str]]: + if arguments is None: + return {}, None + if isinstance(arguments, dict): + return arguments, None + if isinstance(arguments, str): + if not arguments.strip(): + return {}, None + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + return {}, arguments + if not isinstance(parsed, dict): + raise ValueError("Kimi K3 tool call arguments must be a JSON object.") + return parsed, None + raise TypeError( + "Kimi K3 tool call arguments must be a dict or a JSON object string." + ) + + +def normalize_message(message: Any) -> Any: + if not isinstance(message, dict): + return message + + normalized = dict(message) + + tools = normalized.get("tools") + if tools is not None: + normalized["tools"] = deep_sort_dict(tools) + + tool_calls = normalized.get("tool_calls") + if not tool_calls: + return normalized + + normalized_calls = [] + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + normalized_calls.append(tool_call) + continue + + tc = dict(tool_call) + function = tc.get("function") + if isinstance(function, dict): + fn = dict(function) + arguments, json_block = normalize_tool_arguments(fn.get("arguments")) + fn["arguments"] = arguments + if json_block is None: + fn.pop("_xtml_json_block", None) + else: + fn["_xtml_json_block"] = json_block + tc["function"] = fn + else: + arguments, json_block = normalize_tool_arguments(tc.get("arguments")) + tc["arguments"] = arguments + if json_block is None: + tc.pop("_xtml_json_block", None) + else: + tc["_xtml_json_block"] = json_block + normalized_calls.append(tc) + + normalized["tool_calls"] = normalized_calls + return normalized + + +def normalize_conversation(conversation: Any) -> Any: + if not isinstance(conversation, list): + return conversation + + def normalize_messages(messages: list[Any]) -> list[Any]: + return [normalize_message(message) for message in messages] + + if conversation and isinstance(conversation[0], list): + return [normalize_messages(messages) for messages in conversation] + return normalize_messages(conversation) + + +def _tool_call_id_index(tool_calls: Any) -> dict: + """Map assistant ``tool_calls[].id`` to ``(1-based position, function name)``. + + The position mirrors the chat template's enumeration over ``tool_calls`` + (every entry advances the position, even an id-less one). Duplicate ids keep + their first occurrence. + """ + index: dict = {} + if not isinstance(tool_calls, list): + return index + for position, tool_call in enumerate(tool_calls, start=1): + if not isinstance(tool_call, dict): + continue + call_id = tool_call.get("id") + if call_id is None: + continue + key = str(call_id) + if key in index: + continue + function = tool_call.get("function") + name = ( + function.get("name") if isinstance(function, dict) else tool_call.get("name") + ) + index[key] = (position, name) + return index + + +def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]: + """Re-sort K3 XTML tool results into assistant ``tool_calls`` order. + + Serving frameworks generally deliver tool results already in call order. A + direct Transformers caller, however, may pass OpenAI-style tool messages in any + order, so each run of consecutive tool messages is matched against the most + recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` == + ``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and + sorted by the matched 1-based position. The matched call is authoritative, + so each matched message's ``tool`` is set to that call's function name -- + this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting + out of sync with the reordered position. ``index`` is still derived from the + rendered position by the chat template. A run that cannot be fully matched is + left untouched. Re-running is idempotent. + + This function is side-effect free: matched tool messages are shallow-copied + before their ``tool``/``name`` is rewritten, and every other message is + appended to the output as-is. The input list and its message objects are + never mutated. + """ + if not isinstance(messages, list): + return messages + + output: list[Any] = [] + current_index: dict = {} + i = 0 + n = len(messages) + + while i < n: + message = messages[i] + + if isinstance(message, dict) and message.get("role") == "assistant": + tool_calls = message.get("tool_calls") + current_index = _tool_call_id_index(tool_calls) if tool_calls else {} + output.append(message) + i += 1 + continue + + if not isinstance(message, dict) or message.get("role") != "tool": + output.append(message) + i += 1 + continue + + run: list[tuple] = [] # (position, original_offset, message, name) + unresolved = False + offset = 0 + while ( + i < n and isinstance(messages[i], dict) and messages[i].get("role") == "tool" + ): + tool_message = messages[i] + call_id = tool_message.get("tool_call_id", tool_message.get("id")) + matched = current_index.get(str(call_id)) if call_id is not None else None + if matched is None: + unresolved = True + run.append((None, offset, tool_message, None)) + else: + position, name = matched + run.append((position, offset, tool_message, name)) + offset += 1 + i += 1 + + if unresolved: + output.extend(item[2] for item in run) + else: + run.sort(key=lambda item: (item[0], item[1])) + for _, _, tool_message, name in run: + if name is None: + output.append(tool_message) + continue + # The id-matched call is authoritative: align tool (and any + # explicit name) so the rendered XTML tool attribute cannot + # disagree with the reordered position. Copy first so the + # caller's message object is never mutated. + resolved = dict(tool_message) + resolved["tool"] = name + if "name" in resolved: + resolved["name"] = name + output.append(resolved) + + return output + + +def is_batched_conversation(conversation: Any) -> bool: + return ( + isinstance(conversation, list) + and bool(conversation) + and isinstance(conversation[0], list) + ) + + +def _render_content_segments( + content: Any, + image_state: _ImagePromptState, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + if isinstance(content, str): + _append_text(segments, content, image_state) + elif content is not None: + for part in content: + if part["type"] in ["image", "image_url"]: + segments.extend( + _segment(image_state.next_prompt(), allow_special=True)) + else: + _append_text(segments, part["text"], image_state) + return segments + + +def _internal_system_message(message_type: str, body: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_open_tag("message", [("role", "system"), ("type", message_type)])) + segments.extend(_text(body.strip())) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def _render_assistant_segments( + message: dict[str, Any], + image_state: _ImagePromptState, + thinking: bool = True, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + # The channel is structural: in thinking mode every assistant + # message carries the open/close tags even when there is no reasoning + # content to fill in. In non-thinking mode the channel is dropped + # entirely. + if thinking: + reasoning_content = message.get("reasoning_content") or message.get( + "reasoning" + ) + segments.extend(_open_tag("think")) + if reasoning_content is not None and str(reasoning_content).strip(): + _append_text(segments, reasoning_content, image_state) + segments.extend(_close_tag("think")) + + segments.extend(_open_tag("response")) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("response")) + + tool_calls = message.get("tool_calls") + if tool_calls: + segments.extend(_open_tag("tools")) + for index, tool_call in enumerate(tool_calls, start=1): + fn = tool_call.get("function", tool_call) + segments.extend( + _open_tag("call", [("tool", fn["name"]), ("index", index)]) + ) + args = fn.get("arguments", {}) + json_block = fn.get("_xtml_json_block") + if json_block is not None: + segments.extend(_open_tag("json", [("type", "object")])) + _append_text(segments, json_block, image_state) + segments.extend(_close_tag("json")) + elif _is_mapping(args): + for key, value in args.items(): + segments.extend( + _open_tag( + "argument", + [("key", key), ("type", _xtml_type(value))], + ) + ) + _append_text(segments, _xtml_value(value), image_state) + segments.extend(_close_tag("argument")) + segments.extend(_close_tag("call")) + segments.extend(_close_tag("tools")) + + return segments + + +def _render_tool_declare(tools: Any, *, dynamic: bool = False) -> list[EncodeSegment]: + if dynamic: + body = ( + "## New Tools Available\n" + "The system dynamically extends the toolset via lazy-loading.\n" + "You have access to all existing and extended tools.\n" + "Here are the specs for the extended tools.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```" + ) + else: + body = ( + "# Tools\n" + "Here are the available tools, described in JSONSchema.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```" + ) + segments: list[EncodeSegment] = [] + segments.extend(_open_tag("message", [("role", "system"), ("type", "tool-declare")])) + segments.extend(_text(body)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def build_chat_segments( + messages: list[Any], + tools: Optional[list[dict]] = None, + *, + add_generation_prompt: bool = True, + thinking: bool = True, + image_prompts: Optional[list[str]] = None, + **kwargs: Any, +) -> list[EncodeSegment]: + # Re-sort tool results by tool_call_id at the lowest layer so every caller + # (processor or direct tokenizer) gets correctly ordered XTML. The helper is + # side-effect free, so the caller's message objects are left untouched. + messages = normalize_xtml_tool_result_messages(messages) + messages = normalize_conversation(messages) + tools = deep_sort_dict(tools) + + kwargs = dict(kwargs) + response_format = kwargs.get("response_format") + if "response_schema" not in kwargs: + response_schema = extract_response_schema(response_format) + if response_schema is not None: + kwargs["response_schema"] = response_schema + if kwargs.get("response_schema") is not None: + kwargs["response_schema"] = deep_sort_dict(kwargs["response_schema"]) + + image_state = _ImagePromptState(image_prompts) + segments: list[EncodeSegment] = [] + + tool_calls = None + tool_index = 0 + + if tools: + segments.extend(_render_tool_declare(tools)) + + thinking_effort = kwargs.get("thinking_effort") + if thinking and thinking_effort is not None: + assert thinking_effort in _VALID_THINKING_EFFORTS, ( + f"Unsupported thinking_effort={thinking_effort!r}; " + f"supported values are {sorted(_VALID_THINKING_EFFORTS)}." + ) + if thinking and thinking_effort in _VALID_THINKING_EFFORTS: + segments.extend( + _internal_system_message( + "thinking-effort", + "`thinking_effort` guides on how much to think in your " + "thinking channel (not including the response channel), " + "supported values include `low`, `medium`, `high`, and `max`.\n" + f"Now the system is invoked with `thinking_effort={thinking_effort}`.", + ) + ) + + for message_index, message in enumerate(messages): + if not isinstance(message, dict): + continue + + role = message["role"] + if role == "user": + attrs = [("role", "user")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "system" and message.get("tools"): + segments.extend(_render_tool_declare(message["tools"], dynamic=True)) + elif role == "system": + attrs = [("role", "system")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "tool": + tool_index += 1 + tool_name = message.get("tool", message.get("name")) + if ( + tool_name is None + and tool_calls is not None + and tool_index <= len(tool_calls) + ): + tc = tool_calls[tool_index - 1] + fn = tc.get("function", tc) + tool_name = fn["name"] + if tool_name is None: + raise ValueError( + "Kimi K3 tool messages need a resolvable tool name: " + "carry `tool`/`name`, or match a preceding assistant " + "tool_call by order." + ) + segments.extend( + _open_tag( + "message", + [("role", "tool"), ("tool", tool_name), ("index", tool_index)], + ) + ) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "assistant": + tool_calls = message.get("tool_calls") + tool_index = 0 + attrs = [("role", "assistant")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_assistant_segments(message, image_state, thinking)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + + tool_choice = kwargs.get("tool_choice") + if tool_choice == "required": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=required`.\n" + "You MUST call tools in the next message.", + ) + ) + elif tool_choice == "none": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=none`.\n" + "You MUST NOT call any tools in the next message.", + ) + ) + + rf = kwargs.get("response_format") + rf_type = _get_value(rf, "type", rf) if isinstance(rf, dict) else rf + if rf_type == "json_object": + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_object`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.", + ) + ) + elif rf_type == "json_schema": + schema = _json_compact(kwargs.get("response_schema")) + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_schema`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.\n" + "The JSON data must match the following schema:\n" + f"```json\n{schema}\n```", + ) + ) + + if add_generation_prompt: + segments.extend(_open_tag("message", [("role", "assistant")])) + segments.extend(_open_tag("think" if thinking else "response")) + + image_state.assert_consumed() + return segments diff --git a/generation_config.json b/generation_config.json new file mode 100644 index 0000000000000000000000000000000000000000..3908a4a5477122aa4b09386acab760fb7887c0fc --- /dev/null +++ b/generation_config.json @@ -0,0 +1,4 @@ +{ + "max_length": 1048576, + "eos_token_id": 163586 +} \ No newline at end of file diff --git a/kimi_k3_processor.py b/kimi_k3_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..033a5773885c144509788e89a445217b76d6c046 --- /dev/null +++ b/kimi_k3_processor.py @@ -0,0 +1,187 @@ +"""Kimi-K3 processor: wraps vision processor + tokenizer into a single interface. + +Chat rendering (including XTML tool-result ordering) is handled by the +tokenizer's Python encoder; this processor adds multimodal media preprocessing. +""" + +from transformers.feature_extraction_utils import BatchFeature +from transformers.processing_utils import ProcessorMixin +from transformers.utils import logging + +from .media_utils import ensure_media_type + +logger = logging.get_logger(__name__) + +# ── KimiK3Processor ─────────────────────────────────────────────────── + + +class KimiK3Processor(ProcessorMixin): + r""" + Constructs a KimiK3 processor which wraps a KimiK3 image processor + and a tokenizer into a single processor. + + [`KimiK3Processor`] offers all the functionalities of + [`KimiK3VisionProcessor`] and [`TikTokenTokenizer`]. + + Args: + image_processor ([`KimiK3VisionProcessor`], *optional*): + The image processor is a required input. + tokenizer ([`TikTokenTokenizer`], *optional*): + The tokenizer is a required input. + chat_template (`str`, *optional*): Kept for ProcessorMixin + compatibility. Kimi K3 chat encoding is implemented in Python by + the tokenizer. + """ + + attributes = ["image_processor", "tokenizer"] + valid_kwargs = ["chat_template"] + image_processor_class = "AutoImageProcessor" + tokenizer_class = "AutoTokenizer" + + def __init__( + self, + image_processor=None, + tokenizer=None, + chat_template=None, + **kwargs, + ): + super().__init__(image_processor, + tokenizer, + chat_template=chat_template) + self.media_processor = image_processor + self.image_placeholder = "<|kimi_image_placeholder|>" + + # ── Media preprocessing ──────────────────────────────────────────── + + def update_raw_text(self, text: str, image_prompts: list[str]) -> str: + # Replace image placeholders + image_count = text.count(self.image_placeholder) + if image_count > 0: + assert image_count == len(image_prompts), ( + f"image placeholder count {image_count} != " + f"image_prompts count {len(image_prompts)}") + text_parts = text.split(self.image_placeholder) + assert len(text_parts) == len(image_prompts) + 1 + text = "".join([ + text_parts[i] + image_prompts[i] + for i in range(len(image_prompts)) + ]) + text += text_parts[-1] + + return text + + def preprocess_medias(self, + medias: list[dict]) -> tuple[list[dict], list[str]]: + """Process media items and generate corresponding prompts. + + Returns: + A tuple of (updated_medias, image_prompts). + """ + updated_medias = [] + image_prompts = [] + for media in medias: + if media['type'] == 'image': + updated_medias.append(media) + img = ensure_media_type( + media, + transparent_bg_config=self.media_processor. + _transparent_bg_config, + transparent_bg_fill_stage=self.media_processor. + _transparent_bg_fill_stage, + )['image'] + w, h = img.size + image_prompts.append( + self.media_processor.make_image_prompt(w, h)) + else: + raise ValueError(f"unsupported media type: {media['type']}") + return updated_medias, image_prompts + + # ── Main entry points ────────────────────────────────────────────── + + def __call__(self, + messages: list[dict] = None, + medias: list[dict] = None, + text: str = None, + return_tensors: str = "pt", + **kwargs) -> BatchFeature: + """ + Process multimodal inputs for Kimi-K3 model. + + Args: + messages: List of message dicts with 'role' and 'content' fields. + If provided, medias and text will be extracted automatically. + medias: Pre-extracted list of media dicts. + text: Pre-formatted text string. + return_tensors: Format of returned tensors. Default: 'pt'. + **kwargs: Additional arguments passed to apply_chat_template. + + Returns: + BatchFeature with fields: input_ids, attention_mask, + pixel_values, grid_thws. + """ + if messages is None and (medias is None or text is None): + raise ValueError( + "Provide either 'messages' or both 'medias' and 'text'") + + if medias is not None and text is not None: + updated_medias, image_prompts = (self.preprocess_medias(medias)) + preprocessed = self.media_processor.preprocess( + updated_medias, return_tensors=return_tensors) + text = self.update_raw_text(text, image_prompts) + text_inputs = self.tokenizer(text, return_tensors=return_tensors) + return BatchFeature(data={**text_inputs, **preprocessed.data}) + + if medias is None: + medias = self._extract_medias_from_messages(messages) + updated_medias, image_prompts = (self.preprocess_medias(medias)) + preprocessed = self.media_processor.preprocess( + updated_medias, return_tensors=return_tensors) + + if text is None: + text_inputs = self.tokenizer.apply_chat_template( + messages, + tokenize=True, + return_tensors=return_tensors, + return_dict=True, + image_prompts=image_prompts, + **kwargs) + return BatchFeature(data={**text_inputs, **preprocessed.data}) + + text = self.update_raw_text(text, image_prompts) + text_inputs = self.tokenizer(text, return_tensors=return_tensors) + return BatchFeature(data={**text_inputs, **preprocessed.data}) + + @staticmethod + def _extract_medias_from_messages(messages: list[dict]) -> list[dict]: + """Extract media items from messages in a single pass.""" + medias = [] + for msg in messages: + if msg['role'] != 'user' or not msg.get('content'): + continue + + for content_part in msg['content']: + if not isinstance(content_part, dict): + continue + + content_type = content_part.get('type') + if content_type in ['image_url', 'image']: + image_data = content_part.get(content_type) + assert image_data is not None, f"image data is missing for content part: {content_part}" + medias.append({ + 'type': 'image', + 'image': image_data, + }) + return medias + + def apply_chat_template(self, messages, **kwargs): + return self.tokenizer.apply_chat_template(messages, **kwargs) + + def batch_decode(self, *args, **kwargs): + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + return ['input_ids', 'attention_mask', 'pixel_values', 'grid_thws'] diff --git a/kimi_k3_vision_processing.py b/kimi_k3_vision_processing.py new file mode 100644 index 0000000000000000000000000000000000000000..58154424ff951289570c8fb5578e7372bfbfaa8b --- /dev/null +++ b/kimi_k3_vision_processing.py @@ -0,0 +1,179 @@ +"""Image processor class for Kimi-K3. +""" + +import json +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch +from PIL import Image +from transformers.image_processing_utils import (BaseImageProcessor, + BatchFeature) +from transformers.utils import TensorType + +from .media_utils import (MediaInput, TransparentBgConfig, _to_tensor, + ensure_media_type, image_to_np, navit_patchify, + navit_resize_image, normalize) + + +class KimiK3VisionProcessor(BaseImageProcessor): + model_type = "kimi_k3" + + def __init__( + self, + media_proc_cfg: dict, + **kwargs, + ): + super().__init__(**kwargs) + self.media_proc_cfg = media_proc_cfg + + @property + def _transparent_bg_config(self) -> Optional[TransparentBgConfig]: + cfg = self.media_proc_cfg.get("transparent_bg_config") + if cfg is None: + return None + if isinstance(cfg, TransparentBgConfig): + return cfg + return TransparentBgConfig(**cfg) + + @property + def _transparent_bg_fill_stage(self) -> str: + return self.media_proc_cfg.get("transparent_bg_fill_stage", + "before_resize") + + def media_tokens_calculator(self, media: MediaInput): + media = ensure_media_type( + media, + transparent_bg_config=self._transparent_bg_config, + transparent_bg_fill_stage=self._transparent_bg_fill_stage, + ) + ret = self.get_resize_config(media) + return ret['num_tokens'] + + @classmethod + def make_image_prompt(cls, width: int, height: int) -> str: + """Build the K3 image placeholder with resolution info.""" + return (f"<|media_begin|>image {width}x{height}" + f"<|media_content|><|media_pad|><|media_end|>") + + def get_resize_config(self, media_input: MediaInput) -> dict: + if media_input['type'] == 'image': + w, h = media_input['image'].size + ret = navit_resize_image( + w, h, self.media_proc_cfg['patch_size'], + self.media_proc_cfg['merge_kernel_size'], + self.media_proc_cfg['in_patch_limit'], + self.media_proc_cfg['patch_limit_on_one_side'], + self.media_proc_cfg['fixed_output_tokens']) + return ret + else: + raise ValueError("Unsupported type: {}".format( + media_input['type'])) + + def resize_image(self, image: Image.Image, new_width: int, new_height: int, + pad_width: int, pad_height: int) -> np.ndarray: + image_np = image_to_np( + image, + (new_width, new_height), + "resize", + transparent_bg_config=self._transparent_bg_config, + transparent_bg_fill_stage=self._transparent_bg_fill_stage, + ) + image_np = np.pad( + image_np, + ((0, pad_height), (0, pad_width), (0, 0)), + mode="constant", + constant_values=0, + ) + return image_np + + def preprocess( + self, + medias: list[MediaInput], + return_tensors: Optional[Union[str, TensorType]] = None, + ) -> BatchFeature: + """ + Preprocess a atom vision input (images) into model-ready tensors. + + Args: + medias: List of MediaInput. + return_tensors: Desired output format ('pt', 'np', 'tf', or None). + + Returns: + BatchFeature containing 'pixel_values' and 'grid_thws' tensors. + """ + if not isinstance(medias, list): + medias = [medias] + if medias: + pixel_values = [] + for item in medias: + item = ensure_media_type( + item, + transparent_bg_config=self._transparent_bg_config, + transparent_bg_fill_stage=self._transparent_bg_fill_stage, + ) + resize_config = self.get_resize_config(item) + new_width, new_height, pad_width, pad_height = resize_config[ + 'new_width'], resize_config['new_height'], resize_config[ + 'pad_width'], resize_config['pad_height'] + if item['type'] == 'image': + image = item['image'] + image_np = self.resize_image(image, new_width, new_height, + pad_width, pad_height) + pixel_values.append(np.expand_dims(image_np, axis=0)) + else: + raise ValueError("Unsupported type: {}".format( + item['type'])) + normalized_pixel_values = [] + image_std_inv = 1.0 / np.array(self.media_proc_cfg['image_std']) + image_mean = np.array(self.media_proc_cfg['image_mean']) + for pixels in pixel_values: + pixels = normalize(pixels, image_mean, image_std_inv) + pixels_and_thw = navit_patchify( + pixels, + self.media_proc_cfg['patch_size'], + ) + normalized_pixel_values.append(pixels_and_thw) + + pixel_values = torch.cat([ + _to_tensor(pixel_value['pixel_values']) + for pixel_value in normalized_pixel_values + ]) + grid_thws = torch.cat([ + _to_tensor(pixel_value['grid_thw'], + dtype=torch.int64).unsqueeze(0) + for pixel_value in normalized_pixel_values + ]) + + data = { + 'pixel_values': pixel_values, + 'grid_thws': grid_thws, + } + + else: + data = {} + + return BatchFeature(data=data, tensor_type=return_tensors) + + def __repr__(self): + return f"KimiK3VisionProcessor(media_proc_cfg={self.media_proc_cfg})" + + def to_dict(self) -> Dict[str, Any]: + output = super().to_dict() + output["media_proc_cfg"] = self.media_proc_cfg + if "media_processor" in output: + del output["media_processor"] + return output + + @classmethod + def from_dict(cls, config_dict: Dict[str, Any], **kwargs): + config = config_dict.copy() + media_proc_cfg = config.pop("media_proc_cfg", {}) + return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs) + + def to_json_string(self): + dictionary = self.to_dict() + for key, value in dictionary.items(): + if hasattr(value, 'tolist'): + dictionary[key] = value.tolist() + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" diff --git a/media_utils.py b/media_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e8b2ae127e254e364e6b568b1093a7e12ac576e0 --- /dev/null +++ b/media_utils.py @@ -0,0 +1,376 @@ +import base64 +import functools +import io +import math +from dataclasses import dataclass +from typing import Literal, TypedDict + +import numpy as np +from PIL import Image + + +class ImageInput(TypedDict): + type: Literal['image'] + image: Image.Image + + +MediaInput = ImageInput + + +@dataclass +class TransparentBgConfig: + """The config of the transparent background.""" + + pattern: Literal["white", "black", "gray", "chessboard"] = "black" + """The pattern of the transparent background.""" + + chessboard_square_size: int = 16 + """The size of the squares in the chessboard background.""" + + chessboard_square_on_top_left: bool = True + """Whether to start the chessboard with a white square on the top left.""" + + chessboard_white_value: int = 255 + """The value of the white pixels in the background.""" + + chessboard_gray_value: int = 200 + """The value of the gray pixels in the background.""" + + +@functools.lru_cache(maxsize=256) +def _create_chessboard_background( + height: int, + width: int, + square_size: int, + square_on_top_left: bool, + white_value: int, + gray_value: int, +) -> np.ndarray: + """Create a chessboard background.""" + bg = np.ones((height, width, 3), dtype=np.uint8) * white_value + for y in range(0, height, square_size): + for x in range(0, width, square_size): + if (y // square_size + x // square_size) % 2 == ( + 1 if square_on_top_left else 0): + bg[y:y + square_size, x:x + square_size] = gray_value + return bg + + +def fill_transparent_bg_with( + image: Image.Image, + transparent_bg_config: TransparentBgConfig | None = None, +) -> Image.Image: + """Composite a (possibly) transparent image onto a configured background. + + When ``transparent_bg_config`` is ``None``, the image is simply converted + to RGB (preserving the historical behavior). Otherwise the alpha channel + is alpha-composited over a background generated according to the config. + """ + if transparent_bg_config is None: + return image.convert("RGB") + + if image.mode == "RGB": + return image + + has_alpha = "A" in image.getbands() or "transparency" in image.info + if not has_alpha: + return image.convert("RGB") + + img = np.array(image.convert("RGBA")) + height, width = img.shape[:2] + bg_pattern = transparent_bg_config.pattern + if bg_pattern == "white": + bg = np.full((height, width, 3), 255, dtype=np.uint8) + elif bg_pattern == "black": + bg = np.zeros((height, width, 3), dtype=np.uint8) + elif bg_pattern == "gray": + bg = np.full((height, width, 3), 128, dtype=np.uint8) + elif bg_pattern == "chessboard": + bg = _create_chessboard_background( + height, + width, + transparent_bg_config.chessboard_square_size, + transparent_bg_config.chessboard_square_on_top_left, + transparent_bg_config.chessboard_white_value, + transparent_bg_config.chessboard_gray_value, + ) + else: + raise ValueError(f"Invalid background pattern: {bg_pattern}") + + alpha = img[:, :, 3] + img_rgb = img[:, :, :3] + alpha_normalized = alpha.astype(np.float32) / 255.0 + alpha_3d = np.stack([alpha_normalized] * 3, axis=2) + result = alpha_3d * img_rgb + (1 - alpha_3d) * bg + result = result.astype(np.uint8) + return Image.fromarray(result) + + +def navit_resize_image( + width: int, + height: int, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, +): + # Apply the patch limits. + s1 = math.sqrt( + in_patch_limit / + (max(1.0, width // patch_size) * max(1.0, height // patch_size))) + s2 = patch_limit_on_one_side * patch_size / width + s3 = patch_limit_on_one_side * patch_size / height + scale = min(1.0, s1, s2, s3) + new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale)) + new_w = min(new_w, patch_limit_on_one_side * patch_size) + new_h = min(new_h, patch_limit_on_one_side * patch_size) + + # Calculate the padding to make the height and width divisible by the merge kernel size and patch size. + factor = merge_kernel_size * patch_size + + pad_height = (factor - new_h % factor) % factor + pad_width = (factor - new_w % factor) % factor + + if fixed_output_tokens is not None: + num_tokens = fixed_output_tokens + else: + # Calculate new dimensions after padding and patching + token_height = (new_h + pad_height) // factor + token_width = (new_w + pad_width) // factor + + assert token_height * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}" + ) + assert token_width * merge_kernel_size <= patch_limit_on_one_side, ( + f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}" + ) + + num_tokens = token_height * token_width + return { + "num_tokens": num_tokens, + "new_width": new_w, + "new_height": new_h, + "pad_width": pad_width, + "pad_height": pad_height, + "sampled_nframes": 1, + } + + +def _to_pil( + data: str | bytes | Image.Image, + transparent_bg_config: TransparentBgConfig | None = None, + to_rgb: bool = True, +) -> Image.Image: + """Load an image and (optionally) composite its transparent background. + + Args: + data: A PIL Image, a base64 ``data:`` URL, a file path, or raw bytes. + transparent_bg_config: The config used to fill the transparent + background. ``None`` keeps the historical behavior of converting + to RGB without compositing. + to_rgb: If ``False`` the image is returned as-is (the + ``transparent_bg_config`` is ignored). The caller is then + expected to call :func:`fill_transparent_bg_with` later — e.g. + after a resize. + """ + if isinstance(data, Image.Image): + image = data + elif isinstance(data, str): + if data.startswith("data:"): + raw_base64 = data.split(",")[1] + image = Image.open(io.BytesIO(base64.b64decode(raw_base64))) + else: + image = Image.open(data) + elif isinstance(data, bytes): + image = Image.open(io.BytesIO(data)) + else: + raise ValueError(f"Unsupported data type: {type(data)}") + + if not to_rgb: + return image + + return fill_transparent_bg_with(image, transparent_bg_config) + + +def ensure_media_type( + media: MediaInput, + transparent_bg_config: TransparentBgConfig | None = None, + transparent_bg_fill_stage: Literal["before_resize", + "after_resize"] = "before_resize", +) -> MediaInput: + if media['type'] == 'image': + media['image'] = _to_pil( + media['image'], + transparent_bg_config=transparent_bg_config, + to_rgb=transparent_bg_fill_stage == "before_resize", + ) + return media + else: + raise ValueError(f"Unsupported media type: {media['type']}") + + +def image_to_np( + image: Image.Image, + resize_to: tuple[int, int] | None = None, + mode: str = "resize", + raise_error_for_ill_resize: bool = True, + transparent_bg_config: TransparentBgConfig | None = None, + transparent_bg_fill_stage: Literal["before_resize", + "after_resize"] = "before_resize", +) -> np.ndarray: + """Convert an image to a numpy array. + + Args: + content: The image to convert. + resize_to: The size to resize the image to. + mode: The mode to resize the image to. + raise_error_for_ill_resize: Whether to raise an error for ill-sized resize. + transparent_bg_config: The config of the transparent background. Only + used when ``transparent_bg_fill_stage == "after_resize"`` (the + caller is responsible for filling before resize otherwise). + transparent_bg_fill_stage: When to composite the transparent + background — before or after the resize step. + + Returns: + A numpy array. + """ + assert isinstance(image, Image.Image), "image must be a PIL Image" + if resize_to is not None: + if mode == "resize": + image = image.resize(resize_to, resample=Image.Resampling.BICUBIC) + if transparent_bg_fill_stage == "after_resize": + image = fill_transparent_bg_with(image, transparent_bg_config) + + elif mode == "rescale_and_pad_to_center": + scale = min(resize_to[0] / image.width, + resize_to[1] / image.height, 1.0) + new_width = round(image.width * scale) + new_height = round(image.height * scale) + if new_width == 0 or new_height == 0: + if raise_error_for_ill_resize: + raise ValueError( + f"Invalid resize to: {resize_to}, from image size: {image.size}" + ) + else: + return np.zeros((resize_to[1], resize_to[0], 3), + dtype=np.uint8) + + image = image.resize((new_width, new_height), + resample=Image.Resampling.BICUBIC) + if transparent_bg_fill_stage == "after_resize": + image = fill_transparent_bg_with(image, transparent_bg_config) + padding_left = (resize_to[0] - new_width) // 2 + padding_right = resize_to[0] - new_width - padding_left + padding_top = (resize_to[1] - new_height) // 2 + padding_bottom = resize_to[1] - new_height - padding_top + image = np.asarray(image) + image = np.pad( + image, + ((padding_top, padding_bottom), (padding_left, padding_right), + (0, 0)), + mode="constant", + constant_values=0, + ) + assert image.shape == (resize_to[1], resize_to[0], 3) + + elif mode == "rescale_and_pad_to_rightbottom": + scale = min(resize_to[0] / image.width, + resize_to[1] / image.height, 1.0) + new_width = round(image.width * scale) + new_height = round(image.height * scale) + if new_width == 0 or new_height == 0: + if raise_error_for_ill_resize: + raise ValueError( + f"Invalid resize to: {resize_to}, from image size: {image.size}" + ) + else: + return np.zeros((resize_to[1], resize_to[0], 3), + dtype=np.uint8) + + image = image.resize((new_width, new_height), + resample=Image.Resampling.BICUBIC) + if transparent_bg_fill_stage == "after_resize": + image = fill_transparent_bg_with(image, transparent_bg_config) + padding_right = resize_to[0] - new_width + padding_bottom = resize_to[1] - new_height + image = np.asarray(image) + image = np.pad( + image, + ((0, padding_bottom), (0, padding_right), (0, 0)), + mode="constant", + constant_values=0, + ) + assert image.shape == (resize_to[1], resize_to[0], 3) + + else: + raise ValueError(f"Invalid mode: {mode}") + + if isinstance(image, Image.Image): + return np.asarray(image) + else: + return image + + +def navit_patchify(pixel_values: np.ndarray, + patch_size: int) -> dict[str, np.ndarray]: + """Reshape the pixel values to a navit shape. + + Args: + pixel_values: np.ndarray, shape (t, h, w, c) + patch_size: int + + Returns: + dict[str, np.ndarray] + - patches: np.ndarray, shape (t * h//patch_size * w//patch_size, c, patch_size, patch_size) + - grid_thw: np.ndarray, (t, h//patch_size, w//patch_size) + """ + T, H, W, C = pixel_values.shape + assert C == 3, "pixel_values must have 3 channels" + + patches = pixel_values.reshape(T, H // patch_size, patch_size, + W // patch_size, patch_size, C) + # (T, H//patch_size, W//patch_size, C, patch_size, patch_size) + patches = patches.transpose(0, 1, 3, 5, 2, 4) + patches = patches.reshape(-1, C, patch_size, patch_size) + grid_thw = np.array([T, H // patch_size, W // patch_size]) + return {"pixel_values": patches, "grid_thw": grid_thw} + + +def normalize(x: np.ndarray, + mean, + std_inv, + pixels_dtype: np.dtype = np.float32) -> np.ndarray: + """Normalize the image. + + Args: + x: The image to normalize. The shape is (..., 3). The dtype is uint8. The range is [0, 255]. + mean: The mean of the image. + std_inv: The inverse of the std of the image. + pixels_dtype: The dtype of the image. + Returns: + The normalized image. The shape is (..., 3). The dtype is determined by the pixels_dtype. + """ + x = (x / 255.0).astype(pixels_dtype) + x -= mean + x *= std_inv + return x + + +def _to_tensor(data, **kwargs): + import torch + + if isinstance(data, np.ndarray): + return torch.from_numpy(data).to(**kwargs) + elif isinstance(data, torch.Tensor): + return data.to(**kwargs) + elif isinstance(data, list): + return [_to_tensor(item, **kwargs) for item in data] + elif isinstance(data, tuple): + return tuple(_to_tensor(item, **kwargs) for item in data) + elif isinstance(data, dict): + return {k: _to_tensor(v, **kwargs) for k, v in data.items()} + elif data is None: + return None + else: + raise ValueError(f"Unsupported data type: {type(data)}") diff --git a/model-00001-of-000096.safetensors b/model-00001-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f9e9d8d731e99d47d3b8c546fbfe5b4f66422c97 --- /dev/null +++ b/model-00001-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ca3e1f926f0da291c3423ed554e2facfaa1d4121c6a343938f4fab9db2f930b3 +size 2341216112 diff --git a/model-00002-of-000096.safetensors b/model-00002-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..33f53077205bce8804647d08b6c7f27a9b69a1d6 --- /dev/null +++ b/model-00002-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fa4ff4394450e2c9efaea95091b0935ad87b146feb0f7a5aae08c921103390cc +size 16990911504 diff --git a/model-00003-of-000096.safetensors b/model-00003-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..2df90c378e23c302c5b545338a39238a1672dff8 --- /dev/null +++ b/model-00003-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9eb36f76e8df07c3c42f2ed7b2fbb82050557efbaf612b84b15104fc5e28ee8e +size 16990911504 diff --git a/model-00004-of-000096.safetensors b/model-00004-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d76dbd5bf3029985122188787ba3a62f0c49f850 --- /dev/null +++ b/model-00004-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:24c89436d95f03baef3dca9f62e8a0fddbb31ae17b260e025cb40048409b008f +size 16567501776 diff --git a/model-00005-of-000096.safetensors b/model-00005-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a2a370c82f160d8a875422719c061a614236b8fb --- /dev/null +++ b/model-00005-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b600311c5d577340df7d35fb2b978735da595d289ae1aaa6c1f91654cf7fa629 +size 16990911504 diff --git a/model-00006-of-000096.safetensors b/model-00006-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..057e61029b8f51e89b5f5c7143f08f30e6c1781c --- /dev/null +++ b/model-00006-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8db3589e4ea063721ea09e1e13158008f146ad3a61eca208c65cf540cc53ae31 +size 16990911504 diff --git a/model-00007-of-000096.safetensors b/model-00007-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..02cbbf368d7773588711702a08f7eb20ffe2a3dc --- /dev/null +++ b/model-00007-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a23d133889b2485931ea7b64e33c00dc93605f2ddb9593a66af2b86845749d13 +size 16990911504 diff --git a/model-00008-of-000096.safetensors b/model-00008-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0ae754888f29661e70bf13284905b9b1751286fb --- /dev/null +++ b/model-00008-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0ee82aed1bbb5fbacbaccd7c6c20e8e65ea37ca0482d12cca6a73377abf35019 +size 16567501776 diff --git a/model-00009-of-000096.safetensors b/model-00009-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9152d322f895493d5d14ab97b5bb4c5b8f63117b --- /dev/null +++ b/model-00009-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cd0d9c51c586a88b953b28299fd5147f5d172da674bf8f34773c2b6f66971301 +size 16990911504 diff --git a/model-00010-of-000096.safetensors b/model-00010-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5163a305a264f427fcf2721921ababa8081980e4 --- /dev/null +++ b/model-00010-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7b8efeb0b4ef2a633f1d81cf5045d789c75a56a051d2d54ecac15c7654b1bbff +size 16990911504 diff --git a/model-00011-of-000096.safetensors b/model-00011-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a81bd4e0651af72841aa423b93b5987665ecdb64 --- /dev/null +++ b/model-00011-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:019d3c8ba61ec44dd4bc51a93c9b2bf6cc46ce8504950dda1d901d75447645c4 +size 16990916912 diff --git a/model-00012-of-000096.safetensors b/model-00012-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9e25442733e35e8c359c35e60316ff8ec88a76d4 --- /dev/null +++ b/model-00012-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:413ac73803d68d1ed9a65b53352a7f4c48b81a620d4ad827aa986214df09e968 +size 16567507176 diff --git a/model-00013-of-000096.safetensors b/model-00013-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b19058d515d941a9d1d65d2d9ee543d508b978b1 --- /dev/null +++ b/model-00013-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:25643aac8567945b5ec2a8794edc0c48d3efb3cded56ab69108e88b8575105fa +size 16990916912 diff --git a/model-00014-of-000096.safetensors b/model-00014-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..3cb85b4146ddbf567afb52cfada49cebcbcb6a60 --- /dev/null +++ b/model-00014-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d92a11e29ce9bccec1a60d9ed2dbc38ea0a6f6f171fa5c3ef60c9e406f318edc +size 16990916912 diff --git a/model-00015-of-000096.safetensors b/model-00015-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1e57600df70d1ebc0775a35f5ea617d759e444da --- /dev/null +++ b/model-00015-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:842cbe51489f73a25396755b6689e368276192e378b4536d6d7c6c83312c0e4a +size 16990916912 diff --git a/model-00016-of-000096.safetensors b/model-00016-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9efb29985631ffc815a5b21de9bd39e0815c3e5d --- /dev/null +++ b/model-00016-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:635c8221c13ddaf596dc62054f5b08c379c0b81ae87f228cfce08d3cbbb384e8 +size 16567507176 diff --git a/model-00017-of-000096.safetensors b/model-00017-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b5c34d9c70f0dea0dbd09331b388b16bc73d8316 --- /dev/null +++ b/model-00017-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8fd26354232e24ddf1504b4f51511345073ca71db809ca2923a06bb997c93942 +size 16990916912 diff --git a/model-00018-of-000096.safetensors b/model-00018-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..77cd7458387406016ada800be352859819a5b94c --- /dev/null +++ b/model-00018-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a6375a0416d791417627abc87b008c4bd835455fc9609dffa9952e27136fd678 +size 16990916912 diff --git a/model-00019-of-000096.safetensors b/model-00019-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a05ce3af069094aa67c9f17acc4892ae358e9429 --- /dev/null +++ b/model-00019-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b93ccbac9888d939ca94659215d52ea4da53f636515e1bd17088097d7c01f7d5 +size 16990916912 diff --git a/model-00020-of-000096.safetensors b/model-00020-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a165f85ebd5dd02dafc810fc7963a00683cb2c37 --- /dev/null +++ b/model-00020-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9c211ec8bc7ce7c7a4148ebb2f074d6ec225e54f5fa611c046119c33a776b71a +size 16567507176 diff --git a/model-00021-of-000096.safetensors b/model-00021-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4a4afdb4daa884a0ab8b59024f452a2fb072b23e --- /dev/null +++ b/model-00021-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fea2e7f59d76aaa8365055e7f7022623e680137afdb5a0f46a96a5cb52661a23 +size 16990916912 diff --git a/model-00022-of-000096.safetensors b/model-00022-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e4c5ea51efc7bf7308ee51f5d025cb74d99c0e56 --- /dev/null +++ b/model-00022-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5b02b63f0783dc51ce33c5f9ed4deb8695206d6b77995e2a8ff985fd90cdd1b3 +size 16990916912 diff --git a/model-00023-of-000096.safetensors b/model-00023-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..69e166dbd8ee2872a2d058c4ec2b1af7f702536e --- /dev/null +++ b/model-00023-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b1e1b3eaf6d034c9fd33b8296ba723b876174aef5c8335269d9b96b586eb0621 +size 16990916912 diff --git a/model-00024-of-000096.safetensors b/model-00024-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..c6637da4add2cb263f7652e525db2986d3b37c8a --- /dev/null +++ b/model-00024-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7f23243580226c91605b15fe1fd8d7dd5b24d11942fea36c416969c0b15a362f +size 16567507176 diff --git a/model-00025-of-000096.safetensors b/model-00025-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..7ba71c79b6ab9c0c637b3c0f3049b05f63f62cfb --- /dev/null +++ b/model-00025-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:29be09005f802a7351bdec8ee77b42cecfd492e5ffc25f8c91bda5cedb02688a +size 16990916912 diff --git a/model-00026-of-000096.safetensors b/model-00026-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..70accf5eb3bb3690a7bb876db190bb7a92ec6c39 --- /dev/null +++ b/model-00026-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e474b6fc04c8b36cb23cc78f0fd3a00398f26619aa864dbd04e5859f2a39a6bc +size 16990916912 diff --git a/model-00027-of-000096.safetensors b/model-00027-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0f3240136dc82718837fcc1790af5bfce791cf3f --- /dev/null +++ b/model-00027-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eec4d4773a5e25c6f6edf6b664b2aaab050ac6d659da5fd46fffe2172c5c34f0 +size 16990916912 diff --git a/model-00028-of-000096.safetensors b/model-00028-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..eaabdb7f9ae038ef00afa150080133beea508e10 --- /dev/null +++ b/model-00028-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4ffe399002a4982c20fe09ce28762b45e3249fd1d76f8ad61a63b31ebbaa7ebb +size 16567507176 diff --git a/model-00029-of-000096.safetensors b/model-00029-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1f48fc92e127545943fabf67c8da39185a1accaa --- /dev/null +++ b/model-00029-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e384d165847c9dd92e64442d753a8ca38ecd8efda91862abdac99766c4415a47 +size 16990916912 diff --git a/model-00030-of-000096.safetensors b/model-00030-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..7cf88b0d52b30fae47550bbe067f4e26a663ce9f --- /dev/null +++ b/model-00030-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:07b8c5a517cc00b5e59fe1c21888ca86cd8fffb9504007406edb8913379d9a97 +size 16990916912 diff --git a/model-00031-of-000096.safetensors b/model-00031-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..8196d63c02c470a2ac3bcfff5f18c58de2ac4ba6 --- /dev/null +++ b/model-00031-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dcabfea9a0abd0b3e4b2e32807a2ec441d39ce6650988b6c5be657f59db285aa +size 16990916912 diff --git a/model-00032-of-000096.safetensors b/model-00032-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d6086bccbd512cce21ddf84bfb6d87e721eeab47 --- /dev/null +++ b/model-00032-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:271ba5da377eb639efa324c74070d71c9391c310375990847fd63a652c775096 +size 16567507176 diff --git a/model-00033-of-000096.safetensors b/model-00033-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..c723a3826121277bffd211fb1d9e2ae221c34623 --- /dev/null +++ b/model-00033-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85a0abaa9341b7fee090314bf40b3d6f05cb00c26c8438d762dff5445cfbb3cd +size 16990916912 diff --git a/model-00034-of-000096.safetensors b/model-00034-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..2b7e5259ed578ab040b1ebc280510619b0f1522e --- /dev/null +++ b/model-00034-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1f6a1cd3bc7cf3eefc08cc1be94517fd7c5d59969b5916232b6a065becf3ba00 +size 16990916912 diff --git a/model-00035-of-000096.safetensors b/model-00035-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..866f956c176581475dfc8ab530b4126cc4f70536 --- /dev/null +++ b/model-00035-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:835518f58c230eae24ffaa1d6a54746a845843f78b3963be574890cb3bafc600 +size 16990916912 diff --git a/model-00036-of-000096.safetensors b/model-00036-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..41185bce5b9d7d238a551bc94375b14e3e4ad9ac --- /dev/null +++ b/model-00036-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:03e4c43386682c8e07053a96eb841d13282396f0ebc872f80528b22273080825 +size 16567507176 diff --git a/model-00037-of-000096.safetensors b/model-00037-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..ceeb61203abbc30379c7a5ae98f432df1092d452 --- /dev/null +++ b/model-00037-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5087848a65cde0681fa4077f5785e933314b37ebf25e31bc0d00860ab1f0ed8e +size 16990916912 diff --git a/model-00038-of-000096.safetensors b/model-00038-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..92a8d24c0ab39ece1f518f5f201fbc342f8eb2db --- /dev/null +++ b/model-00038-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bbf7b2b58b3d703b48615551bc5709bde0b09f7324dd06e014aca397f90963d8 +size 16990916912 diff --git a/model-00039-of-000096.safetensors b/model-00039-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9990c53cb21b9287b25f50337fbd400695c57ed7 --- /dev/null +++ b/model-00039-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5c1ebb9e62a6697f94d931e27cd7a61d5b386300886b408303a5303823e36e15 +size 16990916912 diff --git a/model-00040-of-000096.safetensors b/model-00040-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4c4e206c57a1e9df89ae9f83a64c80c0168416f8 --- /dev/null +++ b/model-00040-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:27cc21669dc0e3b1f550e2c2c66b1c9601ed79ddeb293e425b46c034b88deae6 +size 16567507176 diff --git a/model-00041-of-000096.safetensors b/model-00041-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..60cbac9c5f3f46e70d6907513493617cbe49d547 --- /dev/null +++ b/model-00041-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a836cbe5d703b50d24e8a49591ce35173a13920133dac0325733b3c608b9948a +size 16990916912 diff --git a/model-00042-of-000096.safetensors b/model-00042-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..eacb1a53bba353b101ad0cb92fce5eb3ad50b90b --- /dev/null +++ b/model-00042-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f723491757fdeb9a11a3278e04824c169ddbc7f76eb4c5af1c1777826c323640 +size 16990916912 diff --git a/model-00043-of-000096.safetensors b/model-00043-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..68f39ce2b99d2cf69f8828ba3bcc323fbd073982 --- /dev/null +++ b/model-00043-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:492bd75434e4a98b5add493b9c48cd3701fce6c5d213f6518c062ea47ea9c9ac +size 16990916912 diff --git a/model-00044-of-000096.safetensors b/model-00044-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..dacf14f540699db7dbe1c899aa17d06d2c4a8c74 --- /dev/null +++ b/model-00044-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:be36ed51ad093b21f7a19ecfa15699a1271ff5164dbc1329814b51d93ac8b657 +size 16567507176 diff --git a/model-00045-of-000096.safetensors b/model-00045-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..424bc818ed60b35739610a1b9845f54e450921cb --- /dev/null +++ b/model-00045-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6a870e05de0ad367300f1aee8dfc506d11fcdd2957c93f1979860cc1e3030310 +size 16990916912 diff --git a/model-00046-of-000096.safetensors b/model-00046-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d8d7c18fd4e7952161f9b140ecd1e027b99422cd --- /dev/null +++ b/model-00046-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fa1b9091bdbd5f4e1a12c3888f253446f901f49e518709b963ac5fd7cdcd7a91 +size 16990916912 diff --git a/model-00047-of-000096.safetensors b/model-00047-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..edd494f83e43d119ccae1d24a01481cbe7668f28 --- /dev/null +++ b/model-00047-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e66b70d813b0baa4a954931b2244a5e7dc04f5de18bb8caaf719896d909001ac +size 16990916912 diff --git a/model-00048-of-000096.safetensors b/model-00048-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a4b42325af66a6c00d44a1e33947f8825bd1dcb6 --- /dev/null +++ b/model-00048-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9fef2805d295e1ff489070b1c91f11cade08d1b2ff38b7347e065291b24a1172 +size 16567507176 diff --git a/model-00049-of-000096.safetensors b/model-00049-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1865f8042289190d7618ff1bf0f2c885174f1161 --- /dev/null +++ b/model-00049-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f72a14d42bc86fb3cb8461eea41f9298d62653e5068c843eb3b450ab53fd943a +size 16990916912 diff --git a/model-00050-of-000096.safetensors b/model-00050-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..49de8969455e11f9d7e7b4bdf6afb7719debee17 --- /dev/null +++ b/model-00050-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:897b56517923bc63ca3dbdf15f6a9ec8a386b23e0d6fc70b282f616636c3871e +size 16990916912 diff --git a/model-00051-of-000096.safetensors b/model-00051-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..08729b0e9c57e44124948ed3dcbecefd14930bab --- /dev/null +++ b/model-00051-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:90a18027c6bf6980f5bff31ba62b68252acc9285b9ac220aef79eaca56e025a5 +size 16990916912 diff --git a/model-00052-of-000096.safetensors b/model-00052-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4eb16f28eeb657a1e9d78e2731a9d8bcabc42352 --- /dev/null +++ b/model-00052-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ecf2eae18212165eeaf576f81534ce21e188c3b28451000415009a4da8a14c72 +size 16567507176 diff --git a/model-00053-of-000096.safetensors b/model-00053-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..922ed7be7a165a509f17c06836886c239bafdfc3 --- /dev/null +++ b/model-00053-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ae8e5ce1655647a0fb33995d8fb21a64d5088bbeb47dab6ae8506fab35e5eeaa +size 16990916912 diff --git a/model-00054-of-000096.safetensors b/model-00054-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..dbad3cf3f1eaee09d849a99e4d371fa5ffa64d04 --- /dev/null +++ b/model-00054-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6f90693f82faea6aa28a26bed8d5f15c5ec54b5bf992c141dab6ce9b668795e0 +size 16990916912 diff --git a/model-00055-of-000096.safetensors b/model-00055-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b5f94696a8854876386aa31947bc985ed82189c2 --- /dev/null +++ b/model-00055-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cc0663717792f2ac512cbadb13112cd006e5e72c7fae894e044ab77ada2a4222 +size 16990916912 diff --git a/model-00056-of-000096.safetensors b/model-00056-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..3f2e6df59b83e5f5665b7fca86301d6debfa3877 --- /dev/null +++ b/model-00056-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c3b1ef3976a0a915f7b0a2d803b62d008b728d3489fae17e71aceeb1774a8c32 +size 16567507176 diff --git a/model-00057-of-000096.safetensors b/model-00057-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b3ce61712edba259fc214c26115a6c1b0f98eee7 --- /dev/null +++ b/model-00057-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f99706201b07b3778c0ea1ff6afcc010f52bf6e65ac19eb1bdc29e994fd09eab +size 16990916912 diff --git a/model-00058-of-000096.safetensors b/model-00058-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f4a55cb3e1a009aa82b351df483ada9a7f353f74 --- /dev/null +++ b/model-00058-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d4d0cc8c01a051f422c79ef702612d6c432a645d5d988c2dc0a771d82d8ea31b +size 16990916912 diff --git a/model-00059-of-000096.safetensors b/model-00059-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..2cd3720b83eb36cb25143974d7678796ff3ac825 --- /dev/null +++ b/model-00059-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e8176ff8dc8bdf3869a25b111d78cf03504c8aaf8aea9e7391f7517baf63b124 +size 16990916912 diff --git a/model-00060-of-000096.safetensors b/model-00060-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5745e8f7cd513714c009ecfe086fe50ebbdee9e4 --- /dev/null +++ b/model-00060-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:88c13ca09dbf9a0fee48894d44c00e4b0994d69e3ae2bb07e192f31f8913ff4b +size 16567507176 diff --git a/model-00061-of-000096.safetensors b/model-00061-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..79accf865a3fcf53fbed36735db006547a3c165a --- /dev/null +++ b/model-00061-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:608b247126c7f80073206d78b8d99f685214a940c882331518f8d2e2b372f6f1 +size 16990916912 diff --git a/model-00062-of-000096.safetensors b/model-00062-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..42ae2e0235f4c685df7d5b63735b5d36cefef03d --- /dev/null +++ b/model-00062-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c827ea899796411560de21b806ebfa592cbfc5a90d7d17ce924700bf0d8bcf19 +size 16990916912 diff --git a/model-00063-of-000096.safetensors b/model-00063-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..74600a9791c6633184e47806ea7dca74618fdcfc --- /dev/null +++ b/model-00063-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8a29ad10c4e26d67ea64c14e73fdd7fe276e0189d4128b3861076d5b662980a7 +size 16990916912 diff --git a/model-00064-of-000096.safetensors b/model-00064-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b6640ea9ed101cac896632ac641a70bbafa93e1c --- /dev/null +++ b/model-00064-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ea2bb9c37bdb28a2b3fe9acae0ab05cff85efa06bffefb2a1205f11dec8c60a6 +size 16567507176 diff --git a/model-00065-of-000096.safetensors b/model-00065-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..25cbd03056995864e8f50a173a2a3f89895e83d9 --- /dev/null +++ b/model-00065-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bdcb8253b41d495e8c3c76a81bee130d4b303b6997f9cfd8ebaa8c1be922817d +size 16990916912 diff --git a/model-00066-of-000096.safetensors b/model-00066-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e6ceadf26eefe8544c4cecf73f8b218d01b3d4d0 --- /dev/null +++ b/model-00066-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3479a3be0b872066b1c0cf52f2df96c0c0f7dbb1f1cc5c2524d2668516f2c23a +size 16990916912 diff --git a/model-00067-of-000096.safetensors b/model-00067-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..42f9199fbf1e6a760a95a89bbd1d624c836e09fe --- /dev/null +++ b/model-00067-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b419d2bea9174ce11ff3ac683978f99ea29441016a464b85f2b9a1144d909e30 +size 16990916912 diff --git a/model-00068-of-000096.safetensors b/model-00068-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..3047e8fbf584a3c5d6ca5484f64fd742c89680a4 --- /dev/null +++ b/model-00068-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a75fb167a6b6799d70119530b7ebe2ec11a0fdaae6ea5e708d547d7ee9e19b5b +size 16567507176 diff --git a/model-00069-of-000096.safetensors b/model-00069-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..99834eb95ae85ae6814433d204e48733ac73ac25 --- /dev/null +++ b/model-00069-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2c581d936845f56af93ec8a3e6a5ae4b233cbd47421cd2d0ab45ca7f188d48fc +size 16990916912 diff --git a/model-00070-of-000096.safetensors b/model-00070-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..ef7c809630dea5fe5da970315e0ad7d5833c52c3 --- /dev/null +++ b/model-00070-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dc4c972b303a8965e60d4c7ba3a7aaa5a8bbb667d6d338f7e45e88a9a8a8144e +size 16990916912 diff --git a/model-00071-of-000096.safetensors b/model-00071-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..85d5d1fa0b96734455ef463ecee1190b83a35dc0 --- /dev/null +++ b/model-00071-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:54068cd012c5661a724c404511601870518b96698c66d79ce8118ffee47cd92e +size 16990916912 diff --git a/model-00072-of-000096.safetensors b/model-00072-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0df89ad8de20c826e2be43e07398a7abdcf18565 --- /dev/null +++ b/model-00072-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0c371b5746d66aae35271f6d4b074137e468c1c92a7d7a8c90dbd4c7a2398bb9 +size 16567507176 diff --git a/model-00073-of-000096.safetensors b/model-00073-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0d04b0757e52a219e68043095f35fa2e90ff2652 --- /dev/null +++ b/model-00073-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2ae258f2c9f590f0df0413a24fed1420c8a3812d6c592335a14ae485dad7f874 +size 16990916912 diff --git a/model-00074-of-000096.safetensors b/model-00074-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6fc349fa7ac8a5fa65bf96f96fc5b62a14a5c5ec --- /dev/null +++ b/model-00074-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cc45b3e6cf51fb18db6a77d4428d5aec1ccbdc243c1486cc747293ff073f3492 +size 16990916912 diff --git a/model-00075-of-000096.safetensors b/model-00075-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..54cb1e48f11e76048c20b31444946d828569fec6 --- /dev/null +++ b/model-00075-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bf3fb0e94a694f274b4de7257d95f22684e9dfae0900bd95c4acd233c05c9f82 +size 16990916912 diff --git a/model-00076-of-000096.safetensors b/model-00076-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e4cb4f6f899cf8062ba08953ca594b7821816e47 --- /dev/null +++ b/model-00076-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9d7ebb90879ea8ff88d8482b23102d067199b863544886161a7c934616a80b51 +size 16567507176 diff --git a/model-00077-of-000096.safetensors b/model-00077-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a37150179515c66b7b4a8d3e66fac7e3c9383aca --- /dev/null +++ b/model-00077-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b74972c08cb5a0b7c24a49b228bd40b4494a5764a0cf944f9339b263f5dc5ce6 +size 16990916912 diff --git a/model-00078-of-000096.safetensors b/model-00078-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..a11f48ec42e883db4408f02cf70e43b6c573121a --- /dev/null +++ b/model-00078-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7433591746769efed5c3234519f47018886ad7a5c2d23c907fc52e06a2e31e6e +size 16990916912 diff --git a/model-00079-of-000096.safetensors b/model-00079-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5aee569f399faddead3825bdb877c2cce2ffd4fc --- /dev/null +++ b/model-00079-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c1e9bc17c33e5ccbcc3284a02744e2c87a462586f70c9b82df39f9dd08cb61e4 +size 16990916912 diff --git a/model-00080-of-000096.safetensors b/model-00080-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9e0e7262ea8097d6be3c81bc3c35d26fdc23f328 --- /dev/null +++ b/model-00080-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:af0ec340d25367f9dd46df33e98793fb15f074c5272192d6fb33bb25695c5654 +size 16567507176 diff --git a/model-00081-of-000096.safetensors b/model-00081-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..937c0848a8a97bb33ef6b4767de05b5cf749a189 --- /dev/null +++ b/model-00081-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:96c4f147b2f0b190ae8f480a383e94c9aaf11f8d109716a7117d0156ecb4a55e +size 16990916912 diff --git a/model-00082-of-000096.safetensors b/model-00082-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1e48693906328e3c3cbf9e8461b18d9232a87d63 --- /dev/null +++ b/model-00082-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0c8bc93c36f8a8699af04ad23de3a62e6581d3cae1275632fa6af1ef3b949c25 +size 16990916912 diff --git a/model-00083-of-000096.safetensors b/model-00083-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fe1930de6f26fa77872a3d330871ee4d2ae22c8f --- /dev/null +++ b/model-00083-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6e753bd257727caf4cff8bcf61e608c3edb6cacdc75e35974e9842b1ba6c1308 +size 16990916912 diff --git a/model-00084-of-000096.safetensors b/model-00084-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..db2a790a16a6a888c65aa834262511c8a6c1c704 --- /dev/null +++ b/model-00084-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cf3958f31aa56a23ff17c9a4ee34457027048822dfa4f7480b1f68bc81702318 +size 16567507176 diff --git a/model-00085-of-000096.safetensors b/model-00085-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..080cd4cf29db9796c1b8dc081999e781bafbe841 --- /dev/null +++ b/model-00085-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d5eb0c1ad6e415893455c3bde89a0617e39af47641cabb6f277359c6a1ffb8fd +size 16990916912 diff --git a/model-00086-of-000096.safetensors b/model-00086-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..66f93183fc75bc8e5010f7ae8dedfe9272f69595 --- /dev/null +++ b/model-00086-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7615a1dbbaec3448a1bbd5f2e5cba21210c717bcd71e20082fa7d324aa59fbc9 +size 16990916912 diff --git a/model-00087-of-000096.safetensors b/model-00087-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..5b77eabd9c1be125f8c7644b266d71d45b39a678 --- /dev/null +++ b/model-00087-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c6c81578d1a496f043133bdf3ab866d7475925d7da1a6181ab0fe9aa1a4472c1 +size 16990916912 diff --git a/model-00088-of-000096.safetensors b/model-00088-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..def9e82a7976ac6fdba806daa9790bc041f55a06 --- /dev/null +++ b/model-00088-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:abc21b10c201b52b6debecf9123d05e8192dff26c7e5edab1e74843e0c35bfa9 +size 16567507176 diff --git a/model-00089-of-000096.safetensors b/model-00089-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..4bbeaa200bdcd1d08a33d7cd6a322cbd703f588c --- /dev/null +++ b/model-00089-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:45d9c198d6de977df43d08bade817b6401acea6f06364f20d2ec38802d3ee981 +size 16990916912 diff --git a/model-00090-of-000096.safetensors b/model-00090-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..1b75ad1ac3e56905d16415168f5f4993705189e3 --- /dev/null +++ b/model-00090-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ee5f33f83ab717b8335853f2c4d7a4787408d83af566510f8c524f7ff8a109e9 +size 16990916912 diff --git a/model-00091-of-000096.safetensors b/model-00091-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0c1f58da614a9b89189b9979a308894809831a48 --- /dev/null +++ b/model-00091-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1b4210640177e2a579a0e5808569eacebc4b2e2d39eec50dfa8b1f46eeda340c +size 16990916912 diff --git a/model-00092-of-000096.safetensors b/model-00092-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..164c0660f6a193c4b86f7caeb29d4a73e1ac4143 --- /dev/null +++ b/model-00092-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b6fb9f9cfb89dbee02aeaca77a60658c2d12512cdbadda00c80d60c921cd2bf6 +size 16567507176 diff --git a/model-00093-of-000096.safetensors b/model-00093-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f8e26b7afab549ac21517e38d00724a71ef620bc --- /dev/null +++ b/model-00093-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e9c537b65bd6e4d2cb5941cc63c1061737633dad9b9bd565c7bb70b913ff51bf +size 16567507176 diff --git a/model-00094-of-000096.safetensors b/model-00094-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..70994b08f8af11b4b77be70e3b21125749ebe2bd --- /dev/null +++ b/model-00094-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5fdf2de358f4cb1070ad0eb6d4bac1a8e0ccdeb18f7deffd5788d5726c0e7f44 +size 4697664072 diff --git a/model-00095-of-000096.safetensors b/model-00095-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..51f85478a3b7a34165adea2210c6bc43add63c61 --- /dev/null +++ b/model-00095-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3e055a0b915b7f6cef3a4866bb45b4dff15c21734710fe181d6b6253ea00326e +size 92289328 diff --git a/model-00096-of-000096.safetensors b/model-00096-of-000096.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..da296080d7800b10d6f98c8b5e093cb35923bd23 --- /dev/null +++ b/model-00096-of-000096.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9d10c74fc10161bef9463a8541a634a97f521f43c99368ea7243ce0c79cdbf7c +size 802448352 diff --git a/model.safetensors.index.json b/model.safetensors.index.json new file mode 100644 index 0000000000000000000000000000000000000000..2e7014d727cf518d32bf21674e3b16687b2ba21b --- /dev/null +++ b/model.safetensors.index.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a1c5210650ce71d2d3ae9ec5a101ac4afd3cf4b10091be589853437eb967febd +size 59764096 diff --git a/modeling_kimi_k3.py b/modeling_kimi_k3.py new file mode 100644 index 0000000000000000000000000000000000000000..be85612710534c2e41065b67933a9639e7817d7f --- /dev/null +++ b/modeling_kimi_k3.py @@ -0,0 +1,1317 @@ +# coding=utf-8 +# Copyright 2025-2026 The Moonshot AI Team and HuggingFace Inc. team. All rights reserved. +# +# The code is based on llava (llava/modeling_llava.py), but modified for Kimi-K3. +# +# Licensing Information: +# - Code derived from llava (llava/modeling_llava.py) is licensed under the Apache License, Version 2.0. +# - Other parts of the code are licensed under the Kimi K3 License (see the LICENSE file in this repository). +# +# Apache License, Version 2.0: +# 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. + + +# NOTE: Reference implementation for model architecture; see the model card for production deployment. +import math +from collections.abc import Sequence +from copy import deepcopy +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import activations + +try: + from transformers.activations import PytorchGELUTanh +except ImportError: + from transformers.activations import GELUTanh + activations.PytorchGELUTanh = GELUTanh + PytorchGELUTanh = GELUTanh +from transformers.activations import PytorchGELUTanh +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_utils import PreTrainedModel +from transformers.models.llava.modeling_llava import \ + LlavaCausalLMOutputWithPast +from transformers.utils import is_flash_attn_2_available + +from .configuration_kimi_k3 import KimiK3Config +from .modeling_kimi_linear import KimiLinearForCausalLM + +# Flash attention imports +if is_flash_attn_2_available(): + from flash_attn import flash_attn_varlen_func +else: + flash_attn_varlen_func = None + + +def multihead_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_cu_seqlens: torch.Tensor | None = None, + k_cu_seqlens: torch.Tensor | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + deterministic: bool = False, +): + """Multi-head attention using flash attention 2. + + Args: + q, k, v: tensor of shape (batch_size, seqlen, num_heads, head_dim), + or (tot_seqlens, num_heads, head_dim) if packing. + q_cu_seqlens (torch.Tensor): cumulative sequence lengths of q. + The first element should be 0 and the last element should be q.shape[0]. + k_cu_seqlens (torch.Tensor): cumulative sequence lengths of k. + The first element should be 0 and the last element should be k.shape[0]. + + Returns: + output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing, + where dim = num_heads * head_dim + """ + attn_out = flash_attn_varlen_func( + q, + k, + v, + q_cu_seqlens, + k_cu_seqlens, + max_seqlen_q, + max_seqlen_k, + causal=False, + deterministic=deterministic, + ) + if isinstance(attn_out, tuple): + attn_out = attn_out[0] + + attn_out = attn_out.flatten(start_dim=-2) + + return attn_out + + +def eager_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_cu_seqlens: Optional[torch.Tensor] = None, + k_cu_seqlens: Optional[torch.Tensor] = None, + **kwargs, +) -> torch.Tensor: + seq_length = q.shape[0] + attention_mask = torch.zeros([1, seq_length, seq_length], + device=q.device, + dtype=torch.bool) + for i in range(1, len(q_cu_seqlens)): + attention_mask[ + ..., + q_cu_seqlens[i - 1]:q_cu_seqlens[i], + q_cu_seqlens[i - 1]:q_cu_seqlens[i], + ] = True + q = q.transpose(0, 1) + k = k.transpose(0, 1) + v = v.transpose(0, 1) + + attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) + attn_weight = attn_weight.masked_fill( + ~attention_mask, torch.finfo(attn_weight.dtype).min) + attn_weight = torch.softmax(attn_weight, dim=-1, + dtype=torch.float32).to(q.dtype) + + attn_output = attn_weight @ v + attn_output = attn_output.transpose(0, 1) + attn_output = attn_output.reshape(seq_length, -1) + return attn_output + + +VL_VISION_ATTENTION_FUNCTIONS = { + "flash_attention_2": multihead_attention, + "eager": eager_attention, +} + + +def _apply_rope_input_validation(x, freqs_cis): + assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape) + assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape) + assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape) + assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype + + +def get_rope_shape_decorate(func): + _get_rope_shape_first_call_flag = set() + + def wrapper(org, interpolation_mode, shape): + key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode) + if key not in _get_rope_shape_first_call_flag: + _get_rope_shape_first_call_flag.add(key) + _ = func(org, interpolation_mode, shape=(64, 64)) + return func(org, interpolation_mode, shape) + + return wrapper + + +@get_rope_shape_decorate +@torch.compile(dynamic=True) +def get_rope_shape(org, interpolation_mode, shape): + return (F.interpolate( + org.permute((2, 0, 1)).unsqueeze(0), + size=shape, + mode=interpolation_mode, + ).squeeze(0).permute((1, 2, 0)).flatten(end_dim=1)) + + +def apply_rope(xq: torch.Tensor, xk: torch.Tensor, + freqs_cis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: (The leading dimensions of all inputs should be the same) + xq: query, tensor of shape (..., num_heads, head_dim) + xk: key, tensor of shape (..., num_heads, head_dim) + freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid. + Returns: + xq_out, xk_out: tensors of shape (..., num_heads, head_dim) + """ + _apply_rope_input_validation(xq, freqs_cis) + _apply_rope_input_validation(xk, freqs_cis) + + freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2 + # ..., num_heads, head_dim/2 + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) + xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2)) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten( + -2) # ..., num_heads, head_dim + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten( + -2) # ..., num_heads, head_dim + return xq_out.type_as(xq), xk_out.type_as(xk) + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + From: + https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86 + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float32) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False): + """ + t_size: int of the temporal size + return: + pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token) + """ + grid_t = np.arange(t_size, dtype=np.float32) + pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], + axis=0) + return pos_embed + + +class Learnable2DInterpPosEmbDivided_fixed(nn.Module): + + def __init__(self, + height: int, + width: int, + num_frames: int, + dim: int, + interpolation_mode: str = 'bicubic') -> None: + super().__init__() + self.height = height + self.width = width + self.num_frames = num_frames + self.dim = dim + self.interpolation_mode = interpolation_mode + self.weight = nn.Parameter(torch.empty(height, width, dim)) + self.register_buffer('time_weight', + torch.from_numpy( + get_1d_sincos_pos_embed( + self.dim, + self.num_frames)).float().unsqueeze(1), + persistent=False) + + self.reset_parameters() + + def reset_parameters(self): + nn.init.normal_(self.weight) + + def forward(self, x: torch.Tensor, + grid_thws: torch.Tensor) -> torch.Tensor: + pos_embs = [] + for t, h, w in grid_thws.tolist(): + assert t <= self.num_frames, f't:{t} > self.num_frames:{self.num_frames}' + if (h, w) == self.weight.shape[:-1]: + pos_emb_2d = self.weight.flatten(end_dim=1) + else: + pos_emb_2d = get_rope_shape( + self.weight, + interpolation_mode=self.interpolation_mode, + shape=(h, w), + ) + + if t == 1: + pos_emb_3d = pos_emb_2d + else: + pos_emb_3d = pos_emb_2d.unsqueeze(0).repeat( + t, 1, 1) + self.time_weight[0:t] + + pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) + + out = x + torch.cat(pos_embs) + return out + + +class MoonVision3dPatchEmbed(nn.Module): + + def __init__(self, + out_dim: int, + in_dim: int = 3, + patch_size: int | tuple[int, int] = (14, 14), + pos_emb_height: int = 14, + pos_emb_width: int = 14, + pos_emb_time: int = 4, + pos_emb_type: str = 'divided_fixed', + patch_embed_proj_bias: bool = True, + pos_emb_interpolation_mode: str = 'bicubic'): + super().__init__() + assert isinstance( + patch_size, + int | Sequence), f'Invalid patch_size type: {type(patch_size)}' + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + assert (len(patch_size) == 2 + ), f'Expected patch_size to be a tuple of 2, got {patch_size}' + self.patch_size = patch_size + + self.proj = nn.Conv2d(in_dim, + out_dim, + kernel_size=patch_size, + stride=patch_size, + bias=patch_embed_proj_bias) + + if pos_emb_type == 'divided_fixed': + self.pos_emb = Learnable2DInterpPosEmbDivided_fixed( + height=pos_emb_height, + width=pos_emb_width, + num_frames=pos_emb_time, + dim=out_dim, + interpolation_mode=pos_emb_interpolation_mode) + else: + raise NotImplementedError( + f'Not support pos_emb_type: {pos_emb_type}') + + def forward(self, x: torch.Tensor, + grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + x (L, Channels): input tensor + grid_hws (N, 3): temporal, height and width + + Returns: + (L, Cout) tensor + """ + x = self.proj(x).view(x.size(0), -1) + # apply positional embedding + x = self.pos_emb(x, grid_thws) + return x + + +class Rope2DPosEmbRepeated(nn.Module): + """2D rotary position embedding with multi-resolution support. + + This class is intended to be used in the following way: + 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis. + 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration. + 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation. + The rope is shared across all attention layers and all heads. + + Refs: + - RoFormer: https://arxiv.org/abs/2104.09864 + - VisionLLaMA: https://arxiv.org/abs/2403.00522 + - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py + + Args: + dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed) + max_height (int): the maximum height of the 2D grid + max_width (int): the maximum width of the 2D grid + theta_base (float): the base of the theta + device (str): the device to store the precomputed cis + """ + + def __init__(self, + dim: int, + max_height: int, + max_width: int, + theta_base=10000): + super().__init__() + self.dim = dim + assert self.dim % 4 == 0, 'dim must be divisible by 4' + self.max_height = max_height + self.max_width = max_width + self.theta_base = theta_base + + def extra_repr(self): + return f'dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}' + + def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: + """Calculate the cis(freqs) for each position in the 2D grid. + + Return: complex tensor of shape (max_height, max_width, dim//2) and value: + height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim)) + weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4)) + note: `cis` is a mathematical notation defined by cis x = cos x + i sin x, + """ + N = self.max_height * self.max_width + flat_pos = torch.arange(0, N).float().to(device) + x_pos = flat_pos % self.max_width + y_pos = flat_pos // self.max_width + dim_range = (torch.arange(0, self.dim, + 4)[:(self.dim // 4)].float().to(device) + ) # C/4 + freqs = 1.0 / (self.theta_base**(dim_range / self.dim)) + x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 + y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 + # N, C/4, 2 + freqs_cis = torch.cat( + [x_cis.unsqueeze(dim=-1), + y_cis.unsqueeze(dim=-1)], dim=-1) + # max_height, max_width, C/2 + freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) + return freqs_cis + + def get_freqs_cis(self, grid_thws: torch.Tensor, + device: torch.device) -> torch.Tensor: + """ + Args: + grid_thws (torch.Tensor): grid time, height and width + + Returns: + freqs_cis: tensor of shape (sum(t * height * width), dim//2) + """ + if not hasattr(self, 'freqs_cis'): + self.register_buffer('freqs_cis', + self._precompute_freqs_cis(device), + persistent=False) + + shapes = grid_thws.tolist() + assert all(1 <= h <= self.max_height and 1 <= w <= self.max_width + for t, h, w in shapes), ( + shapes, + self.max_height, + self.max_width, + ) + freqs_cis = torch.cat( + [ + self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) + for t, h, w in shapes + ], + dim=0, + ) + return freqs_cis + + +class MLP2(nn.Module): + """ + Args: + dims: [in_dim, hidden_dim, out_dim] + bias: whether to use bias in linear layer. + """ + + def __init__(self, dims: list[int], activation, bias=True): + super().__init__() + assert len(dims) == 3 + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + self.activation = activation + for m in [self.fc0, self.fc1]: + nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.fc0(x) + x = self.activation(x) + return self.fc1(x) + + +class MoonViTEncoderLayer(nn.Module): + + def __init__( + self, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + qkv_hidden_size: int | None = None, + norm_type: str = 'layernorm', + mlp_type: str = 'mlp2', + *, + attn_implementation: str = 'flash_attention_2', + activation=F.gelu, + attn_bias: bool = False, + linear_bias: bool = True, + use_deterministic_attn: bool = False, + ): + super().__init__() + self.num_heads = num_heads + self.hidden_dim = hidden_dim + self.qkv_hidden_size = hidden_dim if qkv_hidden_size is None else qkv_hidden_size + self.hidden_size_per_attention_head = self.qkv_hidden_size // self.num_heads + self.attn_implementation = attn_implementation + self.use_deterministic_attn = use_deterministic_attn + + if norm_type == "layernorm": + self.norm0 = nn.LayerNorm(hidden_dim) + self.norm1 = nn.LayerNorm(hidden_dim) + elif norm_type == "rmsnorm": + self.norm0 = nn.RMSNorm(hidden_dim) + self.norm1 = nn.RMSNorm(hidden_dim) + else: + raise NotImplementedError(f"Not support norm_type: {norm_type}") + + if mlp_type == "mlp2": + self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], + activation, + bias=linear_bias) + else: + raise NotImplementedError(f"Not support mlp_type: {mlp_type}") + + self.wqkv = nn.Linear(hidden_dim, + self.qkv_hidden_size * 3, + bias=attn_bias) + self.wo = nn.Linear(self.qkv_hidden_size, hidden_dim, bias=attn_bias) + + def attention_qkvpacked( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + rope_freqs_cis: torch.Tensor | None = None, + ): + """ + Args: + x (torch.Tensor): (batch_size, seqlen, hidden_dim) + cu_seqlens (torch.Tensor): + """ + xqkv = self.wqkv(x) + + qkv_shape = xqkv.size()[:-1] + ( + 3, + self.num_heads, + self.hidden_size_per_attention_head, + ) + # xqkv: (batch_size, seqlen, 3, nheads, headdim) + xqkv = xqkv.view(*qkv_shape) + xq, xk, xv = torch.unbind(xqkv, dim=-3) + + xq, xk = apply_rope(xq, xk, rope_freqs_cis) + + attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation] + attn_out = attn_func(xq, + xk, + xv, + q_cu_seqlens=cu_seqlens, + k_cu_seqlens=cu_seqlens, + max_seqlen_k=max_seqlen, + max_seqlen_q=max_seqlen, + deterministic=self.use_deterministic_attn) + + attn_out = self.wo(attn_out) + return attn_out + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rope_freqs_cis: torch.Tensor | None = None, + ): + residual = hidden_states + hidden_states = self.norm0(hidden_states) + + hidden_states = self.attention_qkvpacked(hidden_states, cu_seqlens, + max_seqlen, rope_freqs_cis) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class MoonViT3dEncoder(nn.Module): + + def __init__(self, + hidden_dim: int, + num_layers: int, + block_cfg: dict, + use_deterministic_attn: bool = False) -> None: + super().__init__() + self.use_deterministic_attn = use_deterministic_attn + + qkv_hidden_size = block_cfg['hidden_dim'] if block_cfg.get( + 'qkv_hidden_size') is None else block_cfg['qkv_hidden_size'] + self.rope_2d = Rope2DPosEmbRepeated( + qkv_hidden_size // block_cfg['num_heads'], 512, 512) + self.blocks = nn.ModuleList([ + MoonViTEncoderLayer( + **block_cfg, + use_deterministic_attn=self.use_deterministic_attn) + for _ in range(num_layers) + ]) + norm_type = block_cfg.get('norm_type', 'layernorm') + if norm_type == "layernorm": + self.final_layernorm = nn.LayerNorm(hidden_dim) + elif norm_type == "rmsnorm": + self.final_layernorm = nn.RMSNorm(hidden_dim) + else: + raise NotImplementedError(f"Not support norm_type: {norm_type}") + + def forward( + self, + hidden_states: torch.Tensor, + grid_thws: torch.Tensor, + ) -> torch.Tensor: + rope_freqs_cis = self.rope_2d.get_freqs_cis( + grid_thws=grid_thws, device=hidden_states.device) + + lengths = torch.cat(( + torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device), + grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2], + )) + + max_seqlen = lengths.max() + cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, + dtype=torch.int32) + for block in self.blocks: + hidden_states = block(hidden_states, + cu_seqlens, + max_seqlen, + rope_freqs_cis=rope_freqs_cis) + + hidden_states = self.final_layernorm(hidden_states) + return hidden_states + + +def tpool_patch_merger( + x: torch.Tensor, + grid_thws: torch.Tensor, + merge_kernel_size: tuple[int, int] = (2, 2), +) -> list[torch.Tensor]: + d_model = x.size(-1) + + outputs = [] + pre_sum = 0 + for t, h, w in grid_thws.tolist(): + # Get the current sequence + seq = x[pre_sum:pre_sum + t * h * w] + # Reshape along self.merge_kernel_size and concat to the last dimension + kernel_height, kernel_width = merge_kernel_size + new_height, new_width = h // kernel_height, w // kernel_width + reshaped_seq = seq.view(t, new_height, kernel_height, new_width, + kernel_width, d_model) + reshaped_seq = reshaped_seq.permute(0, 1, + 3, 2, 4, 5).contiguous().mean( + dim=0) # temporal pooling + padded_seq = reshaped_seq.view(new_height * new_width, + kernel_height * kernel_width, -1) + outputs.append(padded_seq) + pre_sum += t * h * w + + return outputs + + +class MoonViT3dPretrainedModel(PreTrainedModel): + config_class = None + model_type = 'moonvit3d' + _no_split_modules = ['MoonViTEncoderLayer'] + _supports_flash_attn_2 = True + _supports_sdpa = True + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + config = deepcopy(config) + self.merge_kernel_size = config.merge_kernel_size + self.patch_size = config.patch_size + self.merge_type = config.merge_type + + self.patch_embed = MoonVision3dPatchEmbed( + out_dim=config.hidden_size, + patch_size=config.patch_size, + pos_emb_height=config.init_pos_emb_height, + pos_emb_width=config.init_pos_emb_width, + pos_emb_time=config.init_pos_emb_time, + pos_emb_type=config.pos_emb_type, + patch_embed_proj_bias=getattr(config, 'patch_embed_proj_bias', + True), + pos_emb_interpolation_mode=getattr( + config, 'pos_emb_interpolation_mode', 'bicubic'), + ) + + self.encoder = MoonViT3dEncoder( + hidden_dim=config.hidden_size, + num_layers=config.num_hidden_layers, + block_cfg={ + 'num_heads': config.num_attention_heads, + 'hidden_dim': config.hidden_size, + 'qkv_hidden_size': getattr(config, 'qkv_hidden_size', None), + 'mlp_dim': config.intermediate_size, + 'norm_type': getattr(config, 'norm_type', 'layernorm'), + 'mlp_type': getattr(config, 'mlp_type', 'mlp2'), + 'activation': PytorchGELUTanh(), + 'attn_bias': getattr(config, 'attn_bias', True), + 'linear_bias': getattr(config, 'linear_bias', True), + 'attn_implementation': config._attn_implementation, + }, + use_deterministic_attn=getattr(self, 'use_deterministic_attn', + False)) + + def forward(self, pixel_values: torch.Tensor, + grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + pixel_values (torch.Tensor): The input pixel values. + grid_thws (torch.Tensor): Temporal, height and width. + + Returns: + torch.Tensor: The output tokens. + """ + # grid_thws = grid_thws.to('cpu') + assert grid_thws.ndim == 2, f'grid_thws should be 2D, got {grid_thws.ndim}' + assert grid_thws.size(1) == 3, f'No support for thw: {grid_thws}' + hidden_states = self.patch_embed(pixel_values, grid_thws) + hidden_states = self.encoder(hidden_states, grid_thws) + if self.merge_type == 'sd2_tpool': # spatial downsampling 2x with temporal pooling all + hidden_states = tpool_patch_merger( + hidden_states, + grid_thws, + merge_kernel_size=self.merge_kernel_size) + else: + raise NotImplementedError(f'Not support {self.merge_type}') + + return hidden_states + + +# ============================================================================ +# MM Projector Helper Classes (from mm_projector/modeling_mm_projectors.py) +# ============================================================================ + + +class IdentityMap(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x, *args, **kwargs): + return x + + +class MLP(nn.Module): + + def __init__(self, config): + super().__init__() + # TODO, use faster LayerNorm + self.pre_norm = nn.LayerNorm(config.mm_hidden_size) + self.proj = nn.Sequential( + nn.Linear(config.mm_hidden_size, config.hidden_size), nn.GELU(), + nn.Linear(config.hidden_size, config.hidden_size)) + + def forward(self, x, *args, **kwargs): + assert isinstance(x, + list | tuple), f'x is not a list or tuple: {type(x)}' + lengths = [item.shape[0] for item in x] + x = torch.cat(x, dim=0) + x = self.pre_norm(x) + x = self.proj(x) + x = torch.split(x, lengths, dim=0) + + return x + + +class PatchMergerMLP(nn.Module): + + def __init__(self, config): + super().__init__() + eps = config.projector_ln_eps + self.hidden_size = config.mm_hidden_size * ( + config.merge_kernel_size[0] * config.merge_kernel_size[1]) + self.pre_norm = nn.LayerNorm(config.mm_hidden_size, eps=eps) + self.proj = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, config.hidden_size), + ) + + def forward(self, x, *args, **kwargs): + if isinstance(x, list) or isinstance(x, tuple): + x = [ + self.proj(self.pre_norm(item).view(item.shape[0], -1)) + for item in x + ] + else: + # B, N, N_k, C = x.shape + B = x.shape[0] + x = self.proj(self.pre_norm(x).view(B, -1, self.hidden_size)) + return x + + +class PatchMergerMLPV2(nn.Module): + + def __init__(self, config): + super().__init__() + eps = config.projector_ln_eps + self.hidden_size = config.mm_hidden_size * ( + config.merge_kernel_size[0] * config.merge_kernel_size[1]) + self.proj = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size, bias=False), + nn.GELU(), + nn.Linear(self.hidden_size, config.hidden_size, bias=False), + ) + self.post_norm = nn.RMSNorm(config.hidden_size, eps=eps) + for m in self.proj.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, + std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x, *args, **kwargs): + if isinstance(x, list) or isinstance(x, tuple): + lengths = [item.shape[0] for item in x] + x = torch.concat([item.view(item.shape[0], -1) for item in x], + dim=0) + x = self.post_norm(self.proj(x)) + x = torch.split(x, lengths, dim=0) + else: + # B, N, N_k, C = x.shape + B = x.shape[0] + x = self.proj(x.view(B, -1, self.hidden_size)) + x = self.post_norm(x) + return x + + +class KimiK3PreTrainedModel(PreTrainedModel): + config_class = KimiK3Config + base_model_prefix = "model" + _no_split_modules = [ + "MoonViT3dPretrainedModel", + "MoonViTEncoderLayer", + "KimiDecoderLayer", + "PatchMergerMLP", + "PatchMergerMLPV2", + ] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = False + + def _init_weights(self, module): + # important: this ported version of Llava isn't meant for training from scratch - only + # inference and fine-tuning - so the proper init weights code has been removed - the original codebase + # https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose + std = (self.config.initializer_range if hasattr( + self.config, "initializer_range") else + self.config.text_config.initializer_range) + + if hasattr(module, "class_embedding"): + module.class_embedding.data.normal_(mean=0.0, std=std) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class VisionTowerConfig(PretrainedConfig): + model_type = 'moonvit3d' + + def __init__(self, config: KimiK3Config, **kwargs): + super().__init__(**kwargs) + self.patch_size = config.patch_size + self.init_pos_emb_height = config.init_pos_emb_height + self.init_pos_emb_width = config.init_pos_emb_width + self.init_pos_emb_time = config.init_pos_emb_time + self.pos_emb_type = config.pos_emb_type + self.num_attention_heads = config.vt_num_attention_heads + self.num_hidden_layers = config.vt_num_hidden_layers + self.hidden_size = config.vt_hidden_size + self.intermediate_size = config.vt_intermediate_size + self.merge_kernel_size = config.merge_kernel_size + self.merge_type = config.merge_type + self._attn_implementation = config._attn_implementation + self.qkv_hidden_size = getattr(config, 'qkv_hidden_size', None) + self.norm_type = getattr(config, 'norm_type', 'layernorm') + self.attn_bias = getattr(config, 'attn_bias', True) + self.patch_embed_proj_bias = getattr(config, 'patch_embed_proj_bias', + True) + self.mlp_type = getattr(config, 'mlp_type', 'mlp2') + self.linear_bias = getattr(config, 'linear_bias', True) + self.pos_emb_interpolation_mode = getattr( + config, 'pos_emb_interpolation_mode', 'bilinear') + + +class ProjectorConfig: + + def __init__(self, config: KimiK3Config): + self.mm_projector_type = config.mm_projector_type + self.mm_hidden_size = config.mm_hidden_size + self.hidden_size = config.text_hidden_size + self.merge_kernel_size = config.merge_kernel_size + self.projector_hidden_act = config.projector_hidden_act + self.projector_ln_eps = config.projector_ln_eps + + +# ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 +class KimiK3ForConditionalGeneration(KimiK3PreTrainedModel): + + @classmethod + def _supports_default_dynamic_cache(cls) -> bool: + return False + + def __init__(self, config: KimiK3Config): + super().__init__(config) + + vt_config = VisionTowerConfig(config.vision_config) + self.vision_tower = MoonViT3dPretrainedModel(vt_config) + + proj_config = ProjectorConfig(config.vision_config) + if proj_config.mm_projector_type == 'identity': + self.mm_projector = IdentityMap() + elif proj_config.mm_projector_type == 'mlp': + self.mm_projector = MLP(proj_config) + elif proj_config.mm_projector_type == 'patchmerger': + self.mm_projector = PatchMergerMLP(proj_config) + elif proj_config.mm_projector_type == 'patchmergerv2': + self.mm_projector = PatchMergerMLPV2(proj_config) + else: + raise ValueError( + f"Unsupported mm_projector_type: {proj_config.mm_projector_type}" + ) + + self.language_model = KimiLinearForCausalLM(config.text_config) + self.post_init() + + if hasattr(self.language_model, 'dtype'): + target_dtype = self.language_model.dtype + self.vision_tower = self.vision_tower.to(dtype=target_dtype) + self.mm_projector = self.mm_projector.to(dtype=target_dtype) + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.language_model.get_output_embeddings() + + def set_output_embeddings(self, new_embeddings): + self.language_model.set_output_embeddings(new_embeddings) + + def set_decoder(self, decoder): + self.language_model.set_decoder(decoder) + + def get_decoder(self): + return self.language_model.get_decoder() + + def tie_weights(self): + return self.language_model.tie_weights() + + def resize_token_embeddings(self, + new_num_tokens: int | None = None, + pad_to_multiple_of=None) -> nn.Embedding: + model_embeds = self.language_model.resize_token_embeddings( + new_num_tokens, pad_to_multiple_of) + # update vocab size + self.config.text_config.vocab_size = model_embeds.num_embeddings + self.vocab_size = model_embeds.num_embeddings + return model_embeds + + def _merge_input_ids_with_image_features( + self, + image_features: list[torch.Tensor], + inputs_embeds: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + labels: torch.Tensor | None = None, + ): + """ + Args: + image_features (:obj:`torch.Tensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The image features to merge with the input embeddings. + inputs_embeds (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length, embed_dim)`): + The input embeddings. + input_ids (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The input ids. + attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The attention mask. + labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*): + The labels. + """ + _, embed_dim = image_features[0].shape + feature_lengths = [x.shape[0] for x in image_features] + image_features = torch.cat(image_features, dim=0) + + image_token_index: int = self.config.media_placeholder_token_id + pad_token_id: int = self.config.pad_token_id + ignore_index: int = self.config.ignore_index + + batch_size, sequence_length = input_ids.shape + left_padding = not torch.sum( + input_ids[:, -1] == torch.tensor(pad_token_id)) + + # 1. Create a mask to know where special image tokens are + _token_occupation_table = torch.ones_like(input_ids.flatten()) + _token_occupation_table[input_ids.flatten() == + image_token_index] = torch.tensor( + feature_lengths, + dtype=torch.long, + device=input_ids.device) + _token_occupation_table = _token_occupation_table.reshape( + input_ids.shape) + + max_embed_dim = _token_occupation_table.sum(-1).max().item() + assert ( + max_embed_dim >= sequence_length + ), f"The maximum embedding dimension ({max_embed_dim}) is less than the sequence length ({sequence_length})" + batch_indices, non_image_indices = torch.where( + input_ids != image_token_index) + + # 2. Compute the positions where text should be written + # Calculate new positions for text tokens in merged image-text sequence. + new_token_positions = torch.cumsum(_token_occupation_table, -1) - 1 + nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] + if left_padding: + new_token_positions += nb_image_pad[:, + None] # offset for left padding + text_to_overwrite = new_token_positions[batch_indices, + non_image_indices] + + # 3. Create the full embedding, already padded to the maximum position + final_embedding = torch.zeros( + batch_size, + max_embed_dim, + embed_dim, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + final_attention_mask = torch.zeros(batch_size, + max_embed_dim, + dtype=attention_mask.dtype, + device=inputs_embeds.device) + if labels is not None: + final_labels = torch.full( + (batch_size, max_embed_dim), + ignore_index, + dtype=input_ids.dtype, + device=input_ids.device, + ) + # In case the Vision model or the Language model has been offloaded to CPU, we need to manually + # set the corresponding tensors into their correct target device. + target_device = inputs_embeds.device + batch_indices, non_image_indices, text_to_overwrite = ( + batch_indices.to(target_device), + non_image_indices.to(target_device), + text_to_overwrite.to(target_device), + ) + attention_mask = attention_mask.to(target_device) + + # 4. Fill the embeddings based on the mask. + final_embedding[batch_indices, + text_to_overwrite] = inputs_embeds[batch_indices, + non_image_indices] + final_attention_mask[batch_indices, + text_to_overwrite] = attention_mask[ + batch_indices, non_image_indices] + if labels is not None: + final_labels[batch_indices, + text_to_overwrite] = labels[batch_indices, + non_image_indices] + + # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) + image_to_overwrite = torch.full((batch_size, max_embed_dim), + True, + dtype=torch.bool, + device=inputs_embeds.device) + image_to_overwrite[batch_indices, text_to_overwrite] = False + image_to_overwrite &= image_to_overwrite.cumsum( + -1) - 1 >= nb_image_pad[:, None].to(target_device) + + if image_to_overwrite.sum() != image_features.shape[:-1].numel(): + raise ValueError( + f"The input provided to the model are wrong. The number of image tokens is {image_to_overwrite.sum()} while" + f" the number of image features given to the model is {image_features.shape[:-1].numel()}. " + "This prevents correct indexing and breaks batch generation.") + + final_embedding[image_to_overwrite] = ( + image_features.contiguous().reshape(-1, + embed_dim).to(target_device)) + final_attention_mask |= image_to_overwrite + position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_( + (final_attention_mask == 0), 1) + + # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. + batch_indices, pad_indices = torch.where(input_ids == pad_token_id) + indices_to_mask = new_token_positions[batch_indices, pad_indices] + + final_embedding[batch_indices, indices_to_mask] = 0 + + if labels is None: + final_labels = None + + return final_embedding, final_attention_mask, final_labels, position_ids + + def _extract_image_features(self, pixel_values: torch.Tensor, + grid_thws: torch.Tensor) -> list[torch.Tensor]: + """ + Args: + pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`): + The pixel values of the images processed by image processor. + grid_thws (:obj:`torch.Tensor` of shape :obj:`(batch_size, 3)`): + The grid, height, width of the images. + + Returns: + selected_image_feature (:obj:`torch.FloatTensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The selected image features to use as input to the projector head. + + """ + + target_dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + + image_features = self.vision_tower(pixel_values, grid_thws) + return image_features + + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | list[torch.FloatTensor] + | None = None, + grid_thws: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | LlavaCausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + ```""" + assert self.vision_tower is not None, "vision_tower is not loaded" + output_attentions = (output_attentions if output_attentions is not None + else self.config.output_attentions) + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else + self.config.output_hidden_states) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + # 1. Extra the input embeddings + inputs_embeds = self.get_input_embeddings()(input_ids) + + # 2. Merge text and images + if pixel_values is not None and len( + pixel_values) > 0 and input_ids.shape[1] != 1: + image_features = self._extract_image_features( + pixel_values, grid_thws) + if self.mm_projector: + image_features = self.mm_projector(image_features) + + inputs_embeds = inputs_embeds.to( + image_features[0].dtype) # num_tokens, embed_dim + inputs_embeds, attention_mask, labels, position_ids = ( + self._merge_input_ids_with_image_features( + image_features, + inputs_embeds, + input_ids, + attention_mask, + labels, + )) + + # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of + # 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 + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + logits = outputs[0] + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + if attention_mask is not None: + shift_attention_mask = attention_mask[..., 1:] + shift_logits = logits[..., :-1, :][shift_attention_mask.to( + logits.device) != 0].contiguous() + shift_labels = labels[..., 1:][shift_attention_mask.to( + labels.device) != 0].contiguous() + else: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1).to(shift_logits.device), + ) + + if not return_dict: + output = (logits, ) + outputs[1:] + return (loss, ) + output if loss is not None else output + + return LlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + grid_thws=None, + attention_mask=None, + **kwargs, + ): + if past_key_values is not None: + if hasattr(past_key_values, "get_seq_length"): + cache_length = past_key_values.get_seq_length() + past_length = getattr(past_key_values, 'seen_tokens', + cache_length) + else: + cache_length = past_length = past_key_values[0][0].shape[2] + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as + # input) + if attention_mask is not None and attention_mask.shape[ + 1] > input_ids.shape[1]: + input_ids = input_ids[:, -(attention_mask.shape[1] - + past_length):] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + elif self.config.media_placeholder_token_id in input_ids: + input_ids = input_ids[:, input_ids.shape[1] - 1:] + # If the cache has seen more tokens than it can hold, then the cache has a size limit. Let's discard the + # older attention values, as their corresponding values are not part of the input. + if cache_length < past_length and attention_mask is not None: + attention_mask = attention_mask[:, -(cache_length + + input_ids.shape[1]):] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1]:] + + # 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} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update({ + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "grid_thws": grid_thws, + }) + return model_inputs + + def _reorder_cache(self, *args, **kwargs): + return self.language_model._reorder_cache(*args, **kwargs) diff --git a/modeling_kimi_linear.py b/modeling_kimi_linear.py new file mode 100644 index 0000000000000000000000000000000000000000..b8c41e8bfce768d74d8da3a37e693f5ee43876a0 --- /dev/null +++ b/modeling_kimi_linear.py @@ -0,0 +1,1314 @@ +# coding=utf-8 +# Copyright 2025-2026 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved. +# +# The multi-head latent attention, MoE gating and sparse MoE block in this file are +# adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py). They have been +# extensively modified and extended for the Kimi-Linear architecture. +# +# Licensing Information: +# - Code adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py) is licensed under the Apache License, Version 2.0. +# - Other parts of the code are licensed under the Kimi K3 License (see the LICENSE file in this repository). +# +# Apache License, Version 2.0: +# 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 +from collections.abc import Callable +from typing import Any + +import torch +import torch.nn.functional as F +import transformers +from einops import rearrange +from packaging import version +from torch import nn +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin +from transformers.masking_utils import create_causal_mask +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS +from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from transformers.utils.generic import OutputRecorder, check_model_inputs + +try: + from fla.modules import FusedRMSNormGated, ShortConvolution + from fla.ops.kda import chunk_kda, fused_recurrent_kda + # from fla.ops.kda.gate import fused_kda_gate # deprecated, gate is now computed inside chunk_kda/fused_recurrent_kda + from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask + from fla.utils import tensor_cache +except ImportError: + raise ImportError("Plese run `pip install -U fla-core`") + +from .configuration_kimi_k3 import KimiLinearConfig + +assert version.parse(transformers.__version__) >= version.parse("4.56.0"), \ + "Please upgrade transformers to >= 4.56.0" + +logger = logging.get_logger(__name__) + + +# Register Moonshot-specific activation functions +class SituAndMul(nn.Module): + """ + SituAndMul activation: beta * tanh(gate / beta) * sigmoid(gate) * up + When linear_beta is set, up is also transformed by linear_beta * tanh(up / linear_beta). + """ + + def __init__(self, beta: float = 1.0, linear_beta: float | None = None): + super().__init__() + self.beta = beta + self.linear_beta = linear_beta + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d].to(torch.float32) + up = x[..., d:].to(torch.float32) + situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (situ_a * up).to(x.dtype) + + +ACT2FN["situ"] = SituAndMul + + +def _get_situ_activation_params(config: KimiLinearConfig): + beta = getattr(config, "activation_situ_beta", None) + linear_beta = getattr(config, "activation_situ_linear_beta", None) + return beta or 1.0, linear_beta + + +def index_first_axis(x, indices): + return x[indices] + + +@tensor_cache +def get_unpad_data( + attention_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, int]: + lens = prepare_lens_from_mask(attention_mask) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = lens.max().item() + cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask) + return indices, cu_seqlens, max_seqlen_in_batch + + +def pad_input( + hidden_states: torch.Tensor, + indices: torch.LongTensor, + batch_size: int, + seq_len: int, +) -> torch.Tensor: + out = hidden_states.new_zeros((batch_size * seq_len, *hidden_states.shape[1:])) + out[indices] = hidden_states + return out.view(batch_size, seq_len, *hidden_states.shape[1:]) + + +class KimiDynamicCache: + """ + Dynamic cache for Kimi model. + Inspired by Qwen3-Next + """ + is_compileable = False + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + + if config.linear_attn_config is not None: + self.layer_types = [] + for i in range(config.num_hidden_layers): + if config.is_kda_layer(i): + self.layer_types.append("linear_attention") + else: + self.layer_types.append("full_attention") + else: + self.layer_types = ["full_attention"] * config.num_hidden_layers + + self.transformer_layers = [ + i for i in range(config.num_hidden_layers) if self.layer_types[i] == "full_attention" + ] + + linear_layers = [i for i in range( + config.num_hidden_layers) if self.layer_types[i] == "linear_attention"] + self.last_linear_layer = linear_layers[-1] if linear_layers else -1 + + self.conv_states = [None for _ in range(config.num_hidden_layers)] + self.recurrent_states = [None for _ in range(config.num_hidden_layers)] + self.key_cache = [None for _ in range(config.num_hidden_layers)] + self.value_cache = [None for _ in range(config.num_hidden_layers)] + + def __len__(self): + return len(self.layer_types) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.key_cache[layer_idx] is None: + self.key_cache[layer_idx] = key_states + self.value_cache[layer_idx] = value_states + else: + self.key_cache[layer_idx] = torch.cat( + [self.key_cache[layer_idx], key_states], dim=2) + self.value_cache[layer_idx] = torch.cat( + [self.value_cache[layer_idx], value_states], dim=2) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx] is not None: + device = self.key_cache[layer_idx].device + beam_idx = beam_idx.to(device) + self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select( + 0, beam_idx) + self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select( + 0, beam_idx) + + if self.conv_states[layer_idx] is not None: + device = self.conv_states[layer_idx][0].device + beam_idx = beam_idx.to(device) + q_conv, k_conv, v_conv = self.conv_states[layer_idx] + self.conv_states[layer_idx] = ( + q_conv.index_select(0, beam_idx), + k_conv.index_select(0, beam_idx), + v_conv.index_select(0, beam_idx), + ) + self.recurrent_states[layer_idx] = self.recurrent_states[layer_idx].index_select( + 0, beam_idx) + + def get_seq_length(self, layer_idx: int | None = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # take any layer that contains cache and not empty tensor + layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx + if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None: + return 0 + return self.key_cache[layer_idx].shape[-2] + + def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int) -> tuple[int, int]: + """ + Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for + the given layer at `layer_idx`. + The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer. + """ + kv_offset = 0 + query_length = cache_position.shape[0] + past_seen_tokens = self.get_seq_length(layer_idx) + kv_length = query_length + past_seen_tokens + return kv_length, kv_offset + + @property + def has_previous_state(self): + """We have a previous state if the last linear (conv) layer was already updated.""" + if self.last_linear_layer == -1: + return False + return self.conv_states[self.last_linear_layer] is not None + + +class KimiRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + dtype = hidden_states.dtype + x = hidden_states.float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) + return self.weight * x.to(dtype) + + +ALL_LAYERNORM_LAYERS.append(KimiRMSNorm) + + +class KimiBlockSparseMLP(nn.Module): + def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None): + super().__init__() + self.config = config + self.ffn_dim = config.intermediate_size if intermediate_size is None else intermediate_size + self.hidden_dim = config.hidden_size if hidden_size is None else hidden_size + + self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # gate + self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) # down + self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # up + + if config.hidden_act == "situ": + beta, linear_beta = _get_situ_activation_params(config) + self.act_fn = SituAndMul( + beta=beta, + linear_beta=linear_beta, + ) + else: + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states): + if self.config.hidden_act == "situ": + gate_up = torch.cat([self.w1(hidden_states), self.w3(hidden_states)], dim=-1) + current_hidden_states = self.act_fn(gate_up) + else: + current_hidden_states = self.act_fn( + self.w1(hidden_states)) * self.w3(hidden_states) + current_hidden_states = self.w2(current_hidden_states) + return current_hidden_states + + +class KimiMLP(nn.Module): + def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else 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) + if config.hidden_act == "situ": + beta, linear_beta = _get_situ_activation_params(config) + self.act_fn = SituAndMul( + beta=beta, + linear_beta=linear_beta, + ) + else: + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + if self.config.hidden_act == "situ": + gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1) + down_proj = self.down_proj(self.act_fn(gate_up)) + else: + down_proj = self.down_proj(self.act_fn( + self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand the key/value heads from `num_key_value_heads` to `num_attention_heads`.""" + if n_rep == 1: + return hidden_states + return torch.repeat_interleave(hidden_states, dim=1, repeats=n_rep) + + +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 = repeat_kv(key, module.num_key_value_groups) + value = repeat_kv(value, module.num_key_value_groups) + + scores = torch.einsum("bhqd,bhkd->bhqk", query, key) * scaling + if attention_mask is not None: + scores = scores + attention_mask[:, :, :, : key.shape[-2]] + + probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype) + probs = F.dropout(probs, p=dropout, training=module.training) + out = torch.einsum("bhqk,bhkd->bhqd", probs, value).transpose(1, 2).contiguous() + + return out, probs + + +class KimiMLAAttention(nn.Module): + """ + Multi-Latent Attention adapted from deepseek-v3 + """ + + def __init__(self, config: KimiLinearConfig, layer_idx: int): + nn.Module.__init__(self) + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_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.attention_dropout = getattr(config, "attention_dropout", 0.0) + + try: + self.q_lora_rank = config.q_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.use_nope = config.mla_use_nope + self.scaling = self.q_head_dim ** (-0.5) + except Exception as e: + raise ValueError( + f"Kimi MLA config is not found or not properly formatted: {e}") + + if self.q_lora_rank is not None: + self.q_a_proj = nn.Linear( + self.hidden_size, self.q_lora_rank, bias=False, + ) + self.q_a_layernorm = KimiRMSNorm(self.q_lora_rank) + self.q_b_proj = nn.Linear( + self.q_lora_rank, + self.num_heads * self.q_head_dim, + bias=False, + ) + else: + self.q_proj = nn.Linear( + self.hidden_size, self.num_heads * self.q_head_dim, bias=False, + ) + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + ) + self.kv_a_layernorm = KimiRMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads + * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + ) + self.is_causal = True + assert self.use_nope + + self.use_output_gate = getattr(config, "mla_use_output_gate", False) + if self.use_output_gate: + projection_size = self.num_heads * self.v_head_dim + self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + + self.rotary_emb = None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, -1, self.q_head_dim) + key_shape = (batch_size, seq_length, -1, + self.qk_nope_head_dim + self.v_head_dim) + + if self.q_lora_rank is not None: + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + else: + q_states = self.q_proj(hidden_states) + q_states = q_states.view(query_shape).transpose(1, 2) + q_pass, q_rot = torch.split( + q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_pass, k_rot = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + + k_pass = self.kv_b_proj(self.kv_a_layernorm( + k_pass)).view(key_shape).transpose(1, 2) + k_pass, value_states = torch.split( + k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) + + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) + + if past_key_values is not None: + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx) + + if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: + value_states = F.pad( + value_states, [0, self.q_head_dim - self.v_head_dim]) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, _ = 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, + **kwargs, + ) + + if self.config._attn_implementation == "flash_attention_2" and self.q_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape( + batch_size, seq_length, -1).contiguous() + if self.use_output_gate: + g = self.g_proj(hidden_states).sigmoid() + attn_output = attn_output * g + attn_output = self.o_proj(attn_output) + return attn_output + + +class KimiDeltaAttention(nn.Module): + def __init__(self, config: KimiLinearConfig, layer_idx: int): + super().__init__() + self.config = config + self.mode = "chunk" + + self.hidden_size = config.hidden_size + self.conv_size = config.linear_attn_config["short_conv_kernel_size"] + self.head_dim = config.linear_attn_config["head_dim"] + self.num_heads = config.linear_attn_config["num_heads"] + self.head_k_dim = self.head_dim + self.num_k_heads = self.num_heads + + self.layer_idx = layer_idx + + assert self.mode in [ + 'chunk', 'fused_recurrent'], f"Not supported mode `{self.mode}`." + + projection_k_size = self.head_k_dim * self.num_k_heads + projection_size = self.head_dim * self.num_heads + + self.q_proj = nn.Linear( + self.hidden_size, projection_k_size, bias=False) + self.k_proj = nn.Linear( + self.hidden_size, projection_k_size, bias=False) + self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + + self.q_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=self.conv_size, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=self.conv_size, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=projection_size, + kernel_size=self.conv_size, + activation='silu', + ) + + self.A_log = torch.nn.Parameter(torch.log(torch.empty( + self.num_heads, dtype=torch.float32).uniform_(1, 16))) + + self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) + + self.dt_bias = nn.Parameter( + torch.empty(projection_size, dtype=torch.float32)) + + self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False) + + self.use_full_rank_gate = config.linear_attn_config.get("use_full_rank_gate", False) + self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None) + if self.use_full_rank_gate: + self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + else: + self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) + + self.o_norm = FusedRMSNormGated( + self.head_dim, eps=config.rms_norm_eps, activation='sigmoid') + self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + cache_params: KimiDynamicCache | None = None, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + if attention_mask.dim() != 2: + attention_mask = kwargs.get("padding_mask") + + if attention_mask is not None and attention_mask.dim() != 2: + raise ValueError( + "attention_mask must be a 0-1 matrix of shape [batch_size, seq_len] " + "(0 = padding). 3D masks are not supported here.", + ) + use_cache = cache_params is not None + batch_size, q_len, _ = hidden_states.shape + mode = 'fused_recurrent' if use_cache and q_len == 1 else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + cu_seqlens = kwargs.get('cu_seqlens') + indices = None + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis( + rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + conv_state_q, conv_state_k, conv_state_v = None, None, None + recurrent_state = None + if cache_params is not None: + if cache_params.conv_states[self.layer_idx] is not None: + conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[ + self.layer_idx] + recurrent_state = cache_params.recurrent_states[self.layer_idx] + + q_proj_states = self.q_proj(hidden_states) + k_proj_states = self.k_proj(hidden_states) + v_proj_states = self.v_proj(hidden_states) + q, conv_state_q = self.q_conv1d( + x=q_proj_states, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=k_proj_states, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=v_proj_states, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + g = self.f_b_proj(self.f_a_proj(hidden_states)) + g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim) + beta = self.b_proj(hidden_states).float() + + q, k = map(lambda x: rearrange( + x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim) + + if mode == 'chunk': + o, recurrent_state = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=self.gate_lower_bound is not None, + lower_bound=self.gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + else: + o, recurrent_state = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=self.gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + if cache_params is not None: + cache_params.recurrent_states[self.layer_idx] = recurrent_state + cache_params.conv_states[self.layer_idx] = ( + conv_state_q, conv_state_k, conv_state_v) + + if self.use_full_rank_gate: + g = self.g_proj(hidden_states) + else: + g = self.g_b_proj(self.g_a_proj(hidden_states)) + g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim) + o = self.o_norm(o, g) + + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o + + +class KimiMoEGate(nn.Module): + """ + MoEGate adapted from Deepseek-V3. + Parameter correspondences: + num_experts -> n_routed_experts + num_experts_per_token -> num_experts_per_tok + num_expert_group -> n_group + moe_router_activation_func -> scoring_func + """ + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + + # topk selection algorithm + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim)), + ) + + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts), + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + import torch.nn.init as init + + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + # compute gating score + hidden_states = hidden_states.view(-1, h) + logits = F.linear( + hidden_states.type(torch.float32), self.weight.type( + torch.float32), None, + ) + if self.moe_router_activation_func == "sigmoid": + scores = logits.sigmoid() + elif self.moe_router_activation_func == "softmax": + scores = logits.softmax(dim=1) + else: + raise NotImplementedError( + f"insupportable scoring function for MoE gating: {self.moe_router_activation_func}", + ) + + # select top-k experts + assert not self.training + scores = scores.view(bsz * seq_len, -1) + scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) + if self.num_expert_group > 1 and self.num_expert_group > self.topk_group: + group_scores = ( + scores_for_choice.view( + bsz * seq_len, self.num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1) + ) # [n, num_expert_group] + group_idx = torch.topk( + group_scores, k=self.topk_group, dim=-1, sorted=False, + )[ + 1 + ] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, num_expert_group] + group_mask.scatter_(1, group_idx, 1) # [n, num_expert_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand( + bsz * seq_len, self.num_expert_group, self.num_experts // self.num_expert_group, + ) + .reshape(bsz * seq_len, -1) + ) # [n, e] + tmp_scores = scores_for_choice.masked_fill( + ~score_mask.bool(), float("-inf")) # [n, e] + else: + tmp_scores = scores_for_choice + _, topk_idx = torch.topk( + tmp_scores, k=self.top_k, dim=-1, sorted=False, + ) + topk_weight = scores.gather(1, topk_idx) + + # norm gate to sum 1 + if self.top_k > 1 and self.moe_renormalize: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + # must multiply the scaling factor + topk_weight = topk_weight * self.routed_scaling_factor + + return topk_idx, topk_weight + + +class KimiSparseMoeBlock(nn.Module): + """ + Adapted from Deepseek-V3's MOE implementation + The namings are consistent with Kimi's version. + """ + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + self.hidden_dim = config.hidden_size + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_token + self.moe_renormalize = config.moe_renormalize + + self.use_latent_moe = getattr(config, "routed_expert_hidden_size", None) is not None + self.moe_hidden_size = ( + config.routed_expert_hidden_size + if self.use_latent_moe else config.hidden_size + ) + self.latent_moe_use_norm = getattr(config, "latent_moe_use_norm", False) + + self.ep_size = 1 + self.experts_per_rank = config.num_experts + self.ep_rank = 0 + self.experts = nn.ModuleList( + [ + KimiBlockSparseMLP( + config, + hidden_size=self.moe_hidden_size, + intermediate_size=config.moe_intermediate_size, + ) + for _ in range(config.num_experts) + ], + ) + self.gate = KimiMoEGate(config) + if config.num_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.num_shared_experts + self.shared_experts = KimiMLP( + config=config, intermediate_size=intermediate_size, + ) + + if self.use_latent_moe: + self.routed_expert_down_proj = nn.Linear( + config.hidden_size, self.moe_hidden_size, bias=False, + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, config.hidden_size, bias=False, + ) + if self.latent_moe_use_norm: + self.routed_expert_norm = KimiRMSNorm( + self.moe_hidden_size, eps=config.rms_norm_eps, + ) + + def forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + if self.use_latent_moe: + hidden_states = self.routed_expert_down_proj(hidden_states) + + if not self.training: + y = self.moe_infer(hidden_states, topk_idx, topk_weight) + else: + raise NotImplementedError("Training mode is not supported in KimiSparseMoeBlock") + + if self.use_latent_moe: + if self.latent_moe_use_norm: + y = self.routed_expert_norm(y) + y = self.routed_expert_up_proj(y) + + y = y.view(*orig_shape) + + if self.config.num_shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + @torch.no_grad() + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len( + outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + +class KimiDecoderLayer(nn.Module): + def __init__(self, config: KimiLinearConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.config = config + self.layer_idx = layer_idx + if config.is_kda_layer(layer_idx): + self.is_linear_attn = True + self.self_attn = KimiDeltaAttention( + config=config, layer_idx=layer_idx) + elif config.is_mla: + self.is_linear_attn = False + self.self_attn = KimiMLAAttention( + config=config, layer_idx=layer_idx) + else: + raise NotImplementedError + if ( + config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % getattr(config, "moe_layer_freq", 1) == 0 + ): + self.block_sparse_moe = KimiSparseMoeBlock(config) + else: + self.mlp = KimiMLP(config) + self.input_layernorm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + + # Attention residual + self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None + if self.use_attn_residuals: + self.attn_res_block_size = config.attn_res_block_size + self.self_attention_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + self.mlp_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = nn.Linear( + config.hidden_size, 1, bias=False) + self.mlp_res_proj = nn.Linear( + config.hidden_size, 1, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + block_residual: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ): + if self.use_attn_residuals: + return self._forward_attn_residual( + hidden_states, attention_mask, position_ids, + past_key_values, output_attentions, use_cache, + block_residual, **kwargs) + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + if self.is_linear_attn is False: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + cache_params=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + if hasattr(self, "block_sparse_moe"): + hidden_states = self.block_sparse_moe(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + def _forward_attn_residual( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + block_residual: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ): + batch_size, seq_len, hidden_size = hidden_states.shape + prefix_sum = hidden_states + + if block_residual is not None and block_residual.shape[1] > 0: + hidden_states = _apply_attn_res( + prefix_sum.view(-1, hidden_size), + block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ).view(batch_size, seq_len, hidden_size) + + if self.layer_idx % self.attn_res_block_size == 0: + block_residual = torch.cat( + [block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1) + prefix_sum = None + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + if self.is_linear_attn is False: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + cache_params=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + + if prefix_sum is not None: + prefix_sum = prefix_sum + hidden_states + else: + prefix_sum = hidden_states + + hidden_states = _apply_attn_res( + prefix_sum.view(-1, hidden_size), + block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + ).view(batch_size, seq_len, hidden_size) + + hidden_states = self.post_attention_layernorm(hidden_states) + if hasattr(self, "block_sparse_moe"): + hidden_states = self.block_sparse_moe(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + + if prefix_sum is None: + prefix_sum = hidden_states + else: + prefix_sum = prefix_sum + hidden_states + + return prefix_sum, block_residual + + +class KimiPreTrainedModel(PreTrainedModel): + config_class = KimiLinearConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["KimiDecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _can_record_outputs = { + "router_logits": OutputRecorder(KimiBlockSparseMLP, index=1), + "hidden_states": KimiDecoderLayer, + "attentions": KimiMLAAttention, + } + _is_stateful = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +def _apply_attn_res(prefix_sum, block_residual, proj, norm): + """ + prefix_sum: (num_tokens, hidden_size) + block_residual: (num_tokens, num_blocks, hidden_size) + """ + v = torch.cat((block_residual, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.variance_epsilon) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + +class KimiLinearModel(KimiPreTrainedModel): + def __init__(self, config: KimiLinearConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([KimiDecoderLayer( + config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + + self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None + if self.use_attn_residuals: + self.output_attn_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps) + self.output_attn_res_proj = nn.Linear( + config.hidden_size, 1, bias=False) + + if getattr(config, "_attn_implementation", None) is not None: + if config._attn_implementation != "flash_attention_2": + logger.warning_once( + f"Ignoring the provided attention implementation {config._attn_implementation}") + logger.warning_once("Using flash_attention_2 backend instead.") + config._attn_implementation = "flash_attention_2" + else: + config._attn_implementation = "flash_attention_2" + + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def _update_linear_attn_mask(self, attention_mask, cache_position): + """ + NOTE: Left-padding is used for linear attention mask. + No need for zeroing states when + 1. Cached forward + 2. Attending to all inputs + """ + linear_attn_mask = attention_mask + if cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1)): + linear_attn_mask = None + return linear_attn_mask + + @check_model_inputs + # @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + cache_position: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPast: + + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if (input_ids is None) and (inputs_embeds is None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds") + + # Get inputs_embeds + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = KimiDynamicCache(config=self.config) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length( + ) if past_key_values is not None else 0 + cache_position: torch.Tensor = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device, + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=position_ids, + ) + linear_attn_mask = self._update_linear_attn_mask( + attention_mask, cache_position) + + hidden_states = inputs_embeds + if past_key_values is not None: + assert isinstance(past_key_values, KimiDynamicCache) + + block_residual = None + if self.use_attn_residuals: + block_residual = hidden_states.new_zeros( + hidden_states.shape[0] * hidden_states.shape[1], 0, + hidden_states.shape[2]) + + for decoder_layer in self.layers: + layer_mask = linear_attn_mask if decoder_layer.is_linear_attn else causal_mask + + if self.use_attn_residuals: + hidden_states, block_residual = decoder_layer( + hidden_states, + attention_mask=layer_mask, + past_key_values=past_key_values, + cache_position=cache_position, + block_residual=block_residual, + **kwargs, + ) + else: + hidden_states = decoder_layer( + hidden_states, + attention_mask=layer_mask, + past_key_values=past_key_values, + cache_position=cache_position, + **kwargs, + ) + + if self.use_attn_residuals: + hidden_states = self._apply_output_attn_res( + hidden_states, block_residual) + + hidden_states = self.norm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + def _apply_output_attn_res(self, hidden_states, block_residual): + batch_size, seq_len, hidden_size = hidden_states.shape + return _apply_attn_res( + hidden_states.view(-1, hidden_size), + block_residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + ).view(batch_size, seq_len, hidden_size) + + +class KimiLinearForCausalLM(KimiPreTrainedModel, GenerationMixin): + @classmethod + def _supports_default_dynamic_cache(cls) -> bool: + return False + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = KimiLinearModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear( + config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + # @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + generation_mode: bool | None = None, + return_dict: bool | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + """ + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + 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, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + logits = outputs[0] + if generation_mode: + logits = logits[:, -1:] + logits = self.lm_head(logits) + + loss = None + if labels is not None: + loss = self.loss_function( + logits, labels, self.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/preprocessor_config.json b/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..8be28c13c93fb8982f81928d3f4fe80568102286 --- /dev/null +++ b/preprocessor_config.json @@ -0,0 +1,38 @@ +{ + "media_proc_cfg": { + "in_patch_limit": 65536, + "patch_size": 14, + "image_mean": [ + 0.5, + 0.5, + 0.5 + ], + "image_std": [ + 0.5, + 0.5, + 0.5 + ], + "merge_kernel_size": 2, + "fixed_output_tokens": null, + "patch_limit_on_one_side": 512, + "in_patch_limit_each_frame": 16384, + "in_patch_limit_video": 655360, + "sample_fps": 8.0, + "max_num_frames_each_video": null, + "temporal_merge_kernel_size": 4, + "timestamp_mode": "hh:mm:ss.fff", + "transparent_bg_config": { + "pattern": "chessboard", + "chessboard_square_size": 8, + "chessboard_square_on_top_left": true, + "chessboard_white_value": 255, + "chessboard_gray_value": 180 + }, + "transparent_bg_fill_stage": "after_resize", + "config_type": "media_proc.processors.moonvit.MoonViTMediaProcessorConfig" + }, + "auto_map": { + "AutoProcessor": "kimi_k3_processor.KimiK3Processor", + "AutoImageProcessor": "kimi_k3_vision_processing.KimiK3VisionProcessor" + } +} \ No newline at end of file diff --git a/tiktoken.model b/tiktoken.model new file mode 100644 index 0000000000000000000000000000000000000000..b4149a6e17a01b6442187f39890f89bc2fe8d309 --- /dev/null +++ b/tiktoken.model @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103 +size 2795286 diff --git a/tokenization_kimi.py b/tokenization_kimi.py new file mode 100644 index 0000000000000000000000000000000000000000..5df9d462b43295245084fbab46fd956020d4b84e --- /dev/null +++ b/tokenization_kimi.py @@ -0,0 +1,408 @@ +import os +from logging import getLogger +from pathlib import Path +from shutil import copyfile +from typing import Dict, Iterator, List, Optional, Tuple, Union, cast + +import tiktoken +from tiktoken.load import load_tiktoken_bpe +from tokenizers import AddedToken +from transformers.convert_slow_tokenizer import bytes_to_unicode +from transformers.tokenization_utils import PreTrainedTokenizer + +try: + from .encoding_k3 import build_chat_segments, is_batched_conversation +except ImportError: # pragma: no cover - supports direct file execution/import. + from encoding_k3 import build_chat_segments, is_batched_conversation + +logger = getLogger(__name__) +VOCAB_FILES_NAMES = {"vocab_file": "tiktoken.model"} + + +class TikTokenTokenizer(PreTrainedTokenizer): + """ + Tokenizing and encoding/decoding text using the Tiktoken tokenizer. See megatron/tokenizer/tiktoken_tokenizer.py. + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to + this superclass for more information regarding those methods. + + Args: + vocab_file (`str`): + The path to the Tiktoken model file. + bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|begin_of_text|>",`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|end_of_text|>"`): + The end of sequence token. + unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_249|>"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. The second to last item in special_tokens. + pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_250|>"`): + The token used for padding, for example when batching sequences of different lengths. + additional_special_tokens (list of `str`, *optional*): + A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be + skipped when decoding if `skip_special_tokens` is set to `True`. + """ + + vocab_files_names = VOCAB_FILES_NAMES + + model_input_names = ["input_ids", "attention_mask"] + + special_tokens: Dict[str, int] + + num_reserved_special_tokens = 256 + + pat_str = "|".join([ + r"""[\p{Han}]+""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""\p{N}{1,3}""", + r""" ?[^\s\p{L}\p{N}]+[\r\n]*""", + r"""\s*[\r\n]+""", + r"""\s+(?!\S)""", + r"""\s+""", + ]) + + def __init__( + self, + vocab_file, + bos_token: Union[str, AddedToken] = "[BOS]", + eos_token: Union[str, AddedToken] = "[EOS]", + unk_token: Union[str, AddedToken, None] = None, + pad_token: Union[str, AddedToken, None] = None, + additional_special_tokens: List[str] = None, + added_tokens_decoder: Optional[dict] = None, + **kwargs, + ): + assert os.path.isfile(vocab_file), vocab_file + + if additional_special_tokens is None: + additional_special_tokens = [ + "<|im_end|>", + "<|im_user|>", + "<|im_assistant|>", + "<|start_header_id|>", + "<|end_header_id|>", + "[EOT]", + "<|im_system|>", + "<|im_middle|>", + ] + + if added_tokens_decoder: + special_tokens_mapping = { + i: added_tokens_decoder[i].content + for i in added_tokens_decoder + } + else: + special_tokens_mapping = {} + + self.vocab_file = vocab_file + mergeable_ranks = load_tiktoken_bpe(vocab_file) + num_base_tokens = len(mergeable_ranks) + self.special_tokens = { + special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i + for i in range(num_base_tokens, num_base_tokens + + self.num_reserved_special_tokens) + } + + self.model = tiktoken.Encoding( + name=Path(vocab_file).name, + pat_str=self.pat_str, + mergeable_ranks=mergeable_ranks, + special_tokens=self.special_tokens, + ) + logger.info(f"Reloaded tiktoken model from {vocab_file}") + + self.n_words: int = self.model.n_vocab + # BOS / EOS token IDs + self.bos_id: int = self.special_tokens[str(bos_token)] + self.eos_id: int = self.special_tokens[str(eos_token)] + logger.info( + f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}" + ) + + self.pad_id: int = self.special_tokens[str(pad_token)] + self.unk_id: int = self.special_tokens[str(unk_token)] + + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + + self.decoder = {} + for i in range(self.n_words): + # Taken from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee + decoding = ''.join([ + self.byte_encoder[ord(char)] for char in + self.model.decode_single_token_bytes(i).decode('latin-1') + ]) + self.decoder[i] = decoding + + self.encoder = {} + for i in range(self.n_words): + if i in self.decoder: + self.encoder[self.decoder[i]] = i + + super().__init__( + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + additional_special_tokens=additional_special_tokens, + added_tokens_decoder=added_tokens_decoder, + **kwargs, + ) + self.all_special_ids_set = set(self.all_special_ids) + + def _encode_text_piece(self, text: str, + allow_special_tokens: bool = True) -> List[int]: + # The tiktoken tokenizer can handle <=400k chars without + # pyo3_runtime.PanicException. + TIKTOKEN_MAX_ENCODE_CHARS = 400_000 + + # https://github.com/openai/tiktoken/issues/195 + # Here we iterate over subsequences and split if we exceed the limit + # of max consecutive non-whitespace or whitespace characters. + MAX_NO_WHITESPACES_CHARS = 25_000 + + t: List[int] = [] + for i in range(0, len(text), TIKTOKEN_MAX_ENCODE_CHARS): + for substr in self._split_whitespaces_or_nonwhitespaces( + text[i:i + TIKTOKEN_MAX_ENCODE_CHARS], + MAX_NO_WHITESPACES_CHARS, + ): + if allow_special_tokens: + t.extend( + # structural markers: encode <|...|> as their special token IDs + self.model.encode( + substr, + allowed_special="all", + )) + else: + t.extend( + # user/tool text: encode any <|...|> as ordinary BPE tokens (never as control tokens) + self.model.encode( + substr, + disallowed_special=(), + )) + + return t + + def encode(self, + text: str, + allow_special_tokens: bool = True, + **kwargs) -> List[int]: + """ + Encodes a string into a list of token IDs. + + Args: + text (str): The input string to be encoded. + + Returns: + list[int]: A list of token IDs. + """ + # If there are other args, we should call super().encode because there are a lot of code + # to handle those args. supper().encode finally will call _tokenize and _convert_token_to_id. + # NOTE: our encode method is not compatible with the super().encode method, + # e.g. split_special_tokens' default is True in our encode method. + if len(kwargs) > 0: + logger.warning(f"Calling super().encode with {kwargs}") + return super().encode(text, **kwargs) + + assert type(text) is str + return self._encode_text_piece(text, + allow_special_tokens=allow_special_tokens) + + def decode(self, token_ids: Union[int, List[int]], **kwargs) -> str: + """ + Decodes a list of token IDs into a string. + + Args: + token_ids (List[int]): The list of token IDs to be decoded. + + Returns: + str: The decoded string. + """ + # If there are other args, we should call super().decode because there are a lot of code + # to handle those args. supper().encode finally will call convert_tokens_to_string and _convert_id_to_token. + if len(kwargs) > 0: + return super().decode(token_ids, **kwargs) + + if type(token_ids) is int: + token_ids = [token_ids] + + return self.model.decode(cast(List[int], token_ids)) + + @staticmethod + def _split_whitespaces_or_nonwhitespaces( + s: str, max_consecutive_slice_len: int) -> Iterator[str]: + """ + Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len` + consecutive whitespaces or consecutive non-whitespaces. + """ + current_slice_len = 0 + current_slice_is_space = s[0].isspace() if len(s) > 0 else False + slice_start = 0 + + for i in range(len(s)): + is_now_space = s[i].isspace() + + if current_slice_is_space ^ is_now_space: + current_slice_len = 1 + current_slice_is_space = is_now_space + else: + current_slice_len += 1 + if current_slice_len > max_consecutive_slice_len: + yield s[slice_start:i] + slice_start = i + current_slice_len = 1 + yield s[slice_start:] + + def _encode_chat_segments(self, segments) -> List[int]: + token_ids: List[int] = [] + for segment in segments: + token_ids.extend( + self._encode_text_piece( + segment.text, + allow_special_tokens=segment.allow_special, + )) + return token_ids + + @staticmethod + def _truncate(ids: List[int], + truncation: bool = False, + max_length: Optional[int] = None) -> List[int]: + if truncation and max_length is not None: + return ids[:max_length] + return ids + + def _format_chat_token_output(self, + encoded_inputs: List[List[int]], + *, + is_batched: bool, + padding=False, + truncation: bool = False, + max_length: Optional[int] = None, + return_tensors=None, + return_dict: bool = False): + encoded_inputs = [ + self._truncate(ids, truncation=truncation, max_length=max_length) + for ids in encoded_inputs + ] + + needs_batch_encoding = ( + is_batched or padding or return_tensors is not None or return_dict) + if not needs_batch_encoding: + return encoded_inputs[0] + + features = [{ + "input_ids": ids, + "attention_mask": [1] * len(ids) + } for ids in encoded_inputs] + batch = self.pad(features, + padding=padding, + max_length=max_length if padding else None, + return_attention_mask=True, + return_tensors=return_tensors) + + if return_dict: + return batch + if is_batched: + return batch["input_ids"] + return batch["input_ids"][0] if return_tensors is None else batch[ + "input_ids"] + + """ ----- Below are the abstract methods required by PreTrainedTokenizer ----- """ + + @property + def vocab_size(self) -> int: + return self.n_words + + def get_vocab(self) -> Dict[str, int]: + return self.encoder + + def _tokenize(self, text: str, **kwargs) -> List[str]: + return [self.decoder[t] for t in self.encode(text)] + + def _convert_token_to_id(self, token: str) -> int: + return self.encoder.get(token, self.unk_id) + + def _convert_id_to_token(self, index: int) -> str: + return self.decoder.get(index) + + @staticmethod + def clean_up_tokenization(out_string: str) -> str: + return out_string + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + text = ''.join(tokens) + text = bytearray([self.byte_decoder[c] + for c in text]).decode('utf-8', 'replace') + return text + + def save_vocabulary(self, + save_directory: str, + filename_prefix: Optional[str] = None) -> Tuple[str]: + if not os.path.isdir(save_directory): + raise ValueError( + f"vocabulary path ({save_directory}) should be a directory") + out_vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"]) + + if os.path.abspath(self.vocab_file) != os.path.abspath( + out_vocab_file) and os.path.isfile(self.vocab_file): + copyfile(self.vocab_file, out_vocab_file) + + return (out_vocab_file, ) + + def apply_chat_template(self, + conversation, + tools: Optional[list[dict]] = None, + tokenize: bool = False, + add_generation_prompt: bool = True, + thinking: bool = True, + padding=False, + truncation: bool = False, + max_length: Optional[int] = None, + return_tensors=None, + return_dict: bool = False, + **kwargs): + # Tokenizer-level rendering reorders tool result messages to match + # assistant tool_calls, normalizes per-call arguments and response + # schema, then encodes the resulting XTML structure segment-by-segment. + is_batched = is_batched_conversation(conversation) + conversations = conversation if is_batched else [conversation] + image_prompts = kwargs.pop("image_prompts", None) + if is_batched and image_prompts is not None: + raise ValueError("image_prompts is only supported for one chat.") + + # by default set thinking effort to max + kwargs.setdefault("thinking_effort", "max") + + segment_batches = [ + build_chat_segments( + messages, + tools=tools, + add_generation_prompt=add_generation_prompt, + thinking=thinking, + image_prompts=image_prompts, + **kwargs, + ) for messages in conversations + ] + + if not tokenize: + rendered = ["".join(segment.text for segment in segments) + for segments in segment_batches] + return rendered if is_batched else rendered[0] + + encoded_inputs = [ + self._encode_chat_segments(segments) for segments in segment_batches + ] + return self._format_chat_token_output( + encoded_inputs, + is_batched=is_batched, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors, + return_dict=return_dict, + ) diff --git a/tokenizer_config.json b/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..418aaacbd996a6674efa89d26be06dae11433e1c --- /dev/null +++ b/tokenizer_config.json @@ -0,0 +1,157 @@ +{ + "added_tokens_decoder": { + "163584": { + "content": "[BOS]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163585": { + "content": "[EOS]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163586": { + "content": "<|end_of_msg|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163587": { + "content": "<|open|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "163588": { + "content": "<|close|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "163589": { + "content": "<|sep|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "163590": { + "content": "[start_header_id]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163591": { + "content": "[end_header_id]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163593": { + "content": "[EOT]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163602": { + "content": "<|media_begin|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163603": { + "content": "<|media_content|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163604": { + "content": "<|media_end|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163605": { + "content": "<|media_pad|>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163649": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163838": { + "content": "[UNK]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "163839": { + "content": "[PAD]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + } + }, + "additional_special_tokens": [ + "<|end_of_msg|>", + "[start_header_id]", + "[end_header_id]", + "[EOT]", + "<|media_begin|>", + "<|media_content|>", + "<|media_end|>", + "<|media_pad|>", + "" + ], + "bos_token": "[BOS]", + "clean_up_tokenization_spaces": false, + "eos_token": "[EOS]", + "extra_special_tokens": {}, + "model_max_length": 1000000000000000019884624838656, + "pad_token": "[PAD]", + "tokenizer_class": "TikTokenTokenizer", + "unk_token": "[UNK]", + "auto_map": { + "AutoTokenizer": [ + "tokenization_kimi.TikTokenTokenizer", + null + ] + } +} \ No newline at end of file