diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..495a7175bb7b0b48bd0ac67d86f3e0c900127a71 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,4 @@ saved_model/**/* 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
diff --git a/chat_template.jinja b/chat_template.jinja
new file mode 100644
index 0000000000000000000000000000000000000000..f88cb41422626beb4cf1e7826b079448a4663cd0
--- /dev/null
+++ b/chat_template.jinja
@@ -0,0 +1,108 @@
+{%- macro render_content(msg) -%}
+ {%- set c = msg.get('content') -%}
+ {%- if c is string -%}
+ {{ c }}
+ {%- elif c is not none -%}
+ {% for content in c -%}
+ {% if content['type'] == 'image' or content['type'] == 'image_url' -%}
+ <|media_begin|>image<|media_content|><|media_pad|><|media_end|>
+ {% elif content['type'] == 'video' or content['type']== 'video_url'-%}
+ <|kimi_k25_video_placeholder|>
+ {% else -%}
+ {{ content['text'] }}
+ {%- endif -%}
+ {%- endfor -%}
+ {%- endif -%}
+{%- endmacro -%}
+
+{% macro set_roles(message) -%}
+ {%- set role_name = message.get('name') or message['role'] -%}
+ {%- if message['role'] == 'user' -%}
+ <|im_user|>{{role_name}}<|im_middle|>
+ {%- elif message['role'] == 'assistant' -%}
+ <|im_assistant|>{{role_name}}<|im_middle|>
+ {%- else -%}
+ <|im_system|>{{role_name}}<|im_middle|>
+ {%- endif -%}
+{%- endmacro -%}
+
+
+{%- macro render_toolcalls(message) -%}
+ <|tool_calls_section_begin|>
+ {%- for tool_call in message['tool_calls'] -%}
+ {%- set formatted_id = tool_call['id'] -%}
+ <|tool_call_begin|>{{ formatted_id }}<|tool_call_argument_begin|>{% if tool_call['function']['arguments'] is string %}{{ tool_call['function']['arguments'] }}{% else %}{{ tool_call['function']['arguments'] | tojson }}{% endif %}<|tool_call_end|>
+ {%- endfor -%}
+ <|tool_calls_section_end|>
+{%- endmacro -%}
+
+
+{# Find last non-tool-call assisitant message #}
+{%- set ns = namespace(last_non_tool_call_assistant_msg=-1) -%}
+{%- for idx in range(messages|length-1, -1, -1) -%}
+ {%- if messages[idx]['role'] == 'assistant' and not messages[idx].get('tool_calls') -%}
+ {%- set ns.last_non_tool_call_assistant_msg = idx -%}
+ {%- break -%}
+ {%- endif -%}
+{%- endfor -%}
+
+{# split all messages into history & suffix, reasoning_content in suffix should be reserved.#}
+{%- set hist_msgs = messages[:ns.last_non_tool_call_assistant_msg+1] -%}
+{%- set suffix_msgs = messages[ns.last_non_tool_call_assistant_msg+1:] -%}
+
+{%- if tools -%}
+ {%- if tools_ts_str -%}
+ <|im_system|>tool_declare<|im_middle|>{{ tools_ts_str }}<|im_end|>
+ {%- else -%}
+ <|im_system|>tool_declare<|im_middle|>{{ tools | tojson(separators=(',', ':')) }}<|im_end|>
+ {%- endif -%}
+{%- endif -%}
+
+{%- for message in hist_msgs -%}
+ {{set_roles(message)}}
+ {%- if message['role'] == 'assistant' -%}
+ {{render_content(message)}}
+ {%- if message.get('tool_calls') -%}
+ {{render_toolcalls(message)}}
+ {%- endif -%}
+ {%- elif message['role'] == 'tool' -%}
+ {%- set tool_call_id = message.tool_call_id -%}
+ ## Return of {{ tool_call_id }}
+{{render_content(message)}}
+ {%- elif message['content'] is not none -%}
+ {{render_content(message)}}
+ {%- endif -%}
+ <|im_end|>
+{%- endfor -%}
+
+{%- for message in suffix_msgs -%}
+ {{set_roles(message)}}
+ {%- if message['role'] == 'assistant' -%}
+ {%- if thinking is defined and thinking is false -%}
+ {{render_content(message)}}
+ {%- else -%}
+ {%- set rc = message.get('reasoning_content', '') -%}
+ {{rc}}{{render_content(message)}}
+ {%- endif -%}
+ {%- if message.get('tool_calls') -%}
+ {{render_toolcalls(message)}}
+ {%- endif -%}
+ {%- elif message['role'] == 'tool' -%}
+ {%- set tool_call_id = message.tool_call_id -%}
+ ## Return of {{ tool_call_id }}
+{{render_content(message)}}
+ {%- elif message['content'] is not none -%}
+ {{render_content(message)}}
+ {%- endif -%}
+ <|im_end|>
+{%- endfor -%}
+
+
+{%- if add_generation_prompt -%}
+ <|im_assistant|>assistant<|im_middle|>
+ {%- if thinking is defined and thinking is false -%}
+
+ {%- else -%}
+
+ {%- endif -%}
+{%- endif -%}
\ No newline at end of file
diff --git a/config.json b/config.json
new file mode 100644
index 0000000000000000000000000000000000000000..a360507c84898d753ae1f10f7a50ffa0c8d25572
--- /dev/null
+++ b/config.json
@@ -0,0 +1,440 @@
+{
+ "architectures": [
+ "KimiK25ForConditionalGeneration"
+ ],
+ "auto_map": {
+ "AutoConfig": "configuration_kimi_k25.KimiK25Config",
+ "AutoModel": "modeling_kimi_k25.KimiK25ForConditionalGeneration",
+ "AutoModelForCausalLM": "modeling_kimi_k25.KimiK25ForConditionalGeneration"
+ },
+ "bos_token_id": 163584,
+ "dtype": "bfloat16",
+ "eos_token_id": 163585,
+ "ignore_index": -100,
+ "media_placeholder_token_id": 163605,
+ "model_type": "kimi_k25",
+ "pad_token_id": 163839,
+ "text_config": {
+ "_name_or_path": "",
+ "add_cross_attention": false,
+ "architectures": [
+ "DeepseekV3ForCausalLM"
+ ],
+ "attention_bias": false,
+ "attention_dropout": 0.0,
+ "auto_map": {
+ "AutoConfig": "configuration_deepseek.DeepseekV3Config",
+ "AutoModel": "modeling_deepseek.DeepseekV3Model",
+ "AutoModelForCausalLM": "modeling_deepseek.DeepseekV3ForCausalLM"
+ },
+ "aux_loss_alpha": 0.001,
+ "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": 163585,
+ "ep_size": 1,
+ "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": "silu",
+ "hidden_size": 7168,
+ "id2label": {
+ "0": "LABEL_0",
+ "1": "LABEL_1"
+ },
+ "initializer_range": 0.02,
+ "intermediate_size": 18432,
+ "is_decoder": false,
+ "is_encoder_decoder": false,
+ "kv_lora_rank": 512,
+ "label2id": {
+ "LABEL_0": 0,
+ "LABEL_1": 1
+ },
+ "length_penalty": 1.0,
+ "max_length": 20,
+ "max_position_embeddings": 262144,
+ "min_length": 0,
+ "model_type": "deepseek_v3",
+ "moe_intermediate_size": 2048,
+ "moe_layer_freq": 1,
+ "n_group": 1,
+ "n_routed_experts": 384,
+ "n_shared_experts": 1,
+ "no_repeat_ngram_size": 0,
+ "norm_topk_prob": true,
+ "num_attention_heads": 64,
+ "num_beam_groups": 1,
+ "num_beams": 1,
+ "num_experts_per_tok": 8,
+ "num_hidden_layers": 61,
+ "num_key_value_heads": 64,
+ "num_nextn_predict_layers": 0,
+ "num_return_sequences": 1,
+ "output_attentions": false,
+ "output_hidden_states": false,
+ "output_scores": false,
+ "pad_token_id": 163839,
+ "prefix": null,
+ "pretraining_tp": 1,
+ "problem_type": null,
+ "pruned_heads": {},
+ "q_lora_rank": 1536,
+ "qk_nope_head_dim": 128,
+ "qk_rope_head_dim": 64,
+ "remove_invalid_values": false,
+ "repetition_penalty": 1.0,
+ "return_dict": true,
+ "return_dict_in_generate": false,
+ "rms_norm_eps": 1e-05,
+ "rope_scaling": {
+ "beta_fast": 32.0,
+ "beta_slow": 1.0,
+ "factor": 64.0,
+ "mscale": 1.0,
+ "mscale_all_dim": 1.0,
+ "original_max_position_embeddings": 4096,
+ "type": "yarn"
+ },
+ "rope_theta": 50000.0,
+ "routed_scaling_factor": 2.827,
+ "scoring_func": "sigmoid",
+ "sep_token_id": null,
+ "seq_aux": true,
+ "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,
+ "typical_p": 1.0,
+ "use_bfloat16": false,
+ "use_cache": true,
+ "v_head_dim": 128,
+ "vocab_size": 163840
+ },
+ "tie_word_embeddings": false,
+ "transformers_version": "4.57.1",
+ "use_unified_vision_chunk": true,
+ "video_placeholder": "<|kimi_k25_video_placeholder|>",
+ "vision_config": {
+ "init_pos_emb_height": 64,
+ "init_pos_emb_time": 4,
+ "init_pos_emb_width": 64,
+ "merge_kernel_size": [
+ 2,
+ 2
+ ],
+ "merge_type": "sd2_tpool",
+ "mm_hidden_size": 1152,
+ "mm_projector_type": "patchmerger",
+ "model_type": "",
+ "patch_size": 14,
+ "pos_emb_type": "divided_fixed",
+ "projector_hidden_act": "gelu",
+ "projector_ln_eps": 1e-05,
+ "text_hidden_size": 7168,
+ "video_attn_type": "spatial_temporal",
+ "vt_hidden_size": 1152,
+ "vt_intermediate_size": 4304,
+ "vt_num_attention_heads": 16,
+ "vt_num_hidden_layers": 27
+ },
+ "quantization_config": {
+ "config_groups": {
+ "group_0": {
+ "input_activations": {
+ "dynamic": false,
+ "num_bits": 4,
+ "type": "float",
+ "group_size": 16
+ },
+ "weights": {
+ "dynamic": false,
+ "num_bits": 4,
+ "type": "float",
+ "group_size": 16
+ },
+ "targets": [
+ "Linear"
+ ]
+ }
+ },
+ "ignore": [
+ "language_model.lm_head",
+ "language_model.model.layers.0.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.0.self_attn.kv_b_proj",
+ "language_model.model.layers.0.self_attn.q_a_proj",
+ "language_model.model.layers.0.self_attn.q_b_proj",
+ "language_model.model.layers.1.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.1.self_attn.kv_b_proj",
+ "language_model.model.layers.1.self_attn.q_a_proj",
+ "language_model.model.layers.1.self_attn.q_b_proj",
+ "language_model.model.layers.10.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.10.self_attn.kv_b_proj",
+ "language_model.model.layers.10.self_attn.q_a_proj",
+ "language_model.model.layers.10.self_attn.q_b_proj",
+ "language_model.model.layers.11.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.11.self_attn.kv_b_proj",
+ "language_model.model.layers.11.self_attn.q_a_proj",
+ "language_model.model.layers.11.self_attn.q_b_proj",
+ "language_model.model.layers.12.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.12.self_attn.kv_b_proj",
+ "language_model.model.layers.12.self_attn.q_a_proj",
+ "language_model.model.layers.12.self_attn.q_b_proj",
+ "language_model.model.layers.13.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.13.self_attn.kv_b_proj",
+ "language_model.model.layers.13.self_attn.q_a_proj",
+ "language_model.model.layers.13.self_attn.q_b_proj",
+ "language_model.model.layers.14.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.14.self_attn.kv_b_proj",
+ "language_model.model.layers.14.self_attn.q_a_proj",
+ "language_model.model.layers.14.self_attn.q_b_proj",
+ "language_model.model.layers.15.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.15.self_attn.kv_b_proj",
+ "language_model.model.layers.15.self_attn.q_a_proj",
+ "language_model.model.layers.15.self_attn.q_b_proj",
+ "language_model.model.layers.16.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.16.self_attn.kv_b_proj",
+ "language_model.model.layers.16.self_attn.q_a_proj",
+ "language_model.model.layers.16.self_attn.q_b_proj",
+ "language_model.model.layers.17.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.17.self_attn.kv_b_proj",
+ "language_model.model.layers.17.self_attn.q_a_proj",
+ "language_model.model.layers.17.self_attn.q_b_proj",
+ "language_model.model.layers.18.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.18.self_attn.kv_b_proj",
+ "language_model.model.layers.18.self_attn.q_a_proj",
+ "language_model.model.layers.18.self_attn.q_b_proj",
+ "language_model.model.layers.19.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.19.self_attn.kv_b_proj",
+ "language_model.model.layers.19.self_attn.q_a_proj",
+ "language_model.model.layers.19.self_attn.q_b_proj",
+ "language_model.model.layers.2.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.2.self_attn.kv_b_proj",
+ "language_model.model.layers.2.self_attn.q_a_proj",
+ "language_model.model.layers.2.self_attn.q_b_proj",
+ "language_model.model.layers.20.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.20.self_attn.kv_b_proj",
+ "language_model.model.layers.20.self_attn.q_a_proj",
+ "language_model.model.layers.20.self_attn.q_b_proj",
+ "language_model.model.layers.21.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.21.self_attn.kv_b_proj",
+ "language_model.model.layers.21.self_attn.q_a_proj",
+ "language_model.model.layers.21.self_attn.q_b_proj",
+ "language_model.model.layers.22.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.22.self_attn.kv_b_proj",
+ "language_model.model.layers.22.self_attn.q_a_proj",
+ "language_model.model.layers.22.self_attn.q_b_proj",
+ "language_model.model.layers.23.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.23.self_attn.kv_b_proj",
+ "language_model.model.layers.23.self_attn.q_a_proj",
+ "language_model.model.layers.23.self_attn.q_b_proj",
+ "language_model.model.layers.24.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.24.self_attn.kv_b_proj",
+ "language_model.model.layers.24.self_attn.q_a_proj",
+ "language_model.model.layers.24.self_attn.q_b_proj",
+ "language_model.model.layers.25.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.25.self_attn.kv_b_proj",
+ "language_model.model.layers.25.self_attn.q_a_proj",
+ "language_model.model.layers.25.self_attn.q_b_proj",
+ "language_model.model.layers.26.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.26.self_attn.kv_b_proj",
+ "language_model.model.layers.26.self_attn.q_a_proj",
+ "language_model.model.layers.26.self_attn.q_b_proj",
+ "language_model.model.layers.27.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.27.self_attn.kv_b_proj",
+ "language_model.model.layers.27.self_attn.q_a_proj",
+ "language_model.model.layers.27.self_attn.q_b_proj",
+ "language_model.model.layers.28.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.28.self_attn.kv_b_proj",
+ "language_model.model.layers.28.self_attn.q_a_proj",
+ "language_model.model.layers.28.self_attn.q_b_proj",
+ "language_model.model.layers.29.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.29.self_attn.kv_b_proj",
+ "language_model.model.layers.29.self_attn.q_a_proj",
+ "language_model.model.layers.29.self_attn.q_b_proj",
+ "language_model.model.layers.3.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.3.self_attn.kv_b_proj",
+ "language_model.model.layers.3.self_attn.q_a_proj",
+ "language_model.model.layers.3.self_attn.q_b_proj",
+ "language_model.model.layers.30.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.30.self_attn.kv_b_proj",
+ "language_model.model.layers.30.self_attn.q_a_proj",
+ "language_model.model.layers.30.self_attn.q_b_proj",
+ "language_model.model.layers.31.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.31.self_attn.kv_b_proj",
+ "language_model.model.layers.31.self_attn.q_a_proj",
+ "language_model.model.layers.31.self_attn.q_b_proj",
+ "language_model.model.layers.32.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.32.self_attn.kv_b_proj",
+ "language_model.model.layers.32.self_attn.q_a_proj",
+ "language_model.model.layers.32.self_attn.q_b_proj",
+ "language_model.model.layers.33.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.33.self_attn.kv_b_proj",
+ "language_model.model.layers.33.self_attn.q_a_proj",
+ "language_model.model.layers.33.self_attn.q_b_proj",
+ "language_model.model.layers.34.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.34.self_attn.kv_b_proj",
+ "language_model.model.layers.34.self_attn.q_a_proj",
+ "language_model.model.layers.34.self_attn.q_b_proj",
+ "language_model.model.layers.35.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.35.self_attn.kv_b_proj",
+ "language_model.model.layers.35.self_attn.q_a_proj",
+ "language_model.model.layers.35.self_attn.q_b_proj",
+ "language_model.model.layers.36.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.36.self_attn.kv_b_proj",
+ "language_model.model.layers.36.self_attn.q_a_proj",
+ "language_model.model.layers.36.self_attn.q_b_proj",
+ "language_model.model.layers.37.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.37.self_attn.kv_b_proj",
+ "language_model.model.layers.37.self_attn.q_a_proj",
+ "language_model.model.layers.37.self_attn.q_b_proj",
+ "language_model.model.layers.38.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.38.self_attn.kv_b_proj",
+ "language_model.model.layers.38.self_attn.q_a_proj",
+ "language_model.model.layers.38.self_attn.q_b_proj",
+ "language_model.model.layers.39.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.39.self_attn.kv_b_proj",
+ "language_model.model.layers.39.self_attn.q_a_proj",
+ "language_model.model.layers.39.self_attn.q_b_proj",
+ "language_model.model.layers.4.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.4.self_attn.kv_b_proj",
+ "language_model.model.layers.4.self_attn.q_a_proj",
+ "language_model.model.layers.4.self_attn.q_b_proj",
+ "language_model.model.layers.40.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.40.self_attn.kv_b_proj",
+ "language_model.model.layers.40.self_attn.q_a_proj",
+ "language_model.model.layers.40.self_attn.q_b_proj",
+ "language_model.model.layers.41.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.41.self_attn.kv_b_proj",
+ "language_model.model.layers.41.self_attn.q_a_proj",
+ "language_model.model.layers.41.self_attn.q_b_proj",
+ "language_model.model.layers.42.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.42.self_attn.kv_b_proj",
+ "language_model.model.layers.42.self_attn.q_a_proj",
+ "language_model.model.layers.42.self_attn.q_b_proj",
+ "language_model.model.layers.43.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.43.self_attn.kv_b_proj",
+ "language_model.model.layers.43.self_attn.q_a_proj",
+ "language_model.model.layers.43.self_attn.q_b_proj",
+ "language_model.model.layers.44.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.44.self_attn.kv_b_proj",
+ "language_model.model.layers.44.self_attn.q_a_proj",
+ "language_model.model.layers.44.self_attn.q_b_proj",
+ "language_model.model.layers.45.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.45.self_attn.kv_b_proj",
+ "language_model.model.layers.45.self_attn.q_a_proj",
+ "language_model.model.layers.45.self_attn.q_b_proj",
+ "language_model.model.layers.46.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.46.self_attn.kv_b_proj",
+ "language_model.model.layers.46.self_attn.q_a_proj",
+ "language_model.model.layers.46.self_attn.q_b_proj",
+ "language_model.model.layers.47.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.47.self_attn.kv_b_proj",
+ "language_model.model.layers.47.self_attn.q_a_proj",
+ "language_model.model.layers.47.self_attn.q_b_proj",
+ "language_model.model.layers.48.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.48.self_attn.kv_b_proj",
+ "language_model.model.layers.48.self_attn.q_a_proj",
+ "language_model.model.layers.48.self_attn.q_b_proj",
+ "language_model.model.layers.49.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.49.self_attn.kv_b_proj",
+ "language_model.model.layers.49.self_attn.q_a_proj",
+ "language_model.model.layers.49.self_attn.q_b_proj",
+ "language_model.model.layers.5.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.5.self_attn.kv_b_proj",
+ "language_model.model.layers.5.self_attn.q_a_proj",
+ "language_model.model.layers.5.self_attn.q_b_proj",
+ "language_model.model.layers.50.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.50.self_attn.kv_b_proj",
+ "language_model.model.layers.50.self_attn.q_a_proj",
+ "language_model.model.layers.50.self_attn.q_b_proj",
+ "language_model.model.layers.51.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.51.self_attn.kv_b_proj",
+ "language_model.model.layers.51.self_attn.q_a_proj",
+ "language_model.model.layers.51.self_attn.q_b_proj",
+ "language_model.model.layers.52.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.52.self_attn.kv_b_proj",
+ "language_model.model.layers.52.self_attn.q_a_proj",
+ "language_model.model.layers.52.self_attn.q_b_proj",
+ "language_model.model.layers.53.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.53.self_attn.kv_b_proj",
+ "language_model.model.layers.53.self_attn.q_a_proj",
+ "language_model.model.layers.53.self_attn.q_b_proj",
+ "language_model.model.layers.54.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.54.self_attn.kv_b_proj",
+ "language_model.model.layers.54.self_attn.q_a_proj",
+ "language_model.model.layers.54.self_attn.q_b_proj",
+ "language_model.model.layers.55.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.55.self_attn.kv_b_proj",
+ "language_model.model.layers.55.self_attn.q_a_proj",
+ "language_model.model.layers.55.self_attn.q_b_proj",
+ "language_model.model.layers.56.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.56.self_attn.kv_b_proj",
+ "language_model.model.layers.56.self_attn.q_a_proj",
+ "language_model.model.layers.56.self_attn.q_b_proj",
+ "language_model.model.layers.57.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.57.self_attn.kv_b_proj",
+ "language_model.model.layers.57.self_attn.q_a_proj",
+ "language_model.model.layers.57.self_attn.q_b_proj",
+ "language_model.model.layers.58.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.58.self_attn.kv_b_proj",
+ "language_model.model.layers.58.self_attn.q_a_proj",
+ "language_model.model.layers.58.self_attn.q_b_proj",
+ "language_model.model.layers.59.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.59.self_attn.kv_b_proj",
+ "language_model.model.layers.59.self_attn.q_a_proj",
+ "language_model.model.layers.59.self_attn.q_b_proj",
+ "language_model.model.layers.6.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.6.self_attn.kv_b_proj",
+ "language_model.model.layers.6.self_attn.q_a_proj",
+ "language_model.model.layers.6.self_attn.q_b_proj",
+ "language_model.model.layers.60.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.60.self_attn.kv_b_proj",
+ "language_model.model.layers.60.self_attn.q_a_proj",
+ "language_model.model.layers.60.self_attn.q_b_proj",
+ "language_model.model.layers.7.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.7.self_attn.kv_b_proj",
+ "language_model.model.layers.7.self_attn.q_a_proj",
+ "language_model.model.layers.7.self_attn.q_b_proj",
+ "language_model.model.layers.8.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.8.self_attn.kv_b_proj",
+ "language_model.model.layers.8.self_attn.q_a_proj",
+ "language_model.model.layers.8.self_attn.q_b_proj",
+ "language_model.model.layers.9.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.9.self_attn.kv_b_proj",
+ "language_model.model.layers.9.self_attn.q_a_proj",
+ "language_model.model.layers.9.self_attn.q_b_proj",
+ "mm_projector*",
+ "vision_tower*"
+ ],
+ "quant_algo": "NVFP4",
+ "kv_cache_scheme": {
+ "dynamic": false,
+ "num_bits": 8,
+ "type": "float"
+ },
+ "producer": {
+ "name": "modelopt",
+ "version": "0.41.0rc2.dev54+g99912fbdf"
+ },
+ "quant_method": "modelopt"
+ }
+}
\ No newline at end of file
diff --git a/configuration_deepseek.py b/configuration_deepseek.py
new file mode 100644
index 0000000000000000000000000000000000000000..b3152dd7c3e53d223d561848dc967f487daf32ef
--- /dev/null
+++ b/configuration_deepseek.py
@@ -0,0 +1,214 @@
+# Copy from https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/configuration_deepseek.py
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+
+class DeepseekV3Config(PretrainedConfig):
+ r"""
+ This is the configuration class to store the configuration of a [`DeepseekV3Model`]. It is used to instantiate an DeepSeek
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+ defaults will yield a similar configuration to that of the DeepSeek-V3.
+
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+ documentation from [`PretrainedConfig`] for more information.
+
+
+ Args:
+ vocab_size (`int`, *optional*, defaults to 129280):
+ Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`DeepseekV3Model`]
+ hidden_size (`int`, *optional*, defaults to 4096):
+ Dimension of the hidden representations.
+ intermediate_size (`int`, *optional*, defaults to 11008):
+ Dimension of the MLP representations.
+ moe_intermediate_size (`int`, *optional*, defaults to 1407):
+ Dimension of the MoE representations.
+ num_hidden_layers (`int`, *optional*, defaults to 32):
+ Number of hidden layers in the Transformer decoder.
+ num_nextn_predict_layers (`int`, *optional*, defaults to 1):
+ Number of nextn predict layers in the DeepSeekV3 Model.
+ num_attention_heads (`int`, *optional*, defaults to 32):
+ Number of attention heads for each attention layer in the Transformer decoder.
+ n_shared_experts (`int`, *optional*, defaults to None):
+ Number of shared experts, None means dense model.
+ n_routed_experts (`int`, *optional*, defaults to None):
+ Number of routed experts, None means dense model.
+ routed_scaling_factor (`float`, *optional*, defaults to 1.0):
+ Scaling factor or routed experts.
+ topk_method (`str`, *optional*, defaults to `gready`):
+ Topk method used in routed gate.
+ n_group (`int`, *optional*, defaults to None):
+ Number of groups for routed experts.
+ topk_group (`int`, *optional*, defaults to None):
+ Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups).
+ num_experts_per_tok (`int`, *optional*, defaults to None):
+ Number of selected experts, None means dense model.
+ moe_layer_freq (`int`, *optional*, defaults to 1):
+ The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers.
+ first_k_dense_replace (`int`, *optional*, defaults to 0):
+ Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head).
+ \--k dense layers--/
+ norm_topk_prob (`bool`, *optional*, defaults to False):
+ Whether to normalize the weights of the routed experts.
+ scoring_func (`str`, *optional*, defaults to 'softmax'):
+ Method of computing expert weights.
+ aux_loss_alpha (`float`, *optional*, defaults to 0.001):
+ Auxiliary loss weight coefficient.
+ seq_aux = (`bool`, *optional*, defaults to True):
+ Whether to compute the auxiliary loss for each individual sample.
+ num_key_value_heads (`int`, *optional*):
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+ by meanpooling all the original heads within that group. For more details checkout [this
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+ `num_attention_heads`.
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+ The non-linear activation function (function or string) in the decoder.
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
+ The maximum sequence length that this model might ever be used with.
+ initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+ The epsilon used by the rms normalization layers.
+ use_cache (`bool`, *optional*, defaults to `True`):
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
+ relevant if `config.is_decoder=True`.
+ pad_token_id (`int`, *optional*):
+ Padding token id.
+ bos_token_id (`int`, *optional*, defaults to 1):
+ Beginning of stream token id.
+ eos_token_id (`int`, *optional*, defaults to 2):
+ End of stream token id.
+ pretraining_tp (`int`, *optional*, defaults to 1):
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
+ document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
+ necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
+ issue](https://github.com/pytorch/pytorch/issues/76232).
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+ Whether to tie weight embeddings
+ rope_theta (`float`, *optional*, defaults to 10000.0):
+ The base period of the RoPE embeddings.
+ rope_scaling (`Dict`, *optional*):
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
+ `max_position_embeddings` to the expected new maximum.
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
+ attention_dropout (`float`, *optional*, defaults to 0.0):
+ The dropout ratio for the attention probabilities.
+
+ ```python
+ >>> from transformers import DeepseekV3Model, DeepseekV3Config
+
+ >>> # Initializing a Deepseek-V3 style configuration
+ >>> configuration = DeepseekV3Config()
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "deepseek_v3"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ vocab_size=129280,
+ hidden_size=7168,
+ intermediate_size=18432,
+ moe_intermediate_size=2048,
+ num_hidden_layers=61,
+ num_nextn_predict_layers=1,
+ num_attention_heads=128,
+ num_key_value_heads=128,
+ n_shared_experts=1,
+ n_routed_experts=256,
+ ep_size=1,
+ routed_scaling_factor=2.5,
+ kv_lora_rank=512,
+ q_lora_rank=1536,
+ qk_rope_head_dim=64,
+ v_head_dim=128,
+ qk_nope_head_dim=128,
+ topk_method='noaux_tc',
+ n_group=8,
+ topk_group=4,
+ num_experts_per_tok=8,
+ moe_layer_freq=1,
+ first_k_dense_replace=3,
+ norm_topk_prob=True,
+ scoring_func='sigmoid',
+ aux_loss_alpha=0.001,
+ seq_aux=True,
+ hidden_act="silu",
+ max_position_embeddings=4096,
+ initializer_range=0.02,
+ rms_norm_eps=1e-6,
+ use_cache=True,
+ pad_token_id=None,
+ bos_token_id=0,
+ eos_token_id=1,
+ pretraining_tp=1,
+ tie_word_embeddings=False,
+ rope_theta=10000.0,
+ rope_scaling=None,
+ attention_bias=False,
+ attention_dropout=0.0,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.moe_intermediate_size = moe_intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_nextn_predict_layers = num_nextn_predict_layers
+ self.num_attention_heads = num_attention_heads
+ self.n_shared_experts = n_shared_experts
+ self.n_routed_experts = n_routed_experts
+ self.ep_size = ep_size
+ self.routed_scaling_factor = routed_scaling_factor
+ self.kv_lora_rank = kv_lora_rank
+ self.q_lora_rank = q_lora_rank
+ self.qk_rope_head_dim = qk_rope_head_dim
+ self.v_head_dim = v_head_dim
+ self.qk_nope_head_dim = qk_nope_head_dim
+ self.topk_method = topk_method
+ self.n_group = n_group
+ self.topk_group = topk_group
+ self.num_experts_per_tok = num_experts_per_tok
+ self.moe_layer_freq = moe_layer_freq
+ self.first_k_dense_replace = first_k_dense_replace
+ self.norm_topk_prob = norm_topk_prob
+ self.scoring_func = scoring_func
+ self.aux_loss_alpha = aux_loss_alpha
+ self.seq_aux = seq_aux
+ # 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.pretraining_tp = pretraining_tp
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.rope_scaling = rope_scaling
+ self.attention_bias = attention_bias
+ self.attention_dropout = attention_dropout
+
+ 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,
+ )
diff --git a/configuration_kimi_k25.py b/configuration_kimi_k25.py
new file mode 100644
index 0000000000000000000000000000000000000000..d52985edc830ac38140b0a7046a6df10379fed73
--- /dev/null
+++ b/configuration_kimi_k25.py
@@ -0,0 +1,117 @@
+import os
+import sys
+import importlib.util
+from transformers.configuration_utils import PretrainedConfig
+
+# Dynamic import to avoid transformers static import checker
+# Transformers copies only files in auto_map to cache, so we need to try multiple locations
+_current_dir = os.path.dirname(os.path.abspath(__file__))
+
+def _find_and_import_module(module_name, file_name):
+ """Find module file in multiple locations and import it."""
+ # Try current directory first (works when loading directly from model dir)
+ candidates = [
+ os.path.join(_current_dir, file_name),
+ # Fallback: original model path (when running from HF cache)
+ "/workspace/Model-Optimizer/examples/llm_ptq/kimi25_hf-bf16/" + file_name,
+ ]
+
+ for path in candidates:
+ if os.path.exists(path):
+ spec = importlib.util.spec_from_file_location(module_name, path)
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[module_name] = module
+ spec.loader.exec_module(module)
+ return module
+
+ raise FileNotFoundError(f"Could not find {file_name} in any of: {candidates}")
+
+_deepseek_module = _find_and_import_module("configuration_deepseek", "configuration_deepseek.py")
+DeepseekV3Config = _deepseek_module.DeepseekV3Config
+
+
+class KimiK25VisionConfig(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 = 16,
+ vt_num_hidden_layers: int = 27,
+ vt_hidden_size: int = 1152,
+ vt_intermediate_size: int = 4304,
+ merge_kernel_size: tuple = (2, 2),
+ video_attn_type: str = 'spatial_temporal',
+ merge_type: str = 'sd2_tpool',
+ _attn_implementation: str = 'flash_attention_2',
+ # MM Projector parameters
+ mm_projector_type: str = 'patchmerger',
+ mm_hidden_size: int | None = None,
+ projector_hidden_act: str = "gelu",
+ projector_ln_eps: float = 1e-5,
+ # Other parameters
+ ignore_index: int = -100,
+ media_placeholder_token_id: int = 163605,
+ pad_token_id: int = 0,
+ use_unified_vision_chunk: bool = True,
+ video_placeholder="<|kimi_k25_video_placeholder|>",
+ text_hidden_size=7168,
+ **vision_config_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.video_attn_type = video_attn_type
+ 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
+
+
+class KimiK25Config(PretrainedConfig):
+ """Kimi-K2.5 model configuration."""
+
+ model_type = "kimi_k25"
+
+ def __init__(
+ self,
+ text_config: dict | DeepseekV3Config = None,
+ vision_config: dict | KimiK25VisionConfig = None,
+ # Other parameters
+ ignore_index: int = -100,
+ media_placeholder_token_id: int = 163605,
+ pad_token_id: int = 0,
+ use_unified_vision_chunk: bool = True,
+ video_placeholder="<|kimi_k25_video_placeholder|>",
+ **kwargs,
+ ):
+ if isinstance(text_config, dict):
+ text_config = DeepseekV3Config(**text_config)
+ if isinstance(vision_config, dict):
+ vision_config = KimiK25VisionConfig(**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
+ self.use_unified_vision_chunk = use_unified_vision_chunk
+ self.video_placeholder = video_placeholder
+ 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/generation_config.json b/generation_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..4bb3f8ec4a1d604598b7ffb4621b955e995bda92
--- /dev/null
+++ b/generation_config.json
@@ -0,0 +1,4 @@
+{
+ "max_length": 262144,
+ "eos_token_id": 163586
+}
\ No newline at end of file
diff --git a/hf_quant_config.json b/hf_quant_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..5b3ee5aa99242b3f998f84e85385912072583b8f
--- /dev/null
+++ b/hf_quant_config.json
@@ -0,0 +1,260 @@
+{
+ "producer": {
+ "name": "modelopt",
+ "version": "0.41.0rc2.dev54+g99912fbdf"
+ },
+ "quantization": {
+ "quant_algo": "NVFP4",
+ "kv_cache_quant_algo": "FP8",
+ "group_size": 16,
+ "exclude_modules": [
+ "language_model.lm_head",
+ "language_model.model.layers.0.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.0.self_attn.kv_b_proj",
+ "language_model.model.layers.0.self_attn.q_a_proj",
+ "language_model.model.layers.0.self_attn.q_b_proj",
+ "language_model.model.layers.1.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.1.self_attn.kv_b_proj",
+ "language_model.model.layers.1.self_attn.q_a_proj",
+ "language_model.model.layers.1.self_attn.q_b_proj",
+ "language_model.model.layers.10.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.10.self_attn.kv_b_proj",
+ "language_model.model.layers.10.self_attn.q_a_proj",
+ "language_model.model.layers.10.self_attn.q_b_proj",
+ "language_model.model.layers.11.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.11.self_attn.kv_b_proj",
+ "language_model.model.layers.11.self_attn.q_a_proj",
+ "language_model.model.layers.11.self_attn.q_b_proj",
+ "language_model.model.layers.12.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.12.self_attn.kv_b_proj",
+ "language_model.model.layers.12.self_attn.q_a_proj",
+ "language_model.model.layers.12.self_attn.q_b_proj",
+ "language_model.model.layers.13.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.13.self_attn.kv_b_proj",
+ "language_model.model.layers.13.self_attn.q_a_proj",
+ "language_model.model.layers.13.self_attn.q_b_proj",
+ "language_model.model.layers.14.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.14.self_attn.kv_b_proj",
+ "language_model.model.layers.14.self_attn.q_a_proj",
+ "language_model.model.layers.14.self_attn.q_b_proj",
+ "language_model.model.layers.15.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.15.self_attn.kv_b_proj",
+ "language_model.model.layers.15.self_attn.q_a_proj",
+ "language_model.model.layers.15.self_attn.q_b_proj",
+ "language_model.model.layers.16.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.16.self_attn.kv_b_proj",
+ "language_model.model.layers.16.self_attn.q_a_proj",
+ "language_model.model.layers.16.self_attn.q_b_proj",
+ "language_model.model.layers.17.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.17.self_attn.kv_b_proj",
+ "language_model.model.layers.17.self_attn.q_a_proj",
+ "language_model.model.layers.17.self_attn.q_b_proj",
+ "language_model.model.layers.18.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.18.self_attn.kv_b_proj",
+ "language_model.model.layers.18.self_attn.q_a_proj",
+ "language_model.model.layers.18.self_attn.q_b_proj",
+ "language_model.model.layers.19.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.19.self_attn.kv_b_proj",
+ "language_model.model.layers.19.self_attn.q_a_proj",
+ "language_model.model.layers.19.self_attn.q_b_proj",
+ "language_model.model.layers.2.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.2.self_attn.kv_b_proj",
+ "language_model.model.layers.2.self_attn.q_a_proj",
+ "language_model.model.layers.2.self_attn.q_b_proj",
+ "language_model.model.layers.20.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.20.self_attn.kv_b_proj",
+ "language_model.model.layers.20.self_attn.q_a_proj",
+ "language_model.model.layers.20.self_attn.q_b_proj",
+ "language_model.model.layers.21.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.21.self_attn.kv_b_proj",
+ "language_model.model.layers.21.self_attn.q_a_proj",
+ "language_model.model.layers.21.self_attn.q_b_proj",
+ "language_model.model.layers.22.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.22.self_attn.kv_b_proj",
+ "language_model.model.layers.22.self_attn.q_a_proj",
+ "language_model.model.layers.22.self_attn.q_b_proj",
+ "language_model.model.layers.23.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.23.self_attn.kv_b_proj",
+ "language_model.model.layers.23.self_attn.q_a_proj",
+ "language_model.model.layers.23.self_attn.q_b_proj",
+ "language_model.model.layers.24.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.24.self_attn.kv_b_proj",
+ "language_model.model.layers.24.self_attn.q_a_proj",
+ "language_model.model.layers.24.self_attn.q_b_proj",
+ "language_model.model.layers.25.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.25.self_attn.kv_b_proj",
+ "language_model.model.layers.25.self_attn.q_a_proj",
+ "language_model.model.layers.25.self_attn.q_b_proj",
+ "language_model.model.layers.26.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.26.self_attn.kv_b_proj",
+ "language_model.model.layers.26.self_attn.q_a_proj",
+ "language_model.model.layers.26.self_attn.q_b_proj",
+ "language_model.model.layers.27.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.27.self_attn.kv_b_proj",
+ "language_model.model.layers.27.self_attn.q_a_proj",
+ "language_model.model.layers.27.self_attn.q_b_proj",
+ "language_model.model.layers.28.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.28.self_attn.kv_b_proj",
+ "language_model.model.layers.28.self_attn.q_a_proj",
+ "language_model.model.layers.28.self_attn.q_b_proj",
+ "language_model.model.layers.29.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.29.self_attn.kv_b_proj",
+ "language_model.model.layers.29.self_attn.q_a_proj",
+ "language_model.model.layers.29.self_attn.q_b_proj",
+ "language_model.model.layers.3.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.3.self_attn.kv_b_proj",
+ "language_model.model.layers.3.self_attn.q_a_proj",
+ "language_model.model.layers.3.self_attn.q_b_proj",
+ "language_model.model.layers.30.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.30.self_attn.kv_b_proj",
+ "language_model.model.layers.30.self_attn.q_a_proj",
+ "language_model.model.layers.30.self_attn.q_b_proj",
+ "language_model.model.layers.31.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.31.self_attn.kv_b_proj",
+ "language_model.model.layers.31.self_attn.q_a_proj",
+ "language_model.model.layers.31.self_attn.q_b_proj",
+ "language_model.model.layers.32.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.32.self_attn.kv_b_proj",
+ "language_model.model.layers.32.self_attn.q_a_proj",
+ "language_model.model.layers.32.self_attn.q_b_proj",
+ "language_model.model.layers.33.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.33.self_attn.kv_b_proj",
+ "language_model.model.layers.33.self_attn.q_a_proj",
+ "language_model.model.layers.33.self_attn.q_b_proj",
+ "language_model.model.layers.34.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.34.self_attn.kv_b_proj",
+ "language_model.model.layers.34.self_attn.q_a_proj",
+ "language_model.model.layers.34.self_attn.q_b_proj",
+ "language_model.model.layers.35.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.35.self_attn.kv_b_proj",
+ "language_model.model.layers.35.self_attn.q_a_proj",
+ "language_model.model.layers.35.self_attn.q_b_proj",
+ "language_model.model.layers.36.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.36.self_attn.kv_b_proj",
+ "language_model.model.layers.36.self_attn.q_a_proj",
+ "language_model.model.layers.36.self_attn.q_b_proj",
+ "language_model.model.layers.37.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.37.self_attn.kv_b_proj",
+ "language_model.model.layers.37.self_attn.q_a_proj",
+ "language_model.model.layers.37.self_attn.q_b_proj",
+ "language_model.model.layers.38.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.38.self_attn.kv_b_proj",
+ "language_model.model.layers.38.self_attn.q_a_proj",
+ "language_model.model.layers.38.self_attn.q_b_proj",
+ "language_model.model.layers.39.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.39.self_attn.kv_b_proj",
+ "language_model.model.layers.39.self_attn.q_a_proj",
+ "language_model.model.layers.39.self_attn.q_b_proj",
+ "language_model.model.layers.4.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.4.self_attn.kv_b_proj",
+ "language_model.model.layers.4.self_attn.q_a_proj",
+ "language_model.model.layers.4.self_attn.q_b_proj",
+ "language_model.model.layers.40.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.40.self_attn.kv_b_proj",
+ "language_model.model.layers.40.self_attn.q_a_proj",
+ "language_model.model.layers.40.self_attn.q_b_proj",
+ "language_model.model.layers.41.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.41.self_attn.kv_b_proj",
+ "language_model.model.layers.41.self_attn.q_a_proj",
+ "language_model.model.layers.41.self_attn.q_b_proj",
+ "language_model.model.layers.42.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.42.self_attn.kv_b_proj",
+ "language_model.model.layers.42.self_attn.q_a_proj",
+ "language_model.model.layers.42.self_attn.q_b_proj",
+ "language_model.model.layers.43.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.43.self_attn.kv_b_proj",
+ "language_model.model.layers.43.self_attn.q_a_proj",
+ "language_model.model.layers.43.self_attn.q_b_proj",
+ "language_model.model.layers.44.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.44.self_attn.kv_b_proj",
+ "language_model.model.layers.44.self_attn.q_a_proj",
+ "language_model.model.layers.44.self_attn.q_b_proj",
+ "language_model.model.layers.45.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.45.self_attn.kv_b_proj",
+ "language_model.model.layers.45.self_attn.q_a_proj",
+ "language_model.model.layers.45.self_attn.q_b_proj",
+ "language_model.model.layers.46.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.46.self_attn.kv_b_proj",
+ "language_model.model.layers.46.self_attn.q_a_proj",
+ "language_model.model.layers.46.self_attn.q_b_proj",
+ "language_model.model.layers.47.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.47.self_attn.kv_b_proj",
+ "language_model.model.layers.47.self_attn.q_a_proj",
+ "language_model.model.layers.47.self_attn.q_b_proj",
+ "language_model.model.layers.48.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.48.self_attn.kv_b_proj",
+ "language_model.model.layers.48.self_attn.q_a_proj",
+ "language_model.model.layers.48.self_attn.q_b_proj",
+ "language_model.model.layers.49.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.49.self_attn.kv_b_proj",
+ "language_model.model.layers.49.self_attn.q_a_proj",
+ "language_model.model.layers.49.self_attn.q_b_proj",
+ "language_model.model.layers.5.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.5.self_attn.kv_b_proj",
+ "language_model.model.layers.5.self_attn.q_a_proj",
+ "language_model.model.layers.5.self_attn.q_b_proj",
+ "language_model.model.layers.50.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.50.self_attn.kv_b_proj",
+ "language_model.model.layers.50.self_attn.q_a_proj",
+ "language_model.model.layers.50.self_attn.q_b_proj",
+ "language_model.model.layers.51.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.51.self_attn.kv_b_proj",
+ "language_model.model.layers.51.self_attn.q_a_proj",
+ "language_model.model.layers.51.self_attn.q_b_proj",
+ "language_model.model.layers.52.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.52.self_attn.kv_b_proj",
+ "language_model.model.layers.52.self_attn.q_a_proj",
+ "language_model.model.layers.52.self_attn.q_b_proj",
+ "language_model.model.layers.53.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.53.self_attn.kv_b_proj",
+ "language_model.model.layers.53.self_attn.q_a_proj",
+ "language_model.model.layers.53.self_attn.q_b_proj",
+ "language_model.model.layers.54.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.54.self_attn.kv_b_proj",
+ "language_model.model.layers.54.self_attn.q_a_proj",
+ "language_model.model.layers.54.self_attn.q_b_proj",
+ "language_model.model.layers.55.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.55.self_attn.kv_b_proj",
+ "language_model.model.layers.55.self_attn.q_a_proj",
+ "language_model.model.layers.55.self_attn.q_b_proj",
+ "language_model.model.layers.56.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.56.self_attn.kv_b_proj",
+ "language_model.model.layers.56.self_attn.q_a_proj",
+ "language_model.model.layers.56.self_attn.q_b_proj",
+ "language_model.model.layers.57.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.57.self_attn.kv_b_proj",
+ "language_model.model.layers.57.self_attn.q_a_proj",
+ "language_model.model.layers.57.self_attn.q_b_proj",
+ "language_model.model.layers.58.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.58.self_attn.kv_b_proj",
+ "language_model.model.layers.58.self_attn.q_a_proj",
+ "language_model.model.layers.58.self_attn.q_b_proj",
+ "language_model.model.layers.59.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.59.self_attn.kv_b_proj",
+ "language_model.model.layers.59.self_attn.q_a_proj",
+ "language_model.model.layers.59.self_attn.q_b_proj",
+ "language_model.model.layers.6.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.6.self_attn.kv_b_proj",
+ "language_model.model.layers.6.self_attn.q_a_proj",
+ "language_model.model.layers.6.self_attn.q_b_proj",
+ "language_model.model.layers.60.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.60.self_attn.kv_b_proj",
+ "language_model.model.layers.60.self_attn.q_a_proj",
+ "language_model.model.layers.60.self_attn.q_b_proj",
+ "language_model.model.layers.7.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.7.self_attn.kv_b_proj",
+ "language_model.model.layers.7.self_attn.q_a_proj",
+ "language_model.model.layers.7.self_attn.q_b_proj",
+ "language_model.model.layers.8.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.8.self_attn.kv_b_proj",
+ "language_model.model.layers.8.self_attn.q_a_proj",
+ "language_model.model.layers.8.self_attn.q_b_proj",
+ "language_model.model.layers.9.self_attn.kv_a_proj_with_mqa",
+ "language_model.model.layers.9.self_attn.kv_b_proj",
+ "language_model.model.layers.9.self_attn.q_a_proj",
+ "language_model.model.layers.9.self_attn.q_b_proj",
+ "mm_projector*",
+ "vision_tower*"
+ ]
+ }
+}
\ No newline at end of file
diff --git a/kimi_k25_processor.py b/kimi_k25_processor.py
new file mode 100644
index 0000000000000000000000000000000000000000..d526032f91036de5f3d226b866acf449553b986d
--- /dev/null
+++ b/kimi_k25_processor.py
@@ -0,0 +1,165 @@
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.processing_utils import ProcessorMixin
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class KimiK25Processor(ProcessorMixin):
+ r"""
+ Constructs a KimiK25 processor which wraps a KimiK25 image processor and a tokenizer into a single processor.
+
+ [`KimiK25Processor`] offers all the functionalities of [`KimiK25ImageProcessor`] and [`TikTokenTokenizer`]. See the
+ [`~KimiK25Processor.__call__`] and [`~KimiK25Processor.decode`] for more information.
+
+ Args:
+ image_processor ([`KimiK25ImageProcessor`], *optional*):
+ The image processor is a required input.
+ tokenizer ([`TikTokenTokenizer`], *optional*):
+ The tokenizer is a required input.
+ chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+ in a chat into a tokenizable string.
+ """
+
+ 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
+ # A special temporal placeholder to be replaced by actual video placeholders
+ self.video_placeholder = "<|kimi_k25_video_placeholder|>"
+
+ def update_raw_text(self, text: str, video_prompts: list[str]) -> str:
+ # replace video prompt in text with video chunk prompts
+ video_count = text.count(self.video_placeholder)
+ if video_count == 0:
+ return text
+ assert video_count == len(video_prompts)
+ text_parts = text.split(self.video_placeholder)
+ assert len(text_parts) == len(video_prompts) + 1
+ text = "".join([
+ text_parts[i] + video_prompts[i] for i in range(len(video_prompts))
+ ])
+ text += text_parts[-1]
+ return text
+
+ def preprocess_medias(self, medias: list[dict]) -> list[dict]:
+ updated_medias = []
+ video_prompts = []
+ for media in medias:
+ if media['type'] == 'image':
+ updated_medias.append(media)
+ elif media['type'] == 'video':
+ video_chunks = self.media_processor.split_video_chunks(
+ media['video'])
+ updated_medias.extend(video_chunks)
+ video_prompts.append("".join(
+ [vc['prompt'] for vc in video_chunks]))
+ else:
+ raise ValueError(f"unsupported media type: {media['type']}")
+ return updated_medias, video_prompts
+
+ 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-K2.5 model.
+
+ This processor accepts ordered messages and extracts both media and text in a single pass.
+ text will be automatically updated if video input detected in messages
+
+ 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. If None, extracted from messages.
+ text: Pre-formatted text string. If None, generated via apply_chat_template.
+ return_tensors: Format of returned tensors ('pt', 'np', 'tf'). Default: 'pt'.
+ **kwargs: Additional arguments passed to tokenizer.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, video_prompts = self.preprocess_medias(medias)
+ preprocessed = self.media_processor.preprocess(
+ updated_medias, return_tensors=return_tensors)
+ text = self.update_raw_text(text, video_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, video_prompts = self.preprocess_medias(medias)
+ preprocessed = self.media_processor.preprocess(
+ updated_medias, return_tensors=return_tensors)
+
+ # Generate text if not provided
+ if text is None:
+ text = self.tokenizer.apply_chat_template(messages, **kwargs)
+
+ text = self.update_raw_text(text, video_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.
+
+ This is an optimized version that processes messages only once.
+ Kept as internal method since external callers should use __call__.
+ """
+ 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 ['video_url', 'video']:
+ medias.append({
+ 'type': 'video',
+ 'video': content_part['video_url']['url'],
+ 'first_frame_timestamp': 0.0
+ })
+ elif content_type in ['image_url', 'image']:
+ medias.append({
+ 'type': 'image',
+ 'image': content_part['image_url'],
+ })
+ 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_k25_vision_processing.py b/kimi_k25_vision_processing.py
new file mode 100644
index 0000000000000000000000000000000000000000..fdf3ab2f100f7c28a1f1e7295297e54b515d0b53
--- /dev/null
+++ b/kimi_k25_vision_processing.py
@@ -0,0 +1,251 @@
+"""Image processor class for Kimi-K2.5.
+"""
+
+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, VideoChunkInput, _to_tensor,
+ ensure_media_type, get_video_meta, image_to_np,
+ navit_patchify, navit_resize_image,
+ navit_resize_video, normalize,
+ real_sample_fps_and_max_num_frames, timestamp_as_str)
+
+try:
+ from mecord import VideoReader
+except ImportError:
+ VideoReader = None
+
+
+def resampling(video_bytes: bytes,
+ sample_indices: list[int],
+ key_indices=None,
+ frame_time_info=None,
+ num_threads=4) -> str:
+ video = VideoReader(video_bytes,
+ num_threads=num_threads,
+ frame_time_info=frame_time_info,
+ key_indices=key_indices)
+ # extract target frames
+ frames = video[sample_indices]
+ frames = [Image.fromarray(frame) for frame in frames]
+ return frames
+
+
+class KimiK25VisionProcessor(BaseImageProcessor):
+ model_type = "kimi_k25"
+
+ def __init__(
+ self,
+ media_proc_cfg: dict,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.media_proc_cfg = media_proc_cfg
+ self.num_frames_per_chunk = media_proc_cfg[
+ 'temporal_merge_kernel_size']
+
+ def media_tokens_calculator(self, media: MediaInput):
+ media = ensure_media_type(media)
+ ret = self.get_resize_config(media)
+ return ret['num_tokens']
+
+ @classmethod
+ def make_chunk_prompt(cls, timestamp_text: str) -> str:
+ return f"{timestamp_text}<|media_begin|>video<|media_content|><|media_pad|><|media_end|>"
+
+ def split_video_chunks(self,
+ video_url: str | bytes) -> list[list[Image.Image]]:
+ # video_url should be base64 str or bytes
+ video_spec = get_video_meta(video_url)
+ sample_fps = min(self.media_proc_cfg['sample_fps'], video_spec.fps)
+ sampled_nframes = max(
+ round(video_spec.num_frames * sample_fps / video_spec.fps), 1)
+ frame_inds = np.linspace(0, video_spec.num_frames - 1,
+ sampled_nframes).round().astype(int)
+ frame_inds = frame_inds.tolist()
+ sampled_frame_ids = []
+ temporal_merge_kernel_size = self.media_proc_cfg[
+ "temporal_merge_kernel_size"]
+ num_chunks = 0
+ chunk_timestamp = []
+ for i in range(0, len(frame_inds), temporal_merge_kernel_size):
+ sampled_frame_ids.extend(frame_inds[i:i +
+ temporal_merge_kernel_size])
+ start_time = frame_inds[i] / float(video_spec.fps)
+ timestamp_text = timestamp_as_str(
+ start_time, self.media_proc_cfg["timestamp_mode"])
+ chunk_timestamp.append(timestamp_text)
+ num_chunks += 1
+
+ sampled_frames = resampling(video_url, sampled_frame_ids)
+ chunks = []
+ for chunk_id in range(num_chunks):
+ chunk = sampled_frames[chunk_id *
+ temporal_merge_kernel_size:(chunk_id + 1) *
+ temporal_merge_kernel_size]
+ chunks.append(
+ VideoChunkInput(type="video_chunk",
+ video_chunk=chunk,
+ prompt=self.make_chunk_prompt(
+ chunk_timestamp[chunk_id])))
+ return chunks
+
+ 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
+ elif media_input['type'] == 'video_chunk':
+ frame = media_input['video_chunk'][0]
+ width, height = frame.size
+ num_frames = len(media_input["video_chunk"])
+ fps = 1.0
+
+ sample_fps, max_num_frames_each_video = real_sample_fps_and_max_num_frames(
+ media_input["type"],
+ self.media_proc_cfg['sample_fps'],
+ self.media_proc_cfg['max_num_frames_each_video'],
+ )
+
+ in_patch_limit_each_frame = self.media_proc_cfg[
+ 'in_patch_limit_each_frame']
+ if in_patch_limit_each_frame is None:
+ in_patch_limit_each_frame = self.media_proc_cfg[
+ 'in_patch_limit']
+
+ ret = navit_resize_video(
+ width,
+ height,
+ num_frames,
+ fps,
+ sample_fps,
+ self.media_proc_cfg['patch_size'],
+ self.media_proc_cfg['merge_kernel_size'],
+ in_patch_limit_each_frame,
+ self.media_proc_cfg['patch_limit_on_one_side'],
+ self.media_proc_cfg['in_patch_limit_video'],
+ max_num_frames_each_video,
+ 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")
+ 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/video_chunk) 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)
+ 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))
+ elif item['type'] == 'video_chunk':
+ pixels = []
+ for frame in item['video_chunk']:
+ frame_np = self.resize_image(frame, new_width,
+ new_height, pad_width,
+ pad_height)
+ pixels.append(frame_np)
+ pixel_values.append(np.stack(pixels, 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"KimiK25VisionProcessor(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..8795e06f381700d6420798f82174e3f9647e9f89
--- /dev/null
+++ b/media_utils.py
@@ -0,0 +1,368 @@
+import base64
+import io
+import math
+import os
+from datetime import datetime, timezone
+from typing import List, Literal, Optional, TypedDict
+
+import numpy as np
+from PIL import Image
+from pydantic import BaseModel, Field
+
+try:
+ from mecord import VideoReader
+except ImportError:
+ VideoReader = None
+
+
+class VideoSpec(BaseModel):
+ media_type: str = Literal['video']
+ height: int = Field(..., gt=0, description="video frame height")
+ width: int = Field(..., gt=0, description="video frame width")
+ num_frames: int = Field(..., gt=0, description="num frames")
+ fps: float = Field(..., gt=0, description="average fps")
+
+ # optional, help to accelerate video reading
+ key_indices: list[int] = Field(None, description="key indices")
+ frame_time_info: dict = Field(None, description="frame time info")
+
+
+class ImageInput(TypedDict):
+ type: Literal['image']
+ image: Image.Image
+
+
+class VideoChunkInput(TypedDict):
+ type: Literal['video_chunk']
+ video_chunk: List[Image.Image]
+ prompt: Optional[str] = None
+
+
+MediaInput = ImageInput | VideoChunkInput
+
+
+def get_video_meta(video_src: bytes | str | os.PathLike,
+ accurate: bool = True) -> dict:
+ """Get the dimensions of a video."""
+ if isinstance(video_src, os.PathLike):
+ video_src = str(video_src)
+ # if b64 string, decode to bytes
+ if isinstance(video_src,
+ str) and video_src.startswith('data:video/mp4;base64,'):
+ video_src = base64.b64decode(video_src.split(',')[1])
+ video = VideoReader(video_src, auto_init=accurate, num_threads=1)
+ assert video.num_frames > 0, "Invalid video format."
+ assert video.original_width > 0 and video.original_height > 0, (
+ "Invalid video format.")
+ assert video.avg_fps > 0, "Invalid video format."
+ return VideoSpec(media_type='video',
+ height=video.original_height,
+ width=video.original_width,
+ num_frames=video.num_frames,
+ fps=video.avg_fps,
+ key_indices=video.key_indices,
+ frame_time_info=video.frame_time_info)
+
+
+def timestamp_as_str(timestamp: float,
+ timestamp_mode: str = "hh:mm:ss.fff") -> str:
+ """Convert a timestamp to a string in the format of HH:MM:SS.mmm."""
+ if timestamp_mode == "hh:mm:ss.fff":
+ return (datetime.fromtimestamp(timestamp,
+ tz=timezone.utc).strftime("%H:%M:%S") +
+ f".{int((timestamp % 1) * 1000):03d}")
+ elif timestamp_mode == "mm:ss.fff":
+ return (datetime.fromtimestamp(timestamp,
+ tz=timezone.utc).strftime("%M:%S") +
+ f".{int((timestamp % 1) * 1000):03d}")
+ elif timestamp_mode == "mm:ss":
+ return datetime.fromtimestamp(timestamp,
+ tz=timezone.utc).strftime("%M:%S")
+ else:
+ raise ValueError(f"Invalid timestamp mode: {timestamp_mode}")
+
+
+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 navit_resize_video(
+ width: int,
+ height: int,
+ nframes: int,
+ avg_fps: float,
+ sample_fps: float,
+ patch_size: int,
+ merge_kernel_size: int,
+ in_patch_limit_each_frame: int,
+ patch_limit_on_one_side: int,
+ in_patch_limit_total: int | None,
+ max_num_frames_each_video: int | None,
+ fixed_output_tokens_each_frame: int | None,
+):
+ sample_fps = min(sample_fps, avg_fps)
+ # Calculate the number of frames to sample based on target FPS
+ sampled_nframes = max(round(nframes * sample_fps / avg_fps), 1)
+ if max_num_frames_each_video is not None:
+ sampled_nframes = min(sampled_nframes, max_num_frames_each_video)
+
+ if in_patch_limit_total is not None:
+ in_patch_limit_each_frame = min(
+ round(in_patch_limit_total / sampled_nframes),
+ in_patch_limit_each_frame)
+
+ ret = navit_resize_image(
+ width,
+ height,
+ patch_size,
+ merge_kernel_size,
+ in_patch_limit_each_frame,
+ patch_limit_on_one_side,
+ fixed_output_tokens_each_frame,
+ )
+ ret["sampled_nframes"] = sampled_nframes
+ return ret
+
+
+def real_sample_fps_and_max_num_frames(
+ type_name: Literal["video", "video_chunk"],
+ sample_fps: float,
+ max_num_frames_each_video: int | None,
+) -> tuple[int, int | None]:
+ if type_name == "video":
+ return sample_fps, max_num_frames_each_video
+ elif type_name == "video_chunk":
+ max_num_frames_each_video = None
+ sample_fps = math.inf
+ return sample_fps, max_num_frames_each_video
+ else:
+ return math.inf, None
+
+
+def _to_pil(data: str | bytes):
+ if isinstance(data, Image.Image):
+
+ return data.convert("RGB")
+ elif isinstance(data, str):
+ if data.startswith("data:"):
+ raw_base64 = data.split(",")[1]
+ return Image.open(io.BytesIO(
+ base64.b64decode(raw_base64))).convert("RGB")
+ else:
+ return Image.open(data).convert("RGB")
+ elif isinstance(data, bytes):
+ return Image.open(io.BytesIO(data)).convert("RGB")
+ else:
+ raise ValueError(f"Unsupported data type: {type(data)}")
+
+
+def ensure_media_type(media: MediaInput) -> MediaInput:
+ if media['type'] == 'image':
+ media['image'] = _to_pil(media['image'])
+ return media
+ elif media['type'] == 'video_chunk':
+ media['video_chunk'] = [
+ _to_pil(frame) for frame in media['video_chunk']
+ ]
+ 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,
+) -> 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.
+
+ 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)
+
+ 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)
+ 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)
+ 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-00118.safetensors b/model-00001-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..423dcff21307fd22d3e92e8291c998a3612cddb1
--- /dev/null
+++ b/model-00001-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f65c41dd357ea8e3fc400b639fe61faebe68106f90d16437869e3b6df803f93
+size 4996768944
diff --git a/model-00002-of-00118.safetensors b/model-00002-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..38780ed2c429d9f0e02d471f918c5ee3c183392c
--- /dev/null
+++ b/model-00002-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c06cea2f66f1baaf9141e7057609a82172e89acc707ec925d99918c2b319289e
+size 4996135920
diff --git a/model-00003-of-00118.safetensors b/model-00003-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..a370868ff390b268f5f7512b36c82d87a80dc289
--- /dev/null
+++ b/model-00003-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d3cc3ed2153110a3125e20b345591813e7cde2e8eba90a1d463945b95f57e4a
+size 4995637900
diff --git a/model-00004-of-00118.safetensors b/model-00004-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..c707a210c6040476c508aafb54aa818e19116e41
--- /dev/null
+++ b/model-00004-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f820b78ce72228737cc74134b8ed94f1ba0a2bb767fb587f828d1de672d2ac6e
+size 4996136088
diff --git a/model-00005-of-00118.safetensors b/model-00005-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..317ee738129c637fcd719e289c3baa17a9555a3f
--- /dev/null
+++ b/model-00005-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05c6e242542e6868b175ac5614cab172299e77309b6e3003b1d38353f309b44c
+size 4995634336
diff --git a/model-00006-of-00118.safetensors b/model-00006-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..dc2c78c88f176bc2da271c4048096effabece4d5
--- /dev/null
+++ b/model-00006-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e44f0358307af2e74935796bc9bf0eea43ab632759af5c74574424d2600d373
+size 4996124024
diff --git a/model-00007-of-00118.safetensors b/model-00007-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..14ffdd502f241ec84a4de34ac07874e8b57133e0
--- /dev/null
+++ b/model-00007-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07b007910f4144da8ea41c765995ad653f8e16a6139e6b440e8eec8747be1172
+size 4995612028
diff --git a/model-00008-of-00118.safetensors b/model-00008-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..8c756b14f66aea4e5fa364d9711d2b5eec8ecc69
--- /dev/null
+++ b/model-00008-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8923ad3049aaaa152350bc76ab8de20406ce3fcf608ba2aeb6038eaf2f3f5fc
+size 4996089948
diff --git a/model-00009-of-00118.safetensors b/model-00009-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..a4eb86b58611ef29c0c20359e4f466f9ad0722de
--- /dev/null
+++ b/model-00009-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e6a5e40f69e2a88a9e07e28319c652d9ee82395fc66a9b300ca2d41e667caf4
+size 4995586300
diff --git a/model-00010-of-00118.safetensors b/model-00010-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..d5013b75fb4f1041f3d2e4ba48c27b868a888526
--- /dev/null
+++ b/model-00010-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3157533fb9df05c5c6cc6ca564c55d06b120888e2a6525b96b3b585bee738676
+size 4996071708
diff --git a/model-00011-of-00118.safetensors b/model-00011-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..6803033a63a49f654fedd7ffd823c891da1fcdf0
--- /dev/null
+++ b/model-00011-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5f794b2a2027cd2de0969451bd6ade7c03a833b2f6f3f188124523b798a556b
+size 4995571416
diff --git a/model-00012-of-00118.safetensors b/model-00012-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..f73db21f1d81d5e1731a60cebb0cb22dd20fb2be
--- /dev/null
+++ b/model-00012-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a35041c0fa0f8438b6bfaf92568be94a8332828a803dfd774117cdd46c66629d
+size 4996061004
diff --git a/model-00013-of-00118.safetensors b/model-00013-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..2608c011b62803323aa51739b9d38337af8c9e07
--- /dev/null
+++ b/model-00013-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23605de5c9b7dd04c36175d102c4f5276c9aefe859fdafe161e52900d57b15c6
+size 4995566856
diff --git a/model-00014-of-00118.safetensors b/model-00014-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..2a787829bab48c8cf7e4525ad6f9d6c3545f2d21
--- /dev/null
+++ b/model-00014-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:489dd7848b7034086607928ffecbd94ffacf2a7fce2a8298581f25b9bb2d9cde
+size 4996061772
diff --git a/model-00015-of-00118.safetensors b/model-00015-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..1c0f1ce84ae3e1f7b95dc2176881eb11d7acadfe
--- /dev/null
+++ b/model-00015-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4bcfe5b23c0900e401864f594d0ea03f5ec970af7e005c740a01ff33fcc50fa3
+size 4995564940
diff --git a/model-00016-of-00118.safetensors b/model-00016-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..ea62a8497ace66220ca994985c25a2ccd1f9b32c
--- /dev/null
+++ b/model-00016-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb777f7c9c643c0570e9638cf64a5766cb9e36b2583faea51713b945381a4333
+size 4996061780
diff --git a/model-00017-of-00118.safetensors b/model-00017-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..17d30097b5b53fa88d92d9450db2774dcf5eddfd
--- /dev/null
+++ b/model-00017-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bcfce1cf1bc174910d59017854da649c1db1fb10a95131791ff81907755dcd7b
+size 4995565708
diff --git a/model-00018-of-00118.safetensors b/model-00018-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..0fb9ba132b59608150ad473fc269497412e999d3
--- /dev/null
+++ b/model-00018-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b6bb13b08c3880771ab4d1e90a4e69e6967609d33e3be200e4c608d5921ba44
+size 4996061004
diff --git a/model-00019-of-00118.safetensors b/model-00019-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..7f7828929e38591c9250e87ec159e7ae2591f0ca
--- /dev/null
+++ b/model-00019-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bcf52b39ef593c5bb0e860c7205f9f8ba74ac9d87c86c01a22094029757baeb5
+size 4995567656
diff --git a/model-00020-of-00118.safetensors b/model-00020-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..82a0734444da3dc733294d4d71aff38d53a8c229
--- /dev/null
+++ b/model-00020-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d40fba2e973b5d6c71c4a6a67845c5aa1037e70a6b3a586101446ac8529ae2e
+size 4996062448
diff --git a/model-00021-of-00118.safetensors b/model-00021-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..3bb3434eace3b92be2da8059bb3462177c511443
--- /dev/null
+++ b/model-00021-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c11d96f27acf8416dcbd33e9d9a73233c96006dcd8036afd089eddeca36246b7
+size 4995567128
diff --git a/model-00022-of-00118.safetensors b/model-00022-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..8ff1296647beff974b3539f1135f22fed81836e2
--- /dev/null
+++ b/model-00022-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91a0514306e0a53bb320ccc1ba51d7c257604bf9c33333247f879d219cfaf1aa
+size 4993341060
diff --git a/model-00023-of-00118.safetensors b/model-00023-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..a1541e93eb0d9570b016f8ec2dbf83faeb32be79
--- /dev/null
+++ b/model-00023-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bfe9c5468b2ad3cdcce7fb1484b028fb0be3199149d5d1b3dfd5ebf97807813b
+size 4998290456
diff --git a/model-00024-of-00118.safetensors b/model-00024-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..aff40784c45df82fb0221db5a81bc5d89af0d90b
--- /dev/null
+++ b/model-00024-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:707039bd2ef5a4e3ceffd7408a3db7414ae5fc3c63abd2a9f84b768813e4c4fb
+size 4995568348
diff --git a/model-00025-of-00118.safetensors b/model-00025-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..d46a97d1daeeeb6aff296cd8908e6e214bace63c
--- /dev/null
+++ b/model-00025-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a70d9503c8032a81a441083abf0e0a9ab2d7675fd5cf49d47dda351fda4744e9
+size 4996062004
diff --git a/model-00026-of-00118.safetensors b/model-00026-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..76598921b2aa02b4b88c8eface7792d67f5e6933
--- /dev/null
+++ b/model-00026-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d07ed1b644aef63a42b012f6a106ccb6ff8a8d7258a973b496ed1176e052ade0
+size 4995568592
diff --git a/model-00027-of-00118.safetensors b/model-00027-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..17b590fcf854eff4a0eed96f7f595bf4e6bce172
--- /dev/null
+++ b/model-00027-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0961da7c8dce68a865d4577710c800e9ea3417d6f3c7d3891162e70355496efa
+size 4996062132
diff --git a/model-00028-of-00118.safetensors b/model-00028-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..9cf48bd7b207136f6036772adc0a243fdec92239
--- /dev/null
+++ b/model-00028-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cdac5ef6c4267dc8536e313700098a7df5c59182a38aeeae44bc052c203b6f5
+size 4995568100
diff --git a/model-00029-of-00118.safetensors b/model-00029-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..21c0714395c814bb89fa96a6fdeb0ade4deb52bc
--- /dev/null
+++ b/model-00029-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17d53dc085cf53a22dd62bcbd7353c4496ae9851a5d45a68293bccf34fbcbf16
+size 4996062252
diff --git a/model-00030-of-00118.safetensors b/model-00030-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..aa3437111c7053c2d112c5fbe2955ec4c62ad1fc
--- /dev/null
+++ b/model-00030-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:888e2a1e96172455964c2a4297102852c3ff5bdd3dd8dbf0210c8fc8224f86ab
+size 4995567980
diff --git a/model-00031-of-00118.safetensors b/model-00031-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..8d84b58ffbaec77ddbcd2b46750f0123499b8414
--- /dev/null
+++ b/model-00031-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed052a52cd3e61f8c37fbc3ece0c373a5cd22b891478ff3986fcc3675e06aec8
+size 4996062752
diff --git a/model-00032-of-00118.safetensors b/model-00032-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..6747d46e9ed2be2b7a69bb5d962dd0ef37893c5d
--- /dev/null
+++ b/model-00032-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:946618ec542bc64a2c5878fe35a24e7a41f47ebe6e7d19a4b4d5d13126c267f3
+size 4995568232
diff --git a/model-00033-of-00118.safetensors b/model-00033-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..fa9377a734d3685ae8921f1360e13968c2d34193
--- /dev/null
+++ b/model-00033-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f53c031a4ffb23e00254c698d038d6de16f40a4ab80fa7e5410351523b673a6
+size 4996062104
diff --git a/model-00034-of-00118.safetensors b/model-00034-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..e6fb7b7286151fd862501a555c1ca61157c72a45
--- /dev/null
+++ b/model-00034-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed72929dd31ee6aed49694763d59ef61f2ebe446ee8912c1b340c09ea62c7949
+size 4995568888
diff --git a/model-00035-of-00118.safetensors b/model-00035-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..30fa054383fac6e47a477e509677b3176a459135
--- /dev/null
+++ b/model-00035-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0952f3f1e89591828ab8cae349abcb9b3bcdde06cd24d69aa97524e822f9743
+size 4996062224
diff --git a/model-00036-of-00118.safetensors b/model-00036-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..fc87fd82b75d204a831f18273d8b56a8ca481f86
--- /dev/null
+++ b/model-00036-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89856b169ebe5845a084e4249bf7c602ecc0c740404a6aaf07cc052f2a12ddb9
+size 4995566844
diff --git a/model-00037-of-00118.safetensors b/model-00037-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..4646d0e5687e7fcafbc6267b462d68e587345be5
--- /dev/null
+++ b/model-00037-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71a3bcfcb0e57b3199a944384310d97faf26bacd2e4023b7cb1c3a34c7df72d3
+size 4996063120
diff --git a/model-00038-of-00118.safetensors b/model-00038-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..c0d76342bf316a35ae5d80e24e3416d0bbf966dd
--- /dev/null
+++ b/model-00038-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b52539bcdd6ebd02396a73924fd3d2cf34c346b599ab80c009a0f3599cd7bc3
+size 4995567888
diff --git a/model-00039-of-00118.safetensors b/model-00039-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..c2b77fd62f17307c1779fd02be0eabd6bf4fc505
--- /dev/null
+++ b/model-00039-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f9bc3e9fa515e43d1c2e727e1c5a615edee9142a33090543812d25cb31c3317
+size 4996062836
diff --git a/model-00040-of-00118.safetensors b/model-00040-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..1f9a59a51e42dcf1460d8ae45cb23cae8585e86f
--- /dev/null
+++ b/model-00040-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57a7bd1fdda13d730c7bd51f392069b2e07a8403742b110f6ac478d43170ba70
+size 4995566368
diff --git a/model-00041-of-00118.safetensors b/model-00041-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..a5ad7fd809a62c1be668b8f2268ea6984861b094
--- /dev/null
+++ b/model-00041-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7b11d58d9e9520b2b176b0007f2d5a58689c7b7cc6a945b4b94222a139c5520
+size 4996063604
diff --git a/model-00042-of-00118.safetensors b/model-00042-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..2c060c89d681ee8e6c9b45897e3310fd704b0702
--- /dev/null
+++ b/model-00042-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:223fa7dbcd2c2ca12ce7d276e1aa657dffef79ca0561ec697718fa7ec8cb82f8
+size 4995567516
diff --git a/model-00043-of-00118.safetensors b/model-00043-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..7d9f2d28ab8e95ccde0097ea79def753f34b190b
--- /dev/null
+++ b/model-00043-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a484a2d7ff5e2f156560eaecc3e00ac474e53328976bc168d15ccffd09ae506f
+size 4996063612
diff --git a/model-00044-of-00118.safetensors b/model-00044-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..151b0b3ec13835e6f44cf13341e8f66f5db31e3b
--- /dev/null
+++ b/model-00044-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9105bc9ad8c7a9a6e2d2bacc167a06894788387e66e3d5390fd2058511bf8cd2
+size 4995568648
diff --git a/model-00045-of-00118.safetensors b/model-00045-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..9f1a335918c747f5f43de67c0054b28a0c81672a
--- /dev/null
+++ b/model-00045-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd7f4db13360f03012fa3a9932d491060d20e1a2df65e52ed4b90b4e393bbfc4
+size 4996062068
diff --git a/model-00046-of-00118.safetensors b/model-00046-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..c0e4b86022393e72781b942ca38526c1624905b9
--- /dev/null
+++ b/model-00046-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f13c517cc5385a2936dcfa1cf5ce7dc7049e98a50b3719441517cd679fd6979f
+size 4995566748
diff --git a/model-00047-of-00118.safetensors b/model-00047-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..82e48e6ac2ba5293f72935515999e94fe9e0b834
--- /dev/null
+++ b/model-00047-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:793273b59cdd96aeab6507fa2a518cc6db5f01b235dfaab6fabf20c6f52e1caf
+size 4996063604
diff --git a/model-00048-of-00118.safetensors b/model-00048-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..c5d5a868e7fe0220c2bbc84953d0f2696babccfc
--- /dev/null
+++ b/model-00048-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c93fb24a1e31f84bee833e53c3a5a679dd1fd9155c6a2eedf656b0d97bbd58b6
+size 4995567508
diff --git a/model-00049-of-00118.safetensors b/model-00049-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..bb5fbc5206807ead8d68f60d642668359c6e0cea
--- /dev/null
+++ b/model-00049-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27dc0d9240df8198a86cba5b6ceec8bf4add41105ddd64b9dcf17b58dffb9bca
+size 4996063224
diff --git a/model-00050-of-00118.safetensors b/model-00050-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..ec1cb73b3962f72813080f5d51f66c25166a2c97
--- /dev/null
+++ b/model-00050-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c47e3ce388cc3b1fd75e0be2601aa4f9160fd21ceb93f9124198288fd6dcb472
+size 4995567516
diff --git a/model-00051-of-00118.safetensors b/model-00051-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..5e7324ff8cdaffbbee76fe7cd46280dadef577a7
--- /dev/null
+++ b/model-00051-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:580ffe418a4df8d2cade1f7583ef2ab42c31a94414803429a7021f31f3f0b8e8
+size 4996062836
diff --git a/model-00052-of-00118.safetensors b/model-00052-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..5086298dafcbbb04434cad224fd9810c5faac1a4
--- /dev/null
+++ b/model-00052-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2c508b285106c964e640d0016725aad452a650c68f9c2d146daa879a2d26946
+size 4995567508
diff --git a/model-00053-of-00118.safetensors b/model-00053-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..09bdd8abe635a6f5ac1a78bea722ef1e3d4c2c06
--- /dev/null
+++ b/model-00053-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97689c9d30b52fa824540b41ed7378584098ec67b4149f38ba9040ce52d72d9c
+size 4995568380
diff --git a/model-00054-of-00118.safetensors b/model-00054-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..e04f05e98b127edfdcea05357d09a8b817466959
--- /dev/null
+++ b/model-00054-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e5a032d4825315b729658490c3904bd0c608912e0b11203c935602e23723a79
+size 4996063508
diff --git a/model-00055-of-00118.safetensors b/model-00055-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..084d149403a7673b755e018e041877c818e4610d
--- /dev/null
+++ b/model-00055-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75cb1e54579a9f26d62ced74c2b40834956fedef812cfe0dfd7c7ecd3526bedf
+size 4995567080
diff --git a/model-00056-of-00118.safetensors b/model-00056-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..4ca9e536d41e8cb6ffd024ba2664d4ee658d7fdd
--- /dev/null
+++ b/model-00056-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfcb94d30e31721ce3202e8c68e4732804b4d7ebad766dc0d68ddd69d0d98226
+size 4996062884
diff --git a/model-00057-of-00118.safetensors b/model-00057-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..26f127f1e97df5e95112bbd0fa7471201bd08b7a
--- /dev/null
+++ b/model-00057-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15e2bfe98803119aae1a74513a0f77981bed12bdc3534e1a168dbb809ee8a50b
+size 4995568108
diff --git a/model-00058-of-00118.safetensors b/model-00058-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..224d14dfe34122702716e9ce592df721c328dc48
--- /dev/null
+++ b/model-00058-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c12612ffb78ea828791ad13d3db6ebf3ecc3fe72523f96d0b0de37f23b43220
+size 4996062236
diff --git a/model-00059-of-00118.safetensors b/model-00059-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..f68301f31a71554fd3be913b903eae0bfb9d1c61
--- /dev/null
+++ b/model-00059-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cea519049a8311c07719b8d3ab1f1c0dc203a13bd196fda0169b8a13f69a50f
+size 4995568376
diff --git a/model-00060-of-00118.safetensors b/model-00060-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..ebebbec8c4c2f6a88363103e44ec8dc797290ce1
--- /dev/null
+++ b/model-00060-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e181576217ecfca660ef7c2abfb12c61e864729e34b064bced99bff6863cbc9
+size 4996061976
diff --git a/model-00061-of-00118.safetensors b/model-00061-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..07c98a8e44931c7c78704a36c46117bc00535f3e
--- /dev/null
+++ b/model-00061-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f35b8cb945600e22eb3028af28ae449e907f7de2615f03a4526b31164f59d23
+size 4995567876
diff --git a/model-00062-of-00118.safetensors b/model-00062-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..873267a26f757b8a630009c34c01ffc8287ad0ff
--- /dev/null
+++ b/model-00062-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62bf83a7168505164cad66040db8bfdcbf37b173454ae4fd3728b534952010d5
+size 4996062088
diff --git a/model-00063-of-00118.safetensors b/model-00063-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..3e89f4172189a0adf990deadff8816d3620b4f89
--- /dev/null
+++ b/model-00063-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f785339f2f600eb4fb949fe1670fbf8818e17b779175d07c73d115f214c74fa0
+size 4995569276
diff --git a/model-00064-of-00118.safetensors b/model-00064-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..85bbb9fd2b8ea04dede7325a6347ff93a7a0a723
--- /dev/null
+++ b/model-00064-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aac2ee8ae0aee7f08da5b5d7d9905647f57bea7b00cdc9448364a13583f0c47e
+size 4996062216
diff --git a/model-00065-of-00118.safetensors b/model-00065-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..a142710aa76cafeafb30d2b2e6ae0bb3d988ec78
--- /dev/null
+++ b/model-00065-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24c8bac09fbd96e7ca58596b3209f18efc5f18eec62ac04fa7bcf6e1def45380
+size 4995566860
diff --git a/model-00066-of-00118.safetensors b/model-00066-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..cfb8b6d853b5a4791317e12005eea8ca5dd15482
--- /dev/null
+++ b/model-00066-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1691b15a8e325c4e4e8fce6c06b9f5c96758c8b84ad66e3d666bb335a7ec1fa
+size 4996064260
diff --git a/model-00067-of-00118.safetensors b/model-00067-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..59d292a3d84a6fbe4d034818327df88929d0c921
--- /dev/null
+++ b/model-00067-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c52dabc8b13a067e49203e9ead7f679104eccd716bfd29ae61a6f72097796c0
+size 4995567508
diff --git a/model-00068-of-00118.safetensors b/model-00068-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..1e827eb8285186f6c87b2c9dcfc65417032207ed
--- /dev/null
+++ b/model-00068-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27d9947e5781cbbb5134b82149a051cb94e6fa02ab99da711ab7e7ae1be25178
+size 4996062836
diff --git a/model-00069-of-00118.safetensors b/model-00069-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..3aee6df17c7f07fcec9bd840943a97e0ca6c99df
--- /dev/null
+++ b/model-00069-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:68d69e3c2fe1af9db695e4d4ed9fbdef7d70d3d734ab2eda374e434c61496e1b
+size 4995566360
diff --git a/model-00070-of-00118.safetensors b/model-00070-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..62280c2003c9ded834066ae266061fa377a8f6d0
--- /dev/null
+++ b/model-00070-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1c23682f3ba7da436d395e16e7733f9105700437ff6cae337d08a320aeaf6f5
+size 4996063612
diff --git a/model-00071-of-00118.safetensors b/model-00071-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..1b6e7149b2d182f4ae037420cc70c1917db3eb5b
--- /dev/null
+++ b/model-00071-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e08ae34cbbec0f1f859d348be12453cad0bf728bd3b134617fce3e7f3cda0566
+size 4995566748
diff --git a/model-00072-of-00118.safetensors b/model-00072-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..7a35eccd9d4a29058d8a2a8278b39d6ad680f80f
--- /dev/null
+++ b/model-00072-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28ec4b1e70dee559ad45bc1272f191b5e8bf6555740c6ca156309c34ec105e0a
+size 4996063992
diff --git a/model-00073-of-00118.safetensors b/model-00073-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..62d1014e3e9823bac056d3c9593fa98a638895fd
--- /dev/null
+++ b/model-00073-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db5a7f9c68a98666c9194b12f5bd4fd05773cd041494b6dea50ee616ebe02952
+size 4995567516
diff --git a/model-00074-of-00118.safetensors b/model-00074-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..49f17d5fc4aa5c42111d40c2969a5bd8c49cb08e
--- /dev/null
+++ b/model-00074-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c03a88480661852f36b8b7b3cd9cf37642d553e8a5665ed869f571c9bd40ec2
+size 4996062836
diff --git a/model-00075-of-00118.safetensors b/model-00075-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..513d8e1c010c1bd6339717619836d4ee870b7700
--- /dev/null
+++ b/model-00075-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ba794aa37eadedd6ba6640d3e9e8daa9673f45dedb5eca11463e0bcb0c8c5b4
+size 4995566748
diff --git a/model-00076-of-00118.safetensors b/model-00076-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..7834fc2ed78a676624fbfbf327fde7221421b43d
--- /dev/null
+++ b/model-00076-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2fb1be356177c666578adda97872dd13267f8a63d17d6632a2b13bd5801c8526
+size 4996063612
diff --git a/model-00077-of-00118.safetensors b/model-00077-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..a5a817ed6faea4b4a8a180b551744277fe0ceff8
--- /dev/null
+++ b/model-00077-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61568e19921ee548a3c6a30a8af9872fba5f2d1cb8342ccd18183fa189b46a2e
+size 4995567888
diff --git a/model-00078-of-00118.safetensors b/model-00078-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..33a700575f3ff539351693453baa9ae9d507da5e
--- /dev/null
+++ b/model-00078-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7dbabba3784c7e088064c42b55ca74a58676f9eb42808708c26a30886b236547
+size 4996063216
diff --git a/model-00079-of-00118.safetensors b/model-00079-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..56379ea34f456133952229ec06edb1131f082659
--- /dev/null
+++ b/model-00079-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb1505da786e1edb49c96ba6d2cb185a3025dd3af4196adf78f744e5b4947398
+size 4995568284
diff --git a/model-00080-of-00118.safetensors b/model-00080-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..97c1bd7e5c347fde66d457b83e34b51f4e25d8dd
--- /dev/null
+++ b/model-00080-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a7e82ecd6500ecdac9ea5c41881945cc742b72f3b1dedce1d4a6b94b05f7689
+size 4996062068
diff --git a/model-00081-of-00118.safetensors b/model-00081-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..d98e8b920e9fa2c82e344cc319fc94972a698c91
--- /dev/null
+++ b/model-00081-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be836109110a1e82c28ba3294a1caf123ee8f7983399907b164c2680b0c2f641
+size 4995568276
diff --git a/model-00082-of-00118.safetensors b/model-00082-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..8eff7652b8342ff88860590afe73f8be9884d914
--- /dev/null
+++ b/model-00082-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d087a325e13ef0a1f31c4336f1c1d35d28445f9ed5995de0d36307ff88de442b
+size 4995568032
diff --git a/model-00083-of-00118.safetensors b/model-00083-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..f30c80543fc034b42f595035bff42b32917be62d
--- /dev/null
+++ b/model-00083-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3dd7e8cb1fd61028dc39f9f6448d1727e85370ef2c218129c11fd02bb5f880f3
+size 4996061932
diff --git a/model-00084-of-00118.safetensors b/model-00084-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..63489ae8219ad372e74a2f704a49803b65af65a3
--- /dev/null
+++ b/model-00084-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5378ca2273bc3884cdf4273cf49dba1c0a9b729e5a26d72c3d92bb291c285aed
+size 4995568244
diff --git a/model-00085-of-00118.safetensors b/model-00085-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..2d2a0d53ea7d3c417aabd75ce7cf53319ae1a170
--- /dev/null
+++ b/model-00085-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08c8a424a6091ec40ecb00df8f4ef0b64d851a0d85c5c38f1477bf461049458a
+size 4996062488
diff --git a/model-00086-of-00118.safetensors b/model-00086-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..b952ff3d20e316bf7c31698418593498fbcb0d64
--- /dev/null
+++ b/model-00086-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bd81e4dc0693027955f51739a4e904d08fb54659887db253b9a64c250c0bf5d
+size 4995568124
diff --git a/model-00087-of-00118.safetensors b/model-00087-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..26a020ecab64d65dbaa8253226f268395c9f8291
--- /dev/null
+++ b/model-00087-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e78c02a0121435b94289efb0d4977b22a2e009c214b79ae0a09f80f13194e33a
+size 4996062980
diff --git a/model-00088-of-00118.safetensors b/model-00088-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..e18f2758b360ea1d43d772e892e00c68f7cbce96
--- /dev/null
+++ b/model-00088-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85ae98cd2dcaa0e8f6e12dc1887fb6e7c46ae093d156f21310fb979ea4c18091
+size 4995568392
diff --git a/model-00089-of-00118.safetensors b/model-00089-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..28ed30f8ee9beed7e1736a035dfd67f840534d29
--- /dev/null
+++ b/model-00089-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1faf76f707c171d0d596078c76c7367dc47dbbf3aec20a54eaa509e4f499e504
+size 4996062340
diff --git a/model-00090-of-00118.safetensors b/model-00090-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..452b68e633626198ca08c67ea425c393ba3c6876
--- /dev/null
+++ b/model-00090-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0a0dd423703e406cb5e9ca322066b3eff501b20a14b375d79f67acf9a8dfd32
+size 4995567116
diff --git a/model-00091-of-00118.safetensors b/model-00091-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..264b39330ef36ceca7fb655a34c556d9d9402569
--- /dev/null
+++ b/model-00091-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2fbf35cbbc4e57191e920a125732d2b86f47f2fd28209a67987e546f9db1227
+size 4996063996
diff --git a/model-00092-of-00118.safetensors b/model-00092-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..5fca675d5be54f431cefb4ce631647751e5e29d1
--- /dev/null
+++ b/model-00092-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c879442d691a34a42192312c597f977f461defa6770b0190152855f0fd42913d
+size 4995566996
diff --git a/model-00093-of-00118.safetensors b/model-00093-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..c8ea33c0c86fd5c258dbecc620781a2f237dedbb
--- /dev/null
+++ b/model-00093-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8c1de082077ce8aa3668c588d09972073965a2300fcbc86c70794bf32efb15a
+size 4996062580
diff --git a/model-00094-of-00118.safetensors b/model-00094-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..f30273951c12774264960bda9ca99527a3e633df
--- /dev/null
+++ b/model-00094-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d736915a38ef7140d345f06f81317b2fca375958e5c3e2fa1565197cf69dcf95
+size 4995568032
diff --git a/model-00095-of-00118.safetensors b/model-00095-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..79ba1beacd1bde274cdc7a2925d646b8baed79b4
--- /dev/null
+++ b/model-00095-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2833c7c311c0a3770eac5fc28749d654864e744f740ac279ceb5ad1402e11cdb
+size 4996063088
diff --git a/model-00096-of-00118.safetensors b/model-00096-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..1884166bfca28e1fd8f5baf3eb9c1ab4a6e0a48a
--- /dev/null
+++ b/model-00096-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:440e878d70e4f167884627c5778d52a4429b5458b25a2420e7fd740ff27caf92
+size 4995567904
diff --git a/model-00097-of-00118.safetensors b/model-00097-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..bdad310cb5560903a6200bf67191de49013b326f
--- /dev/null
+++ b/model-00097-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fc6a5c84ff97a98523b664337194f22bd89e584b7887a2a18774a0a45708f3f
+size 4996063596
diff --git a/model-00098-of-00118.safetensors b/model-00098-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..158f8660269a903ff0a45b34dc92459e9e77df88
--- /dev/null
+++ b/model-00098-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33a688362cee8e64748c385dd820596b234db4c9002fb94754d57d04e3284504
+size 4995566352
diff --git a/model-00099-of-00118.safetensors b/model-00099-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..8a50c6311ccd906e6adf0605b21ccbf03a371ad6
--- /dev/null
+++ b/model-00099-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31cb81aee26984029fedcda08e3719a57bc9545edc5ad21c54472356769e4c57
+size 4996062836
diff --git a/model-00100-of-00118.safetensors b/model-00100-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..30e92298024b691b761ecdbfba779a8592762917
--- /dev/null
+++ b/model-00100-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec3c4bcbe5d31978d5f21abd5d4bdec25d9c1e28213aac6e214cdf8277bdb570
+size 4995568672
diff --git a/model-00101-of-00118.safetensors b/model-00101-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..b4a47037580db049678cf72fcfdf18d245a2f683
--- /dev/null
+++ b/model-00101-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23881a96b23a5235f57b3809ba9f26129e9414f1e19e33b1bb01ad7bc94270e4
+size 4996062448
diff --git a/model-00102-of-00118.safetensors b/model-00102-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..28fd00c55b7dd8509c0d1e4447e583b2e184570a
--- /dev/null
+++ b/model-00102-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c16834f0fcedacc5edffaded9b32ae1a0acf9193e2eaae1fbe1ebbc8d1103ec
+size 4995567888
diff --git a/model-00103-of-00118.safetensors b/model-00103-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..1eef8e6eb300b30136b2f6e051aad23b829cc0d2
--- /dev/null
+++ b/model-00103-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:561f29ccfe0243acd27659ce9db225e11e2117a37da49c29a9f4d1c32323049a
+size 4996062844
diff --git a/model-00104-of-00118.safetensors b/model-00104-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..84cc67d5f605a7bf46f370ba6283dd442ff59087
--- /dev/null
+++ b/model-00104-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f6b575c0ab4a22704e28ec74eed4511da2bf6ba434f17e9da62f659a766045e
+size 4995567508
diff --git a/model-00105-of-00118.safetensors b/model-00105-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..e5e5bbcb90c779b783fb3515ed7b3623b6e98b80
--- /dev/null
+++ b/model-00105-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac7bd18f9d859c17fac675dc01195177e330613c61637fe62cc858f6e8f41155
+size 4996062448
diff --git a/model-00106-of-00118.safetensors b/model-00106-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..a8a6fd308b6c0d5192045049d19add602249dc2c
--- /dev/null
+++ b/model-00106-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec38ea0f8fbd7ee82d90fe380f547d868e53b67d9cbf64006d48c63f02b38d3a
+size 4995566880
diff --git a/model-00107-of-00118.safetensors b/model-00107-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..5977f347aa790429a58075a3acb5008b33e9e9ec
--- /dev/null
+++ b/model-00107-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0bf21b93061d1c5fb76538351f35c618431c700796ef8918cccb59a63a995a67
+size 4996063472
diff --git a/model-00108-of-00118.safetensors b/model-00108-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..7436cc17541100d22f9e9eb8e8cea3e5075b2148
--- /dev/null
+++ b/model-00108-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd5ccc08421516a97226703407281006ccf0368433e41da72c793b1731455b5c
+size 4995568276
diff --git a/model-00109-of-00118.safetensors b/model-00109-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..b6d9bc6aa5529168dc9d7de58038ec29d486e74b
--- /dev/null
+++ b/model-00109-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0358a30491aab4c7b588b68312043b09bee0b07af624ba7486f394a120e4adc5
+size 4996062456
diff --git a/model-00110-of-00118.safetensors b/model-00110-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..9dde015e4d85362b847d090a309b38c4a5e27b67
--- /dev/null
+++ b/model-00110-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93a6f27833440db39f66925cacbe66eae289aa6c633d872bad83a2b9f6560969
+size 4995566740
diff --git a/model-00111-of-00118.safetensors b/model-00111-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..4ce4d2e74cb9ca7cba7311d8b5d821fb288c6709
--- /dev/null
+++ b/model-00111-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55b9d6d1d2394413b443d14652a29bf615b5f7dcb2cab32ca8ac072a11e469c5
+size 4995569212
diff --git a/model-00112-of-00118.safetensors b/model-00112-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..f2250b570a3d8aafc0635a12b2bc18e391af9d33
--- /dev/null
+++ b/model-00112-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc738d94740ce1f128ddbdb80c958739ce9657f49caacd114416014aa2b1aaf1
+size 4996061528
diff --git a/model-00113-of-00118.safetensors b/model-00113-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..94970f1d1fade4f53c3c0c5a0e447de986e616c1
--- /dev/null
+++ b/model-00113-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ee04b8bae60a34728894ab5d85955f03b3f3a20d701e88758e5e4067dd0b3cb
+size 4995568648
diff --git a/model-00114-of-00118.safetensors b/model-00114-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..39ff5334216c5a4a3b1753b94463c22df98d80d5
--- /dev/null
+++ b/model-00114-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f518f672bd28b6eb1c73306dfa71faae9f7ab418223b3fa67b5e8200b0e5e30b
+size 4996062472
diff --git a/model-00115-of-00118.safetensors b/model-00115-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..ae8a09512c08c58c62fc8b50da78a48d86a39eb5
--- /dev/null
+++ b/model-00115-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d405625a68210277243bda4033bd43316fbca10fc106cb27758ed9a86907ba72
+size 4995567760
diff --git a/model-00116-of-00118.safetensors b/model-00116-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..d41be4a2044853b9b64e2822a1215830f217a892
--- /dev/null
+++ b/model-00116-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44182673b5f728b1c5ba2b46bb715f8dffa74e2e888e40de9d0a3d1547ad945f
+size 4996062972
diff --git a/model-00117-of-00118.safetensors b/model-00117-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..e4e18c0dc04d97d04266c1b9c08a5f5f03e4fcc8
--- /dev/null
+++ b/model-00117-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e4094027f275ac717d116b1f2910505fae634786f8d73bccb7b46bba90c3fca
+size 3795950584
diff --git a/model-00118-of-00118.safetensors b/model-00118-of-00118.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..19c2b42e79d6c7329cdf8a3833b3d2a3d3aff685
--- /dev/null
+++ b/model-00118-of-00118.safetensors
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f608d1f0a9b8468d5f6d420246b7342e331ea20e03aeb3e5d41ad24d74a626fb
+size 2348810384
diff --git a/model.safetensors.index.json b/model.safetensors.index.json
new file mode 100644
index 0000000000000000000000000000000000000000..46f2671c658a1f1c722fe9cb9e1effec9af82711
--- /dev/null
+++ b/model.safetensors.index.json
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b2d068df729805c486e7c735885b682294fa4abe4e69c78b06a0de380417c30
+size 23723724
diff --git a/modeling_deepseek.py b/modeling_deepseek.py
new file mode 100644
index 0000000000000000000000000000000000000000..d67cf312e8987d7aa70bf51ba296d4ffd3c092dd
--- /dev/null
+++ b/modeling_deepseek.py
@@ -0,0 +1,1825 @@
+# coding=utf-8
+# Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+""" PyTorch DeepSeek model."""
+import os
+import sys
+import math
+import importlib.util
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import \
+ _prepare_4d_causal_attention_mask
+from transformers.modeling_outputs import (BaseModelOutputWithPast,
+ CausalLMOutputWithPast,
+ SequenceClassifierOutputWithPast)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import (ALL_LAYERNORM_LAYERS,
+ is_torch_greater_or_equal_than_1_13)
+from transformers.utils import (add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ is_flash_attn_2_available,
+ is_flash_attn_greater_or_equal_2_10, logging,
+ replace_return_docstrings)
+from transformers.utils.import_utils import is_torch_fx_available
+
+# Dynamic import to handle both direct loading and HF cache loading
+_current_dir = os.path.dirname(os.path.abspath(__file__))
+_MODEL_DIR = "/workspace/Model-Optimizer/examples/llm_ptq/kimi25_hf-bf16"
+
+if "configuration_deepseek" in sys.modules:
+ DeepseekV3Config = sys.modules["configuration_deepseek"].DeepseekV3Config
+else:
+ _config_path = os.path.join(_current_dir, "configuration_deepseek.py")
+ if not os.path.exists(_config_path):
+ _config_path = os.path.join(_MODEL_DIR, "configuration_deepseek.py")
+ _spec = importlib.util.spec_from_file_location("configuration_deepseek", _config_path)
+ _config_module = importlib.util.module_from_spec(_spec)
+ sys.modules["configuration_deepseek"] = _config_module
+ _spec.loader.exec_module(_config_module)
+ DeepseekV3Config = _config_module.DeepseekV3Config
+
+if is_flash_attn_2_available():
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
+ from flash_attn.bert_padding import pad_input # noqa
+ from flash_attn.bert_padding import index_first_axis, unpad_input
+
+# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
+# It means that the function will not be traced through and simply appear as a node in the graph.
+if is_torch_fx_available():
+ if not is_torch_greater_or_equal_than_1_13:
+ import torch.fx
+
+ _prepare_4d_causal_attention_mask = torch.fx.wrap(
+ _prepare_4d_causal_attention_mask)
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "DeepseekV3Config"
+
+
+def _get_unpad_data(attention_mask):
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
+ cu_seqlens = F.pad(
+ torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+ return (
+ indices,
+ cu_seqlens,
+ max_seqlen_in_batch,
+ )
+
+
+# code modified from transformers 4.48.3 to amend breaks in newer transformers versions
+def get_usable_length(past_key_value,
+ new_seq_length: int,
+ layer_idx: Optional[int] = 0) -> int:
+ max_length = past_key_value.get_max_cache_shape()
+ previous_seq_length = past_key_value.get_seq_length(layer_idx)
+ if max_length is not None and max_length > 0 and previous_seq_length + new_seq_length > max_length:
+ return max_length - new_seq_length
+ return previous_seq_length
+
+
+class DeepseekV3RMSNorm(nn.Module):
+
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ DeepseekV3RMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states):
+ input_dtype = hidden_states.dtype
+ hidden_states = hidden_states.to(torch.float32)
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance +
+ self.variance_epsilon)
+ return self.weight * hidden_states.to(input_dtype)
+
+
+ALL_LAYERNORM_LAYERS.append(DeepseekV3RMSNorm)
+
+
+class DeepseekV3RotaryEmbedding(nn.Module):
+
+ def __init__(self,
+ dim,
+ max_position_embeddings=2048,
+ base=10000,
+ device=None):
+ super().__init__()
+
+ self.dim = dim
+ self.max_position_embeddings = max_position_embeddings
+ self.base = base
+ inv_freq = 1.0 / (self.base**(
+ torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ # Build here to make `torch.jit.trace` work.
+ self._set_cos_sin_cache(
+ seq_len=max_position_embeddings,
+ device=self.inv_freq.device,
+ dtype=torch.get_default_dtype(),
+ )
+ self.max_seq_len_cached = None
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+ t = torch.arange(self.max_seq_len_cached,
+ device=device,
+ dtype=self.inv_freq.dtype)
+
+ freqs = torch.outer(t, self.inv_freq.to(t.device))
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer("cos_cached",
+ emb.cos().to(dtype),
+ persistent=False)
+ self.register_buffer("sin_cached",
+ emb.sin().to(dtype),
+ persistent=False)
+
+ def forward(self, x, seq_len=None):
+ # x: [bs, num_attention_heads, seq_len, head_size]
+ if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached:
+ self._set_cos_sin_cache(seq_len=seq_len,
+ device=x.device,
+ dtype=x.dtype)
+
+ return (
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
+ )
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->DeepseekV3
+class DeepseekV3LinearScalingRotaryEmbedding(DeepseekV3RotaryEmbedding):
+ """DeepseekV3RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
+
+ def __init__(
+ self,
+ dim,
+ max_position_embeddings=2048,
+ base=10000,
+ device=None,
+ scaling_factor=1.0,
+ ):
+ self.scaling_factor = scaling_factor
+ super().__init__(dim, max_position_embeddings, base, device)
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+ t = torch.arange(self.max_seq_len_cached,
+ device=device,
+ dtype=self.inv_freq.dtype)
+ t = t / self.scaling_factor
+
+ freqs = torch.outer(t, self.inv_freq)
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer("cos_cached",
+ emb.cos().to(dtype),
+ persistent=False)
+ self.register_buffer("sin_cached",
+ emb.sin().to(dtype),
+ persistent=False)
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->DeepseekV3
+class DeepseekV3DynamicNTKScalingRotaryEmbedding(DeepseekV3RotaryEmbedding):
+ """DeepseekV3RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
+
+ def __init__(
+ self,
+ dim,
+ max_position_embeddings=2048,
+ base=10000,
+ device=None,
+ scaling_factor=1.0,
+ ):
+ self.scaling_factor = scaling_factor
+ super().__init__(dim, max_position_embeddings, base, device)
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+
+ if seq_len > self.max_position_embeddings:
+ base = self.base * ((self.scaling_factor * seq_len /
+ self.max_position_embeddings) -
+ (self.scaling_factor - 1))**(self.dim /
+ (self.dim - 2))
+ inv_freq = 1.0 / (base**(
+ torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ t = torch.arange(self.max_seq_len_cached,
+ device=device,
+ dtype=self.inv_freq.dtype)
+
+ freqs = torch.outer(t, self.inv_freq)
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer("cos_cached",
+ emb.cos().to(dtype),
+ persistent=False)
+ self.register_buffer("sin_cached",
+ emb.sin().to(dtype),
+ persistent=False)
+
+
+# Inverse dim formula to find dim based on number of rotations
+def yarn_find_correction_dim(num_rotations,
+ dim,
+ base=10000,
+ max_position_embeddings=2048):
+ return (dim * math.log(max_position_embeddings /
+ (num_rotations * 2 * math.pi))) / (2 *
+ math.log(base))
+
+
+# Find dim range bounds based on rotations
+def yarn_find_correction_range(low_rot,
+ high_rot,
+ dim,
+ base=10000,
+ max_position_embeddings=2048):
+ low = math.floor(
+ yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings))
+ high = math.ceil(
+ yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings))
+ return max(low, 0), min(high, dim - 1) # Clamp values just in case
+
+
+def yarn_get_mscale(scale=1, mscale=1):
+ if scale <= 1:
+ return 1.0
+ return 0.1 * mscale * math.log(scale) + 1.0
+
+
+def yarn_linear_ramp_mask(min, max, dim):
+ if min == max:
+ max += 0.001 # Prevent singularity
+
+ linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
+ ramp_func = torch.clamp(linear_func, 0, 1)
+ return ramp_func
+
+
+class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding):
+
+ def __init__(
+ self,
+ dim,
+ max_position_embeddings=2048,
+ base=10000,
+ device=None,
+ scaling_factor=1.0,
+ original_max_position_embeddings=4096,
+ beta_fast=32,
+ beta_slow=1,
+ mscale=1,
+ mscale_all_dim=0,
+ ):
+ self.scaling_factor = scaling_factor
+ self.original_max_position_embeddings = original_max_position_embeddings
+ self.beta_fast = beta_fast
+ self.beta_slow = beta_slow
+ self.mscale = mscale
+ self.mscale_all_dim = mscale_all_dim
+ super().__init__(dim, max_position_embeddings, base, device)
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+ dim = self.dim
+
+ freq_extra = 1.0 / (self.base**(
+ torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
+ freq_inter = 1.0 / (self.scaling_factor * self.base**(
+ torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
+
+ low, high = yarn_find_correction_range(
+ self.beta_fast,
+ self.beta_slow,
+ dim,
+ self.base,
+ self.original_max_position_embeddings,
+ )
+ inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to(
+ device=device, dtype=torch.float32)
+ inv_freq = freq_inter * (1 -
+ inv_freq_mask) + freq_extra * inv_freq_mask
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ t = torch.arange(seq_len, device=device, dtype=torch.float32)
+
+ freqs = torch.outer(t, inv_freq)
+
+ _mscale = float(
+ yarn_get_mscale(self.scaling_factor, self.mscale) /
+ yarn_get_mscale(self.scaling_factor, self.mscale_all_dim))
+
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer("cos_cached", (emb.cos() * _mscale).to(dtype),
+ persistent=False)
+ self.register_buffer("sin_cached", (emb.sin() * _mscale).to(dtype),
+ persistent=False)
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., :x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2:]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ position_ids (`torch.Tensor`):
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
+ used to pass offsetted position ids when working with a KV-cache.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+
+ b, h, s, d = q.shape
+ q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+
+ b, h, s, d = k.shape
+ k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class DeepseekV3MLP(nn.Module):
+
+ def __init__(self, config, 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)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(
+ self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class MoEGate(nn.Module):
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.top_k = config.num_experts_per_tok
+ self.n_routed_experts = config.n_routed_experts
+ self.routed_scaling_factor = config.routed_scaling_factor
+ self.scoring_func = config.scoring_func
+ self.seq_aux = config.seq_aux
+ self.topk_method = config.topk_method
+ self.n_group = config.n_group
+ self.topk_group = config.topk_group
+
+ # topk selection algorithm
+ self.norm_topk_prob = config.norm_topk_prob
+ self.gating_dim = config.hidden_size
+ self.weight = nn.Parameter(
+ torch.empty((self.n_routed_experts, self.gating_dim)))
+ if self.topk_method == "noaux_tc":
+ self.e_score_correction_bias = nn.Parameter(
+ torch.empty((self.n_routed_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.scoring_func == "sigmoid":
+ scores = logits.sigmoid()
+ else:
+ raise NotImplementedError(
+ f"insupportable scoring function for MoE gating: {self.scoring_func}"
+ )
+
+ ### select top-k experts
+ if self.topk_method == "noaux_tc":
+ assert not self.training
+ scores_for_choice = scores.view(
+ bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0)
+ group_scores = (scores_for_choice.view(
+ bsz * seq_len, self.n_group,
+ -1).topk(2, dim=-1)[0].sum(dim=-1)) # [n, n_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, n_group]
+ group_mask.scatter_(1, group_idx, 1) # [n, n_group]
+ score_mask = (group_mask.unsqueeze(-1).expand(
+ bsz * seq_len, self.n_group,
+ self.n_routed_experts // self.n_group).reshape(
+ bsz * seq_len, -1)) # [n, e]
+ tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(),
+ 0.0) # [n, e]
+ _, topk_idx = torch.topk(tmp_scores,
+ k=self.top_k,
+ dim=-1,
+ sorted=False)
+ topk_weight = scores.gather(1, topk_idx)
+ else:
+ raise NotImplementedError(
+ f"insupportable TopK function for MoE gating: {self.topk_method}"
+ )
+
+ ### norm gate to sum 1
+ if self.top_k > 1 and self.norm_topk_prob:
+ denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
+ topk_weight = topk_weight / denominator
+ topk_weight = topk_weight * self.routed_scaling_factor # must multiply the scaling factor
+
+ return topk_idx, topk_weight
+
+
+class DeepseekV3MoE(nn.Module):
+ """
+ A mixed expert module containing shared experts.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.num_experts_per_tok = config.num_experts_per_tok
+
+ if hasattr(config, "ep_size") and config.ep_size > 1:
+ assert config.ep_size == dist.get_world_size()
+ self.ep_size = config.ep_size
+ self.experts_per_rank = config.n_routed_experts // config.ep_size
+ self.ep_rank = dist.get_rank()
+ self.experts = nn.ModuleList([
+ (DeepseekV3MLP(config,
+ intermediate_size=config.moe_intermediate_size)
+ if i >= self.ep_rank * self.experts_per_rank
+ and i < (self.ep_rank + 1) * self.experts_per_rank else None)
+ for i in range(config.n_routed_experts)
+ ])
+ else:
+ self.ep_size = 1
+ self.experts_per_rank = config.n_routed_experts
+ self.ep_rank = 0
+ self.experts = nn.ModuleList([
+ DeepseekV3MLP(config,
+ intermediate_size=config.moe_intermediate_size)
+ for i in range(config.n_routed_experts)
+ ])
+ self.gate = MoEGate(config)
+ if config.n_shared_experts is not None:
+ intermediate_size = config.moe_intermediate_size * config.n_shared_experts
+ self.shared_experts = DeepseekV3MLP(
+ config=config, intermediate_size=intermediate_size)
+
+ 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])
+ flat_topk_idx = topk_idx.view(-1)
+ if not self.training:
+ y = self.moe_infer(hidden_states, topk_idx,
+ topk_weight).view(*orig_shape)
+ if self.config.n_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]]
+ sorted_tokens_shape = sorted_tokens.shape
+ if self.ep_size > 1:
+ tokens_per_ep_rank = tokens_per_expert.view(self.ep_size,
+ -1).sum(dim=1)
+ tokens_per_expert_group = tokens_per_expert.new_empty(
+ tokens_per_expert.shape[0])
+ dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert)
+ output_splits = (tokens_per_expert_group.view(
+ self.ep_size, -1).sum(1).cpu().numpy().tolist())
+ gathered_tokens = sorted_tokens.new_empty(
+ tokens_per_expert_group.sum(dim=0).cpu().item(),
+ sorted_tokens.shape[1])
+ input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist()
+ dist.all_to_all(
+ list(gathered_tokens.split(output_splits)),
+ list(sorted_tokens.split(input_split_sizes)),
+ )
+ tokens_per_expert_post_gather = tokens_per_expert_group.view(
+ self.ep_size, self.experts_per_rank).sum(dim=0)
+ gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0], ),
+ dtype=np.int32)
+ s = 0
+ for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
+ gatherd_idxs[s:s + k] = i % self.experts_per_rank
+ s += k
+ gatherd_idxs = gatherd_idxs.argsort()
+ sorted_tokens = gathered_tokens[gatherd_idxs]
+ tokens_per_expert = tokens_per_expert_post_gather
+ 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)
+ if self.ep_size > 1:
+ new_x = torch.empty_like(outs)
+ new_x[gatherd_idxs] = outs
+ gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
+ dist.all_to_all(
+ list(gathered_tokens.split(input_split_sizes)),
+ list(new_x.split(output_splits)),
+ )
+ outs = gathered_tokens
+
+ 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
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :,
+ None, :, :].expand(batch,
+ num_key_value_heads,
+ n_rep, slen, head_dim)
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen,
+ head_dim)
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->DeepseekV3
+class DeepseekV3Attention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self,
+ config: DeepseekV3Config,
+ layer_idx: Optional[int] = None):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ if layer_idx is None:
+ logger.warning_once(
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class.")
+
+ self.attention_dropout = config.attention_dropout
+ self.hidden_size = config.hidden_size
+ self.num_heads = config.num_attention_heads
+
+ self.max_position_embeddings = config.max_position_embeddings
+ self.rope_theta = config.rope_theta
+ 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 = config.qk_nope_head_dim + config.qk_rope_head_dim
+
+ self.is_causal = True
+
+ if self.q_lora_rank is None:
+ self.q_proj = nn.Linear(self.hidden_size,
+ self.num_heads * self.q_head_dim,
+ bias=False)
+ else:
+ self.q_a_proj = nn.Linear(self.hidden_size,
+ config.q_lora_rank,
+ bias=config.attention_bias)
+ self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank)
+ self.q_b_proj = nn.Linear(config.q_lora_rank,
+ self.num_heads * self.q_head_dim,
+ bias=False)
+
+ self.kv_a_proj_with_mqa = nn.Linear(
+ self.hidden_size,
+ config.kv_lora_rank + config.qk_rope_head_dim,
+ bias=config.attention_bias,
+ )
+ self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank)
+ self.kv_b_proj = nn.Linear(
+ config.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=config.attention_bias,
+ )
+ self._init_rope()
+
+ self.softmax_scale = self.q_head_dim**(-0.5)
+ if self.config.rope_scaling is not None:
+ mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
+ scaling_factor = self.config.rope_scaling["factor"]
+ if mscale_all_dim:
+ mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
+ self.softmax_scale = self.softmax_scale * mscale * mscale
+
+ def _init_rope(self):
+ if self.config.rope_scaling is None:
+ self.rotary_emb = DeepseekV3RotaryEmbedding(
+ self.qk_rope_head_dim,
+ max_position_embeddings=self.max_position_embeddings,
+ base=self.rope_theta,
+ )
+ else:
+ scaling_type = self.config.rope_scaling["type"]
+ scaling_factor = self.config.rope_scaling["factor"]
+ if scaling_type == "linear":
+ self.rotary_emb = DeepseekV3LinearScalingRotaryEmbedding(
+ self.qk_rope_head_dim,
+ max_position_embeddings=self.max_position_embeddings,
+ scaling_factor=scaling_factor,
+ base=self.rope_theta,
+ )
+ elif scaling_type == "dynamic":
+ self.rotary_emb = DeepseekV3DynamicNTKScalingRotaryEmbedding(
+ self.qk_rope_head_dim,
+ max_position_embeddings=self.max_position_embeddings,
+ scaling_factor=scaling_factor,
+ base=self.rope_theta,
+ )
+ elif scaling_type == "yarn":
+ kwargs = {
+ key: self.config.rope_scaling[key]
+ for key in [
+ "original_max_position_embeddings",
+ "beta_fast",
+ "beta_slow",
+ "mscale",
+ "mscale_all_dim",
+ ] if key in self.config.rope_scaling
+ }
+ self.rotary_emb = DeepseekV3YarnRotaryEmbedding(
+ self.qk_rope_head_dim,
+ max_position_embeddings=self.max_position_embeddings,
+ scaling_factor=scaling_factor,
+ base=self.rope_theta,
+ **kwargs,
+ )
+ else:
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+ return (tensor.view(bsz, seq_len, self.num_heads,
+ self.v_head_dim).transpose(1, 2).contiguous())
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: bool = False,
+ use_cache: bool = False,
+ **kwargs,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
+ Optional[Tuple[torch.Tensor]]]:
+ if "padding_mask" in kwargs:
+ warnings.warn(
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+ )
+ bsz, q_len, _ = hidden_states.size()
+
+ if self.q_lora_rank is None:
+ q = self.q_proj(hidden_states)
+ else:
+ q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
+ q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
+ q_nope, q_pe = torch.split(
+ q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+ compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
+ compressed_kv, k_pe = torch.split(
+ compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+ k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
+ kv = (self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(
+ bsz, q_len, self.num_heads,
+ self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2))
+
+ k_nope, value_states = torch.split(
+ kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+ kv_seq_len = value_states.shape[-2]
+ if past_key_value is not None:
+ if self.layer_idx is None:
+ raise ValueError(
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+ "with a layer index.")
+ kv_seq_len += get_usable_length(past_key_value, kv_seq_len,
+ self.layer_idx)
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+ q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)
+
+ query_states = k_pe.new_empty(bsz, self.num_heads, q_len,
+ self.q_head_dim)
+ query_states[:, :, :, :self.qk_nope_head_dim] = q_nope
+ query_states[:, :, :, self.qk_nope_head_dim:] = q_pe
+
+ key_states = k_pe.new_empty(bsz, self.num_heads, q_len,
+ self.q_head_dim)
+ key_states[:, :, :, :self.qk_nope_head_dim] = k_nope
+ key_states[:, :, :, self.qk_nope_head_dim:] = k_pe
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
+ key_states, value_states = past_key_value.update(
+ key_states, value_states, self.layer_idx, cache_kwargs)
+
+ attn_weights = (
+ torch.matmul(query_states, key_states.transpose(2, 3)) *
+ self.softmax_scale)
+
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+ raise ValueError(
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
+ f" {attn_weights.size()}")
+ assert attention_mask is not None
+ if attention_mask is not None:
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+ raise ValueError(
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+ )
+ attn_weights = attn_weights + attention_mask
+
+ # upcast attention to fp32
+ attn_weights = nn.functional.softmax(attn_weights,
+ dim=-1,
+ dtype=torch.float32).to(
+ query_states.dtype)
+ attn_weights = nn.functional.dropout(attn_weights,
+ p=self.attention_dropout,
+ training=self.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim):
+ raise ValueError(
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is"
+ f" {attn_output.size()}")
+
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ attn_output = attn_output.reshape(bsz, q_len,
+ self.num_heads * self.v_head_dim)
+
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->DeepseekV3
+class DeepseekV3FlashAttention2(DeepseekV3Attention):
+ """
+ DeepseekV3 flash attention module. This module inherits from `DeepseekV3Attention` as the weights of the module stays
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+ flash attention and deal with padding tokens in case the input contains any of them.
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10(
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.LongTensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: bool = False,
+ use_cache: bool = False,
+ **kwargs,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
+ Optional[Tuple[torch.Tensor]]]:
+ # DeepseekV3FlashAttention2 attention does not support output_attentions
+ if "padding_mask" in kwargs:
+ warnings.warn(
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+ )
+
+ # overwrite attention_mask with padding_mask
+ attention_mask = kwargs.pop("padding_mask")
+
+ output_attentions = False
+
+ bsz, q_len, _ = hidden_states.size()
+
+ if self.q_lora_rank is None:
+ q = self.q_proj(hidden_states)
+ else:
+ q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
+ q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
+ q_nope, q_pe = torch.split(
+ q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+ # Flash attention requires the input to have the shape
+ # batch_size x seq_length x head_dim x hidden_dim
+ # therefore we just need to keep the original shape
+ compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
+ compressed_kv, k_pe = torch.split(
+ compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+ k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)
+ kv = (self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(
+ bsz, q_len, self.num_heads,
+ self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2))
+
+ k_nope, value_states = torch.split(
+ kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+ kv_seq_len = value_states.shape[-2]
+
+ kv_seq_len = value_states.shape[-2]
+ if past_key_value is not None:
+ kv_seq_len += get_usable_length(past_key_value, kv_seq_len,
+ self.layer_idx)
+
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+ q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)
+
+ query_states = k_pe.new_empty(bsz, self.num_heads, q_len,
+ self.q_head_dim)
+ query_states[:, :, :, :self.qk_nope_head_dim] = q_nope
+ query_states[:, :, :, self.qk_nope_head_dim:] = q_pe
+
+ key_states = k_pe.new_empty(bsz, self.num_heads, q_len,
+ self.q_head_dim)
+ key_states[:, :, :, :self.qk_nope_head_dim] = k_nope
+ key_states[:, :, :, self.qk_nope_head_dim:] = k_pe
+
+ if self.q_head_dim != self.v_head_dim:
+ value_states = F.pad(value_states,
+ [0, self.q_head_dim - self.v_head_dim])
+
+ if past_key_value is not None:
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
+ key_states, value_states = past_key_value.update(
+ key_states, value_states, self.layer_idx, cache_kwargs)
+
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+ # to be able to avoid many of these transpose/reshape/view.
+ query_states = query_states.transpose(1, 2)
+ key_states = key_states.transpose(1, 2)
+ value_states = value_states.transpose(1, 2)
+
+ dropout_rate = self.attention_dropout if self.training else 0.0
+
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
+ # cast them back in the correct dtype just to be sure everything works as expected.
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+ # in fp32. (DeepseekV3RMSNorm handles it correctly)
+
+ input_dtype = query_states.dtype
+ if input_dtype == torch.float32:
+ # Handle the case where the model is quantized
+ if hasattr(self.config, "_pre_quantization_dtype"):
+ target_dtype = self.config._pre_quantization_dtype
+ elif torch.is_autocast_enabled():
+ target_dtype = torch.get_autocast_gpu_dtype()
+ else:
+ target_dtype = (self.q_proj.weight.dtype if self.q_lora_rank
+ is None else self.q_a_proj.weight.dtype)
+
+ logger.warning_once(
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+ f" {target_dtype}.")
+
+ query_states = query_states.to(target_dtype)
+ key_states = key_states.to(target_dtype)
+ value_states = value_states.to(target_dtype)
+
+ attn_output = self._flash_attention_forward(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ q_len,
+ dropout=dropout_rate,
+ softmax_scale=self.softmax_scale,
+ )
+ if self.q_head_dim != self.v_head_dim:
+ attn_output = attn_output[:, :, :, :self.v_head_dim]
+
+ attn_output = attn_output.reshape(bsz, q_len, self.num_heads *
+ self.v_head_dim).contiguous()
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+ def _flash_attention_forward(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ query_length,
+ dropout=0.0,
+ softmax_scale=None,
+ ):
+ """
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+ first unpad the input, then computes the attention scores and pad the final attention scores.
+
+ Args:
+ query_states (`torch.Tensor`):
+ Input query states to be passed to Flash Attention API
+ key_states (`torch.Tensor`):
+ Input key states to be passed to Flash Attention API
+ value_states (`torch.Tensor`):
+ Input value states to be passed to Flash Attention API
+ attention_mask (`torch.Tensor`):
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+ position of padding tokens and 1 for the position of non-padding tokens.
+ dropout (`int`, *optional*):
+ Attention dropout
+ softmax_scale (`float`, *optional*):
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+ """
+ if not self._flash_attn_uses_top_left_mask:
+ causal = self.is_causal
+ else:
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in DeepseekV3FlashAttention2 __init__.
+ causal = self.is_causal and query_length != 1
+
+ # Contains at least one padding token in the sequence
+ if attention_mask is not None:
+ batch_size = query_states.shape[0]
+ (
+ query_states,
+ key_states,
+ value_states,
+ indices_q,
+ cu_seq_lens,
+ max_seq_lens,
+ ) = self._upad_input(query_states, key_states, value_states,
+ attention_mask, query_length)
+
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+ attn_output_unpad = flash_attn_varlen_func(
+ query_states,
+ key_states,
+ value_states,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_k=cu_seqlens_k,
+ max_seqlen_q=max_seqlen_in_batch_q,
+ max_seqlen_k=max_seqlen_in_batch_k,
+ dropout_p=dropout,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ )
+
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size,
+ query_length)
+ else:
+ attn_output = flash_attn_func(
+ query_states,
+ key_states,
+ value_states,
+ dropout,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ )
+
+ return attn_output
+
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask,
+ query_length):
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(
+ attention_mask)
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+ key_layer = index_first_axis(
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,
+ head_dim),
+ indices_k,
+ )
+ value_layer = index_first_axis(
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,
+ head_dim),
+ indices_k,
+ )
+ if query_length == kv_seq_len:
+ query_layer = index_first_axis(
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads,
+ head_dim),
+ indices_k,
+ )
+ cu_seqlens_q = cu_seqlens_k
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
+ indices_q = indices_k
+ elif query_length == 1:
+ max_seqlen_in_batch_q = 1
+ cu_seqlens_q = torch.arange(
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
+ ) # There is a memcpy here, that is very bad.
+ indices_q = cu_seqlens_q[:-1]
+ query_layer = query_layer.squeeze(1)
+ else:
+ # The -q_len: slice assumes left padding.
+ attention_mask = attention_mask[:, -query_length:]
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
+ query_layer, attention_mask)
+
+ return (
+ query_layer,
+ key_layer,
+ value_layer,
+ indices_q,
+ (cu_seqlens_q, cu_seqlens_k),
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+ )
+
+
+ATTENTION_CLASSES = {
+ "eager": DeepseekV3Attention,
+ "flash_attention_2": DeepseekV3FlashAttention2,
+}
+
+
+class DeepseekV3DecoderLayer(nn.Module):
+
+ def __init__(self, config: DeepseekV3Config, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = ATTENTION_CLASSES[config._attn_implementation](
+ config=config, layer_idx=layer_idx)
+
+ self.mlp = (DeepseekV3MoE(config) if
+ (config.n_routed_experts is not None
+ and layer_idx >= config.first_k_dense_replace
+ and layer_idx % config.moe_layer_freq == 0) else
+ DeepseekV3MLP(config))
+ self.input_layernorm = DeepseekV3RMSNorm(config.hidden_size,
+ eps=config.rms_norm_eps)
+ self.post_attention_layernorm = DeepseekV3RMSNorm(
+ config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
+ output_attentions: Optional[bool] = False,
+ use_cache: Optional[bool] = False,
+ **kwargs,
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor,
+ torch.FloatTensor]]]:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`, *optional*):
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+ query_sequence_length, key_sequence_length)` if default attention is used.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+ returned tensors for more detail.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+ (see `past_key_values`).
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+ """
+ if "padding_mask" in kwargs:
+ warnings.warn(
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
+ )
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_value,
+ 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)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ outputs = (hidden_states, )
+
+ if output_attentions:
+ outputs += (self_attn_weights, )
+
+ if use_cache:
+ outputs += (present_key_value, )
+
+ return outputs
+
+
+DeepseekV3_START_DOCSTRING = r"""
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+ etc.)
+
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+ and behavior.
+
+ Parameters:
+ config ([`DeepseekV3Config`]):
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
+ load the weights associated with the model, only the configuration. Check out the
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+@add_start_docstrings(
+ "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.",
+ DeepseekV3_START_DOCSTRING,
+)
+class DeepseekV3PreTrainedModel(PreTrainedModel):
+ config_class = DeepseekV3Config
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["DeepseekV3DecoderLayer"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn_2 = True
+ _supports_cache_class = 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_()
+
+
+DeepseekV3_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+ it.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+
+ [What are attention masks?](../glossary#attention-mask)
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
+ `past_key_values`).
+
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+ information on the default strategy.
+
+ - 1 indicates the head is **not masked**,
+ - 0 indicates the head is **masked**.
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+ config.n_positions - 1]`.
+
+ [What are position IDs?](../glossary#position-ids)
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+ Two formats are allowed:
+ - a [`~cache_utils.Cache`] instance;
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+ cache format.
+
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+ legacy cache format will be returned.
+
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+ of shape `(batch_size, sequence_length)`.
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+ model's internal embedding lookup matrix.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+ `past_key_values`).
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+ tensors for more detail.
+ output_hidden_states (`bool`, *optional*):
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+ more detail.
+ return_dict (`bool`, *optional*):
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+ "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.",
+ DeepseekV3_START_DOCSTRING,
+)
+class DeepseekV3Model(DeepseekV3PreTrainedModel):
+ """
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeepseekV3DecoderLayer`]
+
+ Args:
+ config: DeepseekV3Config
+ """
+
+ def __init__(self, config: DeepseekV3Config):
+ 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([
+ DeepseekV3DecoderLayer(config, layer_idx)
+ for layer_idx in range(config.num_hidden_layers)
+ ])
+ self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+ self.norm = DeepseekV3RMSNorm(config.hidden_size,
+ eps=config.rms_norm_eps)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.embed_tokens = value
+
+ @add_start_docstrings_to_model_forward(DeepseekV3_INPUTS_DOCSTRING)
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
+ 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)
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+ return_dict = (return_dict if return_dict is not None else
+ self.config.use_return_dict)
+
+ # retrieve input_ids and inputs_embeds
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError(
+ "You cannot specify both input_ids and inputs_embeds at the same time"
+ )
+ elif input_ids is not None:
+ batch_size, seq_length = input_ids.shape[:2]
+ elif inputs_embeds is not None:
+ batch_size, seq_length = inputs_embeds.shape[:2]
+ else:
+ raise ValueError(
+ "You have to specify either input_ids or inputs_embeds")
+
+ past_key_values_length = 0
+ if use_cache:
+ use_legacy_cache = not isinstance(past_key_values, Cache)
+ if use_legacy_cache:
+ past_key_values = DynamicCache.from_legacy_cache(
+ past_key_values)
+ past_key_values_length = get_usable_length(past_key_values,
+ seq_length)
+
+ if position_ids is None:
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ position_ids = torch.arange(
+ past_key_values_length,
+ seq_length + past_key_values_length,
+ dtype=torch.long,
+ device=device,
+ )
+ position_ids = position_ids.unsqueeze(0)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if self._use_flash_attention_2:
+ # 2d mask is passed through the layers
+ attention_mask = (attention_mask if
+ (attention_mask is not None
+ and 0 in attention_mask) else None)
+ else:
+ # 4d mask is passed through the layers
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask,
+ (batch_size, seq_length),
+ inputs_embeds,
+ past_key_values_length,
+ )
+
+ # embed positions
+ hidden_states = inputs_embeds
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+ next_decoder_cache = None
+
+ for decoder_layer in self.layers:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states, )
+
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_values,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ if use_cache:
+ next_decoder_cache = layer_outputs[
+ 2 if output_attentions else 1]
+
+ if output_attentions:
+ all_self_attns += (layer_outputs[1], )
+
+ hidden_states = self.norm(hidden_states)
+
+ # add hidden states from the last decoder layer
+ if output_hidden_states:
+ all_hidden_states += (hidden_states, )
+
+ next_cache = None
+ if use_cache:
+ next_cache = (next_decoder_cache.to_legacy_cache()
+ if use_legacy_cache else next_decoder_cache)
+ if not return_dict:
+ return tuple(
+ v for v in
+ [hidden_states, next_cache, all_hidden_states, all_self_attns]
+ if v is not None)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+
+
+class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel):
+ _tied_weights_keys = ["lm_head.weight"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = DeepseekV3Model(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()
+
+ def get_input_embeddings(self):
+ return self.model.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.embed_tokens = value
+
+ def get_output_embeddings(self):
+ return self.lm_head
+
+ def set_output_embeddings(self, new_embeddings):
+ self.lm_head = new_embeddings
+
+ def set_decoder(self, decoder):
+ self.model = decoder
+
+ def get_decoder(self):
+ return self.model
+
+ @add_start_docstrings_to_model_forward(DeepseekV3_INPUTS_DOCSTRING)
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast,
+ config_class=_CONFIG_FOR_DOC)
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[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, transformers.,
+ 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, transformers., config.vocab_size]`.
+
+ Returns:
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM
+
+ >>> model = DeepseekV3ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ 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)
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ 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,
+ )
+
+ hidden_states = outputs[0]
+ logits = self.lm_head(hidden_states)
+ logits = logits.float()
+
+ loss = None
+ if labels is not None:
+ # Shift so that tokens < n predict n
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = labels[..., 1:].contiguous()
+ # Flatten the tokens
+ loss_fct = CrossEntropyLoss()
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
+ shift_labels = shift_labels.view(-1)
+ # Enable model parallelism
+ shift_labels = shift_labels.to(shift_logits.device)
+ loss = loss_fct(shift_logits, shift_labels)
+
+ if not return_dict:
+ output = (logits, ) + outputs[1:]
+ return (loss, ) + output if loss is not None else output
+
+ return CausalLMOutputWithPast(
+ 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,
+ attention_mask=None,
+ inputs_embeds=None,
+ **kwargs,
+ ):
+ if past_key_values is not None:
+ if isinstance(past_key_values, Cache):
+ cache_length = past_key_values.get_seq_length()
+ # seen_tokens 可能在某些 transformers 版本中不存在，使用 getattr 安全访问
+ past_length = getattr(past_key_values, 'seen_tokens',
+ cache_length)
+ max_cache_length = past_key_values.get_max_length()
+ else:
+ cache_length = past_length = past_key_values[0][0].shape[2]
+ max_cache_length = None
+
+ # 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.
+
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+ if (max_cache_length is not None and attention_mask is not None
+ and cache_length + input_ids.shape[1] > max_cache_length):
+ attention_mask = attention_mask[:, -max_cache_length:]
+
+ 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,
+ })
+ return model_inputs
+
+ @staticmethod
+ def _reorder_cache(past_key_values, beam_idx):
+ reordered_past = ()
+ for layer_past in past_key_values:
+ reordered_past += (tuple(
+ past_state.index_select(0, beam_idx.to(past_state.device))
+ for past_state in layer_past), )
+ return reordered_past
+
+
+@add_start_docstrings(
+ """
+ The DeepseekV3 Model transformer with a sequence classification head on top (linear layer).
+
+ [`DeepseekV3ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+ (e.g. GPT-2) do.
+
+ Since it does classification on the last token, it requires to know the position of the last token. If a
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+ each row of the batch).
+ """,
+ DeepseekV3_START_DOCSTRING,
+)
+class DeepseekV3ForSequenceClassification(DeepseekV3PreTrainedModel):
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.model = DeepseekV3Model(config)
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.embed_tokens = value
+
+ @add_start_docstrings_to_model_forward(DeepseekV3_INPUTS_DOCSTRING)
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.LongTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, transformers.,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ return_dict = (return_dict if return_dict is not None else
+ self.config.use_return_dict)
+
+ transformer_outputs = self.model(
+ 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,
+ )
+ hidden_states = transformer_outputs[0]
+ logits = self.score(hidden_states)
+
+ if input_ids is not None:
+ batch_size = input_ids.shape[0]
+ else:
+ batch_size = inputs_embeds.shape[0]
+
+ if self.config.pad_token_id is None and batch_size != 1:
+ raise ValueError(
+ "Cannot handle batch sizes > 1 if no padding token is defined."
+ )
+ if self.config.pad_token_id is None:
+ sequence_lengths = -1
+ else:
+ if input_ids is not None:
+ sequence_lengths = (torch.eq(
+ input_ids, self.config.pad_token_id).int().argmax(-1) -
+ 1).to(logits.device)
+ else:
+ sequence_lengths = -1
+
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device),
+ sequence_lengths]
+
+ loss = None
+ if labels is not None:
+ labels = labels.to(logits.device)
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long
+ or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(pooled_logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels),
+ labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(pooled_logits, labels)
+ if not return_dict:
+ output = (pooled_logits, ) + transformer_outputs[1:]
+ return ((loss, ) + output) if loss is not None else output
+
+ return SequenceClassifierOutputWithPast(
+ loss=loss,
+ logits=pooled_logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
diff --git a/modeling_kimi_k25.py b/modeling_kimi_k25.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d20eaecfaf361b41db086c48f8f572ccce36054
--- /dev/null
+++ b/modeling_kimi_k25.py
@@ -0,0 +1,1297 @@
+# coding=utf-8
+# Copyright 2025-2026 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved.
+#
+# The code is based on llava (llava/modeling_llava.py) and DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py), but modified for Kimi-K2.5.
+#
+# Licensing Information:
+# - Code derived from llava (llava/modeling_llava.py) and 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 MIT License.
+#
+# 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.
+#
+# MIT License:
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", 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 THE
+# AUTHORS 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.
+import os
+import sys
+import math
+import importlib.util
+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.cache_utils import Cache
+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
+
+# Dynamic imports to avoid transformers static import checker
+# Transformers copies only files in auto_map to cache, so we need to try multiple locations
+_current_dir = os.path.dirname(os.path.abspath(__file__))
+_MODEL_DIR = "/workspace/Model-Optimizer/examples/llm_ptq/kimi25_hf-bf16"
+
+def _find_and_import_module(module_name, file_name):
+ """Find module file in multiple locations and import it."""
+ # Check if already imported
+ if module_name in sys.modules:
+ return sys.modules[module_name]
+
+ # Try current directory first (works when loading directly from model dir)
+ candidates = [
+ os.path.join(_current_dir, file_name),
+ # Fallback: original model path (when running from HF cache)
+ os.path.join(_MODEL_DIR, file_name),
+ ]
+
+ for path in candidates:
+ if os.path.exists(path):
+ spec = importlib.util.spec_from_file_location(module_name, path)
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[module_name] = module
+ spec.loader.exec_module(module)
+ return module
+
+ raise FileNotFoundError(f"Could not find {file_name} in any of: {candidates}")
+
+# Import local modules dynamically
+_config_module = _find_and_import_module("configuration_kimi_k25", "configuration_kimi_k25.py")
+KimiK25Config = _config_module.KimiK25Config
+
+_deepseek_module = _find_and_import_module("modeling_deepseek", "modeling_deepseek.py")
+DeepseekV3ForCausalLM = _deepseek_module.DeepseekV3ForCausalLM
+
+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.cache_utils import Cache
+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
+
+# 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 += attention_mask
+ 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'):
+ 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)
+
+ 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)
+ 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,
+ *,
+ attn_implementation: str = 'flash_attention_2',
+ activation=F.gelu,
+ attn_bias: bool = False,
+ use_deterministic_attn: bool = False,
+ ):
+ super().__init__()
+ self.num_heads = num_heads
+ self.hidden_dim = hidden_dim
+ self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
+ self.attn_implementation = attn_implementation
+ self.use_deterministic_attn = use_deterministic_attn
+
+ self.norm0 = nn.LayerNorm(hidden_dim)
+ self.norm1 = nn.LayerNorm(hidden_dim)
+ self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation)
+ self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias)
+ self.wo = nn.Linear(hidden_dim, 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,
+ video_attn_type: str = 'spatial_temporal') -> None:
+ super().__init__()
+ self.use_deterministic_attn = False
+ assert video_attn_type == 'spatial_temporal', f'video_attn_type must be "spatial_temporal", got {video_attn_type}'
+ self.video_attn_type = video_attn_type
+ self.rope_2d = Rope2DPosEmbRepeated(
+ block_cfg['hidden_dim'] // 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)
+ ])
+ self.final_layernorm = nn.LayerNorm(hidden_dim)
+
+ 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 = ['PackingTransformer']
+ _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,
+ )
+
+ 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,
+ 'mlp_dim':
+ config.intermediate_size,
+ 'activation':
+ PytorchGELUTanh(),
+ 'attn_bias':
+ True,
+ 'attn_implementation':
+ config._attn_implementation,
+ },
+ video_attn_type=config.video_attn_type)
+
+ 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 KimiK25PreTrainedModel(PreTrainedModel):
+ config_class = KimiK25Config
+ base_model_prefix = "model"
+ _no_split_modules = [
+ "MoonViT3dPretrainedModel",
+ "MoonViTEncoderLayer",
+ "DeepseekDecoderLayer",
+ "PatchMergerMLP",
+ ]
+ _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: KimiK25Config, **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.video_attn_type = config.video_attn_type
+ self.merge_type = config.merge_type
+ self._attn_implementation = config._attn_implementation
+
+
+class ProjectorConfig:
+
+ def __init__(self, config: KimiK25Config):
+ 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 KimiK25ForConditionalGeneration(KimiK25PreTrainedModel):
+
+ def __init__(self, config: KimiK25Config):
+ 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)
+ else:
+ raise ValueError(
+ f"Unsupported mm_projector_type: {proj_config.mm_projector_type}"
+ )
+
+ self.language_model = DeepseekV3ForCausalLM(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 isinstance(past_key_values, Cache):
+ 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/preprocessor_config.json b/preprocessor_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..b0accb4b335998b847e0093bf4a63e1f32447cc9
--- /dev/null
+++ b/preprocessor_config.json
@@ -0,0 +1,30 @@
+{
+ "auto_map": {
+ "AutoProcessor": "kimi_k25_processor.KimiK25Processor",
+ "AutoImageProcessor": "kimi_k25_vision_processing.KimiK25VisionProcessor"
+ },
+ "media_proc_cfg": {
+ "in_patch_limit": 16384,
+ "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": 4096,
+ "in_patch_limit_video": null,
+ "sample_fps": 2.0,
+ "max_num_frames_each_video": null,
+ "temporal_merge_kernel_size": 4,
+ "timestamp_mode": "hh:mm:ss.fff",
+ "config_type": "media_proc.processors.moonvit.MoonViTMediaProcessorConfig"
+ }
+}
\ No newline at end of file
diff --git a/processor_config.json b/processor_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..31838a3332c7aad5652f4fa910b2f5c3cae3636c
--- /dev/null
+++ b/processor_config.json
@@ -0,0 +1,6 @@
+{
+ "auto_map": {
+ "AutoProcessor": "kimi_k25_processor.KimiK25Processor"
+ },
+ "processor_class": "KimiK25Processor"
+}
diff --git a/special_tokens_map.json b/special_tokens_map.json
new file mode 100644
index 0000000000000000000000000000000000000000..a7815a0bd492c62d5f88efca347cb04bcf99b408
--- /dev/null
+++ b/special_tokens_map.json
@@ -0,0 +1,38 @@
+{
+ "additional_special_tokens": [
+ "<|im_end|>",
+ "<|im_user|>",
+ "<|im_assistant|>",
+ "<|start_header_id|>",
+ "<|end_header_id|>",
+ "[EOT]",
+ "<|im_system|>",
+ "<|im_middle|>",
+ "<|media_begin|>",
+ "<|media_content|>",
+ "<|media_end|>",
+ "<|media_pad|>"
+ ],
+ "bos_token": {
+ "content": "[BOS]",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false
+ },
+ "eos_token": {
+ "content": "[EOS]",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false
+ },
+ "pad_token": "[EOS]",
+ "unk_token": {
+ "content": "[UNK]",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false
+ }
+}
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..7868ea7598734c24cbbbbf904799c76b29af9803
--- /dev/null
+++ b/tokenization_kimi.py
@@ -0,0 +1,352 @@
+import os
+from collections import OrderedDict
+from logging import getLogger
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, 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
+
+from .tool_declaration_ts import encode_tools_to_typescript_style
+
+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
+
+ self._token_config_cache = OrderedDict()
+ self._cache_max_size = 128
+
+ 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(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
+
+ # 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
+
+ texts = self.pre_tokenizer_process(text)
+
+ all_substrs = []
+ for text in texts:
+ substrs = (
+ substr 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))
+ all_substrs.extend(substrs)
+
+ t: List[int] = []
+ for substr in all_substrs:
+ if allow_special_tokens:
+ t.extend(
+ # we should consider special token as a common token
+ self.model.encode(
+ substr,
+ allowed_special="all",
+ ))
+ else:
+ t.extend(
+ # we should consider special token as a common token
+ self.model.encode(
+ substr,
+ disallowed_special=(),
+ ))
+
+ return t
+
+ 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 pre_tokenizer_process(self, text: str) -> List[str]:
+ """
+ pre-tokenizes the input text into a list of tokens.
+ This method is used to split the input text into smaller chunks for internal processing.
+ """
+ return [text]
+
+ """ ----- 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,
+ **kwargs):
+
+ tools = deep_sort_dict(tools)
+
+ # Convert tools to TypeScript style string if tools are provided
+ tools_ts_str = None
+ if tools:
+ try:
+ tools_ts_str = encode_tools_to_typescript_style(tools)
+
+ except Exception as e:
+ print(f"Failed to convert tools to TypeScript style: {e}")
+ tools_ts_str = None
+
+ # Store the TypeScript string in kwargs so it can be accessed by the template
+ if tools_ts_str is not None:
+ kwargs['tools_ts_str'] = tools_ts_str
+ return super().apply_chat_template(
+ conversation,
+ tools=tools,
+ tokenize=tokenize,
+ add_generation_prompt=add_generation_prompt,
+ thinking=thinking,
+ **kwargs)
+
+
+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
diff --git a/tokenizer_config.json b/tokenizer_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..2a1e4b6eadcea8e6ea6f3a227d48e603e19b7140
--- /dev/null
+++ b/tokenizer_config.json
@@ -0,0 +1,216 @@
+{
+ "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": "<|im_end|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": true
+ },
+ "163587": {
+ "content": "<|im_user|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": true
+ },
+ "163588": {
+ "content": "<|im_assistant|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": true
+ },
+ "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
+ },
+ "163594": {
+ "content": "<|im_system|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": true
+ },
+ "163595": {
+ "content": "<|tool_calls_section_begin|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": false
+ },
+ "163596": {
+ "content": "<|tool_calls_section_end|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": false
+ },
+ "163597": {
+ "content": "<|tool_call_begin|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": false
+ },
+ "163598": {
+ "content": "<|tool_call_argument_begin|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": false
+ },
+ "163599": {
+ "content": "<|tool_call_end|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": false
+ },
+ "163601": {
+ "content": "<|im_middle|>",
+ "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
+ },
+ "163606": {
+ "content": "",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": false
+ },
+ "163607": {
+ "content": "",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false,
+ "special": false
+ },
+ "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": [
+ "<|im_end|>",
+ "<|im_user|>",
+ "<|im_assistant|>",
+ "<|start_header_id|>",
+ "<|end_header_id|>",
+ "[EOT]",
+ "<|im_system|>",
+ "<|im_middle|>",
+ "<|media_begin|>",
+ "<|media_content|>",
+ "<|media_end|>",
+ "<|media_pad|>"
+ ],
+ "auto_map": {
+ "AutoTokenizer": [
+ "tokenization_kimi.TikTokenTokenizer",
+ null
+ ]
+ },
+ "bos_token": "[BOS]",
+ "clean_up_tokenization_spaces": false,
+ "eos_token": "[EOS]",
+ "extra_special_tokens": {},
+ "model_max_length": 1000000000000000019884624838656,
+ "pad_token": "[EOS]",
+ "tokenizer_class": "TikTokenTokenizer",
+ "unk_token": "[UNK]"
+}
diff --git a/tool_declaration_ts.py b/tool_declaration_ts.py
new file mode 100644
index 0000000000000000000000000000000000000000..3cc7727ddbeabef16b90ce0219446bc2d4ea9032
--- /dev/null
+++ b/tool_declaration_ts.py
@@ -0,0 +1,479 @@
+"""
+Encode structured tool declaration to typescript style string.
+"""
+import dataclasses
+import json
+import logging
+from collections.abc import Sequence
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_TS_INDENT = " "
+_TS_FIELD_DELIMITER = ",\n"
+
+
+class _SchemaRegistry:
+ """Registry for schema definitions to handle $ref resolution"""
+
+ def __init__(self):
+ self.definitions = {}
+ self.has_self_ref = False
+
+ def register_definitions(self, defs: dict[str, Any]):
+ """Register schema definitions from $defs section"""
+ if not defs:
+ return
+ for def_name, def_schema in defs.items():
+ self.definitions[def_name] = def_schema
+
+ def resolve_ref(self, ref: str) -> dict[str, Any]:
+ """Resolve a reference to its schema definition"""
+ if ref == "#":
+ self.has_self_ref = True
+ return {"$self_ref": True}
+ elif ref.startswith("#/$defs/"):
+ def_name = ref.split("/")[-1]
+ if def_name not in self.definitions:
+ raise ValueError(f"Reference not found: {ref}")
+ return self.definitions[def_name]
+ else:
+ raise ValueError(f"Unsupported reference format: {ref}")
+
+
+def _format_description(description: str, indent: str = "") -> str:
+ return "\n".join([
+ f"{indent}// {line}" if line else ""
+ for line in description.split("\n")
+ ])
+
+
+class _BaseType:
+ description: str
+ constraints: dict[str, Any]
+
+ def __init__(
+ self,
+ extra_props: dict[str, Any],
+ *,
+ allowed_constraint_keys: Sequence[str] = (),
+ ):
+ self.description = extra_props.get("description", "")
+ self.constraints = {
+ k: v
+ for k, v in extra_props.items() if k in allowed_constraint_keys
+ }
+
+ def to_typescript_style(self, indent: str = "") -> str:
+ raise NotImplementedError
+
+ def format_docstring(self, indent: str) -> str:
+ lines = []
+ if self.description:
+ lines.append(_format_description(self.description, indent))
+ if self.constraints:
+ constraints_str = ", ".join(f"{k}: {v}" for k, v in sorted(
+ self.constraints.items(), key=lambda kv: kv[0]))
+ lines.append(f"{indent}// {constraints_str}")
+
+ return "".join(x + "\n" for x in lines)
+
+
+class _ParameterTypeScalar(_BaseType):
+ type: str
+
+ def __init__(self, type: str, extra_props: dict[str, Any] | None = None):
+ self.type = type
+
+ allowed_constraint_keys: list[str] = []
+ if self.type == "string":
+ allowed_constraint_keys = ["maxLength", "minLength", "pattern"]
+ elif self.type in ("number", "integer"):
+ allowed_constraint_keys = ["maximum", "minimum"]
+
+ super().__init__(extra_props or {},
+ allowed_constraint_keys=allowed_constraint_keys)
+
+ def to_typescript_style(self, indent: str = "") -> str:
+ # Map integer to number in TypeScript
+ if self.type == "integer":
+ return "number"
+ return self.type
+
+
+class _ParameterTypeObject(_BaseType):
+ properties: list["_Parameter"]
+ additional_properties: Any | None = None
+
+ def __init__(self,
+ json_schema_object: dict[str, Any],
+ registry: _SchemaRegistry | None = None):
+ super().__init__(json_schema_object)
+
+ self.properties = []
+ self.additional_properties = None
+
+ if not json_schema_object:
+ return
+
+ if "$defs" in json_schema_object and registry:
+ registry.register_definitions(json_schema_object["$defs"])
+
+ self.additional_properties = json_schema_object.get(
+ "additionalProperties")
+ if isinstance(self.additional_properties, dict):
+ self.additional_properties = _parse_parameter_type(
+ self.additional_properties, registry)
+
+ if "properties" not in json_schema_object:
+ return
+
+ required_parameters = json_schema_object.get("required", [])
+ optional_parameters = set(
+ json_schema_object["properties"].keys()) - set(required_parameters)
+
+ self.properties = [
+ _Parameter(
+ name=name,
+ type=_parse_parameter_type(prop, registry),
+ optional=name in optional_parameters,
+ default=prop.get("default")
+ if isinstance(prop, dict) else None,
+ ) for name, prop in json_schema_object["properties"].items()
+ ]
+
+ def to_typescript_style(self, indent: str = "") -> str:
+ # sort by optional, make the required parameters first
+ parameters = [p for p in self.properties if not p.optional]
+ opt_params = [p for p in self.properties if p.optional]
+
+ parameters = sorted(parameters, key=lambda p: p.name)
+ parameters.extend(sorted(opt_params, key=lambda p: p.name))
+
+ param_strs = []
+ for p in parameters:
+ one = p.to_typescript_style(indent=indent + _TS_INDENT)
+ param_strs.append(one)
+
+ if self.additional_properties is not None:
+ ap_type_str = "any"
+ if self.additional_properties is True:
+ ap_type_str = "any"
+ elif self.additional_properties is False:
+ ap_type_str = "never"
+ elif isinstance(self.additional_properties, _ParameterType):
+ ap_type_str = self.additional_properties.to_typescript_style(
+ indent=indent + _TS_INDENT)
+ else:
+ raise ValueError(
+ f"Unknown additionalProperties: {self.additional_properties}"
+ )
+ param_strs.append(
+ f"{indent + _TS_INDENT}[k: string]: {ap_type_str}")
+
+ if not param_strs:
+ return "{}"
+
+ params_str = _TS_FIELD_DELIMITER.join(param_strs)
+ if params_str:
+ # add new line before and after
+ params_str = f"\n{params_str}\n"
+ # always wrap with object
+ return f"{{{params_str}{indent}}}"
+
+
+class _ParameterTypeArray(_BaseType):
+ item: "_ParameterType"
+
+ def __init__(self,
+ json_schema_object: dict[str, Any],
+ registry: _SchemaRegistry | None = None):
+ super().__init__(json_schema_object,
+ allowed_constraint_keys=("minItems", "maxItems"))
+ if json_schema_object.get("items"):
+ self.item = _parse_parameter_type(json_schema_object["items"],
+ registry)
+ else:
+ self.item = _ParameterTypeScalar(type="any")
+
+ def to_typescript_style(self, indent: str = "") -> str:
+ item_docstring = self.item.format_docstring(indent + _TS_INDENT)
+ if item_docstring:
+ return ("Array<\n" + item_docstring + indent + _TS_INDENT +
+ self.item.to_typescript_style(indent=indent + _TS_INDENT) +
+ "\n" + indent + ">")
+ else:
+ return f"Array<{self.item.to_typescript_style(indent=indent)}>"
+
+
+class _ParameterTypeEnum(_BaseType):
+ # support scalar types only
+ enum: list[str | int | float | bool | None]
+
+ def __init__(self, json_schema_object: dict[str, Any]):
+ super().__init__(json_schema_object)
+ self.enum = json_schema_object["enum"]
+
+ # Validate enum values against declared type if present
+ if "type" in json_schema_object:
+ typ = json_schema_object["type"]
+ if isinstance(typ, list):
+ if len(typ) == 1:
+ typ = typ[0]
+ elif len(typ) == 2:
+ if "null" not in typ:
+ raise ValueError(f"Enum type {typ} is not supported")
+ else:
+ typ = typ[0] if typ[0] != "null" else typ[1]
+ else:
+ raise ValueError(f"Enum type {typ} is not supported")
+ for val in self.enum:
+ if val is None:
+ continue
+ if typ == "string" and not isinstance(val, str):
+ raise ValueError(f"Enum value {val} is not a string")
+ elif typ == "number" and not isinstance(val, (int, float)):
+ raise ValueError(f"Enum value {val} is not a number")
+ elif typ == "integer" and not isinstance(val, int):
+ raise ValueError(f"Enum value {val} is not an integer")
+ elif typ == "boolean" and not isinstance(val, bool):
+ raise ValueError(f"Enum value {val} is not a boolean")
+
+ def to_typescript_style(self, indent: str = "") -> str:
+ return " | ".join(
+ [f'"{e}"' if isinstance(e, str) else str(e) for e in self.enum])
+
+
+class _ParameterTypeAnyOf(_BaseType):
+ types: list["_ParameterType"]
+
+ def __init__(
+ self,
+ json_schema_object: dict[str, Any],
+ registry: _SchemaRegistry | None = None,
+ ):
+ super().__init__(json_schema_object)
+ self.types = [
+ _parse_parameter_type(t, registry)
+ for t in json_schema_object["anyOf"]
+ ]
+
+ def to_typescript_style(self, indent: str = "") -> str:
+ return " | ".join(
+ [t.to_typescript_style(indent=indent) for t in self.types])
+
+
+class _ParameterTypeUnion(_BaseType):
+ types: list[str]
+
+ def __init__(self, json_schema_object: dict[str, Any]):
+ super().__init__(json_schema_object)
+
+ mapping = {
+ "string": "string",
+ "number": "number",
+ "integer": "number",
+ "boolean": "boolean",
+ "null": "null",
+ "object": "{}",
+ "array": "Array",
+ }
+ self.types = [mapping[t] for t in json_schema_object["type"]]
+
+ def to_typescript_style(self, indent: str = "") -> str:
+ return " | ".join(self.types)
+
+
+class _ParameterTypeRef(_BaseType):
+ ref_name: str
+ is_self_ref: bool = False
+
+ def __init__(self, json_schema_object: dict[str, Any],
+ registry: _SchemaRegistry):
+ super().__init__(json_schema_object)
+
+ ref = json_schema_object["$ref"]
+ resolved_schema = registry.resolve_ref(ref)
+
+ if resolved_schema.get("$self_ref", False):
+ self.ref_name = "parameters"
+ self.is_self_ref = True
+ else:
+ self.ref_name = ref.split("/")[-1]
+
+ def to_typescript_style(self, indent: str = "") -> str:
+ return self.ref_name
+
+
+_ParameterType = (_ParameterTypeScalar
+ | _ParameterTypeObject
+ | _ParameterTypeArray
+ | _ParameterTypeEnum
+ | _ParameterTypeAnyOf
+ | _ParameterTypeUnion
+ | _ParameterTypeRef)
+
+
+@dataclasses.dataclass
+class _Parameter:
+ """
+ A parameter in a function, or a field in a object.
+ It consists of the type as well as the name.
+ """
+
+ type: _ParameterType
+ name: str = "_"
+ optional: bool = True
+ default: Any | None = None
+
+ @classmethod
+ def parse_extended(cls, attributes: dict[str, Any]) -> "_Parameter":
+ if not attributes:
+ raise ValueError("attributes is empty")
+
+ return cls(
+ name=attributes.get("name", "_"),
+ type=_parse_parameter_type(attributes),
+ optional=attributes.get("optional", False),
+ default=attributes.get("default"),
+ )
+
+ def to_typescript_style(self, indent: str = "") -> str:
+ comments = self.type.format_docstring(indent)
+
+ if self.default is not None:
+ default_repr = (json.dumps(self.default, ensure_ascii=False)
+ if not isinstance(self.default, (int, float, bool))
+ else repr(self.default))
+ comments += f"{indent}// Default: {default_repr}\n"
+
+ return (
+ comments +
+ f"{indent}{self.name}{'?' if self.optional else ''}: {self.type.to_typescript_style(indent=indent)}"
+ )
+
+
+def _parse_parameter_type(
+ json_schema_object: dict[str, Any] | bool,
+ registry: _SchemaRegistry | None = None) -> _ParameterType:
+ if isinstance(json_schema_object, bool):
+ if json_schema_object:
+ return _ParameterTypeScalar(type="any")
+ else:
+ logger.warning(
+ f"Warning: Boolean value {json_schema_object} is not supported, use null instead."
+ )
+ return _ParameterTypeScalar(type="null")
+
+ if "$ref" in json_schema_object and registry:
+ return _ParameterTypeRef(json_schema_object, registry)
+
+ if "anyOf" in json_schema_object:
+ return _ParameterTypeAnyOf(json_schema_object, registry)
+ elif "enum" in json_schema_object:
+ return _ParameterTypeEnum(json_schema_object)
+ elif "type" in json_schema_object:
+ typ = json_schema_object["type"]
+ if isinstance(typ, list):
+ return _ParameterTypeUnion(json_schema_object)
+ elif typ == "object":
+ return _ParameterTypeObject(json_schema_object, registry)
+ elif typ == "array":
+ return _ParameterTypeArray(json_schema_object, registry)
+ else:
+ return _ParameterTypeScalar(typ, json_schema_object)
+ elif json_schema_object == {}:
+ return _ParameterTypeScalar(type="any")
+ else:
+ raise ValueError(f"Invalid JSON Schema object: {json_schema_object}")
+
+
+def _openai_function_to_typescript_style(function: dict[str, Any], ) -> str:
+ """Convert OpenAI function definition (dict) to TypeScript style string."""
+ registry = _SchemaRegistry()
+ parameters = function.get("parameters") or {}
+ parsed = _ParameterTypeObject(parameters, registry)
+
+ interfaces = []
+ root_interface_name = None
+ if registry.has_self_ref:
+ root_interface_name = "parameters"
+ params_str = _TS_FIELD_DELIMITER.join([
+ p.to_typescript_style(indent=_TS_INDENT) for p in parsed.properties
+ ])
+ params_str = f"\n{params_str}\n" if params_str else ""
+ interface_def = f"interface {root_interface_name} {{{params_str}}}"
+ interfaces.append(interface_def)
+
+ definitions_copy = dict(registry.definitions)
+ for def_name, def_schema in definitions_copy.items():
+ obj_type = _parse_parameter_type(def_schema, registry)
+ params_str = obj_type.to_typescript_style()
+
+ description_part = ""
+ if obj_description := def_schema.get("description", ""):
+ description_part = _format_description(obj_description) + "\n"
+
+ interface_def = f"{description_part}interface {def_name} {params_str}"
+ interfaces.append(interface_def)
+
+ interface_str = "\n".join(interfaces)
+ function_name = function.get("name", "function")
+ if root_interface_name:
+ type_def = f"type {function_name} = (_: {root_interface_name}) => any;"
+ else:
+ params_str = parsed.to_typescript_style()
+ type_def = f"type {function_name} = (_: {params_str}) => any;"
+
+ description = function.get("description")
+ return "\n".join(
+ filter(
+ bool,
+ [
+ interface_str,
+ ((description and _format_description(description)) or ""),
+ type_def,
+ ],
+ ))
+
+
+def encode_tools_to_typescript_style(tools: list[dict[str, Any]], ) -> str:
+ """
+ Convert tools (list of dict) to TypeScript style string.
+
+ Supports OpenAI format: {"type": "function", "function": {...}}
+
+ Args:
+ tools: List of tool definitions in dict format
+
+ Returns:
+ TypeScript style string representation of the tools
+ """
+ if not tools:
+ return ""
+
+ functions = []
+
+ for tool in tools:
+ tool_type = tool.get("type")
+ if tool_type == "function":
+ func_def = tool.get("function", {})
+ if func_def:
+ functions.append(
+ _openai_function_to_typescript_style(func_def))
+ else:
+ # Skip unsupported tool types (like "_plugin")
+ continue
+
+ if not functions:
+ return ""
+
+ functions_str = "\n".join(functions)
+ result = "# Tools\n\n"
+
+ if functions_str:
+ result += "## functions\nnamespace functions {\n"
+ result += functions_str + "\n"
+ result += "}\n"
+
+ return result