Duplicate from moonshotai/Kimi-Audio-7B-Instruct

Co-authored-by: moyanwang <[email protected]>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,161 @@
+---
+license: mit
+language:
+- en
+- zh
+tags:
+- audio
+- audio-language-model
+- speech-recognition
+- audio-understanding
+- text-to-speech
+- audio-generation
+- chat
+- kimi-audio
+---
+
+# Kimi-Audio
+
+<p align="center">
+    <img src="https://raw.githubusercontent.com/MoonshotAI/Kimi-Audio/master/assets/kimia_logo.png" width="400"/> <!-- TODO: Replace with actual raw image URL from your repo -->
+<p>
+
+<p align="center">
+<a href="https://huggingface.co/moonshotai/Kimi-Audio-7B">🤗 Kimi-Audio-7B</a>&nbsp; | <a href="https://huggingface.co/moonshotai/Kimi-Audio-7B-Instruct">🤗 Kimi-Audio-7B-Instruct </a>&nbsp; | <a href="https://raw.githubusercontent.com/MoonshotAI/Kimi-Audio/master/assets/kimia_report.pdf">📑 Paper</a>
+</p>
+
+## Introduction
+
+We present Kimi-Audio, an open-source audio foundation model excelling in **audio understanding, generation, and conversation**. This repository hosts the model checkpoints for Kimi-Audio-7B-Instruct.
+
+Kimi-Audio is designed as a universal audio foundation model capable of handling a wide variety of audio processing tasks within a single unified framework. Key features include:
+
+*   **Universal Capabilities:** Handles diverse tasks like speech recognition (ASR), audio question answering (AQA), audio captioning (AAC), speech emotion recognition (SER), sound event/scene classification (SEC/ASC) and end-to-end speech conversation.
+*   **State-of-the-Art Performance:** Achieves SOTA results on numerous audio benchmarks (see our [Technical Report](https://raw.githubusercontent.com/MoonshotAI/Kimi-Audio/master/assets/kimia_report.pdf)).
+*   **Large-Scale Pre-training:** Pre-trained on over 13 million hours of diverse audio data (speech, music, sounds) and text data.
+*   **Novel Architecture:** Employs a hybrid audio input (continuous acoustic + discrete semantic tokens) and an LLM core with parallel heads for text and audio token generation.
+*   **Efficient Inference:** Features a chunk-wise streaming detokenizer based on flow matching for low-latency audio generation.
+
+For more details, please refer to our [GitHub Repository](https://github.com/MoonshotAI/Kimi-Audio) and [Technical Report](https://raw.githubusercontent.com/MoonshotAI/Kimi-Audio/master/assets/kimia_report.pdf).
+
+## Requirements
+
+We recommend that you build a Docker image to run the inference. After cloning the inference code, you can construct the image using the `docker build` command.
+```bash
+git clone https://github.com/MoonshotAI/Kimi-Audio
+git submodule update --init
+cd Kimi-Audio
+docker build -t kimi-audio:v0.1 .
+```
+Alternatively, You can also use our pre-built image:
+```bash
+docker pull moonshotai/kimi-audio:v0.1
+```
+
+Or, you can install requirments by:
+```bash
+pip install -r requirements.txt
+```
+
+You may refer to the Dockerfile in case of any environment issues.
+
+## Quickstart
+
+This example demonstrates basic usage for generating text from audio (ASR) and generating both text and speech in a conversational turn using the `Kimi-Audio-7B-Instruct` model.
+
+```python
+import soundfile as sf
+# Assuming the KimiAudio class is available after installation
+from kimia_infer.api.kimia import KimiAudio
+import torch # Ensure torch is imported if needed for device placement
+
+# --- 1. Load Model ---
+# Load the model from Hugging Face Hub
+# Make sure you are logged in (`huggingface-cli login`) if the repo is private.
+model_id = "moonshotai/Kimi-Audio-7B-Instruct" # Or "Kimi/Kimi-Audio-7B"
+device = "cuda" if torch.cuda.is_available() else "cpu" # Example device placement
+# Note: The KimiAudio class might handle model loading differently.
+# You might need to pass the model_id directly or download checkpoints manually
+# and provide the local path as shown in the original readme_kimia.md.
+# Please refer to the main Kimi-Audio repository for precise loading instructions.
+# Example assuming KimiAudio takes the HF ID or a local path:
+try:
+    model = KimiAudio(model_path=model_id, load_detokenizer=True) # May need device argument
+    model.to(device) # Example device placement
+except Exception as e:
+    print(f"Automatic loading from HF Hub might require specific setup.")
+    print(f"Refer to Kimi-Audio docs. Trying local path example (update path!). Error: {e}")
+    # Fallback example:
+    # model_path = "/path/to/your/downloaded/kimia-hf-ckpt" # IMPORTANT: Update this path if loading locally
+    # model = KimiAudio(model_path=model_path, load_detokenizer=True)
+    # model.to(device) # Example device placement
+
+# --- 2. Define Sampling Parameters ---
+sampling_params = {
+    "audio_temperature": 0.8,
+    "audio_top_k": 10,
+    "text_temperature": 0.0,
+    "text_top_k": 5,
+    "audio_repetition_penalty": 1.0,
+    "audio_repetition_window_size": 64,
+    "text_repetition_penalty": 1.0,
+    "text_repetition_window_size": 16,
+}
+
+# --- 3. Example 1: Audio-to-Text (ASR) ---
+# TODO: Provide actual example audio files or URLs accessible to users
+# E.g., download sample files first or use URLs
+# wget https://path/to/your/asr_example.wav -O asr_example.wav
+# wget https://path/to/your/qa_example.wav -O qa_example.wav
+asr_audio_path = "asr_example.wav" # IMPORTANT: Make sure this file exists
+qa_audio_path = "qa_example.wav" # IMPORTANT: Make sure this file exists
+
+messages_asr = [
+    {"role": "user", "message_type": "text", "content": "Please transcribe the following audio:"},
+    {"role": "user", "message_type": "audio", "content": asr_audio_path}
+]
+
+# Generate only text output
+# Note: Ensure the model object and generate method accept device placement if needed
+_, text_output = model.generate(messages_asr, **sampling_params, output_type="text")
+print(">>> ASR Output Text: ", text_output)
+# Expected output: "这并不是告别，这是一个篇章的结束，也是新篇章的开始。" (Example)
+
+# --- 4. Example 2: Audio-to-Audio/Text Conversation ---
+messages_conversation = [
+    {"role": "user", "message_type": "audio", "content": qa_audio_path}
+]
+
+# Generate both audio and text output
+wav_output, text_output = model.generate(messages_conversation, **sampling_params, output_type="both")
+
+# Save the generated audio
+output_audio_path = "output_audio.wav"
+# Ensure wav_output is on CPU and flattened before saving
+sf.write(output_audio_path, wav_output.detach().cpu().view(-1).numpy(), 24000) # Assuming 24kHz output
+print(f">>> Conversational Output Audio saved to: {output_audio_path}")
+print(">>> Conversational Output Text: ", text_output)
+# Expected output: "A." (Example)
+
+print("Kimi-Audio inference examples complete.")
+
+```
+
+## Citation
+
+If you find Kimi-Audio useful in your research or applications, please cite our technical report:
+
+```bibtex
+@misc{kimi_audio_2024,
+      title={Kimi-Audio Technical Report},
+      author={Kimi Team},
+      year={2024},
+      eprint={arXiv:placeholder},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+```
+
+## License
+
+The model is based and modified from [Qwen 2.5-7B](https://github.com/QwenLM/Qwen2.5). Code derived from Qwen2.5-7B is licensed under the [Apache 2.0 License](https://www.apache.org/licenses/LICENSE-2.0). Other parts of the code are licensed under the [MIT License](https://opensource.org/licenses/MIT).

audio_detokenizer/config.yaml

@@ -0,0 +1,123 @@
+accumulate_grad_batches: 1
+base_config: config/config_base.yaml
+batch_max_tokens: 12000
+batch_size: 2
+cfg_init: 1.0
+cfg_scale: 4.0
+cfg_schedule: linear
+check_val_every_n_epoch: 10
+clip_grad_norm: 0
+data_dir: ''
+debug: false
+deep_speed_strategy_stage: 2
+drop_last: true
+dynamic_cfg: false
+endless_ds: false
+filter_args:
+  lang:
+  - zh
+  - en
+  max_spk_num: 6
+  speech_ratio: 0.6
+gradient_clip_val: 1.0
+indexed_ds: true
+infer: false
+infer_exp_name: ''
+infer_json_path: ''
+inference_ckpt: ''
+inference_mode: nonstreaming
+learning_rate: 1e-4
+limit_val_batches: 100
+load_opt: false
+log_interval: 10
+logger_type: tensorboard
+loss:
+  lambda_fm: 1.0
+  lambda_phone: 0.0
+  mel_loss: l1
+max_epochs: 1000
+max_eval_sentences: -1
+max_eval_tokens: -1
+max_prompt_ratio: 0.5
+max_segment_cnt: 20000
+max_sentences: -1
+max_speech_duration: 20
+max_tokens: 31250
+max_training_steps: 100000
+max_updates: 160000
+mel_mean: -4.479605
+mel_std: 3.4584913
+meta_dir: null
+min_prompt_duration: 0.5
+min_speech_duration: -1
+model:
+  condition_prenet_depth: 6
+  dit:
+    chunk_params:
+      hz: 50
+      max_chunk: 3.0
+      max_chunk_history: 50000000
+      min_chunk: 0.5
+      need_block_shift: false
+    condition_input_dim: 1280
+    condition_type: discrete_codes
+    depth: 16
+    ffn_act_layer: gleu_tanh
+    ffn_conv_kernel_size: 5
+    ffn_gated_glu: false
+    ffn_type: vanilla_mlp
+    hidden_size: 2304
+    input_size: 80
+    max_seq_len: 4096
+    mlp_ratio: 4.0
+    num_heads: 18
+    position_embedding_type: skip
+    prompt_cfg_dropout: 0.2
+    rope_params:
+      max_position_embeddings: 4096
+      rope_base: 10000.0
+      rope_interpolation_factor: 1.0
+    semantic_cfg_dropout: 0.2
+    semantic_vocab_size: 16384
+    use_chunk_setting: true
+    use_rope: true
+  phone_predictor:
+    blank_id: 4
+    phone_vocab_size: 5000
+  position_id_start_from: 0
+  random_position_start: true
+  restart_position_ids: false
+  use_condition_prenet: false
+need_merge_same_speaker: true
+need_precise_phones: false
+no_verlap: true
+normalize_mel: true
+num_nodes: 1
+num_sanity_val_steps: 0
+num_workers: 1
+ode_steps: 150
+optimizer_adam_beta1: 0.9
+optimizer_adam_beta2: 0.98
+optimizer_class: adamw
+pin_memory: true
+precision: bf16-mixed
+save_interval: 2000
+save_topk: 10
+seed: 1234
+shuffle: true
+sort_by_len: true
+src_sample_rate: 16000
+strategy: ddp
+tensorboard_dir: tb_logs
+test_num: 100
+tgt_sample_rate: 24000
+timescale: 80000
+use_cfg: false
+use_cfg_rescale: false
+use_distributed_sampler: false
+use_uncondition: false
+val_check_interval: 2000000
+vocoder_ckpt: ''
+wandb_name: glm4_semantic_cfm_v2_debug
+warmup_updates: 100
+weight_decay: 0.0001

audio_detokenizer/model.pt

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

config.json

@@ -0,0 +1,44 @@
+{
+  "architectures": [
+    "MoonshotKimiaForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_moonshot_kimia.KimiAudioConfig",
+    "AutoModel": "modeling_moonshot_kimia.MoonshotKimiaModel",
+    "AutoModelForCausalLM": "modeling_moonshot_kimia.MoonshotKimiaForCausalLM"
+  },
+  "bos_token_id": 151643,
+  "eos_token_ids": [
+    151644,
+    151645
+  ],
+  "hidden_act": "silu",
+  "hidden_size": 3584,
+  "initializer_range": 0.02,
+  "intermediate_size": 18944,
+  "kimia_adaptor_input_dim": 5120,
+  "kimia_audio_output_vocab": 16896,
+  "kimia_media_begin": 151661,
+  "kimia_media_end": 151663,
+  "kimia_mimo_audiodelaytokens": 5,
+  "kimia_mimo_layers": 6,
+  "kimia_mimo_transformer_from_layer_index": 21,
+  "kimia_text_output_vocab": 152064,
+  "kimia_token_offset": 152064,
+  "num_attention_heads": 28,
+  "num_audio_special_tokens": 512,
+  "num_base_tokens": 151643,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 4,
+  "pad_token_id": 152063,
+  "max_position_embeddings": 8192,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.44.1",
+  "use_cache": true,
+  "use_whisper_feature": true,
+  "vocab_size": 168448
+}

configuration_moonshot_kimia.py

@@ -0,0 +1,66 @@
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+
+class KimiAudioConfig(Qwen2Config):
+    def __init__(
+        self,
+        vocab_size=163840,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        hidden_act="silu",
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        tie_word_embeddings=False,
+        kimia_mimo_layers: int = 6,
+        kimia_mimo_audiodelaytokens: int = 5,
+        kimia_mimo_transformer_from_layer_index: int = 21,
+        kimia_audio_output_vocab: int = 16896,
+        kimia_text_output_vocab: int = 152064,
+        num_audio_special_tokens: int = 512,
+        num_base_tokens: int = 151643,
+        kimia_token_offset: int = 152064,
+        use_whisper_feature: bool = True,
+        kimia_adaptor_input_dim: int = 5120,
+        kimia_media_begin: int = 151661,
+        kimia_media_end: int = 151663,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_hidden_layers=num_hidden_layers,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            hidden_act=hidden_act,
+            initializer_range=initializer_range,
+            rms_norm_eps=rms_norm_eps,
+            use_cache=use_cache,
+            tie_word_embeddings=tie_word_embeddings,
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            **kwargs,
+        )
+
+        self.kimia_mimo_layers = kimia_mimo_layers
+        self.kimia_mimo_audiodelaytokens = kimia_mimo_audiodelaytokens
+        # vocab
+        self.kimia_mimo_transformer_from_layer_index = (
+            kimia_mimo_transformer_from_layer_index
+        )
+        self.kimia_audio_output_vocab = kimia_audio_output_vocab
+        self.kimia_text_output_vocab = kimia_text_output_vocab
+        self.num_audio_special_tokens = num_audio_special_tokens
+        self.num_base_tokens = num_base_tokens
+        self.kimia_token_offset = kimia_token_offset
+        self.use_whisper_feature = use_whisper_feature
+        self.kimia_adaptor_input_dim = kimia_adaptor_input_dim
+        # special tokens
+        self.kimia_media_begin = kimia_media_begin
+        self.kimia_media_end = kimia_media_end

generation_config.json

@@ -0,0 +1,3 @@
+{
+  "max_length": 8192
+}

model-1-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:462878caab2cc405a9665569e9c4a72191b070f73df5530f6e9c419108a24fe2
+size 466117192

model-10-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:835a9fe443d1118f5ff46d1d35814e6080ff2a9ed578654887638ee853f47ae2
+size 466117192

model-11-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c9de9a823c1b63a4813b4e1d0f6d9c522382981f1889f9423aa077cc91a8f0f
+size 466117208

model-12-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e38dd5adcf335dfe2da28613ab0098570ae9f5fb0f0d137720b5b2458d68d45
+size 466117208

model-13-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:895e02180f74fd811874d200229a549344e6a84a12943aa78d7bf03c3ffb6140
+size 466117208

model-14-of-35.safetensors

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

model-15-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76c3fb75b201ccbbfc1f9b909b75924445efe3e09199ac6d13db0857356975fa
+size 466117208

model-16-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9af4933a1ef591f45d2e61fd7cc5b4e00969e59339a74a10a625de85ccb26362
+size 466117208

model-17-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39545d953a1db3ffb63f981e57c2f7d4fd178131c515c58955ece020be96241d
+size 466117208

model-18-of-35.safetensors

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

model-19-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51e850d87b4ece3566404d66c1e27d474549a160b6bc216607bd11fc336e25ee
+size 466117208

model-2-of-35.safetensors

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

model-20-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:434af95fd5e5a9e38dfe37c2123de7e55e8676c45ecc3f11ab94c0984588ab40
+size 466117208

model-21-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f1dd7e0d633053d61bdd14f1ff96f3dc27e7e4ee66c707ce4ff4f4bb1490040
+size 466117208

model-22-of-35.safetensors

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

model-23-of-35.safetensors

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

model-24-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17c2348936c30f4715c6ac406d88d3fcd2ede9b860051b0e43357398b6b2c392
+size 466117208

model-25-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f0f0d3852637374b66bd250abab868a48ca4120d01e70acff617ce751e2f9ce
+size 466117208

model-26-of-35.safetensors

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

model-27-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95c4714d295631d3562b25c5168a6551a92450a4c63bb9592eadde18347e2a3c
+size 466117208

model-28-of-35.safetensors

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

model-29-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02c5b5edc547ab8ff6ac83fec6f10a018ce111cc1ce6f6e3ced4db432b60903f
+size 466117264

model-3-of-35.safetensors

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

model-30-of-35.safetensors

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

model-31-of-35.safetensors

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

model-32-of-35.safetensors

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

model-33-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e7ab0dbed449d1ba7ed1f2582e279af6585c505158b1af66a7a77521677b70a
+size 466117264

model-34-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46b117a2fe7dd2c2eb3d2c6fb367bc3813a2d3b2a01bd330d983c52112a32ffd
+size 466117264

model-35-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7dbbe9894d818f2751c32a36fe6db1f7e5d0f7d94e9752ee2ba81535ec777e33
+size 62419592

model-36-of-36.safetensors

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

model-4-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d5fe7491bac40d3b5c5eec23ba449f15e3248d12bdf582881fc65dafebe5a7a
+size 466117192

model-5-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:193972da11808be8be48b74f439c5cd0567d825c84505565d30e208bf3dff0ed
+size 466117192

model-6-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6aabe4e1e987bfc7ecc46933ae4a5f37c6e262411c1d86c663709abe633f3756
+size 466117192

model-7-of-35.safetensors

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

model-8-of-35.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54f14bd698a405430c529895e167d39827e349a48181544e7e12ff41637f28e8
+size 466117192

model-9-of-35.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,460 @@
+{
+  "metadata": {
+    "total_size": 19532673280
+  },
+  "weight_map": {
+    "model.layers.0.self_attn.q_proj.weight": "model-1-of-35.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-1-of-35.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-1-of-35.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-1-of-35.safetensors",
+    "model.layers.0.self_attn.q_proj.bias": "model-1-of-35.safetensors",
+    "model.layers.0.self_attn.k_proj.bias": "model-1-of-35.safetensors",
+    "model.layers.0.self_attn.v_proj.bias": "model-1-of-35.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-1-of-35.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-1-of-35.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-1-of-35.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-1-of-35.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-1-of-35.safetensors",
+    "model.layers.0.self_attn.rotary_emb.inv_freq": "model-1-of-35.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-2-of-35.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-2-of-35.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-2-of-35.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-2-of-35.safetensors",
+    "model.layers.1.self_attn.q_proj.bias": "model-2-of-35.safetensors",
+    "model.layers.1.self_attn.k_proj.bias": "model-2-of-35.safetensors",
+    "model.layers.1.self_attn.v_proj.bias": "model-2-of-35.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-2-of-35.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-2-of-35.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-2-of-35.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-2-of-35.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-2-of-35.safetensors",
+    "model.layers.1.self_attn.rotary_emb.inv_freq": "model-2-of-35.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-3-of-35.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-3-of-35.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-3-of-35.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-3-of-35.safetensors",
+    "model.layers.2.self_attn.q_proj.bias": "model-3-of-35.safetensors",
+    "model.layers.2.self_attn.k_proj.bias": "model-3-of-35.safetensors",
+    "model.layers.2.self_attn.v_proj.bias": "model-3-of-35.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-3-of-35.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-3-of-35.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-3-of-35.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-3-of-35.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-3-of-35.safetensors",
+    "model.layers.2.self_attn.rotary_emb.inv_freq": "model-3-of-35.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-4-of-35.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-4-of-35.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-4-of-35.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-4-of-35.safetensors",
+    "model.layers.3.self_attn.q_proj.bias": "model-4-of-35.safetensors",
+    "model.layers.3.self_attn.k_proj.bias": "model-4-of-35.safetensors",
+    "model.layers.3.self_attn.v_proj.bias": "model-4-of-35.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-4-of-35.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-4-of-35.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-4-of-35.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-4-of-35.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-4-of-35.safetensors",
+    "model.layers.3.self_attn.rotary_emb.inv_freq": "model-4-of-35.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-5-of-35.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-5-of-35.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-5-of-35.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-5-of-35.safetensors",
+    "model.layers.4.self_attn.q_proj.bias": "model-5-of-35.safetensors",
+    "model.layers.4.self_attn.k_proj.bias": "model-5-of-35.safetensors",
+    "model.layers.4.self_attn.v_proj.bias": "model-5-of-35.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-5-of-35.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-5-of-35.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-5-of-35.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-5-of-35.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-5-of-35.safetensors",
+    "model.layers.4.self_attn.rotary_emb.inv_freq": "model-5-of-35.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-6-of-35.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-6-of-35.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-6-of-35.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-6-of-35.safetensors",
+    "model.layers.5.self_attn.q_proj.bias": "model-6-of-35.safetensors",
+    "model.layers.5.self_attn.k_proj.bias": "model-6-of-35.safetensors",
+    "model.layers.5.self_attn.v_proj.bias": "model-6-of-35.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-6-of-35.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-6-of-35.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-6-of-35.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-6-of-35.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-6-of-35.safetensors",
+    "model.layers.5.self_attn.rotary_emb.inv_freq": "model-6-of-35.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-7-of-35.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-7-of-35.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-7-of-35.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-7-of-35.safetensors",
+    "model.layers.6.self_attn.q_proj.bias": "model-7-of-35.safetensors",
+    "model.layers.6.self_attn.k_proj.bias": "model-7-of-35.safetensors",
+    "model.layers.6.self_attn.v_proj.bias": "model-7-of-35.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-7-of-35.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-7-of-35.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-7-of-35.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-7-of-35.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-7-of-35.safetensors",
+    "model.layers.6.self_attn.rotary_emb.inv_freq": "model-7-of-35.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-8-of-35.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-8-of-35.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-8-of-35.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-8-of-35.safetensors",
+    "model.layers.7.self_attn.q_proj.bias": "model-8-of-35.safetensors",
+    "model.layers.7.self_attn.k_proj.bias": "model-8-of-35.safetensors",
+    "model.layers.7.self_attn.v_proj.bias": "model-8-of-35.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-8-of-35.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-8-of-35.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-8-of-35.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-8-of-35.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-8-of-35.safetensors",
+    "model.layers.7.self_attn.rotary_emb.inv_freq": "model-8-of-35.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-9-of-35.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-9-of-35.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-9-of-35.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-9-of-35.safetensors",
+    "model.layers.8.self_attn.q_proj.bias": "model-9-of-35.safetensors",
+    "model.layers.8.self_attn.k_proj.bias": "model-9-of-35.safetensors",
+    "model.layers.8.self_attn.v_proj.bias": "model-9-of-35.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-9-of-35.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-9-of-35.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-9-of-35.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-9-of-35.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-9-of-35.safetensors",
+    "model.layers.8.self_attn.rotary_emb.inv_freq": "model-9-of-35.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-10-of-35.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-10-of-35.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-10-of-35.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-10-of-35.safetensors",
+    "model.layers.9.self_attn.q_proj.bias": "model-10-of-35.safetensors",
+    "model.layers.9.self_attn.k_proj.bias": "model-10-of-35.safetensors",
+    "model.layers.9.self_attn.v_proj.bias": "model-10-of-35.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-10-of-35.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-10-of-35.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-10-of-35.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-10-of-35.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-10-of-35.safetensors",
+    "model.layers.9.self_attn.rotary_emb.inv_freq": "model-10-of-35.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-11-of-35.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-11-of-35.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-11-of-35.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-11-of-35.safetensors",
+    "model.layers.10.self_attn.q_proj.bias": "model-11-of-35.safetensors",
+    "model.layers.10.self_attn.k_proj.bias": "model-11-of-35.safetensors",
+    "model.layers.10.self_attn.v_proj.bias": "model-11-of-35.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-11-of-35.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-11-of-35.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-11-of-35.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-11-of-35.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-11-of-35.safetensors",
+    "model.layers.10.self_attn.rotary_emb.inv_freq": "model-11-of-35.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-12-of-35.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-12-of-35.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-12-of-35.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-12-of-35.safetensors",
+    "model.layers.11.self_attn.q_proj.bias": "model-12-of-35.safetensors",
+    "model.layers.11.self_attn.k_proj.bias": "model-12-of-35.safetensors",
+    "model.layers.11.self_attn.v_proj.bias": "model-12-of-35.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-12-of-35.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-12-of-35.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-12-of-35.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-12-of-35.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-12-of-35.safetensors",
+    "model.layers.11.self_attn.rotary_emb.inv_freq": "model-12-of-35.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-13-of-35.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-13-of-35.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-13-of-35.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-13-of-35.safetensors",
+    "model.layers.12.self_attn.q_proj.bias": "model-13-of-35.safetensors",
+    "model.layers.12.self_attn.k_proj.bias": "model-13-of-35.safetensors",
+    "model.layers.12.self_attn.v_proj.bias": "model-13-of-35.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-13-of-35.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-13-of-35.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-13-of-35.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-13-of-35.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-13-of-35.safetensors",
+    "model.layers.12.self_attn.rotary_emb.inv_freq": "model-13-of-35.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-14-of-35.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-14-of-35.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-14-of-35.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-14-of-35.safetensors",
+    "model.layers.13.self_attn.q_proj.bias": "model-14-of-35.safetensors",
+    "model.layers.13.self_attn.k_proj.bias": "model-14-of-35.safetensors",
+    "model.layers.13.self_attn.v_proj.bias": "model-14-of-35.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-14-of-35.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-14-of-35.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-14-of-35.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-14-of-35.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-14-of-35.safetensors",
+    "model.layers.13.self_attn.rotary_emb.inv_freq": "model-14-of-35.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-15-of-35.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-15-of-35.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-15-of-35.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-15-of-35.safetensors",
+    "model.layers.14.self_attn.q_proj.bias": "model-15-of-35.safetensors",
+    "model.layers.14.self_attn.k_proj.bias": "model-15-of-35.safetensors",
+    "model.layers.14.self_attn.v_proj.bias": "model-15-of-35.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-15-of-35.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-15-of-35.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-15-of-35.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-15-of-35.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-15-of-35.safetensors",
+    "model.layers.14.self_attn.rotary_emb.inv_freq": "model-15-of-35.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-16-of-35.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-16-of-35.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-16-of-35.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-16-of-35.safetensors",
+    "model.layers.15.self_attn.q_proj.bias": "model-16-of-35.safetensors",
+    "model.layers.15.self_attn.k_proj.bias": "model-16-of-35.safetensors",
+    "model.layers.15.self_attn.v_proj.bias": "model-16-of-35.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-16-of-35.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-16-of-35.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-16-of-35.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-16-of-35.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-16-of-35.safetensors",
+    "model.layers.15.self_attn.rotary_emb.inv_freq": "model-16-of-35.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-17-of-35.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-17-of-35.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-17-of-35.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-17-of-35.safetensors",
+    "model.layers.16.self_attn.q_proj.bias": "model-17-of-35.safetensors",
+    "model.layers.16.self_attn.k_proj.bias": "model-17-of-35.safetensors",
+    "model.layers.16.self_attn.v_proj.bias": "model-17-of-35.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-17-of-35.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-17-of-35.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-17-of-35.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-17-of-35.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-17-of-35.safetensors",
+    "model.layers.16.self_attn.rotary_emb.inv_freq": "model-17-of-35.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-18-of-35.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-18-of-35.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-18-of-35.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-18-of-35.safetensors",
+    "model.layers.17.self_attn.q_proj.bias": "model-18-of-35.safetensors",
+    "model.layers.17.self_attn.k_proj.bias": "model-18-of-35.safetensors",
+    "model.layers.17.self_attn.v_proj.bias": "model-18-of-35.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-18-of-35.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-18-of-35.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-18-of-35.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-18-of-35.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-18-of-35.safetensors",
+    "model.layers.17.self_attn.rotary_emb.inv_freq": "model-18-of-35.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-19-of-35.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-19-of-35.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-19-of-35.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-19-of-35.safetensors",
+    "model.layers.18.self_attn.q_proj.bias": "model-19-of-35.safetensors",
+    "model.layers.18.self_attn.k_proj.bias": "model-19-of-35.safetensors",
+    "model.layers.18.self_attn.v_proj.bias": "model-19-of-35.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-19-of-35.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-19-of-35.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-19-of-35.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-19-of-35.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-19-of-35.safetensors",
+    "model.layers.18.self_attn.rotary_emb.inv_freq": "model-19-of-35.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-20-of-35.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-20-of-35.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-20-of-35.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-20-of-35.safetensors",
+    "model.layers.19.self_attn.q_proj.bias": "model-20-of-35.safetensors",
+    "model.layers.19.self_attn.k_proj.bias": "model-20-of-35.safetensors",
+    "model.layers.19.self_attn.v_proj.bias": "model-20-of-35.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-20-of-35.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-20-of-35.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-20-of-35.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-20-of-35.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-20-of-35.safetensors",
+    "model.layers.19.self_attn.rotary_emb.inv_freq": "model-20-of-35.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-21-of-35.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-21-of-35.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-21-of-35.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-21-of-35.safetensors",
+    "model.layers.20.self_attn.q_proj.bias": "model-21-of-35.safetensors",
+    "model.layers.20.self_attn.k_proj.bias": "model-21-of-35.safetensors",
+    "model.layers.20.self_attn.v_proj.bias": "model-21-of-35.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-21-of-35.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-21-of-35.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-21-of-35.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-21-of-35.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-21-of-35.safetensors",
+    "model.layers.20.self_attn.rotary_emb.inv_freq": "model-21-of-35.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-22-of-35.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-22-of-35.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-22-of-35.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-22-of-35.safetensors",
+    "model.layers.21.self_attn.q_proj.bias": "model-22-of-35.safetensors",
+    "model.layers.21.self_attn.k_proj.bias": "model-22-of-35.safetensors",
+    "model.layers.21.self_attn.v_proj.bias": "model-22-of-35.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-22-of-35.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-22-of-35.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-22-of-35.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-22-of-35.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-22-of-35.safetensors",
+    "model.layers.21.self_attn.rotary_emb.inv_freq": "model-22-of-35.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-23-of-35.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-23-of-35.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-23-of-35.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-23-of-35.safetensors",
+    "model.layers.22.self_attn.q_proj.bias": "model-23-of-35.safetensors",
+    "model.layers.22.self_attn.k_proj.bias": "model-23-of-35.safetensors",
+    "model.layers.22.self_attn.v_proj.bias": "model-23-of-35.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-23-of-35.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-23-of-35.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-23-of-35.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-23-of-35.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-23-of-35.safetensors",
+    "model.layers.22.self_attn.rotary_emb.inv_freq": "model-23-of-35.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-24-of-35.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-24-of-35.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-24-of-35.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-24-of-35.safetensors",
+    "model.layers.23.self_attn.q_proj.bias": "model-24-of-35.safetensors",
+    "model.layers.23.self_attn.k_proj.bias": "model-24-of-35.safetensors",
+    "model.layers.23.self_attn.v_proj.bias": "model-24-of-35.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-24-of-35.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-24-of-35.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-24-of-35.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-24-of-35.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-24-of-35.safetensors",
+    "model.layers.23.self_attn.rotary_emb.inv_freq": "model-24-of-35.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-25-of-35.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-25-of-35.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-25-of-35.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-25-of-35.safetensors",
+    "model.layers.24.self_attn.q_proj.bias": "model-25-of-35.safetensors",
+    "model.layers.24.self_attn.k_proj.bias": "model-25-of-35.safetensors",
+    "model.layers.24.self_attn.v_proj.bias": "model-25-of-35.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-25-of-35.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-25-of-35.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-25-of-35.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-25-of-35.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-25-of-35.safetensors",
+    "model.layers.24.self_attn.rotary_emb.inv_freq": "model-25-of-35.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-26-of-35.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-26-of-35.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-26-of-35.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-26-of-35.safetensors",
+    "model.layers.25.self_attn.q_proj.bias": "model-26-of-35.safetensors",
+    "model.layers.25.self_attn.k_proj.bias": "model-26-of-35.safetensors",
+    "model.layers.25.self_attn.v_proj.bias": "model-26-of-35.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-26-of-35.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-26-of-35.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-26-of-35.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-26-of-35.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-26-of-35.safetensors",
+    "model.layers.25.self_attn.rotary_emb.inv_freq": "model-26-of-35.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-27-of-35.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-27-of-35.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-27-of-35.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-27-of-35.safetensors",
+    "model.layers.26.self_attn.q_proj.bias": "model-27-of-35.safetensors",
+    "model.layers.26.self_attn.k_proj.bias": "model-27-of-35.safetensors",
+    "model.layers.26.self_attn.v_proj.bias": "model-27-of-35.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-27-of-35.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-27-of-35.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-27-of-35.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-27-of-35.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-27-of-35.safetensors",
+    "model.layers.26.self_attn.rotary_emb.inv_freq": "model-27-of-35.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-28-of-35.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-28-of-35.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-28-of-35.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-28-of-35.safetensors",
+    "model.layers.27.self_attn.q_proj.bias": "model-28-of-35.safetensors",
+    "model.layers.27.self_attn.k_proj.bias": "model-28-of-35.safetensors",
+    "model.layers.27.self_attn.v_proj.bias": "model-28-of-35.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-28-of-35.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-28-of-35.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-28-of-35.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-28-of-35.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-28-of-35.safetensors",
+    "model.layers.27.self_attn.rotary_emb.inv_freq": "model-28-of-35.safetensors",
+    "model.mimo_layers.0.self_attn.q_proj.weight": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.self_attn.k_proj.weight": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.self_attn.v_proj.weight": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.self_attn.o_proj.weight": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.input_layernorm.weight": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.post_attention_layernorm.weight": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.mlp.gate_proj.weight": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.mlp.down_proj.weight": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.mlp.up_proj.weight": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.self_attn.q_proj.bias": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.self_attn.k_proj.bias": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.self_attn.v_proj.bias": "model-29-of-35.safetensors",
+    "model.mimo_layers.0.self_attn.rotary_emb.inv_freq": "model-29-of-35.safetensors",
+    "model.mimo_layers.1.self_attn.q_proj.weight": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.self_attn.k_proj.weight": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.self_attn.v_proj.weight": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.self_attn.o_proj.weight": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.input_layernorm.weight": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.post_attention_layernorm.weight": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.mlp.gate_proj.weight": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.mlp.down_proj.weight": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.mlp.up_proj.weight": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.self_attn.q_proj.bias": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.self_attn.k_proj.bias": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.self_attn.v_proj.bias": "model-30-of-35.safetensors",
+    "model.mimo_layers.1.self_attn.rotary_emb.inv_freq": "model-30-of-35.safetensors",
+    "model.mimo_layers.2.self_attn.q_proj.weight": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.self_attn.k_proj.weight": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.self_attn.v_proj.weight": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.self_attn.o_proj.weight": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.input_layernorm.weight": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.post_attention_layernorm.weight": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.mlp.gate_proj.weight": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.mlp.down_proj.weight": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.mlp.up_proj.weight": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.self_attn.q_proj.bias": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.self_attn.k_proj.bias": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.self_attn.v_proj.bias": "model-31-of-35.safetensors",
+    "model.mimo_layers.2.self_attn.rotary_emb.inv_freq": "model-31-of-35.safetensors",
+    "model.mimo_layers.3.self_attn.q_proj.weight": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.self_attn.k_proj.weight": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.self_attn.v_proj.weight": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.self_attn.o_proj.weight": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.input_layernorm.weight": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.post_attention_layernorm.weight": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.mlp.gate_proj.weight": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.mlp.down_proj.weight": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.mlp.up_proj.weight": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.self_attn.q_proj.bias": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.self_attn.k_proj.bias": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.self_attn.v_proj.bias": "model-32-of-35.safetensors",
+    "model.mimo_layers.3.self_attn.rotary_emb.inv_freq": "model-32-of-35.safetensors",
+    "model.mimo_layers.4.self_attn.q_proj.weight": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.self_attn.k_proj.weight": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.self_attn.v_proj.weight": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.self_attn.o_proj.weight": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.input_layernorm.weight": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.post_attention_layernorm.weight": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.mlp.gate_proj.weight": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.mlp.down_proj.weight": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.mlp.up_proj.weight": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.self_attn.q_proj.bias": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.self_attn.k_proj.bias": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.self_attn.v_proj.bias": "model-33-of-35.safetensors",
+    "model.mimo_layers.4.self_attn.rotary_emb.inv_freq": "model-33-of-35.safetensors",
+    "model.mimo_layers.5.self_attn.q_proj.weight": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.self_attn.k_proj.weight": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.self_attn.v_proj.weight": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.self_attn.o_proj.weight": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.input_layernorm.weight": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.post_attention_layernorm.weight": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.mlp.gate_proj.weight": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.mlp.down_proj.weight": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.mlp.up_proj.weight": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.self_attn.q_proj.bias": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.self_attn.k_proj.bias": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.self_attn.v_proj.bias": "model-34-of-35.safetensors",
+    "model.mimo_layers.5.self_attn.rotary_emb.inv_freq": "model-34-of-35.safetensors",
+    "model.vq_adaptor.layers.0.weight": "model-35-of-35.safetensors",
+    "model.vq_adaptor.layers.0.bias": "model-35-of-35.safetensors",
+    "model.vq_adaptor.layers.3.weight": "model-35-of-35.safetensors",
+    "model.vq_adaptor.layers.3.bias": "model-35-of-35.safetensors",
+    "model.vq_adaptor.layers.4.weight": "model-35-of-35.safetensors",
+    "model.vq_adaptor.layers.4.bias": "model-35-of-35.safetensors",
+    "model.embed_tokens.weight": "model-36-of-36.safetensors",
+    "model.norm.weight": "model-36-of-36.safetensors",
+    "lm_head.weight": "model-36-of-36.safetensors",
+    "mimo_output.weight": "model-36-of-36.safetensors",
+    "model.mimo_norm.weight": "model-36-of-36.safetensors"
+  }
+}

modeling_moonshot_kimia.py

@@ -0,0 +1,917 @@
+# coding=utf-8
+# Copyright 2025 The Moonshot AI Team, Qwen Team, and HuggingFace Inc. team. All rights reserved.
+#
+# The code is based on Qwen2.5-7B, but modified for KimiAudio.
+#
+# Licensing Information:
+# - Code derived from Qwen2.5-7B 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.
+"""PyTorch KimiAudio model."""
+
+from typing import List, Optional, Tuple, Union
+import torch
+import torch.utils.checkpoint
+from torch import nn
+
+import transformers
+from packaging import version
+
+assert version.parse(transformers.__version__) >= version.parse("4.34.1")
+
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.utils import (
+    logging,
+)
+from .configuration_moonshot_kimia import KimiAudioConfig
+import torch.nn.functional as F
+from transformers.models.qwen2.modeling_qwen2 import (
+    Qwen2RMSNorm,
+    Qwen2MLP,
+    Qwen2PreTrainedModel,
+)
+from transformers.models.qwen2.modeling_qwen2 import apply_rotary_pos_emb
+
+if version.parse(transformers.__version__) >= version.parse("4.35.0"):
+    from transformers.utils import is_flash_attn_2_available as is_flash_attn_available
+else:
+    from transformers.utils import is_flash_attn_available
+
+if is_flash_attn_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+else:
+    raise RuntimeError("flash attention must be installed")
+
+
+logger = logging.get_logger(__name__)
+
+
+def _get_unpad_data(padding_mask):
+    seqlens_in_batch = padding_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(padding_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,
+    )
+
+
+def _upad_input(query_layer, key_layer, value_layer, padding_mask, query_length):
+    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(padding_mask)
+    batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+    num_heads = query_layer.shape[2]
+
+    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, 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.
+        padding_mask = padding_mask[:, -query_length:]
+        query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
+            query_layer, padding_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),
+    )
+
+
+# Copied from transformers.models.bart.modeling_bart._make_causal_mask
+def _make_causal_mask(
+    input_ids_shape: torch.Size,
+    dtype: torch.dtype,
+    device: torch.device,
+    past_key_values_length: int = 0,
+):
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat(
+            [
+                torch.zeros(
+                    tgt_len, past_key_values_length, dtype=dtype, device=device
+                ),
+                mask,
+            ],
+            dim=-1,
+        )
+    return mask[None, None, :, :].expand(
+        bsz, 1, tgt_len, tgt_len + past_key_values_length
+    )
+
+
+# Copied from transformers.models.bart.modeling_bart._expand_mask
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+    bsz, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(
+        inverted_mask.to(torch.bool), torch.finfo(dtype).min
+    )
+
+
+class RotaryEmbedding(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(),
+        )
+
+    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.einsum("i,j->ij", 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)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if 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),
+        )
+
+
+class MoonshotAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: KimiAudioConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=True
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=False
+        )
+
+        self._init_rope()
+
+    def _init_rope(self):
+
+        self.rotary_emb = RotaryEmbedding(
+            self.head_dim,
+            max_position_embeddings=self.max_position_embeddings,
+            base=self.rope_theta,
+        )
+
+    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: bool = False,
+        use_cache: bool = False,
+        padding_mask: Optional[torch.LongTensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # LlamaFlashAttention2 attention does not support output_attentions
+
+        output_attentions = False
+
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dime x hidden_dim
+        # therefore we just need to keep the original shape
+        query_states = query_states.view(
+            bsz, q_len, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        value_states = value_states.view(
+            bsz, q_len, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        cos = cos[position_ids]
+        sin = sin[position_ids]
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin, position_ids
+        )
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        # TODO: llama does not have dropout in the config??
+        # It is recommended to use dropout with FA according to the docs
+        # when training.
+        dropout_rate = 0.0  # if not self.training else self.attn_dropout
+
+        # 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 float16 just to be sure everything works as expected.
+        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+        # in fp32. (LlamaRMSNorm handles it correctly)
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            logger.warning_once(
+                "The input hidden states seems to be silently casted in float32, this might be related to"
+                " the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                " float16."
+            )
+
+            query_states = query_states.to(torch.float16)
+            key_states = key_states.to(torch.float16)
+            value_states = value_states.to(torch.float16)
+
+        attn_output = self._flash_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            padding_mask,
+            q_len,
+            dropout=dropout_rate,
+        )
+
+        if input_dtype == torch.float32:
+            attn_output = attn_output.to(torch.float32)
+
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).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,
+        padding_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
+            padding_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)
+        """
+        # Contains at least one padding token in the sequence
+        if padding_mask is not None:
+            batch_size = query_states.shape[0]
+            (
+                query_states,
+                key_states,
+                value_states,
+                indices_q,
+                cu_seq_lens,
+                max_seq_lens,
+            ) = _upad_input(
+                query_states, key_states, value_states, padding_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=True,
+            )
+
+            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=True,
+            )
+
+        return attn_output
+
+
+class MoonshotDecoderLayer(nn.Module):
+    def __init__(self, config: KimiAudioConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.config = config
+
+        logger.warning_once("using normal flash attention")
+        self.self_attn = MoonshotAttention(config=config)
+
+        self.mlp = Qwen2MLP(config)
+        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen2RMSNorm(
+            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,
+        padding_mask: Optional[torch.LongTensor] = None,
+    ) -> 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, 1, tgt_len, src_len)` where padding elements are indicated by very large negative 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.
+            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
+        """
+
+        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,
+            padding_mask=padding_mask,
+        )
+        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
+
+
+class VQAdaptor(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(config.kimia_adaptor_input_dim, config.hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Dropout(0.0),
+            nn.Linear(config.hidden_size, config.hidden_size, bias=True),
+            nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps, bias=True),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class MoonshotKimiaModel(Qwen2PreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`QwenDecoderLayer`]
+
+    Args:
+        config: KimiAudioConfig
+    """
+
+    config_class = KimiAudioConfig
+
+    def __init__(self, config: KimiAudioConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.kimia_mimo_transformer_from_layer_index = (
+            config.kimia_mimo_transformer_from_layer_index
+        )
+
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [MoonshotDecoderLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # extra 1B audio transformers
+        self.mimo_layers = nn.ModuleList(
+            [MoonshotDecoderLayer(config) for _ in range(config.kimia_mimo_layers)]
+        )
+        self.mimo_norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.use_whisper_feature = config.use_whisper_feature
+        if self.use_whisper_feature:
+            self.vq_adaptor = VQAdaptor(config)
+        self.kimia_media_begin = config.kimia_media_begin
+        self.kimia_media_end = config.kimia_media_end
+
+        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
+
+    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
+    def _prepare_decoder_attention_mask(
+        self, attention_mask, input_shape, inputs_embeds, past_key_values_length
+    ):
+        # create causal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(
+                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+            ).to(inputs_embeds.device)
+            combined_attention_mask = (
+                expanded_attn_mask
+                if combined_attention_mask is None
+                else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        text_input_ids: torch.LongTensor = None,
+        whisper_input_feature: Optional[torch.FloatTensor] = None,
+        is_continuous_mask: Optional[torch.Tensor] = 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
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_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:
+            # shape: batch, seq_len, hidden_size
+            input_ids = input_ids.to(torch.cuda.current_device())
+            text_input_ids = text_input_ids.to(torch.cuda.current_device())
+            audio_emb = self.embed_tokens(input_ids)
+            if self.use_whisper_feature and whisper_input_feature is not None:
+                if not isinstance(whisper_input_feature, list):
+                    whisper_input_feature = whisper_input_feature.squeeze(0)
+                    whisper_input_feature = [whisper_input_feature]
+
+                media_start_idx = (input_ids == self.kimia_media_begin).nonzero()
+                media_end_idx = (input_ids == self.kimia_media_end).nonzero()
+                # shape: batch, seq_len, hidden_size
+                whisper_input_dim = whisper_input_feature[0].shape[-1]
+                whisper_dtype = whisper_input_feature[0].dtype
+                expanded_whisper = (
+                    torch.zeros(audio_emb.shape[1], whisper_input_dim)
+                    .to(torch.cuda.current_device())
+                    .to(whisper_dtype)
+                )
+                for (seg_idx, start_idx), (_, end_idx) in zip(
+                    media_start_idx, media_end_idx
+                ):
+                    # assert whisper_emb.shape[1] == end_idx - (start_idx + 1)
+
+                    feat_len = end_idx - (start_idx + 1)
+                    whisper_input_feature_i = whisper_input_feature[seg_idx].squeeze(0)
+                    assert feat_len == is_continuous_mask[seg_idx].sum()
+                    expanded_whisper[start_idx + 1 : end_idx, :] = (
+                        whisper_input_feature_i[:feat_len, :]
+                    )
+
+                expanded_whisper = expanded_whisper.unsqueeze(0)
+                whisper_emb = self.vq_adaptor(
+                    expanded_whisper.transpose(0, 1)
+                ).transpose(0, 1)
+                is_continuous_mask = is_continuous_mask.to(torch.cuda.current_device())
+                whisper_emb = whisper_emb.to(torch.cuda.current_device())
+                whisper_emb = whisper_emb * is_continuous_mask[:, :, None]
+
+                encoder_input_addwith_discrete_token = (
+                    audio_emb + whisper_emb
+                ) * torch.sqrt(
+                    torch.tensor(
+                        2.0, dtype=whisper_emb.dtype, device=torch.cuda.current_device()
+                    )
+                )
+                audio_emb = (
+                    audio_emb * (~is_continuous_mask[:, :, None])
+                    + encoder_input_addwith_discrete_token
+                    * is_continuous_mask[:, :, None]
+                )
+            if text_input_ids is not None and text_input_ids.sum() != 0:
+                inputs_embeds = audio_emb + self.embed_tokens(text_input_ids)
+            else:
+                inputs_embeds = audio_emb
+        # embed positions
+        # TODO kill attention_mask for prefill
+        padding_mask = attention_mask
+
+        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 = () if use_cache else None
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            past_key_value = (
+                past_key_values[idx] if past_key_values is not None else None
+            )
+            layer_outputs = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                padding_mask=padding_mask,
+            )
+
+            hidden_states = layer_outputs[0]
+            if idx == self.kimia_mimo_transformer_from_layer_index:
+                mimo_hidden_states = hidden_states.clone()
+
+            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)
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        # apply audio transformer layers
+        for idx, decoder_layer in enumerate(self.mimo_layers):
+            if output_hidden_states:
+                all_hidden_states += (mimo_hidden_states,)
+
+            past_key_value = (
+                past_key_values[idx + len(self.layers)]
+                if past_key_values is not None
+                else None
+            )
+            layer_outputs = decoder_layer(
+                mimo_hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                padding_mask=padding_mask,
+            )
+
+            mimo_hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+        mimo_hidden_states = self.mimo_norm(mimo_hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (mimo_hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    mimo_hidden_states,
+                    next_cache,
+                    all_hidden_states,
+                    all_hidden_states,
+                    all_self_attns,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPast(
+            last_hidden_state=(hidden_states, mimo_hidden_states),
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class MoonshotKimiaForCausalLM(Qwen2PreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight", "mimo_output.weight"]
+    config_class = KimiAudioConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MoonshotKimiaModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.mimo_output = 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
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        text_input_ids: torch.LongTensor = None,
+        whisper_input_feature: Optional[torch.FloatTensor] = None,
+        is_continuous_mask: Optional[torch.Tensor] = 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,
+        generation_mode: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+        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,
+            text_input_ids=text_input_ids,
+            whisper_input_feature=whisper_input_feature,
+            is_continuous_mask=is_continuous_mask,
+            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,
+        )
+        if return_dict:
+            hidden_states, mimo_hidden_states = (
+                outputs.last_hidden_state[0],
+                outputs.last_hidden_state[1],
+            )
+        else:
+            hidden_states, mimo_hidden_states = outputs[0], outputs[1]
+
+        audio_logits = self.lm_head(hidden_states)
+        text_logits = self.mimo_output(mimo_hidden_states)
+
+        if not return_dict:
+            output = (text_logits, audio_logits) + outputs[2:]
+            return output
+        return CausalLMOutputWithPast(
+            loss=None,
+            logits=(text_logits, audio_logits),
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

special_tokens_map.json

@@ -0,0 +1,425 @@
+{
+  "additional_special_tokens": [
+      "<|im_msg_end|>",
+      "<|im_user_msg_start|>",
+      "<|im_assistant_msg_start|>",
+      "<|reserved_token_0|>",
+      "<|reserved_token_1|>",
+      "<|reserved_token_2|>",
+      "<|reserved_token_3|>",
+      "[EOT]",
+      "<|reserved_token_4|>",
+      "<|reserved_token_5|>",
+      "<|reserved_token_6|>",
+      "<|reserved_token_7|>",
+      "<|reserved_token_8|>",
+      "<|reserved_token_9|>",
+      "<|reserved_token_10|>",
+      "<|reserved_token_11|>",
+      "<|im_media_begin|>",
+      "<|reserved_token_12|>",
+      "<|im_media_end|>",
+      "<|reserved_token_13|>",
+      "<|reserved_token_14|>",
+      "<|im_kimia_text_blank|>",
+      "<|im_kimia_text_eos|>",
+      "<|reserved_token_15|>",
+      "<|reserved_token_16|>",
+      "<|im_kimia_user_msg_start|>",
+      "<|im_kimia_assistant_msg_start|>",
+      "<|reserved_token_17|>",
+      "<|reserved_token_18|>",
+      "<|reserved_token_19|>",
+      "<|im_kimia_speech_ct_id|>",
+      "<|im_kimia_speech_ctd_id|>",
+      "<|reserved_token_20|>",
+      "<|reserved_token_21|>",
+      "<|reserved_token_22|>",
+      "<|reserved_token_23|>",
+      "<|reserved_token_24|>",
+      "<|reserved_token_25|>",
+      "<|reserved_token_26|>",
+      "<|reserved_token_27|>",
+      "<|reserved_token_28|>",
+      "<|reserved_token_29|>",
+      "<|reserved_token_30|>",
+      "<|reserved_token_31|>",
+      "<|reserved_token_32|>",
+      "<|reserved_token_33|>",
+      "<|reserved_token_34|>",
+      "<|reserved_token_35|>",
+      "<|reserved_token_36|>",
+      "<|reserved_token_37|>",
+      "<|reserved_token_38|>",
+      "<|reserved_token_39|>",
+      "<|reserved_token_40|>",
+      "<|reserved_token_41|>",
+      "<|reserved_token_42|>",
+      "<|reserved_token_43|>",
+      "<|reserved_token_44|>",
+      "<|reserved_token_45|>",
+      "<|reserved_token_46|>",
+      "<|reserved_token_47|>",
+      "<|reserved_token_48|>",
+      "<|reserved_token_49|>",
+      "<|reserved_token_50|>",
+      "<|reserved_token_51|>",
+      "<|reserved_token_52|>",
+      "<|reserved_token_53|>",
+      "<|reserved_token_54|>",
+      "<|reserved_token_55|>",
+      "<|reserved_token_56|>",
+      "<|reserved_token_57|>",
+      "<|reserved_token_58|>",
+      "<|reserved_token_59|>",
+      "<|reserved_token_60|>",
+      "<|reserved_token_61|>",
+      "<|reserved_token_62|>",
+      "<|reserved_token_63|>",
+      "<|reserved_token_64|>",
+      "<|reserved_token_65|>",
+      "<|reserved_token_66|>",
+      "<|reserved_token_67|>",
+      "<|reserved_token_68|>",
+      "<|reserved_token_69|>",
+      "<|reserved_token_70|>",
+      "<|reserved_token_71|>",
+      "<|reserved_token_72|>",
+      "<|reserved_token_73|>",
+      "<|reserved_token_74|>",
+      "<|reserved_token_75|>",
+      "<|reserved_token_76|>",
+      "<|reserved_token_77|>",
+      "<|reserved_token_78|>",
+      "<|reserved_token_79|>",
+      "<|reserved_token_80|>",
+      "<|reserved_token_81|>",
+      "<|reserved_token_82|>",
+      "<|reserved_token_83|>",
+      "<|reserved_token_84|>",
+      "<|reserved_token_85|>",
+      "<|reserved_token_86|>",
+      "<|reserved_token_87|>",
+      "<|reserved_token_88|>",
+      "<|reserved_token_89|>",
+      "<|reserved_token_90|>",
+      "<|reserved_token_91|>",
+      "<|reserved_token_92|>",
+      "<|reserved_token_93|>",
+      "<|reserved_token_94|>",
+      "<|reserved_token_95|>",
+      "<|reserved_token_96|>",
+      "<|reserved_token_97|>",
+      "<|reserved_token_98|>",
+      "<|reserved_token_99|>",
+      "<|reserved_token_100|>",
+      "<|reserved_token_101|>",
+      "<|reserved_token_102|>",
+      "<|reserved_token_103|>",
+      "<|reserved_token_104|>",
+      "<|reserved_token_105|>",
+      "<|reserved_token_106|>",
+      "<|reserved_token_107|>",
+      "<|reserved_token_108|>",
+      "<|reserved_token_109|>",
+      "<|reserved_token_110|>",
+      "<|reserved_token_111|>",
+      "<|reserved_token_112|>",
+      "<|reserved_token_113|>",
+      "<|reserved_token_114|>",
+      "<|reserved_token_115|>",
+      "<|reserved_token_116|>",
+      "<|reserved_token_117|>",
+      "<|reserved_token_118|>",
+      "<|reserved_token_119|>",
+      "<|reserved_token_120|>",
+      "<|reserved_token_121|>",
+      "<|reserved_token_122|>",
+      "<|reserved_token_123|>",
+      "<|reserved_token_124|>",
+      "<|reserved_token_125|>",
+      "<|reserved_token_126|>",
+      "<|reserved_token_127|>",
+      "<|reserved_token_128|>",
+      "<|reserved_token_129|>",
+      "<|reserved_token_130|>",
+      "<|reserved_token_131|>",
+      "<|reserved_token_132|>",
+      "<|reserved_token_133|>",
+      "<|reserved_token_134|>",
+      "<|reserved_token_135|>",
+      "<|reserved_token_136|>",
+      "<|reserved_token_137|>",
+      "<|reserved_token_138|>",
+      "<|reserved_token_139|>",
+      "<|reserved_token_140|>",
+      "<|reserved_token_141|>",
+      "<|reserved_token_142|>",
+      "<|reserved_token_143|>",
+      "<|reserved_token_144|>",
+      "<|reserved_token_145|>",
+      "<|reserved_token_146|>",
+      "<|reserved_token_147|>",
+      "<|reserved_token_148|>",
+      "<|reserved_token_149|>",
+      "<|reserved_token_150|>",
+      "<|reserved_token_151|>",
+      "<|reserved_token_152|>",
+      "<|reserved_token_153|>",
+      "<|reserved_token_154|>",
+      "<|reserved_token_155|>",
+      "<|reserved_token_156|>",
+      "<|reserved_token_157|>",
+      "<|reserved_token_158|>",
+      "<|reserved_token_159|>",
+      "<|reserved_token_160|>",
+      "<|reserved_token_161|>",
+      "<|reserved_token_162|>",
+      "<|reserved_token_163|>",
+      "<|reserved_token_164|>",
+      "<|reserved_token_165|>",
+      "<|reserved_token_166|>",
+      "<|reserved_token_167|>",
+      "<|reserved_token_168|>",
+      "<|reserved_token_169|>",
+      "<|reserved_token_170|>",
+      "<|reserved_token_171|>",
+      "<|reserved_token_172|>",
+      "<|reserved_token_173|>",
+      "<|reserved_token_174|>",
+      "<|reserved_token_175|>",
+      "<|reserved_token_176|>",
+      "<|reserved_token_177|>",
+      "<|reserved_token_178|>",
+      "<|reserved_token_179|>",
+      "<|reserved_token_180|>",
+      "<|reserved_token_181|>",
+      "<|reserved_token_182|>",
+      "<|reserved_token_183|>",
+      "<|reserved_token_184|>",
+      "<|reserved_token_185|>",
+      "<|reserved_token_186|>",
+      "<|reserved_token_187|>",
+      "<|reserved_token_188|>",
+      "<|reserved_token_189|>",
+      "<|reserved_token_190|>",
+      "<|reserved_token_191|>",
+      "<|reserved_token_192|>",
+      "<|reserved_token_193|>",
+      "<|reserved_token_194|>",
+      "<|reserved_token_195|>",
+      "<|reserved_token_196|>",
+      "<|reserved_token_197|>",
+      "<|reserved_token_198|>",
+      "<|reserved_token_199|>",
+      "<|reserved_token_200|>",
+      "<|reserved_token_201|>",
+      "<|reserved_token_202|>",
+      "<|reserved_token_203|>",
+      "<|reserved_token_204|>",
+      "<|reserved_token_205|>",
+      "<|reserved_token_206|>",
+      "<|reserved_token_207|>",
+      "<|reserved_token_208|>",
+      "<|reserved_token_209|>",
+      "<|reserved_token_210|>",
+      "<|reserved_token_211|>",
+      "<|reserved_token_212|>",
+      "<|reserved_token_213|>",
+      "<|reserved_token_214|>",
+      "<|reserved_token_215|>",
+      "<|reserved_token_216|>",
+      "<|reserved_token_217|>",
+      "<|reserved_token_218|>",
+      "<|reserved_token_219|>",
+      "<|reserved_token_220|>",
+      "<|reserved_token_221|>",
+      "<|reserved_token_222|>",
+      "<|reserved_token_223|>",
+      "<|reserved_token_224|>",
+      "<|reserved_token_225|>",
+      "<|reserved_token_226|>",
+      "<|reserved_token_227|>",
+      "<|reserved_token_228|>",
+      "<|reserved_token_229|>",
+      "<|reserved_token_230|>",
+      "<|reserved_token_231|>",
+      "<|reserved_token_232|>",
+      "<|reserved_token_233|>",
+      "<|reserved_token_234|>",
+      "<|reserved_token_235|>",
+      "<|reserved_token_236|>",
+      "<|reserved_token_237|>",
+      "<|reserved_token_238|>",
+      "<|reserved_token_239|>",
+      "<|reserved_token_240|>",
+      "<|reserved_token_241|>",
+      "<|reserved_token_242|>",
+      "<|reserved_token_243|>",
+      "<|reserved_token_244|>",
+      "<|reserved_token_245|>",
+      "<|reserved_token_246|>",
+      "<|reserved_token_247|>",
+      "<|reserved_token_248|>",
+      "<|reserved_token_249|>",
+      "<|reserved_token_250|>",
+      "<|reserved_token_251|>",
+      "<|reserved_token_252|>",
+      "<|reserved_token_253|>",
+      "<|reserved_token_254|>",
+      "<|reserved_token_255|>",
+      "<|reserved_token_256|>",
+      "<|reserved_token_257|>",
+      "<|reserved_token_258|>",
+      "<|reserved_token_259|>",
+      "<|reserved_token_260|>",
+      "<|reserved_token_261|>",
+      "<|reserved_token_262|>",
+      "<|reserved_token_263|>",
+      "<|reserved_token_264|>",
+      "<|reserved_token_265|>",
+      "<|reserved_token_266|>",
+      "<|reserved_token_267|>",
+      "<|reserved_token_268|>",
+      "<|reserved_token_269|>",
+      "<|reserved_token_270|>",
+      "<|reserved_token_271|>",
+      "<|reserved_token_272|>",
+      "<|reserved_token_273|>",
+      "<|reserved_token_274|>",
+      "<|reserved_token_275|>",
+      "<|reserved_token_276|>",
+      "<|reserved_token_277|>",
+      "<|reserved_token_278|>",
+      "<|reserved_token_279|>",
+      "<|reserved_token_280|>",
+      "<|reserved_token_281|>",
+      "<|reserved_token_282|>",
+      "<|reserved_token_283|>",
+      "<|reserved_token_284|>",
+      "<|reserved_token_285|>",
+      "<|reserved_token_286|>",
+      "<|reserved_token_287|>",
+      "<|reserved_token_288|>",
+      "<|reserved_token_289|>",
+      "<|reserved_token_290|>",
+      "<|reserved_token_291|>",
+      "<|reserved_token_292|>",
+      "<|reserved_token_293|>",
+      "<|reserved_token_294|>",
+      "<|reserved_token_295|>",
+      "<|reserved_token_296|>",
+      "<|reserved_token_297|>",
+      "<|reserved_token_298|>",
+      "<|reserved_token_299|>",
+      "<|reserved_token_300|>",
+      "<|reserved_token_301|>",
+      "<|reserved_token_302|>",
+      "<|reserved_token_303|>",
+      "<|reserved_token_304|>",
+      "<|reserved_token_305|>",
+      "<|reserved_token_306|>",
+      "<|reserved_token_307|>",
+      "<|reserved_token_308|>",
+      "<|reserved_token_309|>",
+      "<|reserved_token_310|>",
+      "<|reserved_token_311|>",
+      "<|reserved_token_312|>",
+      "<|reserved_token_313|>",
+      "<|reserved_token_314|>",
+      "<|reserved_token_315|>",
+      "<|reserved_token_316|>",
+      "<|reserved_token_317|>",
+      "<|reserved_token_318|>",
+      "<|reserved_token_319|>",
+      "<|reserved_token_320|>",
+      "<|reserved_token_321|>",
+      "<|reserved_token_322|>",
+      "<|reserved_token_323|>",
+      "<|reserved_token_324|>",
+      "<|reserved_token_325|>",
+      "<|reserved_token_326|>",
+      "<|reserved_token_327|>",
+      "<|reserved_token_328|>",
+      "<|reserved_token_329|>",
+      "<|reserved_token_330|>",
+      "<|reserved_token_331|>",
+      "<|reserved_token_332|>",
+      "<|reserved_token_333|>",
+      "<|reserved_token_334|>",
+      "<|reserved_token_335|>",
+      "<|reserved_token_336|>",
+      "<|reserved_token_337|>",
+      "<|reserved_token_338|>",
+      "<|reserved_token_339|>",
+      "<|reserved_token_340|>",
+      "<|reserved_token_341|>",
+      "<|reserved_token_342|>",
+      "<|reserved_token_343|>",
+      "<|reserved_token_344|>",
+      "<|reserved_token_345|>",
+      "<|reserved_token_346|>",
+      "<|reserved_token_347|>",
+      "<|reserved_token_348|>",
+      "<|reserved_token_349|>",
+      "<|reserved_token_350|>",
+      "<|reserved_token_351|>",
+      "<|reserved_token_352|>",
+      "<|reserved_token_353|>",
+      "<|reserved_token_354|>",
+      "<|reserved_token_355|>",
+      "<|reserved_token_356|>",
+      "<|reserved_token_357|>",
+      "<|reserved_token_358|>",
+      "<|reserved_token_359|>",
+      "<|reserved_token_360|>",
+      "<|reserved_token_361|>",
+      "<|reserved_token_362|>",
+      "<|reserved_token_363|>",
+      "<|reserved_token_364|>",
+      "<|reserved_token_365|>",
+      "<|reserved_token_366|>",
+      "<|reserved_token_367|>",
+      "<|reserved_token_368|>",
+      "<|reserved_token_369|>",
+      "<|reserved_token_370|>",
+      "<|reserved_token_371|>",
+      "<|reserved_token_372|>",
+      "<|reserved_token_373|>",
+      "<|reserved_token_374|>",
+      "<|reserved_token_375|>",
+      "<|reserved_token_376|>",
+      "<|reserved_token_377|>",
+      "<|reserved_token_378|>",
+      "<|reserved_token_379|>",
+      "<|reserved_token_380|>",
+      "<|reserved_token_381|>",
+      "<|reserved_token_382|>",
+      "<|reserved_token_383|>",
+      "<|reserved_token_384|>",
+      "<|reserved_token_385|>",
+      "<|reserved_token_386|>",
+      "<|reserved_token_387|>",
+      "<|reserved_token_388|>",
+      "<|reserved_token_389|>",
+      "<|reserved_token_390|>",
+      "<|reserved_token_391|>",
+      "<|reserved_token_392|>",
+      "<|reserved_token_393|>",
+      "<|reserved_token_394|>",
+      "<|reserved_token_395|>",
+      "<|reserved_token_396|>",
+      "<|reserved_token_397|>",
+      "<|reserved_token_398|>",
+      "<|reserved_token_399|>",
+      "<|reserved_token_400|>",
+      "<|reserved_token_401|>",
+      "<|reserved_token_402|>",
+      "<|reserved_token_403|>",
+      "<|reserved_token_404|>"
+  ],
+  "bos_token": "[BOS]",
+  "eos_token": "[EOS]",
+  "pad_token": "<|reserved_token_406|>",
+  "unk_token": "<|reserved_token_405|>"
+}

tiktoken.model

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

tokenization_kimia.py

@@ -0,0 +1,335 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""Megatron tokenizers."""
+from transformers.tokenization_utils import PreTrainedTokenizer
+from typing import Union
+from typing import (
+    AbstractSet,
+    cast,
+    Collection,
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Sequence,
+    Union,
+    Optional,
+)
+from tiktoken.load import load_tiktoken_bpe
+import tiktoken
+from pathlib import Path
+import os
+import logging
+from tokenizers import AddedToken
+
+logger = logging.getLogger(__name__)
+VOCAB_FILES_NAMES = {"vocab_file": "tiktoken.model"}
+
+
+class TikTokenTokenizer(PreTrainedTokenizer):
+    """
+    Tokenizing and encoding/decoding text using the Tiktoken tokenizer.
+    """
+
+    special_tokens: Dict[str, int]
+
+    num_reserved_special_tokens = 293 + 128
+
+    pat_str = "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
+
+    vocab_files_names = VOCAB_FILES_NAMES
+
+    def __init__(
+        self,
+        vocab_file,
+        bos_token: Union[str, AddedToken] = "[BOS]",
+        eos_token: Union[str, AddedToken] = "[EOS]",
+        unk_token: Union[str, AddedToken] = "[UNK]",
+        pad_token: Union[str, AddedToken] = "[PAD]",
+        additional_special_tokens: Optional[List[str]] = None,
+        added_tokens_decoder: Optional[dict] = None,
+        **kwargs,
+    ):
+        """
+        Initializes the Tokenizer with a Tiktoken model.
+
+        Args:
+            model_path (str): The path to the Tiktoken model file.
+        """
+        assert os.path.isfile(vocab_file), vocab_file
+
+        mergeable_ranks = load_tiktoken_bpe(vocab_file)
+        num_base_tokens = len(mergeable_ranks)
+
+        used_special_tokens = [
+            "[BOS]",
+            "[EOS]",
+            "<|im_msg_end|>",  # 0
+            "<|im_user_msg_start|>",  # 1
+            "<|im_assistant_msg_start|>",  # 2
+            "<|reserved_token_0|>",  # 3
+            "<|reserved_token_1|>",
+            "<|reserved_token_2|>",
+            "<|reserved_token_3|>",  # 4
+            "[EOT]",
+            "<|reserved_token_4|>",  # 5
+            "<|reserved_token_5|>",  # 6
+            "<|reserved_token_6|>",  # 7
+            "<|reserved_token_7|>",  # 8
+            "<|reserved_token_8|>",  # 9
+            "<|reserved_token_9|>",  # 10
+            "<|reserved_token_10|>",  # 11
+            "<|reserved_token_11|>",  # 12
+            "<|im_media_begin|>",  # 13
+            "<|reserved_token_12|>",  # 14
+            "<|im_media_end|>",  # 15
+            "<|reserved_token_13|>",  # 16
+            "<|reserved_token_14|>",  # 17
+            "<|im_kimia_text_blank|>",  # 18
+            "<|im_kimia_text_eos|>",  # 19
+            "<|reserved_token_15|>",  # 20
+            "<|reserved_token_16|>",  # 21
+            "<|im_kimia_user_msg_start|>",  # 22
+            "<|im_kimia_assistant_msg_start|>",  # 23
+            "<|reserved_token_17|>",  # 24
+            "<|reserved_token_18|>",  # 25
+            "<|reserved_token_19|>",  # 26
+            "<|im_kimia_speech_ct_id|>",  # 27
+            "<|im_kimia_speech_ctd_id|>",  # 28
+        ]
+        autoset_special_tokens = [
+            f"<|reserved_token_{i}|>"
+            for i in range(
+                20, self.num_reserved_special_tokens - len(used_special_tokens) + 20
+            )
+        ]
+        special_tokens = used_special_tokens + autoset_special_tokens
+        self.special_tokens = {
+            token: num_base_tokens + i for i, token in enumerate(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_token = "[BOS]"
+        self.bos_id: int = self.special_tokens["[BOS]"]
+        self.eos_token = "[EOS]"
+        self.eos_id: int = self.special_tokens["[EOS]"]
+
+        # use last speical token as pad token, the last - 1 is unk_token
+        self.pad_token: str = special_tokens[-1]
+        self.pad_id: int = self.special_tokens[self.pad_token]
+
+        self.unk_token: str = special_tokens[-2]
+        self.unk_id: int = self.special_tokens[self.pad_token]
+
+        self.stop_tokens = {
+            self.special_tokens["[EOS]"],
+            self.special_tokens["[EOT]"],
+        }
+
+        logger.info(
+            f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
+        )
+
+    def encode(
+        self,
+        s: str,
+        *,
+        bos: bool,
+        eos: bool,
+        allowed_special: Union[Literal["all"], AbstractSet[str]] = set(),
+        disallowed_special: Union[Literal["all"], Collection[str]] = (),
+    ) -> List[int]:
+        """
+        Encodes a string into a list of token IDs.
+
+        Args:
+            s (str): The input string to be encoded.
+            bos (bool): Whether to prepend the beginning-of-sequence token.
+            eos (bool): Whether to append the end-of-sequence token.
+            allowed_tokens ("all"|set[str]): allowed special tokens in string
+            disallowed_tokens ("all"|set[str]): special tokens that raise an error when in string
+
+        Returns:
+            list[int]: A list of token IDs.
+
+        By default, setting disallowed_special=() encodes a string by ignoring
+        special tokens. Specifically:
+        - Setting `disallowed_special` to () will cause all text corresponding
+          to special tokens to be encoded as natural text (insteading of raising
+          an error).
+        - Setting `allowed_special` to "all" will treat all text corresponding
+          to special tokens to be encoded as special tokens.
+        """
+        assert type(s) 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
+
+        substrs = (
+            substr
+            for i in range(0, len(s), TIKTOKEN_MAX_ENCODE_CHARS)
+            for substr in self._split_whitespaces_or_nonwhitespaces(
+                s[i : i + TIKTOKEN_MAX_ENCODE_CHARS], MAX_NO_WHITESPACES_CHARS
+            )
+        )
+        t: List[int] = []
+        for substr in substrs:
+            t.extend(
+                self.model.encode(
+                    substr,
+                    allowed_special=allowed_special,
+                    disallowed_special=disallowed_special,
+                )
+            )
+        if bos:
+            t.insert(0, self.bos_id)
+        if eos:
+            t.append(self.eos_id)
+        return t
+
+    def decode(self, t: Sequence[int]) -> str:
+        """
+        Decodes a list of token IDs into a string.
+
+        Args:
+            t (List[int]): The list of token IDs to be decoded.
+
+        Returns:
+            str: The decoded string.
+        """
+        # Typecast is safe here. Tiktoken doesn't do anything list-related with the sequence.
+        return self.model.decode(cast(List[int], t))
+
+    @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:]
+
+    """ ----- Below are the abstract methods required by megatron ----- """
+
+    @property
+    def vocab_size(self):
+        return self.n_words
+
+    @property
+    def vocab(self):
+        if hasattr(self, "str_vocab"):
+            return self.str_vocab
+        self.str_vocab = {}
+
+        # convert mergeable_ranks from bytes to string
+        utf8_num, unicode_num = 0, 0
+        for byte_key, index in self.model._mergeable_ranks.items():
+            try:
+                str_key = byte_key.decode("utf-8")
+                utf8_num += 1
+            except UnicodeDecodeError:
+                # use backslashreplace so we can get num vocab different tokens
+                # see: https://docs.python.org/3/howto/unicode.html
+                # this vocab is only used for offline processing, so this is fine
+                str_key = byte_key.decode("utf-8", "backslashreplace") + "_unicode_"
+                unicode_num += 1
+
+            self.str_vocab[str_key] = index
+        logger.info(f"num utf8: {utf8_num}, num unicode: {unicode_num}")
+
+        # add all special tokens to the dictionary
+        self.str_vocab.update(self.model._special_tokens)
+
+        assert len(self.str_vocab) == self.vocab_size
+        return self.str_vocab
+
+    @property
+    def inv_vocab(self):
+        return {v: k for k, v in self.vocab.items()}
+
+    def tokenize(self, text, eos=True):
+        # BOS: always add bos token
+        # EOS:
+        #    Most cases should be true when we are tokenizing a full sequence
+        #    Only setting to false when we are running a inference
+        return self.encode(text, bos=True, eos=eos)
+
+    def detokenize(self, tokens):
+        # convert tensor to list if needed...
+        if not isinstance(tokens, list):
+            tokens = tokens.tolist()
+        return self.decode(tokens)
+
+    @property
+    def eod(self):
+        return self.eos_id
+
+    def bod(self):
+        return self.bos_id
+
+    @property
+    def msk_start_id(self):
+        return self.msk_start
+
+    @property
+    def msk_end_id(self):
+        return self.msk_end
+
+    def _get_index_2_bytes(self):
+        if hasattr(self, "index_2_bytes"):
+            return self.index_2_bytes
+
+        # use array rather than dict for faster access
+        self.index_2_bytes = [0] * self.model.n_vocab
+        for byte_key, index in self.model._mergeable_ranks.items():
+            self.index_2_bytes[index] = len(byte_key)
+
+        for _, index in self.model._special_tokens.items():
+            # in total we have 256 special tokens, 2^8 = 256
+            # so the num of bytes of each token is only 1
+            self.index_2_bytes[index] = 1
+
+        return self.index_2_bytes
+
+    def get_array_bytes(self, array):
+        index_2_bytes = self._get_index_2_bytes()
+        return sum(index_2_bytes[i] for i in array)
+
+    @property
+    def eos_token_id(self):
+        return self.eos_id
+
+    @property
+    def pad_token_id(self):
+        return self.pad_id