Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,202 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright (c) 2025 Baidu, Inc. All Rights Reserved.
+
+   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.
+

README.md

@@ -1,22 +1,223 @@
 ---
 license: apache-2.0
+language:
+- en
+- zh
+pipeline_tag: text-generation
+tags:
+- ERNIE4.5
+library_name: transformers
 ---
-EXL3 quants of [ERNIE-4.5-300B-A47B-PT](https://huggingface.co/baidu/ERNIE-4.5-300B-A47B-PT)
-
-[2.00 bits per weight](https://huggingface.co/turboderp/ERNIE-4.5-300B-A47B-PT-exl3/tree/2.0bpw)    
-[2.10 bits per weight](https://huggingface.co/turboderp/ERNIE-4.5-300B-A47B-PT-exl3/tree/2.1bpw) (optimized)   
-[2.25 bits per weight](https://huggingface.co/turboderp/ERNIE-4.5-300B-A47B-PT-exl3/tree/2.25bpw) (optimized)     
-[2.50 bits per weight](https://huggingface.co/turboderp/ERNIE-4.5-300B-A47B-PT-exl3/tree/2.5bpw) (optimized)     
-[3.00 bits per weight](https://huggingface.co/turboderp/ERNIE-4.5-300B-A47B-PT-exl3/tree/3.0bpw)     
-[4.00 bits per weight](https://huggingface.co/turboderp/ERNIE-4.5-300B-A47B-PT-exl3/tree/4.0bpw)    
-
-
-Quant    | Weights/VRAM | Perplexity | KL-div          
----------|--------------|------------|-----------
-2.00 bpw | 70.2 GB      | 7.4131     | 0.5283
-2.10 bpw | 73.4 GB      | 6.7507     | 0.2202
-2.25 bpw | 78.6 GB      | 6.5576     | 0.2074
-2.50 bpw | 87.8 GB      | 6.3504     | 0.1899
-3.00 bpw | 104.9 GB     | 5.8913     | 0.1547
-4.00 bpw | 139.5 GB     | 5.8132     | 0.0717
-Original | 597.1 GB     | 5.4131     | 
+
+<div align="center" style="line-height: 1;">
+  <a href="https://ernie.baidu.com/" target="_blank" style="margin: 2px;">
+    <img alt="Chat" src="https://img.shields.io/badge/🤖_Chat-ERNIE_Bot-blue" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+  <a href="https://huggingface.co/baidu" target="_blank" style="margin: 2px;">
+    <img alt="Hugging Face" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Baidu-ffc107?color=ffc107&logoColor=white" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+  <a href="https://github.com/PaddlePaddle/ERNIE" target="_blank" style="margin: 2px;">
+    <img alt="Github" src="https://img.shields.io/badge/GitHub-ERNIE-000?logo=github&color=0000FF" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+  <a href="https://ernie.baidu.com/blog/ernie4.5" target="_blank" style="margin: 2px;">
+    <img alt="Blog" src="https://img.shields.io/badge/🖖_Blog-ERNIE4.5-A020A0" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+</div>
+
+<div align="center" style="line-height: 1;">
+  <a href="#license" style="margin: 2px;">
+    <img alt="License" src="https://img.shields.io/badge/License-Apache2.0-A5de54" style="display: inline-block; vertical-align: middle;"/>
+  </a>
+</div>
+
+# ERNIE-4.5-300B-A47B
+
+> [!NOTE]
+> Note: "**-Paddle**" models use [PaddlePaddle](https://github.com/PaddlePaddle/Paddle) weights, while "**-PT**" models use Transformer-style PyTorch weights.
+
+
+## ERNIE 4.5 Highlights
+
+The advanced capabilities of the ERNIE 4.5 models, particularly the MoE-based A47B and A3B series, are underpinned by several key technical innovations:
+
+1. **Multimodal Heterogeneous MoE Pre-Training:** Our models are jointly trained on both textual and visual modalities to better capture the nuances of multimodal information and improve performance on tasks involving text understanding and generation, image understanding, and cross-modal reasoning. To achieve this without one modality hindering the learning of another, we designed a *heterogeneous MoE structure*, incorporated *modality-isolated routing*, and employed *router orthogonal loss* and *multimodal token-balanced loss*. These architectural choices ensure that both modalities are effectively represented, allowing for mutual reinforcement during training.
+
+2. **Scaling-Efficient Infrastructure:** We propose a novel heterogeneous hybrid parallelism and hierarchical load balancing strategy for efficient training of ERNIE 4.5 models. By using intra-node expert parallelism, memory-efficient pipeline scheduling, FP8 mixed-precision training and finegrained recomputation methods, we achieve remarkable pre-training throughput. For inference, we propose *multi-expert parallel collaboration* method and *convolutional code quantization* algorithm to achieve 4-bit/2-bit lossless quantization. Furthermore, we introduce PD disaggregation with dynamic role switching for effective resource utilization to enhance inference performance for ERNIE 4.5 MoE models. Built on [PaddlePaddle](https://github.com/PaddlePaddle/Paddle), ERNIE 4.5 delivers high-performance inference across a wide range of hardware platforms.
+
+3. **Modality-Specific Post-Training:** To meet the diverse requirements of real-world applications, we fine-tuned variants of the pre-trained model for specific modalities. Our LLMs are optimized for general-purpose language understanding and generation. The VLMs focuses on visuallanguage understanding and supports both thinking and non-thinking modes. Each model employed a combination of *Supervised Fine-tuning (SFT)*, *Direct Preference Optimization (DPO)* or a modified reinforcement learning method named *Unified Preference Optimization (UPO)* for post-training.
+
+## Model Overview
+
+ERNIE-4.5-300B-A47B is a text MoE Post-trained model, with 300B total parameters and 47B activated parameters for each token. The following are the model configuration details:
+
+|Key|Value|
+|-|-|
+|Modality|Text|
+|Training Stage|Pretraining|
+|Params(Total / Activated)|300B / 47B|
+|Layers|54|
+|Heads(Q/KV)|64 / 8|
+|Text Experts(Total / Activated)|64 / 8|
+|Vision Experts(Total / Activated)|64 / 8|
+|Context Length|131072|
+
+## Quickstart
+
+### Using `transformers` library
+
+**Note**: Before using the model, please ensure you have the `transformers` library installed (version 4.50.0 or higher)
+
+The following contains a code snippet illustrating how to use the model generate content based on given inputs.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "baidu/ERNIE-4.5-300B-A47B-PT"
+
+# load the tokenizer and the model
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+
+# prepare the model input
+prompt = "Give me a short introduction to large language model."
+messages = [
+    {"role": "user", "content": prompt}
+]
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+model_inputs = tokenizer([text], add_special_tokens=False, return_tensors="pt").to(model.device)
+
+# conduct text completion
+generated_ids = model.generate(
+    model_inputs.input_ids,
+    max_new_tokens=1024
+)
+output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
+
+# decode the generated ids
+generate_text = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
+print("generate_text:", generate_text)
+```
+
+### Using vLLM
+
+[vllm](https://github.com/vllm-project/vllm/tree/main) github library. Python-only [build](https://docs.vllm.ai/en/latest/getting_started/installation/gpu.html#set-up-using-python-only-build-without-compilation).
+
+```bash
+# 80G * 16 GPU
+vllm serve baidu/ERNIE-4.5-300B-A47B-PT --trust-remote-code
+```
+
+```bash
+# FP8 online quantification 80G * 16 GPU
+vllm serve baidu/ERNIE-4.5-300B-A47B-PT --trust-remote-code --quantization fp8
+```
+
+## Best Practices
+
+### **Sampling Parameters**
+
+To achieve optimal performance, we suggest using `Temperature=0.8`, `TopP=0.8`.
+
+### Prompts for Web Search
+
+For Web Search, {references}, {date}, and {question} are arguments.
+
+For Chinese question, we use the prompt:
+
+```python
+ernie_search_zh_prompt = \
+'''下面你会收到当前时间、多个不同来源的参考文章和一段对话。你的任务是阅读多个参考文章，并根据参考文章中的信息回答对话中的问题。
+以下是当前时间和参考文章：
+---------
+#当前时间
+{date}
+
+#参考文章
+{references}
+
+---------
+请注意：
+1. 回答必须结合问题需求和当前时间，对参考文章的可用性进行判断，避免在回答中使用错误或过时的信息。
+2. 当参考文章中的信息无法准确地回答问题时，你需要在回答中提供获取相应信息的建议，或承认无法提供相应信息。
+3. 你需要优先根据百科、官网、权威机构、专业网站等高权威性来源的信息来回答问题。
+4. 回复需要综合参考文章中的相关数字、案例、法律条文、公式等信息，使你的答案更专业。
+5. 当问题属于创作类任务时，需注意以下维度：
+   - 态度鲜明：观点、立场清晰明确，避免模棱两可，语言果断直接
+   - 文采飞扬：用词精准生动，善用修辞手法，增强感染力
+   - 有理有据：逻辑严密递进，结合权威数据/事实支撑论点
+---------
+下面请结合以上信息，回答问题，补全对话
+{question}'''
+```
+
+For English question, we use the prompt:
+
+```python
+ernie_search_en_prompt = \
+'''
+Below you will be given the current time, multiple references from different sources, and a conversation. Your task is to read the references and use the information in them to answer the question in the conversation.
+Here are the current time and the references:
+---------
+#Current Time
+{date}
+
+#References
+{references}
+
+---------
+Please note:
+1. Based on the question’s requirements and the current time, assess the usefulness of the references to avoid using inaccurate or outdated information in the answer.  
+2. If the references do not provide enough information to accurately answer the question, you should suggest how to obtain the relevant information or acknowledge that you are unable to provide it.  
+3. Prioritize using information from highly authoritative sources such as encyclopedias, official websites, authoritative institutions, and professional websites when answering questions.
+4. Incorporate relevant numbers, cases, legal provisions, formulas, and other details from the references to make your answer more professional.
+5. For creative tasks, keep these dimensions in mind:
+   - Clear attitude: Clear views and positions, avoid ambiguity, and use decisive and direct language
+   - Brilliant writing: Precise and vivid words, good use of rhetoric, and enhance the appeal
+   - Well-reasoned: Rigorous logic and progressive, combined with authoritative data/facts to support the argument
+
+---------
+Now, using the information above, answer the question and complete the conversation:  
+{question}'''
+```
+
+Parameter notes:
+
+* {question} is the user’s question
+* {date} is the current time, and the recommended format is “YYYY-MM-DD HH:MM:SS, Day of the Week, Beijing/China.”
+* {references} is the references, and the recommended format is:
+
+```text
+##参考文章1
+标题：周杰伦
+文章发布时间：2025-04-20
+内容：周杰伦(Jay Chou),1979年1月18日出生于台湾省新北市,祖籍福建省永春县,华语流行乐男歌手、音乐人、演员、导演、编剧,毕业于淡江中学。2000年,发行个人首张音乐专辑《Jay》。...
+来源网站网址：baike.baidu.com
+来源网站的网站名：百度百科
+
+##参考文章2
+...
+```
+
+## License
+
+The ERNIE 4.5 models are provided under the Apache License 2.0. This license permits commercial use, subject to its terms and conditions. Copyright (c) 2025 Baidu, Inc. All Rights Reserved.
+
+## Citation
+
+If you find ERNIE 4.5 useful or wish to use it in your projects, please kindly cite our technical report:
+
+```bibtex
+@misc{ernie2025technicalreport,
+      title={ERNIE 4.5 Technical Report},
+      author={Baidu ERNIE Team},
+      year={2025},
+      eprint={},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={}
+}
+```

added_tokens.json

@@ -0,0 +1 @@
+{"<|IMAGE_PLACEHOLDER|>": 100295, "<|AUDIO_PLACEHOLDER|>": 100296, "<|LOC_0|>": 100297, "<|LOC_1|>": 100298, "<|LOC_2|>": 100299, "<|LOC_3|>": 100300, "<|LOC_4|>": 100301, "<|LOC_5|>": 100302, "<|LOC_6|>": 100303, "<|LOC_7|>": 100304, "<|LOC_8|>": 100305, "<|LOC_9|>": 100306, "<|LOC_10|>": 100307, "<|LOC_11|>": 100308, "<|LOC_12|>": 100309, "<|LOC_13|>": 100310, "<|LOC_14|>": 100311, "<|LOC_15|>": 100312, "<|LOC_16|>": 100313, "<|LOC_17|>": 100314, "<|LOC_18|>": 100315, "<|LOC_19|>": 100316, "<|LOC_20|>": 100317, "<|LOC_21|>": 100318, "<|LOC_22|>": 100319, "<|LOC_23|>": 100320, "<|LOC_24|>": 100321, "<|LOC_25|>": 100322, "<|LOC_26|>": 100323, "<|LOC_27|>": 100324, "<|LOC_28|>": 100325, "<|LOC_29|>": 100326, "<|LOC_30|>": 100327, "<|LOC_31|>": 100328, "<|LOC_32|>": 100329, "<|LOC_33|>": 100330, "<|LOC_34|>": 100331, "<|LOC_35|>": 100332, "<|LOC_36|>": 100333, "<|LOC_37|>": 100334, "<|LOC_38|>": 100335, "<|LOC_39|>": 100336, "<|LOC_40|>": 100337, "<|LOC_41|>": 100338, "<|LOC_42|>": 100339, "<|LOC_43|>": 100340, "<|LOC_44|>": 100341, "<|LOC_45|>": 100342, "<|LOC_46|>": 100343, "<|LOC_47|>": 100344, "<|LOC_48|>": 100345, "<|LOC_49|>": 100346, "<|LOC_50|>": 100347, "<|LOC_51|>": 100348, "<|LOC_52|>": 100349, "<|LOC_53|>": 100350, "<|LOC_54|>": 100351, "<|LOC_55|>": 100352, "<|LOC_56|>": 100353, "<|LOC_57|>": 100354, "<|LOC_58|>": 100355, "<|LOC_59|>": 100356, "<|LOC_60|>": 100357, "<|LOC_61|>": 100358, "<|LOC_62|>": 100359, "<|LOC_63|>": 100360, "<|LOC_64|>": 100361, "<|LOC_65|>": 100362, "<|LOC_66|>": 100363, "<|LOC_67|>": 100364, "<|LOC_68|>": 100365, "<|LOC_69|>": 100366, "<|LOC_70|>": 100367, "<|LOC_71|>": 100368, "<|LOC_72|>": 100369, "<|LOC_73|>": 100370, "<|LOC_74|>": 100371, "<|LOC_75|>": 100372, "<|LOC_76|>": 100373, "<|LOC_77|>": 100374, "<|LOC_78|>": 100375, "<|LOC_79|>": 100376, "<|LOC_80|>": 100377, "<|LOC_81|>": 100378, "<|LOC_82|>": 100379, "<|LOC_83|>": 100380, "<|LOC_84|>": 100381, "<|LOC_85|>": 100382, "<|LOC_86|>": 100383, "<|LOC_87|>": 100384, "<|LOC_88|>": 100385, "<|LOC_89|>": 100386, "<|LOC_90|>": 100387, "<|LOC_91|>": 100388, "<|LOC_92|>": 100389, "<|LOC_93|>": 100390, "<|LOC_94|>": 100391, "<|LOC_95|>": 100392, "<|LOC_96|>": 100393, "<|LOC_97|>": 100394, "<|LOC_98|>": 100395, "<|LOC_99|>": 100396, "<|LOC_100|>": 100397, "<|LOC_101|>": 100398, "<|LOC_102|>": 100399, "<|LOC_103|>": 100400, "<|LOC_104|>": 100401, "<|LOC_105|>": 100402, "<|LOC_106|>": 100403, "<|LOC_107|>": 100404, "<|LOC_108|>": 100405, "<|LOC_109|>": 100406, "<|LOC_110|>": 100407, "<|LOC_111|>": 100408, "<|LOC_112|>": 100409, "<|LOC_113|>": 100410, "<|LOC_114|>": 100411, "<|LOC_115|>": 100412, "<|LOC_116|>": 100413, "<|LOC_117|>": 100414, "<|LOC_118|>": 100415, "<|LOC_119|>": 100416, "<|LOC_120|>": 100417, "<|LOC_121|>": 100418, "<|LOC_122|>": 100419, "<|LOC_123|>": 100420, "<|LOC_124|>": 100421, "<|LOC_125|>": 100422, "<|LOC_126|>": 100423, "<|LOC_127|>": 100424, "<|LOC_128|>": 100425, "<|LOC_129|>": 100426, "<|LOC_130|>": 100427, "<|LOC_131|>": 100428, "<|LOC_132|>": 100429, "<|LOC_133|>": 100430, "<|LOC_134|>": 100431, "<|LOC_135|>": 100432, "<|LOC_136|>": 100433, "<|LOC_137|>": 100434, "<|LOC_138|>": 100435, "<|LOC_139|>": 100436, "<|LOC_140|>": 100437, "<|LOC_141|>": 100438, "<|LOC_142|>": 100439, "<|LOC_143|>": 100440, "<|LOC_144|>": 100441, "<|LOC_145|>": 100442, "<|LOC_146|>": 100443, "<|LOC_147|>": 100444, "<|LOC_148|>": 100445, "<|LOC_149|>": 100446, "<|LOC_150|>": 100447, "<|LOC_151|>": 100448, "<|LOC_152|>": 100449, "<|LOC_153|>": 100450, "<|LOC_154|>": 100451, "<|LOC_155|>": 100452, "<|LOC_156|>": 100453, "<|LOC_157|>": 100454, "<|LOC_158|>": 100455, "<|LOC_159|>": 100456, "<|LOC_160|>": 100457, "<|LOC_161|>": 100458, "<|LOC_162|>": 100459, "<|LOC_163|>": 100460, "<|LOC_164|>": 100461, "<|LOC_165|>": 100462, "<|LOC_166|>": 100463, "<|LOC_167|>": 100464, "<|LOC_168|>": 100465, "<|LOC_169|>": 100466, "<|LOC_170|>": 100467, "<|LOC_171|>": 100468, "<|LOC_172|>": 100469, "<|LOC_173|>": 100470, "<|LOC_174|>": 100471, "<|LOC_175|>": 100472, "<|LOC_176|>": 100473, "<|LOC_177|>": 100474, "<|LOC_178|>": 100475, "<|LOC_179|>": 100476, "<|LOC_180|>": 100477, "<|LOC_181|>": 100478, "<|LOC_182|>": 100479, "<|LOC_183|>": 100480, "<|LOC_184|>": 100481, "<|LOC_185|>": 100482, "<|LOC_186|>": 100483, "<|LOC_187|>": 100484, "<|LOC_188|>": 100485, "<|LOC_189|>": 100486, "<|LOC_190|>": 100487, "<|LOC_191|>": 100488, "<|LOC_192|>": 100489, "<|LOC_193|>": 100490, "<|LOC_194|>": 100491, "<|LOC_195|>": 100492, "<|LOC_196|>": 100493, "<|LOC_197|>": 100494, "<|LOC_198|>": 100495, "<|LOC_199|>": 100496, "<|LOC_200|>": 100497, "<|LOC_201|>": 100498, "<|LOC_202|>": 100499, "<|LOC_203|>": 100500, "<|LOC_204|>": 100501, "<|LOC_205|>": 100502, "<|LOC_206|>": 100503, "<|LOC_207|>": 100504, "<|LOC_208|>": 100505, "<|LOC_209|>": 100506, "<|LOC_210|>": 100507, "<|LOC_211|>": 100508, "<|LOC_212|>": 100509, "<|LOC_213|>": 100510, "<|LOC_214|>": 100511, "<|LOC_215|>": 100512, "<|LOC_216|>": 100513, "<|LOC_217|>": 100514, "<|LOC_218|>": 100515, "<|LOC_219|>": 100516, "<|LOC_220|>": 100517, "<|LOC_221|>": 100518, "<|LOC_222|>": 100519, "<|LOC_223|>": 100520, "<|LOC_224|>": 100521, "<|LOC_225|>": 100522, "<|LOC_226|>": 100523, "<|LOC_227|>": 100524, "<|LOC_228|>": 100525, "<|LOC_229|>": 100526, "<|LOC_230|>": 100527, "<|LOC_231|>": 100528, "<|LOC_232|>": 100529, "<|LOC_233|>": 100530, "<|LOC_234|>": 100531, "<|LOC_235|>": 100532, "<|LOC_236|>": 100533, "<|LOC_237|>": 100534, "<|LOC_238|>": 100535, "<|LOC_239|>": 100536, "<|LOC_240|>": 100537, "<|LOC_241|>": 100538, "<|LOC_242|>": 100539, "<|LOC_243|>": 100540, "<|LOC_244|>": 100541, "<|LOC_245|>": 100542, "<|LOC_246|>": 100543, "<|LOC_247|>": 100544, "<|LOC_248|>": 100545, "<|LOC_249|>": 100546, "<|LOC_250|>": 100547, "<|LOC_251|>": 100548, "<|LOC_252|>": 100549, "<|LOC_253|>": 100550, "<|LOC_254|>": 100551, "<|LOC_255|>": 100552, "<|LOC_256|>": 100553, "<|LOC_257|>": 100554, "<|LOC_258|>": 100555, "<|LOC_259|>": 100556, "<|LOC_260|>": 100557, "<|LOC_261|>": 100558, "<|LOC_262|>": 100559, "<|LOC_263|>": 100560, "<|LOC_264|>": 100561, "<|LOC_265|>": 100562, "<|LOC_266|>": 100563, "<|LOC_267|>": 100564, "<|LOC_268|>": 100565, "<|LOC_269|>": 100566, "<|LOC_270|>": 100567, "<|LOC_271|>": 100568, "<|LOC_272|>": 100569, "<|LOC_273|>": 100570, "<|LOC_274|>": 100571, "<|LOC_275|>": 100572, "<|LOC_276|>": 100573, "<|LOC_277|>": 100574, "<|LOC_278|>": 100575, "<|LOC_279|>": 100576, "<|LOC_280|>": 100577, "<|LOC_281|>": 100578, "<|LOC_282|>": 100579, "<|LOC_283|>": 100580, "<|LOC_284|>": 100581, "<|LOC_285|>": 100582, "<|LOC_286|>": 100583, "<|LOC_287|>": 100584, "<|LOC_288|>": 100585, "<|LOC_289|>": 100586, "<|LOC_290|>": 100587, "<|LOC_291|>": 100588, "<|LOC_292|>": 100589, "<|LOC_293|>": 100590, "<|LOC_294|>": 100591, "<|LOC_295|>": 100592, "<|LOC_296|>": 100593, "<|LOC_297|>": 100594, "<|LOC_298|>": 100595, "<|LOC_299|>": 100596, "<|LOC_300|>": 100597, "<|LOC_301|>": 100598, "<|LOC_302|>": 100599, "<|LOC_303|>": 100600, "<|LOC_304|>": 100601, "<|LOC_305|>": 100602, "<|LOC_306|>": 100603, "<|LOC_307|>": 100604, "<|LOC_308|>": 100605, "<|LOC_309|>": 100606, "<|LOC_310|>": 100607, "<|LOC_311|>": 100608, "<|LOC_312|>": 100609, "<|LOC_313|>": 100610, "<|LOC_314|>": 100611, "<|LOC_315|>": 100612, "<|LOC_316|>": 100613, "<|LOC_317|>": 100614, "<|LOC_318|>": 100615, "<|LOC_319|>": 100616, "<|LOC_320|>": 100617, "<|LOC_321|>": 100618, "<|LOC_322|>": 100619, "<|LOC_323|>": 100620, "<|LOC_324|>": 100621, "<|LOC_325|>": 100622, "<|LOC_326|>": 100623, "<|LOC_327|>": 100624, "<|LOC_328|>": 100625, "<|LOC_329|>": 100626, "<|LOC_330|>": 100627, "<|LOC_331|>": 100628, "<|LOC_332|>": 100629, "<|LOC_333|>": 100630, "<|LOC_334|>": 100631, "<|LOC_335|>": 100632, "<|LOC_336|>": 100633, "<|LOC_337|>": 100634, "<|LOC_338|>": 100635, "<|LOC_339|>": 100636, "<|LOC_340|>": 100637, "<|LOC_341|>": 100638, "<|LOC_342|>": 100639, "<|LOC_343|>": 100640, "<|LOC_344|>": 100641, "<|LOC_345|>": 100642, "<|LOC_346|>": 100643, "<|LOC_347|>": 100644, "<|LOC_348|>": 100645, "<|LOC_349|>": 100646, "<|LOC_350|>": 100647, "<|LOC_351|>": 100648, "<|LOC_352|>": 100649, "<|LOC_353|>": 100650, "<|LOC_354|>": 100651, "<|LOC_355|>": 100652, "<|LOC_356|>": 100653, "<|LOC_357|>": 100654, "<|LOC_358|>": 100655, "<|LOC_359|>": 100656, "<|LOC_360|>": 100657, "<|LOC_361|>": 100658, "<|LOC_362|>": 100659, "<|LOC_363|>": 100660, "<|LOC_364|>": 100661, "<|LOC_365|>": 100662, "<|LOC_366|>": 100663, "<|LOC_367|>": 100664, "<|LOC_368|>": 100665, "<|LOC_369|>": 100666, "<|LOC_370|>": 100667, "<|LOC_371|>": 100668, "<|LOC_372|>": 100669, "<|LOC_373|>": 100670, "<|LOC_374|>": 100671, "<|LOC_375|>": 100672, "<|LOC_376|>": 100673, "<|LOC_377|>": 100674, "<|LOC_378|>": 100675, "<|LOC_379|>": 100676, "<|LOC_380|>": 100677, "<|LOC_381|>": 100678, "<|LOC_382|>": 100679, "<|LOC_383|>": 100680, "<|LOC_384|>": 100681, "<|LOC_385|>": 100682, "<|LOC_386|>": 100683, "<|LOC_387|>": 100684, "<|LOC_388|>": 100685, "<|LOC_389|>": 100686, "<|LOC_390|>": 100687, "<|LOC_391|>": 100688, "<|LOC_392|>": 100689, "<|LOC_393|>": 100690, "<|LOC_394|>": 100691, "<|LOC_395|>": 100692, "<|LOC_396|>": 100693, "<|LOC_397|>": 100694, "<|LOC_398|>": 100695, "<|LOC_399|>": 100696, "<|LOC_400|>": 100697, "<|LOC_401|>": 100698, "<|LOC_402|>": 100699, "<|LOC_403|>": 100700, "<|LOC_404|>": 100701, "<|LOC_405|>": 100702, "<|LOC_406|>": 100703, "<|LOC_407|>": 100704, "<|LOC_408|>": 100705, "<|LOC_409|>": 100706, "<|LOC_410|>": 100707, "<|LOC_411|>": 100708, "<|LOC_412|>": 100709, "<|LOC_413|>": 100710, "<|LOC_414|>": 100711, "<|LOC_415|>": 100712, "<|LOC_416|>": 100713, "<|LOC_417|>": 100714, "<|LOC_418|>": 100715, "<|LOC_419|>": 100716, "<|LOC_420|>": 100717, "<|LOC_421|>": 100718, "<|LOC_422|>": 100719, "<|LOC_423|>": 100720, "<|LOC_424|>": 100721, "<|LOC_425|>": 100722, "<|LOC_426|>": 100723, "<|LOC_427|>": 100724, "<|LOC_428|>": 100725, "<|LOC_429|>": 100726, "<|LOC_430|>": 100727, "<|LOC_431|>": 100728, "<|LOC_432|>": 100729, "<|LOC_433|>": 100730, "<|LOC_434|>": 100731, "<|LOC_435|>": 100732, "<|LOC_436|>": 100733, "<|LOC_437|>": 100734, "<|LOC_438|>": 100735, "<|LOC_439|>": 100736, "<|LOC_440|>": 100737, "<|LOC_441|>": 100738, "<|LOC_442|>": 100739, "<|LOC_443|>": 100740, "<|LOC_444|>": 100741, "<|LOC_445|>": 100742, "<|LOC_446|>": 100743, "<|LOC_447|>": 100744, "<|LOC_448|>": 100745, "<|LOC_449|>": 100746, "<|LOC_450|>": 100747, "<|LOC_451|>": 100748, "<|LOC_452|>": 100749, "<|LOC_453|>": 100750, "<|LOC_454|>": 100751, "<|LOC_455|>": 100752, "<|LOC_456|>": 100753, "<|LOC_457|>": 100754, "<|LOC_458|>": 100755, "<|LOC_459|>": 100756, "<|LOC_460|>": 100757, "<|LOC_461|>": 100758, "<|LOC_462|>": 100759, "<|LOC_463|>": 100760, "<|LOC_464|>": 100761, "<|LOC_465|>": 100762, "<|LOC_466|>": 100763, "<|LOC_467|>": 100764, "<|LOC_468|>": 100765, "<|LOC_469|>": 100766, "<|LOC_470|>": 100767, "<|LOC_471|>": 100768, "<|LOC_472|>": 100769, "<|LOC_473|>": 100770, "<|LOC_474|>": 100771, "<|LOC_475|>": 100772, "<|LOC_476|>": 100773, "<|LOC_477|>": 100774, "<|LOC_478|>": 100775, "<|LOC_479|>": 100776, "<|LOC_480|>": 100777, "<|LOC_481|>": 100778, "<|LOC_482|>": 100779, "<|LOC_483|>": 100780, "<|LOC_484|>": 100781, "<|LOC_485|>": 100782, "<|LOC_486|>": 100783, "<|LOC_487|>": 100784, "<|LOC_488|>": 100785, "<|LOC_489|>": 100786, "<|LOC_490|>": 100787, "<|LOC_491|>": 100788, "<|LOC_492|>": 100789, "<|LOC_493|>": 100790, "<|LOC_494|>": 100791, "<|LOC_495|>": 100792, "<|LOC_496|>": 100793, "<|LOC_497|>": 100794, "<|LOC_498|>": 100795, "<|LOC_499|>": 100796, "<|LOC_500|>": 100797, "<|LOC_501|>": 100798, "<|LOC_502|>": 100799, "<|LOC_503|>": 100800, "<|LOC_504|>": 100801, "<|LOC_505|>": 100802, "<|LOC_506|>": 100803, "<|LOC_507|>": 100804, "<|LOC_508|>": 100805, "<|LOC_509|>": 100806, "<|LOC_510|>": 100807, "<|LOC_511|>": 100808, "<|LOC_512|>": 100809, "<|LOC_513|>": 100810, "<|LOC_514|>": 100811, "<|LOC_515|>": 100812, "<|LOC_516|>": 100813, "<|LOC_517|>": 100814, "<|LOC_518|>": 100815, "<|LOC_519|>": 100816, "<|LOC_520|>": 100817, "<|LOC_521|>": 100818, "<|LOC_522|>": 100819, "<|LOC_523|>": 100820, "<|LOC_524|>": 100821, "<|LOC_525|>": 100822, "<|LOC_526|>": 100823, "<|LOC_527|>": 100824, "<|LOC_528|>": 100825, "<|LOC_529|>": 100826, "<|LOC_530|>": 100827, "<|LOC_531|>": 100828, "<|LOC_532|>": 100829, "<|LOC_533|>": 100830, "<|LOC_534|>": 100831, "<|LOC_535|>": 100832, "<|LOC_536|>": 100833, "<|LOC_537|>": 100834, "<|LOC_538|>": 100835, "<|LOC_539|>": 100836, "<|LOC_540|>": 100837, "<|LOC_541|>": 100838, "<|LOC_542|>": 100839, "<|LOC_543|>": 100840, "<|LOC_544|>": 100841, "<|LOC_545|>": 100842, "<|LOC_546|>": 100843, "<|LOC_547|>": 100844, "<|LOC_548|>": 100845, "<|LOC_549|>": 100846, "<|LOC_550|>": 100847, "<|LOC_551|>": 100848, "<|LOC_552|>": 100849, "<|LOC_553|>": 100850, "<|LOC_554|>": 100851, "<|LOC_555|>": 100852, "<|LOC_556|>": 100853, "<|LOC_557|>": 100854, "<|LOC_558|>": 100855, "<|LOC_559|>": 100856, "<|LOC_560|>": 100857, "<|LOC_561|>": 100858, "<|LOC_562|>": 100859, "<|LOC_563|>": 100860, "<|LOC_564|>": 100861, "<|LOC_565|>": 100862, "<|LOC_566|>": 100863, "<|LOC_567|>": 100864, "<|LOC_568|>": 100865, "<|LOC_569|>": 100866, "<|LOC_570|>": 100867, "<|LOC_571|>": 100868, "<|LOC_572|>": 100869, "<|LOC_573|>": 100870, "<|LOC_574|>": 100871, "<|LOC_575|>": 100872, "<|LOC_576|>": 100873, "<|LOC_577|>": 100874, "<|LOC_578|>": 100875, "<|LOC_579|>": 100876, "<|LOC_580|>": 100877, "<|LOC_581|>": 100878, "<|LOC_582|>": 100879, "<|LOC_583|>": 100880, "<|LOC_584|>": 100881, "<|LOC_585|>": 100882, "<|LOC_586|>": 100883, "<|LOC_587|>": 100884, "<|LOC_588|>": 100885, "<|LOC_589|>": 100886, "<|LOC_590|>": 100887, "<|LOC_591|>": 100888, "<|LOC_592|>": 100889, "<|LOC_593|>": 100890, "<|LOC_594|>": 100891, "<|LOC_595|>": 100892, "<|LOC_596|>": 100893, "<|LOC_597|>": 100894, "<|LOC_598|>": 100895, "<|LOC_599|>": 100896, "<|LOC_600|>": 100897, "<|LOC_601|>": 100898, "<|LOC_602|>": 100899, "<|LOC_603|>": 100900, "<|LOC_604|>": 100901, "<|LOC_605|>": 100902, "<|LOC_606|>": 100903, "<|LOC_607|>": 100904, "<|LOC_608|>": 100905, "<|LOC_609|>": 100906, "<|LOC_610|>": 100907, "<|LOC_611|>": 100908, "<|LOC_612|>": 100909, "<|LOC_613|>": 100910, "<|LOC_614|>": 100911, "<|LOC_615|>": 100912, "<|LOC_616|>": 100913, "<|LOC_617|>": 100914, "<|LOC_618|>": 100915, "<|LOC_619|>": 100916, "<|LOC_620|>": 100917, "<|LOC_621|>": 100918, "<|LOC_622|>": 100919, "<|LOC_623|>": 100920, "<|LOC_624|>": 100921, "<|LOC_625|>": 100922, "<|LOC_626|>": 100923, "<|LOC_627|>": 100924, "<|LOC_628|>": 100925, "<|LOC_629|>": 100926, "<|LOC_630|>": 100927, "<|LOC_631|>": 100928, "<|LOC_632|>": 100929, "<|LOC_633|>": 100930, "<|LOC_634|>": 100931, "<|LOC_635|>": 100932, "<|LOC_636|>": 100933, "<|LOC_637|>": 100934, "<|LOC_638|>": 100935, "<|LOC_639|>": 100936, "<|LOC_640|>": 100937, "<|LOC_641|>": 100938, "<|LOC_642|>": 100939, "<|LOC_643|>": 100940, "<|LOC_644|>": 100941, "<|LOC_645|>": 100942, "<|LOC_646|>": 100943, "<|LOC_647|>": 100944, "<|LOC_648|>": 100945, "<|LOC_649|>": 100946, "<|LOC_650|>": 100947, "<|LOC_651|>": 100948, "<|LOC_652|>": 100949, "<|LOC_653|>": 100950, "<|LOC_654|>": 100951, "<|LOC_655|>": 100952, "<|LOC_656|>": 100953, "<|LOC_657|>": 100954, "<|LOC_658|>": 100955, "<|LOC_659|>": 100956, "<|LOC_660|>": 100957, "<|LOC_661|>": 100958, "<|LOC_662|>": 100959, "<|LOC_663|>": 100960, "<|LOC_664|>": 100961, "<|LOC_665|>": 100962, "<|LOC_666|>": 100963, "<|LOC_667|>": 100964, "<|LOC_668|>": 100965, "<|LOC_669|>": 100966, "<|LOC_670|>": 100967, "<|LOC_671|>": 100968, "<|LOC_672|>": 100969, "<|LOC_673|>": 100970, "<|LOC_674|>": 100971, "<|LOC_675|>": 100972, "<|LOC_676|>": 100973, "<|LOC_677|>": 100974, "<|LOC_678|>": 100975, "<|LOC_679|>": 100976, "<|LOC_680|>": 100977, "<|LOC_681|>": 100978, "<|LOC_682|>": 100979, "<|LOC_683|>": 100980, "<|LOC_684|>": 100981, "<|LOC_685|>": 100982, "<|LOC_686|>": 100983, "<|LOC_687|>": 100984, "<|LOC_688|>": 100985, "<|LOC_689|>": 100986, "<|LOC_690|>": 100987, "<|LOC_691|>": 100988, "<|LOC_692|>": 100989, "<|LOC_693|>": 100990, "<|LOC_694|>": 100991, "<|LOC_695|>": 100992, "<|LOC_696|>": 100993, "<|LOC_697|>": 100994, "<|LOC_698|>": 100995, "<|LOC_699|>": 100996, "<|LOC_700|>": 100997, "<|LOC_701|>": 100998, "<|LOC_702|>": 100999, "<|LOC_703|>": 101000, "<|LOC_704|>": 101001, "<|LOC_705|>": 101002, "<|LOC_706|>": 101003, "<|LOC_707|>": 101004, "<|LOC_708|>": 101005, "<|LOC_709|>": 101006, "<|LOC_710|>": 101007, "<|LOC_711|>": 101008, "<|LOC_712|>": 101009, "<|LOC_713|>": 101010, "<|LOC_714|>": 101011, "<|LOC_715|>": 101012, "<|LOC_716|>": 101013, "<|LOC_717|>": 101014, "<|LOC_718|>": 101015, "<|LOC_719|>": 101016, "<|LOC_720|>": 101017, "<|LOC_721|>": 101018, "<|LOC_722|>": 101019, "<|LOC_723|>": 101020, "<|LOC_724|>": 101021, "<|LOC_725|>": 101022, "<|LOC_726|>": 101023, "<|LOC_727|>": 101024, "<|LOC_728|>": 101025, "<|LOC_729|>": 101026, "<|LOC_730|>": 101027, "<|LOC_731|>": 101028, "<|LOC_732|>": 101029, "<|LOC_733|>": 101030, "<|LOC_734|>": 101031, "<|LOC_735|>": 101032, "<|LOC_736|>": 101033, "<|LOC_737|>": 101034, "<|LOC_738|>": 101035, "<|LOC_739|>": 101036, "<|LOC_740|>": 101037, "<|LOC_741|>": 101038, "<|LOC_742|>": 101039, "<|LOC_743|>": 101040, "<|LOC_744|>": 101041, "<|LOC_745|>": 101042, "<|LOC_746|>": 101043, "<|LOC_747|>": 101044, "<|LOC_748|>": 101045, "<|LOC_749|>": 101046, "<|LOC_750|>": 101047, "<|LOC_751|>": 101048, "<|LOC_752|>": 101049, "<|LOC_753|>": 101050, "<|LOC_754|>": 101051, "<|LOC_755|>": 101052, "<|LOC_756|>": 101053, "<|LOC_757|>": 101054, "<|LOC_758|>": 101055, "<|LOC_759|>": 101056, "<|LOC_760|>": 101057, "<|LOC_761|>": 101058, "<|LOC_762|>": 101059, "<|LOC_763|>": 101060, "<|LOC_764|>": 101061, "<|LOC_765|>": 101062, "<|LOC_766|>": 101063, "<|LOC_767|>": 101064, "<|LOC_768|>": 101065, "<|LOC_769|>": 101066, "<|LOC_770|>": 101067, "<|LOC_771|>": 101068, "<|LOC_772|>": 101069, "<|LOC_773|>": 101070, "<|LOC_774|>": 101071, "<|LOC_775|>": 101072, "<|LOC_776|>": 101073, "<|LOC_777|>": 101074, "<|LOC_778|>": 101075, "<|LOC_779|>": 101076, "<|LOC_780|>": 101077, "<|LOC_781|>": 101078, "<|LOC_782|>": 101079, "<|LOC_783|>": 101080, "<|LOC_784|>": 101081, "<|LOC_785|>": 101082, "<|LOC_786|>": 101083, "<|LOC_787|>": 101084, "<|LOC_788|>": 101085, "<|LOC_789|>": 101086, "<|LOC_790|>": 101087, "<|LOC_791|>": 101088, "<|LOC_792|>": 101089, "<|LOC_793|>": 101090, "<|LOC_794|>": 101091, "<|LOC_795|>": 101092, "<|LOC_796|>": 101093, "<|LOC_797|>": 101094, "<|LOC_798|>": 101095, "<|LOC_799|>": 101096, "<|LOC_800|>": 101097, "<|LOC_801|>": 101098, "<|LOC_802|>": 101099, "<|LOC_803|>": 101100, "<|LOC_804|>": 101101, "<|LOC_805|>": 101102, "<|LOC_806|>": 101103, "<|LOC_807|>": 101104, "<|LOC_808|>": 101105, "<|LOC_809|>": 101106, "<|LOC_810|>": 101107, "<|LOC_811|>": 101108, "<|LOC_812|>": 101109, "<|LOC_813|>": 101110, "<|LOC_814|>": 101111, "<|LOC_815|>": 101112, "<|LOC_816|>": 101113, "<|LOC_817|>": 101114, "<|LOC_818|>": 101115, "<|LOC_819|>": 101116, "<|LOC_820|>": 101117, "<|LOC_821|>": 101118, "<|LOC_822|>": 101119, "<|LOC_823|>": 101120, "<|LOC_824|>": 101121, "<|LOC_825|>": 101122, "<|LOC_826|>": 101123, "<|LOC_827|>": 101124, "<|LOC_828|>": 101125, "<|LOC_829|>": 101126, "<|LOC_830|>": 101127, "<|LOC_831|>": 101128, "<|LOC_832|>": 101129, "<|LOC_833|>": 101130, "<|LOC_834|>": 101131, "<|LOC_835|>": 101132, "<|LOC_836|>": 101133, "<|LOC_837|>": 101134, "<|LOC_838|>": 101135, "<|LOC_839|>": 101136, "<|LOC_840|>": 101137, "<|LOC_841|>": 101138, "<|LOC_842|>": 101139, "<|LOC_843|>": 101140, "<|LOC_844|>": 101141, "<|LOC_845|>": 101142, "<|LOC_846|>": 101143, "<|LOC_847|>": 101144, "<|LOC_848|>": 101145, "<|LOC_849|>": 101146, "<|LOC_850|>": 101147, "<|LOC_851|>": 101148, "<|LOC_852|>": 101149, "<|LOC_853|>": 101150, "<|LOC_854|>": 101151, "<|LOC_855|>": 101152, "<|LOC_856|>": 101153, "<|LOC_857|>": 101154, "<|LOC_858|>": 101155, "<|LOC_859|>": 101156, "<|LOC_860|>": 101157, "<|LOC_861|>": 101158, "<|LOC_862|>": 101159, "<|LOC_863|>": 101160, "<|LOC_864|>": 101161, "<|LOC_865|>": 101162, "<|LOC_866|>": 101163, "<|LOC_867|>": 101164, "<|LOC_868|>": 101165, "<|LOC_869|>": 101166, "<|LOC_870|>": 101167, "<|LOC_871|>": 101168, "<|LOC_872|>": 101169, "<|LOC_873|>": 101170, "<|LOC_874|>": 101171, "<|LOC_875|>": 101172, "<|LOC_876|>": 101173, "<|LOC_877|>": 101174, "<|LOC_878|>": 101175, "<|LOC_879|>": 101176, "<|LOC_880|>": 101177, "<|LOC_881|>": 101178, "<|LOC_882|>": 101179, "<|LOC_883|>": 101180, "<|LOC_884|>": 101181, "<|LOC_885|>": 101182, "<|LOC_886|>": 101183, "<|LOC_887|>": 101184, "<|LOC_888|>": 101185, "<|LOC_889|>": 101186, "<|LOC_890|>": 101187, "<|LOC_891|>": 101188, "<|LOC_892|>": 101189, "<|LOC_893|>": 101190, "<|LOC_894|>": 101191, "<|LOC_895|>": 101192, "<|LOC_896|>": 101193, "<|LOC_897|>": 101194, "<|LOC_898|>": 101195, "<|LOC_899|>": 101196, "<|LOC_900|>": 101197, "<|LOC_901|>": 101198, "<|LOC_902|>": 101199, "<|LOC_903|>": 101200, "<|LOC_904|>": 101201, "<|LOC_905|>": 101202, "<|LOC_906|>": 101203, "<|LOC_907|>": 101204, "<|LOC_908|>": 101205, "<|LOC_909|>": 101206, "<|LOC_910|>": 101207, "<|LOC_911|>": 101208, "<|LOC_912|>": 101209, "<|LOC_913|>": 101210, "<|LOC_914|>": 101211, "<|LOC_915|>": 101212, "<|LOC_916|>": 101213, "<|LOC_917|>": 101214, "<|LOC_918|>": 101215, "<|LOC_919|>": 101216, "<|LOC_920|>": 101217, "<|LOC_921|>": 101218, "<|LOC_922|>": 101219, "<|LOC_923|>": 101220, "<|LOC_924|>": 101221, "<|LOC_925|>": 101222, "<|LOC_926|>": 101223, "<|LOC_927|>": 101224, "<|LOC_928|>": 101225, "<|LOC_929|>": 101226, "<|LOC_930|>": 101227, "<|LOC_931|>": 101228, "<|LOC_932|>": 101229, "<|LOC_933|>": 101230, "<|LOC_934|>": 101231, "<|LOC_935|>": 101232, "<|LOC_936|>": 101233, "<|LOC_937|>": 101234, "<|LOC_938|>": 101235, "<|LOC_939|>": 101236, "<|LOC_940|>": 101237, "<|LOC_941|>": 101238, "<|LOC_942|>": 101239, "<|LOC_943|>": 101240, "<|LOC_944|>": 101241, "<|LOC_945|>": 101242, "<|LOC_946|>": 101243, "<|LOC_947|>": 101244, "<|LOC_948|>": 101245, "<|LOC_949|>": 101246, "<|LOC_950|>": 101247, "<|LOC_951|>": 101248, "<|LOC_952|>": 101249, "<|LOC_953|>": 101250, "<|LOC_954|>": 101251, "<|LOC_955|>": 101252, "<|LOC_956|>": 101253, "<|LOC_957|>": 101254, "<|LOC_958|>": 101255, "<|LOC_959|>": 101256, "<|LOC_960|>": 101257, "<|LOC_961|>": 101258, "<|LOC_962|>": 101259, "<|LOC_963|>": 101260, "<|LOC_964|>": 101261, "<|LOC_965|>": 101262, "<|LOC_966|>": 101263, "<|LOC_967|>": 101264, "<|LOC_968|>": 101265, "<|LOC_969|>": 101266, "<|LOC_970|>": 101267, "<|LOC_971|>": 101268, "<|LOC_972|>": 101269, "<|LOC_973|>": 101270, "<|LOC_974|>": 101271, "<|LOC_975|>": 101272, "<|LOC_976|>": 101273, "<|LOC_977|>": 101274, "<|LOC_978|>": 101275, "<|LOC_979|>": 101276, "<|LOC_980|>": 101277, "<|LOC_981|>": 101278, "<|LOC_982|>": 101279, "<|LOC_983|>": 101280, "<|LOC_984|>": 101281, "<|LOC_985|>": 101282, "<|LOC_986|>": 101283, "<|LOC_987|>": 101284, "<|LOC_988|>": 101285, "<|LOC_989|>": 101286, "<|LOC_990|>": 101287, "<|LOC_991|>": 101288, "<|LOC_992|>": 101289, "<|LOC_993|>": 101290, "<|LOC_994|>": 101291, "<|LOC_995|>": 101292, "<|LOC_996|>": 101293, "<|LOC_997|>": 101294, "<|LOC_998|>": 101295, "<|LOC_999|>": 101296, "<|LOC_1000|>": 101297, "<|LOC_BEGIN|>": 101298, "<|LOC_END|>": 101299, "<|LOC_SEP|>": 101300, "<|CROP_COL_SEP|>": 101301, "<|CROP_ROW_SEP|>": 101302, "<|IMAGE_SEP|>": 101303}

config.json

@@ -0,0 +1,53 @@
+{
+    "_attn_implementation": "sdpa",
+    "architectures": [
+        "Ernie4_5_MoeForCausalLM"
+    ],
+    "auto_map": {
+        "AutoConfig": "configuration_ernie4_5_moe.Ernie4_5_MoeConfig",
+        "AutoModel": "modeling_ernie4_5_moe.Ernie4_5_Model",
+        "AutoModelForCausalLM": "modeling_ernie4_5_moe.Ernie4_5_MoeForCausalLM"
+    },
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "hidden_act": "silu",
+    "hidden_size": 8192,
+    "intermediate_size": 28672,
+    "max_position_embeddings": 131072,
+    "model_type": "ernie4_5_moe",
+    "moe_capacity": [
+        64,
+        64,
+        64
+    ],
+    "moe_gate": "topk",
+    "moe_intermediate_size": 3584,
+    "moe_k": 8,
+    "moe_layer_interval": 1,
+    "moe_layer_start_index": 3,
+    "moe_num_experts": 64,
+    "moe_use_aux_free": true,
+    "num_attention_heads": 64,
+    "num_hidden_layers": 54,
+    "num_key_value_heads": 8,
+    "num_nextn_predict_layers": 1,
+    "pad_token_id": 0,
+    "rms_norm_eps": 1e-05,
+    "rope_theta": 500000,
+    "tie_word_embeddings": false,
+    "torch_dtype": "bfloat16",
+    "use_bias": false,
+    "use_cache": true,
+    "vocab_size": 103424,
+    "quantization_config": {
+        "quant_method": "exl3",
+        "version": "0.0.4",
+        "bits": 2.5,
+        "head_bits": 6,
+        "calibration": {
+            "rows": 100,
+            "cols": 2048
+        },
+        "out_scales": "auto"
+    }
+}

configuration_ernie4_5_moe.py

@@ -0,0 +1,192 @@
+# Copyright (c) 2025 Baidu, Inc. All Rights Reserved.
+#
+# 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.
+"""Ernie4_5_Moe model configuration"""
+
+from transformers import PretrainedConfig
+
+
+
+class Ernie4_5_MoeConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Ernie4_5_Model`].
+    It is used to instantiate an ERNIE-4.5 model according to the specified arguments,
+    defining the model architecture. Instantiating a configuration with the defaults
+    will yield a similar configuration to that of ERNIE-4.5-300B-A47B-PT [baidu/ERNIE-4.5-300B-A47B-PT].
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (int): Size of the vocabulary (number of unique tokens)
+        hidden_size (int): Dimensionality of the encoder layers and the pooler layer
+        intermediate_size (int): Dimensionality of the "intermediate" (feed-forward) layer
+        max_position_embeddings (int): Maximum sequence length the model can handle
+        num_hidden_layers (int): Number of hidden layers in the Transformer encoder
+        num_attention_heads (int): Number of attention heads for each attention layer
+        rms_norm_eps (float): The epsilon used by the RMS normalization layers
+        use_cache (bool): Whether to use caching for faster generation (decoding)
+        use_flash_attention (bool): Whether to use FlashAttention for optimized attention computation
+        pad_token_id (int): Token ID used for padding sequences
+        bos_token_id (int): Token ID used for beginning-of-sequence
+        eos_token_id (int): Token ID used for end-of-sequence
+        use_bias (bool): Whether to use bias terms in linear layers
+        rope_theta (float): The base period of the RoPE embeddings
+        weight_share_add_bias (bool): Whether to share bias weights in certain layers
+        ignored_index (int): Target value that is ignored during loss computation
+        attention_probs_dropout_prob (float): Dropout probability for attention weights
+        hidden_dropout_prob (float): Dropout probability for hidden layers
+        num_key_value_heads (int): Number of key/value heads (for Grouped Query Attention)
+        max_sequence_length (int): Maximum sequence length for positional embeddings
+        moe_num_experts: Number of experts in MoE layers
+        moe_capacity: Capacity configuration for MoE layers
+        moe_layer_interval: Interval between MoE layers
+        moe_layer_start_index: Starting layer index for MoE
+        moe_layer_end_index: Ending layer index for MoE (-1 means last layer)
+        sinkhorn_2gate: Whether to use sinkhorn 2-gate routing
+        sinkhorn_temp: Temperature for sinkhorn routing
+        moe_dropout_prob: Dropout probability for MoE layers
+        moe_gate: Type of gating mechanism ('top2', etc.)
+        moe_intermediate_size: Intermediate size for MoE layers
+        moe_gate_act: Activation function for gating
+        moe_k: Number of experts to route to
+        **kwargs: Additional base model configuration parameters
+    """
+
+    model_type = "ernie4_5_moe"
+    use_keep_in_fp32_modules = True
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    attribute_map = {
+        "n_positions": "max_position_embeddings",
+        "n_embd": "hidden_size",
+        "n_layer": "num_hidden_layers",
+        "n_head": "num_attention_heads",
+        "n_inner": "intermediate_size",
+        "activation_function": "hidden_act",
+    }
+
+    # Default tensor parallel plan for base model `ernie_4_5_moe`
+    base_model_tp_plan = {
+        "model.layers.*.self_attn.q_proj": "colwise_rep",
+        "model.layers.*.self_attn.k_proj": "colwise_rep",
+        "model.layers.*.self_attn.v_proj": "colwise_rep",
+        "model.layers.*.self_attn.o_proj": "rowwise_rep",
+        "model.layers.*.mlp.experts.*.gate_proj": "colwise",
+        "model.layers.*.mlp.experts.*.up_proj": "colwise",
+        "model.layers.*.mlp.experts.*.down_proj": "rowwise",
+        "model.layers.*.mlp.gate_proj": "colwise",
+        "model.layers.*.mlp.up_proj": "colwise",
+        "model.layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=768,
+        intermediate_size=11008,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_key_value_heads=None,
+        max_position_embeddings=32768,
+        use_sliding_window=None,
+        sliding_window=None,
+        rms_norm_eps=1e-6,
+        use_cache=False,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        attention_probs_dropout_prob=0.0,
+        hidden_dropout_prob=0.0,
+        rope_theta=10000.0,
+        use_flash_attention=False,
+        use_rmsnorm=True,
+        use_bias=False,
+        weight_share_add_bias=True,
+        max_sequence_length=None,
+        ignored_index=-100,
+        use_moe=True,
+        moe_num_experts=64,
+        moe_capacity=(64, 64, 64),
+        moe_layer_interval=2,
+        moe_layer_start_index=0,
+        moe_layer_end_index=-1,
+        sinkhorn_2gate=True,
+        sinkhorn_temp=3e-2,
+        moe_dropout_prob=0.0,
+        moe_gate="top2",
+        moe_intermediate_size=3584,
+        moe_k=2,
+        moe_gate_act="softmax",
+        moe_use_aux_free=False,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.use_rmsnorm = use_rmsnorm
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.max_sequence_length = max_sequence_length
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.ignored_index = ignored_index
+        self.use_cache = use_cache
+        self.use_bias = use_bias
+        self.weight_share_add_bias = weight_share_add_bias
+        self.use_flash_attention = use_flash_attention
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.hidden_dropout_prob = hidden_dropout_prob
+
+        self.use_moe = moe_num_experts > 0 and use_moe
+        self.moe_num_experts = moe_num_experts
+        self.moe_capacity = moe_capacity
+        self.sinkhorn_2gate = sinkhorn_2gate
+        self.sinkhorn_temp = sinkhorn_temp
+        self.moe_layer_interval = moe_layer_interval
+        self.moe_dropout_prob = moe_dropout_prob
+        self.moe_gate = moe_gate
+        self.moe_intermediate_size = moe_intermediate_size
+        self.moe_k = moe_k
+        self.moe_layer_start_index = moe_layer_start_index
+        self.moe_layer_end_index = self.num_hidden_layers - 1 if moe_layer_end_index == -1 else moe_layer_end_index
+        self.moe_gate_act = moe_gate_act
+        self.moe_use_aux_free = moe_use_aux_free
+
+        # Set default for tied embeddings if not specified.
+        if "tie_word_embeddings" not in kwargs:
+            kwargs["tie_word_embeddings"] = False
+        
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "do_sample": true,
+  "top_p": 0.8,
+  "temperature": 0.8,
+  "repetition_penalty": 1.0,
+  "frequency_penalty": 0.0,
+  "presence_penalty": 0.0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "transformers_version": "4.52.4",
+  "use_cache": true
+}

model-00001-of-00014.safetensors

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

model-00002-of-00014.safetensors

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

model-00003-of-00014.safetensors

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

model-00004-of-00014.safetensors

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

model-00005-of-00014.safetensors

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

model-00006-of-00014.safetensors

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

model-00007-of-00014.safetensors

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

model-00008-of-00014.safetensors

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

model-00009-of-00014.safetensors

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

model-00010-of-00014.safetensors

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

model-00011-of-00014.safetensors

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

model-00012-of-00014.safetensors

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

model-00013-of-00014.safetensors

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

model-00014-of-00014.safetensors

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

modeling_ernie4_5_moe.py

@@ -0,0 +1,1590 @@
+# Copyright (c) 2025 Baidu, Inc. All Rights Reserved.
+#
+# 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.
+""" Ernie4_5_Moe model """
+
+from copy import deepcopy
+from dataclasses import dataclass
+from functools import partial
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+from transformers.cache_utils import (
+    Cache,
+    DynamicCache,
+    SlidingWindowCache,
+    StaticCache,
+)
+from transformers.generation import GenerationMixin
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput, MoeCausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    LossKwargs,
+    auto_docstring,
+    can_return_tuple,
+    logging,
+    is_torch_flex_attn_available,
+)
+
+from .configuration_ernie4_5_moe import Ernie4_5_MoeConfig
+
+
+if is_torch_flex_attn_available():
+    from torch.nn.attention.flex_attention import BlockMask
+
+    from transformers.integrations.flex_attention import make_flex_block_causal_mask
+
+logger = logging.get_logger(__name__)
+
+
+class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs):
+    """Kwargs class used during autoregressive generation"""
+
+    ...
+
+
+@dataclass
+class Erine4_5_MoeModelOutputWithPast(ModelOutput):
+    """Class for Ernie4_5_Moe model outputs with past keys."""
+
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+    router_loss: Optional[torch.FloatTensor] = None
+    gate_logits: Optional[tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class Ernie4_5_MoeCausalLMOutputWithPast(MoeCausalLMOutputWithPast):
+    """Class for  Ernie4_5_Moe causal LM output with past keys"""
+
+    router_loss: Optional[torch.FloatTensor] = None
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+
+    x1 = x[..., 0::2]
+    x2 = x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).reshape(x.shape)
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    orig_dtype = q.dtype
+    sin_pos = torch.stack([sin, sin], dim=-1).reshape(*sin.shape[:-1], -1)
+    cos_pos = torch.stack([cos, cos], dim=-1).reshape(*sin.shape[:-1], -1)
+    q_embed = (q.float() * cos_pos) + (rotate_half(q).float() * sin_pos)
+    k_embed = (k.float() * cos_pos) + (rotate_half(k).float() * sin_pos)
+    return q_embed.to(orig_dtype), k_embed.to(orig_dtype)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    """
+    Eager attention for Ernie4_5_Attention forward function.
+    """
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask.to(attn_weights.device)
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query.dtype
+    )
+    attn_weights = nn.functional.dropout(
+        attn_weights, p=dropout, training=module.training
+    )
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+def topk_gate_func(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+):
+    """
+    Topk gate function for Ernie4_5_MoEMlp
+    """
+    capacity = module.get_capacity(hidden_states.shape[0])
+    with torch.autocast(device_type="cuda", dtype=torch.float32):
+        logits = module.gate(hidden_states.float())
+    router_loss = torch.zeros([1], dtype=torch.float32, device=hidden_states.device)
+    router_loss.detach()
+    return logits, capacity, router_loss
+
+
+class Ernie4_5_ResidualWithDropout(nn.Module):
+    """
+    Fused dropout implementation with residual connection support.
+
+    This layer combines dropout and residual addition in a single operation for better performance,
+    particularly on GPU devices. The dropout is conditionally applied based on the probability.
+
+    Args:
+        prob (float): Dropout probability (between 0 and 1)
+
+    Attributes:
+        prob (float): Stores the dropout probability
+        dropout (nn.Dropout): The actual dropout layer instance
+    """
+
+    def __init__(self, prob):
+        """
+        Initialize the fused dropout layer.
+
+        Args:
+            prob (float): Dropout probability (0 means no dropout)
+        """
+        super().__init__()
+        self.prob = prob
+        self.dropout = nn.Dropout(p=prob)
+
+    def forward(self, x, y):
+        """
+        Forward pass of the fused dropout layer.
+
+        Args:
+            x (torch.Tensor): Input tensor to potentially apply dropout on
+            y (torch.Tensor): Residual tensor to add to the (possibly dropped out) x
+
+        Returns:
+            torch.Tensor: Result of x (with optional dropout) + y
+        """
+        if self.prob > 0:
+            x = self.dropout(x)
+        output = x + y
+
+        return output
+
+
+class Ernie4_5_Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config, layer_idx=0):
+        """
+        Args:
+            config (ErnieConfig): Model configuration.
+            layer_idx (int, optional): Index in transformer stack. Defaults to 0.
+        """
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = (
+            config.num_key_value_heads
+            if config.num_key_value_heads is not None
+            else self.nums_head
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.head_dim = self.hidden_size // self.num_heads
+        self.freq_allocation = (
+            config.freq_allocation if hasattr(config, "freq_allocation") else 0
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = getattr(config, "attention_probs_dropout_prob", 0.0)
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            self.hidden_size,
+            self.num_heads * self.head_dim,
+            bias=config.use_bias,
+        )
+
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.use_bias,
+        )
+
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.use_bias,
+        )
+
+        self.o_proj = nn.Linear(
+            self.hidden_size,
+            self.hidden_size,
+            bias=config.use_bias,
+        )
+
+        self.config = config
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Cache] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[
+        torch.Tensor,
+        Optional[torch.Tensor],
+        Optional[Tuple[torch.Tensor, torch.Tensor]],
+    ]:
+        """
+        Ernie4_5_Attention forward function
+        """
+        B, L = hidden_states.shape[:-1]
+
+        query_states = (
+            self.q_proj(hidden_states).view(B, L, self.num_heads, -1).transpose(1, 2)
+        )
+        key_states = (
+            self.k_proj(hidden_states)
+            .view(B, L, self.num_key_value_heads, -1)
+            .transpose(1, 2)
+        )
+        value_states = (
+            self.v_proj(hidden_states)
+            .view(B, L, self.num_key_value_heads, -1)
+            .transpose(1, 2)
+        )
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"cache_position": cache_position}
+            key_states, value_states = past_key_value.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(B, L, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class Ernie4_5_MLP(nn.Module):
+    """
+    Ernie4_5_MLP - Gated Multi-Layer Perceptron module used in Ernie model.
+    """
+
+    def __init__(self, config, intermediate_size=None):
+        """
+        Initialize the MLP module with configuration options.
+
+        Args:
+            config: Model configuration object with attributes:
+                - hidden_size: int
+                - intermediate_size: int
+                - use_bias: bool
+            layer_idx (int): Index of current layer (default: 0)
+        """
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = (
+            intermediate_size
+            if intermediate_size is not None
+            else config.intermediate_size
+        )
+        self.gate_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=config.use_bias
+        )
+        self.up_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=config.use_bias
+        )
+        self.down_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=config.use_bias
+        )
+
+    def forward(self, x):
+        """
+        Args:
+            x (Tensor): shape [batch_size, seq_len, hidden_size]
+
+        Returns:
+            Tensor: shape [batch_size, seq_len, hidden_size]
+        """
+        down_proj = self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class Ernie4_5_MoeStatics(nn.Module):
+    """
+    Stores MoE (Mixture of Experts) statistics
+    and expert usage information.
+    """
+
+    def __init__(self, config):
+        """
+        Initialize MoE statistics tracking.
+
+        Args:
+            config: Model configuration containing MoE parameters
+        """
+        super().__init__()
+
+        num_experts = config.moe_num_experts
+        num_experts_groups = 1
+
+        self.e_score_correction_bias = nn.Parameter(
+            torch.zeros(num_experts_groups, num_experts, dtype=torch.float32),
+            requires_grad=False,
+        )
+
+
+class Ernie4_5_MoeMLP(nn.Module):
+    """Mixture of Experts (MoE) variant of ERNIE's MLP layer."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.k = config.moe_k
+        self.sinkhorn_2gate = config.sinkhorn_2gate
+        self.sinkhorn_temp = config.sinkhorn_temp
+
+        moe_intermediate_size = (
+            config.moe_intermediate_size
+            if config.moe_intermediate_size
+            else config.intermediate_size
+        )
+        self.gate = nn.Linear(
+            config.hidden_size, config.moe_num_experts, bias=False, dtype=torch.float32
+        )
+        if config.moe_gate_act == "softmax":
+            self.gate_act = partial(F.softmax, dim=-1)
+        elif config.moe_gate_act == "sigmoid":
+            self.gate_act = F.sigmoid
+        else:
+            raise ValueError(f"{config.moe_gate_act} is not supported.")
+
+        self.experts = nn.ModuleList(
+            [
+                Ernie4_5_MLP(config, moe_intermediate_size)
+                for i in range(config.moe_num_experts)
+            ]
+        )
+
+        if config.moe_use_aux_free:
+            self.moe_statics = Ernie4_5_MoeStatics(config)
+
+        self.use_correction_bias = config.moe_use_aux_free
+        self.num_local_experts = len(self.experts)
+
+        self.shared_experts = self._init_shared_experts()
+
+    def _init_shared_experts(self):
+        """
+        Initialize the shared expert module.
+
+        Returns:
+            shared_experts: Shared expert module, returns None if no shared experts are needed.
+
+        """
+        cfg = deepcopy(self.config)
+        if getattr(cfg, "moe_num_shared_experts", 0) > 0:
+            if getattr(cfg, "moe_intermediate_size", None):
+                cfg.intermediate_size = (
+                    cfg.moe_intermediate_size * cfg.moe_num_shared_experts
+                )
+            else:
+                cfg.intermediate_size = (
+                    cfg.intermediate_size * cfg.moe_num_shared_experts
+                )
+            shared_experts = Ernie4_5_MLP(cfg, cfg.intermediate_size)
+        else:
+            shared_experts = None
+        return shared_experts
+
+    def forward(
+        self,
+        input: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Forward pass through MoE layer.
+
+        Args:
+            input (Tensor): Input tensor of shape [s, d].
+            token_type_ids: Optional tensor for token types.
+
+        Returns:
+            tuple: (output, combine_weights, router_loss, gate_logits)
+        """
+
+        if input.dim() == 3:
+            orig_shape = input.shape
+            input = input.reshape(-1, input.shape[-1])
+        else:
+            orig_shape = None
+        assert (
+            input.dim() == 2
+        ), f"input Tensor must have dimensions: (s)equence, (d)im, got:{input.shape}"
+
+        assert self.gate is not None
+
+        gate_input = input
+
+        (
+            dispatched_input,
+            combine_weights,
+            dispatch_mask,
+            scatter_index,
+            router_loss,
+            gate_logits,
+            gate_prob,
+        ) = self.gate_and_dispatch(gate_input)
+
+        expert_out = self.forward_experts(dispatched_input)
+
+        combined_output = self.combine_expert_output(
+            expert_out, combine_weights, scatter_index
+        )
+
+        if self.shared_experts is not None:
+            shared_expert_out = self.shared_experts(gate_input)
+            combined_output += shared_expert_out
+
+        if orig_shape:
+            combined_output = combined_output.reshape(
+                orig_shape[:-1] + (combined_output.shape[-1],)
+            )
+
+        return combined_output, combine_weights, router_loss, gate_logits
+
+    def forward_experts(self, dispatched_input: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass through experts sequentially.
+
+        Args:
+            dispatched_input (Tensor): Input tensor of shape [num_experts, capacity, dim].
+
+        Returns:
+            Tensor: Expert outputs of shape [num_experts, capacity, dim].
+        """
+        true_experts = self.experts
+        dispatched_input = dispatched_input.reshape(
+            1, self.num_local_experts, -1, dispatched_input.shape[-1]
+        )
+        expert_outputs = []
+        if isinstance(self.experts, nn.ModuleList):
+            chunks = dispatched_input.permute(1, 0, 2, 3).contiguous().unbind(0)
+            assert len(chunks) == len(
+                true_experts
+            ), f"{len(chunks)}, {len(true_experts)}"
+            for chunk, expert in zip(chunks, true_experts):
+                expert_outputs.append(expert(chunk))
+        else:
+            dispatched_input = dispatched_input.permute(1, 0, 2, 3).contiguous()
+            orig_shape = dispatched_input.shape
+            chunks = dispatched_input.reshape(orig_shape[0], -1, orig_shape[-1])
+            chunks = self.experts(chunks)
+            chunks = chunks.reshape(orig_shape[:-1] + (chunks.shape[-1],)).unbind(0)
+            expert_outputs.extend(chunks)
+
+        expert_output = torch.stack(expert_outputs, dim=1)
+        return expert_output
+
+    def moe_gate_dispatch(
+        self,
+        x: torch.Tensor,
+        gate_logits: torch.Tensor,
+        k: int,
+        capacity: Optional[int],
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Dispatch inputs to experts based on their routing probabilities.
+        """
+        S, H = x.shape
+        E = gate_logits.shape[1]
+        device = x.device
+        topk_prob, topk_idx = torch.topk(gate_logits, k, dim=-1)
+        combine_weights = topk_prob
+        expert_id = topk_idx
+        y = x.new_zeros((E, capacity, H))
+        scatter_index = x.new_full((k, S), -1, dtype=torch.int32)
+
+        # per-expert slot counters
+        slot_counter = torch.zeros(E, dtype=torch.int32, device=device)
+
+        for tok in range(S):
+            for route in range(k):
+                e = expert_id[tok, route].item()
+                slot = slot_counter[e].item()
+                if slot >= capacity:
+                    combine_weights[tok, route] = 0.0
+                    continue
+
+                # record mapping & dispatch activation
+                scatter_index[route, tok] = e * capacity + slot
+                y[e, slot] = x[tok]
+                slot_counter[e] += 1
+
+        expert_offset = torch.cumsum(slot_counter, 0, dtype=torch.int64)
+
+        return y, combine_weights, scatter_index, expert_offset, expert_id
+
+    def combine_expert_output(
+        self,
+        expert_output: torch.Tensor,
+        combine_weights: torch.Tensor,
+        scatter_index: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Combine expert outputs using combination weights.
+
+        Args:
+            expert_output (Tensor): Expert outputs [num_experts, capacity, dim].
+            combine_weights (Tensor): Combination weights.
+            scatter_index (Tensor): Scatter indices.
+
+        Returns:
+            Tensor: Combined output [seqlen, dim].
+        """
+        expert_output = expert_output.reshape(-1, expert_output.shape[-1])
+        combined_output = self.combining(expert_output, combine_weights, scatter_index)
+        return combined_output
+
+    def combining(self, x, combine_weights, scatter_index):
+        """
+        Combines and aggregates input matrix using combination weights.
+
+        Args:
+            x (Tensor): Input tensor of shape [num_experts * capacity, dim]
+            combine_weights (Tensor): Combination weights of shape [seq, 2]
+            scatter_index (Tensor): Scatter indices of shape [seq, 2]
+
+        Returns:
+            Tensor: Combined output tensor of shape [seq, dim]
+        """
+        dim = x.shape[-1]
+
+        scatter_index = scatter_index.reshape([-1])
+        num_k = combine_weights.shape[-1]
+
+        combine_weights = combine_weights.unsqueeze(1)
+
+        x = x[scatter_index].reshape([-1, num_k, dim])
+
+        return torch.matmul(combine_weights, x).squeeze(1)
+
+    def gate_and_dispatch(self, input):
+        """
+        Calculate gate and dispatch inputs.
+
+        Args:
+            input: Input tensor of shape [seq, dim]
+
+        Returns:
+            tuple: (dispatched_input, combine_weights, dispatch_mask,
+            scatter_index, router_loss, gate_logits, gate_prob)
+        """
+        gate_logits, capacity, router_loss = topk_gate_func(self, input)
+
+        # capacity no use
+        prob = self.gate_act(gate_logits)
+        (
+            dispatched_input,
+            combine_weights_unnorm,
+            scatter_index,
+            dispatch_mask,
+            _,
+        ) = self.moe_gate_dispatch(input, prob, k=self.k, capacity=capacity)
+        dispatch_mask = torch.diff(F.pad(dispatch_mask, (1, 0)))
+
+        scatter_index.detach()
+        dispatch_mask.detach()
+
+        scatter_index = scatter_index.transpose(0, 1)  # [k, s] -> [s, k]
+        combine_weights = combine_weights_unnorm / torch.clamp(
+            combine_weights_unnorm.sum(dim=-1, keepdim=True), min=1e-12
+        )
+        combine_weights = combine_weights.to(dtype=dispatched_input.dtype)
+
+        return (
+            dispatched_input,
+            combine_weights,
+            dispatch_mask,
+            scatter_index,
+            router_loss,
+            gate_logits,
+            prob,
+        )
+
+    def get_capacity(self, num_tokens, cap_factor=None):
+        """
+        Calculate capacity based on number of tokens.
+
+        Args:
+            num_tokens: Number of input tokens
+            cap_factor: Optional capacity factor override
+
+        Returns:
+            int: Calculated capacity
+        """
+        num_experts = self.config.moe_num_experts
+        if cap_factor is not None:
+            cap = cap_factor
+        else:
+            if self.training:
+                cap = self.config.moe_capacity[0]
+            elif num_tokens < num_experts:
+                cap = self.config.moe_capacity[2]
+            else:
+                cap = self.config.moe_capacity[1]
+
+        capacity = int(cap * num_tokens // num_experts)
+        assert (
+            capacity > 0
+        ), f"requires capacity to >= 0. cap={cap}, num_tokens={num_tokens}"
+        return capacity
+
+
+class Ernie4_5_RMSNorm(nn.Module):
+    """
+    Ernie Root Mean Square Layer Normalization (Ernie4_5_RMSNorm) implementation.
+
+    Ernie4_5_RMSNorm is a simplified version of LayerNorm that focuses on the root mean square of inputs,
+    omitting the mean-centering operation. This provides computational efficiency while maintaining
+    good performance.
+
+    """
+
+    def __init__(self, config):
+        """
+        Initialize RMSNorm layer.
+
+        Args:
+            config (ErnieConfig): Model configuration.
+        """
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.weight = nn.Parameter(torch.ones(config.hidden_size))
+        self.variance_epsilon = config.rms_norm_eps
+
+    def forward(self, hidden_states):
+        """
+        Apply RMS normalization to input hidden states.
+
+        Args:
+            hidden_states (Tensor): Input tensor of shape [batch_size, seq_len, hidden_size]
+
+        Returns:
+            Tensor: Normalized output tensor of same shape as input
+        """
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class Ernie4_5_RopeEmbedding(nn.Module):
+    """
+    Implements Rotary Position Embedding (RoPE) for Ernie4_5_MoeModel.
+    """
+
+    def __init__(self, config: Ernie4_5_MoeConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get(
+                "rope_type", config.rope_scaling.get("type")
+            )
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, None, :].float()
+        position_ids_expanded = position_ids[..., None].float()
+        freqs = inv_freq_expanded.float() * position_ids_expanded.float()
+        cos = torch.cos(freqs) * self.attention_scaling
+        sin = torch.sin(freqs) * self.attention_scaling
+        return cos, sin
+        # return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class Ernie4_5_DecoderLayer(nn.Module):
+    """A single transformer decoder layer in ERNIE-MoE model.
+
+    Contains self-attention and feed-forward components with optional MoE (Mixture of Experts)
+    support, residual connections, and layer normalization.
+    """
+
+    def __init__(self, config, layer_idx):
+        """Initialize the decoder layer.
+
+        Args:
+            config (ErnieMoEConfig): Model configuration.
+            layer_idx (int): Index of this layer in the transformer stack
+        """
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.layer_idx = layer_idx
+        self.config = config
+        self.use_moe = config.use_moe
+        self.self_attn = Ernie4_5_Attention(config, layer_idx)
+
+        moe_layer_start_index = (
+            min(config.moe_layer_start_index)
+            if isinstance(config.moe_layer_start_index, (tuple, list))
+            else config.moe_layer_start_index
+        )
+        moe_layer_end_index = (
+            max(config.moe_layer_end_index)
+            if isinstance(config.moe_layer_end_index, (tuple, list))
+            else config.moe_layer_end_index
+        )
+
+        if (
+            self.use_moe
+            and ((layer_idx + 1) % config.moe_layer_interval == 0)
+            and layer_idx >= moe_layer_start_index
+            and layer_idx <= moe_layer_end_index
+        ):
+            self.mlp = Ernie4_5_MoeMLP(config)
+        else:
+            self.mlp = Ernie4_5_MLP(config)
+
+        self.input_layernorm = Ernie4_5_RMSNorm(config)
+        self.post_attention_layernorm = Ernie4_5_RMSNorm(config)
+
+        self.residual_add1 = Ernie4_5_ResidualWithDropout(config.hidden_dropout_prob)
+        self.residual_add2 = Ernie4_5_ResidualWithDropout(config.hidden_dropout_prob)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[
+            tuple[torch.Tensor, torch.Tensor]
+        ] = None,  # necessary, but kept here for BC
+        output_router_loss: bool = True,
+        output_gate_logits: bool = True,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[
+        torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        """Forward pass through the decoder layer.
+
+        Args:
+            hidden_states (torch.Tensor): Input tensor [batch_size, seq_len, hidden_size]
+            attention_mask (Optional[torch.Tensor]): Attention mask tensor
+            position_ids (Optional[torch.Tensor]): Position indices for rotary embeddings
+            past_key_value (Optional[Tuple[torch.Tensor]]): Cached key/value states
+            output_attentions (Optional[bool]): Whether to return attention weights
+            use_cache (Optional[bool]): Whether to cache key/value states
+            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+                Indices depicting the position of the input sequence tokens in the sequence.
+            position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+                with `head_dim` being the embedding dimension of each attention head.
+            output_router_loss (bool): Whether to return MoE router loss
+            output_gate_logits (bool): Whether to return MoE gate logits
+
+        Returns:
+            Union: Various output combinations depending on arguments:
+                - Base case: Hidden states tensor
+                - With attention: Tuple of (hidden_states, attention_weights)
+                - With router loss: May include gate logits in output tuple
+                - With MoE gate logits: May include gate logits in output tuple
+        """
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            past_key_value=past_key_value,
+            position_ids=position_ids,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+
+        hidden_states = self.residual_add1(hidden_states, residual)
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+
+        router_loss = None
+        gate_logits = None
+
+        if isinstance(self.mlp, Ernie4_5_MoeMLP):
+            hidden_states, _, router_loss, gate_logits = self.mlp(hidden_states)
+        else:
+            hidden_states = self.mlp(hidden_states)
+
+        hidden_states = self.residual_add2(hidden_states, residual)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if output_router_loss:
+            outputs += (router_loss,)
+
+        if output_gate_logits:
+            outputs += (gate_logits,)
+
+        return outputs
+
+
+@auto_docstring
+class Ernie4_5_PretrainedModel(PreTrainedModel):
+    """Base class for ERNIE pretrained models."""
+
+    config_class = Ernie4_5_MoeConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Ernie4_5_DecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_cache_class = True
+    _supports_quantized_cache = True
+    _supports_static_cache = False  # MoE models don't work with torch.compile (`torch.where(condition)` not supported)
+
+
+def subbatch(f, arg_idx, axis, bs, out_idx, same_arg_idx={}):
+    """
+    Converts a function to one that applies to subbatch of an input dimension.
+    Useful for processing large tensors in smaller chunks to reduce memory usage.
+
+    Args:
+        f (Callable): Function to be subbatched.
+        arg_idx ([int]): Indices of the inputs to be subbatched.
+        axis ([int]): Indices of the dimensions to be subbatched for each input.
+        bs (int): Subbatch size.
+        out_idx (int): Dimension to concatenate outputs along.
+        same_arg_idx (dict): Mapping of argument indices that share the same tensor.
+
+    Returns:
+        Callable: New function that processes inputs in subbatches.
+    """
+
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+
+        assert len(arg_idx) == len(
+            axis
+        ), "Number of batching args and number of batching dims should match."
+
+        inps = [args[i] for i in arg_idx]
+        axis_width = [inp.shape[d] for inp, d in zip(inps, axis)]
+        assert len(set(axis_width)) == 1, "Batch sizes should be kept equal."
+
+        inp_axis = {idx: d for idx, d in zip(arg_idx, axis)}
+
+        axis_width = axis_width[0]
+        if axis_width < bs:
+            return f(*args, **kwargs)
+
+        outs = []
+        for slice_at in range(0, axis_width, bs):
+            _args = []
+            for i, inp in enumerate(args):
+                if i in same_arg_idx:
+                    assert (
+                        i > same_arg_idx[i]
+                    ), f"expect i > same_arg_idx[i], but got i: {i} and same_arg_idx[i]: {same_arg_idx[i]}"
+                    _args.append(_args[same_arg_idx[i]])
+                elif i in arg_idx:
+                    d = inp_axis[i]
+                    start = slice_at
+                    end = min(inp.shape[d], slice_at + bs)
+                    # Build slice for all dims, only slice along axis d
+                    slices = [slice(None)] * inp.ndim
+                    slices[d] = slice(start, end)
+                    _args.append(inp[tuple(slices)])
+                else:
+                    _args.append(inp)
+
+            out = f(*_args, **kwargs)
+            outs.append(out)
+
+        return torch.cat(outs, dim=out_idx)
+
+    return wrapper
+
+
+class ErniePretrainingCriterion(nn.Module):
+    """Criterion for ERNIE pretraining task."""
+
+    def __init__(self, config, return_tuple=True):
+        """Initialize the pretraining criterion.
+
+        Args:
+            config (ErnieConfig): Model configuration.
+            return_tuple (bool): Whether to return loss as tuple (loss, loss_sum). Defaults to True.
+        """
+        super().__init__()
+        self.ignored_index = getattr(config, "ignored_index", -100)
+        self.config = config
+        self.return_tuple = return_tuple
+
+        self.loss_func = nn.CrossEntropyLoss(reduction="none")
+
+    def forward(self, prediction_scores, masked_lm_labels, loss_mask, router_loss=None):
+        """Compute the combined pretraining loss.
+
+        Args:
+            prediction_scores: Prediction scores tensor, [batch_size, seq_len, vocab_size]
+            masked_lm_labels: Target labels tensor [batch_size, seq_len]
+            loss_mask: Optional mask for valid tokens
+            router_loss: Optional MoE router loss tensor
+
+        Returns:
+            Union:
+                - If return_tuple=True: Tuple of (combined_loss, mlm_loss_sum)
+                - If return_tuple=False: Combined loss tensor
+        """
+        res = self.forward_impl(prediction_scores, masked_lm_labels, loss_mask)
+
+        if self.return_tuple:
+            loss, loss_sum = res
+        else:
+            loss, loss_sum = res, None
+
+        if router_loss is not None and isinstance(router_loss, torch.Tensor):
+            loss = loss + router_loss - router_loss.detach()
+
+        return loss, loss_sum
+
+    def loss_impl(
+        self, prediction_scores: torch.Tensor, masked_lm_labels: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Core loss computation without reduction (but per-token).
+
+        Args:
+            prediction_scores (torch.Tensor): Logits tensor [batch_size, seq_len, vocab_size].
+            masked_lm_labels (torch.Tensor): Target labels tensor [batch_size, seq_len].
+
+        Returns:
+            torch.Tensor: Unreduced loss tensor of shape [batch_size, seq_len].
+                          Losses are calculated in float32.
+        """
+        scores_float32 = prediction_scores.to(torch.float32)
+        # prediction_scores: [batch_size, seq_len, vocab_size]
+        # masked_lm_labels: [batch_size, seq_len]
+        # Transpose prediction_scores to [batch_size, vocab_size, seq_len]
+        unreduced_loss = self.loss_func(
+            scores_float32.transpose(1, 2),  # Shape: [batch_size, vocab_size, seq_len]
+            masked_lm_labels.long(),  # Shape: [batch_size, seq_len], ensure long type
+        )
+        # unreduced_loss will be of shape [batch_size, seq_len] and dtype float32
+        return unreduced_loss
+
+    def forward_impl(self, prediction_scores, masked_lm_labels, loss_mask=None):
+        """
+        Loss function forward pass implementation.
+        """
+        prediction_scores_dims = len(prediction_scores.shape)
+
+        loss_subbatch_seqlen_config_key = "loss_subbatch_seqlen"
+        default_loss_subbatch_seqlen = 32768
+
+        current_loss_subbatch_seqlen = self.config.get(
+            loss_subbatch_seqlen_config_key, default_loss_subbatch_seqlen
+        )
+
+        if (
+            prediction_scores_dims == 2
+            and prediction_scores.shape[0] > current_loss_subbatch_seqlen
+        ):
+            sb_loss_func = subbatch(
+                self.loss_impl, [0, 1], [0, 0], current_loss_subbatch_seqlen, 0
+            )
+            masked_lm_loss = sb_loss_func(prediction_scores, masked_lm_labels)
+        elif (
+            prediction_scores_dims == 3
+            and prediction_scores.shape[1] > current_loss_subbatch_seqlen
+        ):
+            sb_loss_func = subbatch(
+                self.loss_impl, [0, 1], [1, 1], current_loss_subbatch_seqlen, 1
+            )
+            masked_lm_loss = sb_loss_func(prediction_scores, masked_lm_labels)
+        else:
+            masked_lm_loss = self.loss_impl(prediction_scores, masked_lm_labels)
+
+        if loss_mask is None:
+            loss_mask = masked_lm_labels != self.ignored_index
+
+        loss_mask = loss_mask.reshape(-1).to(torch.float32)
+
+        masked_lm_loss = torch.sum(
+            masked_lm_loss.to(torch.float32).reshape(-1) * loss_mask
+        )
+
+        # The division will be in float32
+        loss = masked_lm_loss / loss_mask.sum()
+
+        loss_sum = masked_lm_loss.sum().detach()
+
+        if not self.return_tuple:
+            if self.training:
+                return loss
+            return loss_sum
+        return loss, loss_sum
+
+
+@auto_docstring
+class Ernie4_5_Model(Ernie4_5_PretrainedModel):
+    """The core ERNIE transformer model with MoE (Mixture of Experts) support."""
+
+    _keep_in_fp32_modules = ["gate"]
+
+    def __init__(self, config: Ernie4_5_MoeConfig):
+        """Initialize the ERNIE model architecture."""
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.hidden_size = config.hidden_size
+        self.config = config
+
+        self.embed_tokens = nn.Embedding(
+            self.vocab_size,
+            self.hidden_size,
+        )
+
+        self.layers = nn.ModuleList(
+            [Ernie4_5_DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = Ernie4_5_RMSNorm(config)
+        self.rotary_emb = Ernie4_5_RopeEmbedding(config=config)
+
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        """Get the input embedding layer."""
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        """Set new input embeddings."""
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
+    ):
+        """Forward pass through the ERNIE model."""
+        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
+        )
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        inputs_embeds = inputs_embeds.to(self.embed_tokens.weight.dtype)
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(
+            attention_mask,
+            inputs_embeds,
+            cache_position,
+            past_key_values,
+            output_attentions,
+        )
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        all_router_loss = (
+            torch.tensor(0.0, device=inputs_embeds.device)
+            if self.config.use_moe
+            else None
+        )
+        all_gate_logits = ()
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    partial(decoder_layer.__call__, **flash_attn_kwargs),
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                    position_embeddings,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    causal_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    use_cache,
+                    cache_position,
+                    position_embeddings,
+                    **flash_attn_kwargs,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+            if self.config.use_moe:
+                layer_outputs, gate_logits = layer_outputs[:-1], layer_outputs[-1]
+                all_gate_logits = all_gate_logits + (gate_logits,)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        # assert all_router_loss is None, f'moe not support `return-dict`'
+        return Erine4_5_MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            router_loss=all_router_loss,
+            gate_logits=all_gate_logits,
+        )
+
+    def _update_causal_mask(
+        self,
+        attention_mask: Union[torch.Tensor, "BlockMask"],
+        input_tensor: torch.Tensor,
+        cache_position: torch.Tensor,
+        past_key_values: Cache,
+        output_attentions: bool = False,
+    ):
+        if self.config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and past_key_values is not None:
+                is_padding_right = (
+                    attention_mask[:, -1].sum().item() != input_tensor.size()[0]
+                )
+                if is_padding_right:
+                    raise ValueError(
+                        "You are attempting to perform batched generation with padding_side='right'"
+                        " this may lead to unexpected behaviour for Flash Attention version of Qwen3. Make sure to "
+                        " call `tokenizer.padding_side  = 'left'` before tokenizing the input. "
+                    )
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+        if self.config._attn_implementation == "flex_attention":
+            if isinstance(attention_mask, torch.Tensor):
+                attention_mask = make_flex_block_causal_mask(attention_mask)
+            return attention_mask
+
+        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
+        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
+        # to infer the attention mask.
+        past_seen_tokens = (
+            past_key_values.get_seq_length() if past_key_values is not None else 0
+        )
+        using_static_cache = isinstance(past_key_values, StaticCache)
+        using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)
+
+        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
+        if (
+            self.config._attn_implementation == "sdpa"
+            and not (using_static_cache or using_sliding_window_cache)
+            and not output_attentions
+        ):
+            if AttentionMaskConverter._ignore_causal_mask_sdpa(
+                attention_mask,
+                inputs_embeds=input_tensor,
+                past_key_values_length=past_seen_tokens,
+                sliding_window=self.config.sliding_window,
+                is_training=self.training,
+            ):
+                return None
+
+        dtype = input_tensor.dtype
+        min_dtype = torch.finfo(dtype).min
+        sequence_length = input_tensor.shape[1]
+        # SlidingWindowCache or StaticCache
+        if using_sliding_window_cache or using_static_cache:
+            target_length = past_key_values.get_max_cache_shape()
+        # DynamicCache or no cache
+        else:
+            target_length = (
+                attention_mask.shape[-1]
+                if isinstance(attention_mask, torch.Tensor)
+                else past_seen_tokens + sequence_length + 1
+            )
+
+        # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
+        causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
+            attention_mask,
+            sequence_length=sequence_length,
+            target_length=target_length,
+            dtype=dtype,
+            cache_position=cache_position,
+            batch_size=input_tensor.shape[0],
+            config=self.config,
+            past_key_values=past_key_values,
+        )
+
+        if (
+            self.config._attn_implementation == "sdpa"
+            and attention_mask is not None
+            and attention_mask.device.type in ["cuda", "xpu", "npu"]
+            and not output_attentions
+        ):
+            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
+            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
+            # Details: https://github.com/pytorch/pytorch/issues/110213
+            causal_mask = AttentionMaskConverter._unmask_unattended(
+                causal_mask, min_dtype
+            )
+
+        return causal_mask
+
+    @staticmethod
+    def _prepare_4d_causal_attention_mask_with_cache_position(
+        attention_mask: torch.Tensor,
+        sequence_length: int,
+        target_length: int,
+        dtype: torch.dtype,
+        cache_position: torch.Tensor,
+        batch_size: int,
+        config: Ernie4_5_MoeConfig,
+        past_key_values: Cache,
+    ):
+        """
+        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
+        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                A 2D attention mask of shape `(batch_size, key_value_length)`,
+                or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
+            sequence_length (`int`):
+                The sequence length being processed.
+            target_length (`int`):
+                The target length: when generating with static cache, the mask should be as long as the static cache,
+                 to account for the 0 padding, the part of the cache that is not filled yet.
+            dtype (`torch.dtype`):
+                The dtype to use for the 4D attention mask.
+            cache_position (`torch.Tensor`):
+                Indices depicting the position of the input sequence tokens in the sequence.
+            batch_size (`torch.Tensor`):
+                Batch size.
+            config (`Ernie4_5_MoeConfig`):
+                The model's configuration class
+            past_key_values (`Cache`):
+                The cache class that is being used currently to generate
+        """
+        if attention_mask is not None and attention_mask.dim() == 4:
+            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
+            causal_mask = attention_mask
+        else:
+            min_dtype = torch.finfo(dtype).min
+            causal_mask = torch.full(
+                (sequence_length, target_length),
+                fill_value=min_dtype,
+                dtype=dtype,
+                device=cache_position.device,
+            )
+            diagonal_attend_mask = torch.arange(
+                target_length, device=cache_position.device
+            ) > cache_position.reshape(-1, 1)
+            text_config = config.get_text_config()
+            if (
+                getattr(text_config, "use_sliding_window", True)
+                and text_config.sliding_window is not None
+            ):
+                if (
+                    not isinstance(past_key_values, SlidingWindowCache)
+                    or sequence_length > target_length
+                ):
+                    sliding_attend_mask = torch.arange(
+                        target_length, device=cache_position.device
+                    ) <= (cache_position.reshape(-1, 1) - text_config.sliding_window)
+                    diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
+            causal_mask *= diagonal_attend_mask
+            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
+            if attention_mask is not None:
+                causal_mask = (
+                    causal_mask.clone()
+                )  # copy to contiguous memory for in-place edit
+                if attention_mask.shape[-1] > target_length:
+                    attention_mask = attention_mask[:, :target_length]
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[
+                    :, None, None, :
+                ].to(causal_mask.device)
+                padding_mask = padding_mask == 0
+                causal_mask[:, :, :, :mask_length] = causal_mask[
+                    :, :, :, :mask_length
+                ].masked_fill(padding_mask, min_dtype)
+        return causal_mask
+
+
+@auto_docstring
+class Ernie4_5_MoeForCausalLM(Ernie4_5_PretrainedModel, GenerationMixin):
+    """ERNIE Mixture of Experts (MoE) model for causal language modeling."""
+
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        """
+        Initializes the ERNIE MoE model for causal language modeling.
+
+        Args:
+            config (dict): Model configuration.
+        """
+        super().__init__(config)
+        self.config = config
+        self.model = Ernie4_5_Model(config)
+        self.lm_head = nn.Linear(
+            config.hidden_size,
+            config.vocab_size,
+            bias=config.weight_share_add_bias and config.use_bias,
+        )  # TODO
+        self.loss_function = ErniePretrainingCriterion(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        """Returns the input embeddings layer."""
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        """Sets the input embeddings layer."""
+        self.ernie.embed_tokens = value
+
+    def get_output_embeddings(self):
+        """Returns the output embeddings (LM head)."""
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        """Sets the output embeddings layer."""
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        """Sets the ERNIE decoder model."""
+        self.model = decoder
+
+    def get_decoder(self):
+        """Get the transformer decoder."""
+        return self.model
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        position_ids=None,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds=None,
+        labels=None,
+        loss_mask=None,
+        use_cache=False,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs: Unpack[KwargsForCausalLM],
+    ):
+        """
+        Forward pass for causal language modeling.
+        """
+        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
+        )
+
+        outputs = self.model(
+            input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            past_key_values=past_key_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+
+        loss, router_loss = None, None
+        if getattr(self.config, "use_moe", False):
+            router_loss = outputs.router_loss
+
+        if labels is not None:
+            loss, _ = self.loss_function(logits, labels, loss_mask, router_loss)
+
+        return Ernie4_5_MoeCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_loss=router_loss,
+        )
+
+
+__all__ = ["Ernie4_5_Model", "Ernie4_5_MoeForCausalLM", "Ernie4_5_PretrainedModel"]

special_tokens_map.json

@@ -0,0 +1 @@
+{"bos_token": "<s>", "eos_token": "</s>", "pad_token": "<unk>", "unk_token": "<unk>", "cls_token": "<|begin_of_sentence|>", "sep_token": "<|end_of_sentence|>", "mask_token": "<mask:1>", "sys_start_token": "<mask:4>", "sys_end_token": "<mask:5>", "header_start_token": "<mask:6>", "header_end_token": "<mask:7>",  "additional_special_tokens": ["<|IMAGE_PLACEHOLDER|>", "<|AUDIO_PLACEHOLDER|>", "<|LOC_0|>", "<|LOC_1|>", "<|LOC_2|>", "<|LOC_3|>", "<|LOC_4|>", "<|LOC_5|>", "<|LOC_6|>", "<|LOC_7|>", "<|LOC_8|>", "<|LOC_9|>", "<|LOC_10|>", "<|LOC_11|>", "<|LOC_12|>", "<|LOC_13|>", "<|LOC_14|>", "<|LOC_15|>", "<|LOC_16|>", "<|LOC_17|>", "<|LOC_18|>", "<|LOC_19|>", "<|LOC_20|>", "<|LOC_21|>", "<|LOC_22|>", "<|LOC_23|>", "<|LOC_24|>", "<|LOC_25|>", "<|LOC_26|>", "<|LOC_27|>", "<|LOC_28|>", "<|LOC_29|>", "<|LOC_30|>", "<|LOC_31|>", "<|LOC_32|>", "<|LOC_33|>", "<|LOC_34|>", "<|LOC_35|>", "<|LOC_36|>", "<|LOC_37|>", "<|LOC_38|>", "<|LOC_39|>", "<|LOC_40|>", "<|LOC_41|>", "<|LOC_42|>", "<|LOC_43|>", "<|LOC_44|>", "<|LOC_45|>", "<|LOC_46|>", "<|LOC_47|>", "<|LOC_48|>", "<|LOC_49|>", "<|LOC_50|>", "<|LOC_51|>", "<|LOC_52|>", "<|LOC_53|>", "<|LOC_54|>", "<|LOC_55|>", "<|LOC_56|>", "<|LOC_57|>", "<|LOC_58|>", "<|LOC_59|>", "<|LOC_60|>", "<|LOC_61|>", "<|LOC_62|>", "<|LOC_63|>", "<|LOC_64|>", "<|LOC_65|>", "<|LOC_66|>", "<|LOC_67|>", "<|LOC_68|>", "<|LOC_69|>", "<|LOC_70|>", "<|LOC_71|>", "<|LOC_72|>", "<|LOC_73|>", "<|LOC_74|>", "<|LOC_75|>", "<|LOC_76|>", "<|LOC_77|>", "<|LOC_78|>", "<|LOC_79|>", "<|LOC_80|>", "<|LOC_81|>", "<|LOC_82|>", "<|LOC_83|>", "<|LOC_84|>", "<|LOC_85|>", "<|LOC_86|>", "<|LOC_87|>", "<|LOC_88|>", "<|LOC_89|>", "<|LOC_90|>", "<|LOC_91|>", "<|LOC_92|>", "<|LOC_93|>", "<|LOC_94|>", "<|LOC_95|>", "<|LOC_96|>", "<|LOC_97|>", "<|LOC_98|>", "<|LOC_99|>", "<|LOC_100|>", "<|LOC_101|>", "<|LOC_102|>", "<|LOC_103|>", "<|LOC_104|>", "<|LOC_105|>", "<|LOC_106|>", "<|LOC_107|>", "<|LOC_108|>", "<|LOC_109|>", "<|LOC_110|>", "<|LOC_111|>", "<|LOC_112|>", "<|LOC_113|>", "<|LOC_114|>", "<|LOC_115|>", "<|LOC_116|>", "<|LOC_117|>", "<|LOC_118|>", "<|LOC_119|>", "<|LOC_120|>", "<|LOC_121|>", "<|LOC_122|>", "<|LOC_123|>", "<|LOC_124|>", "<|LOC_125|>", "<|LOC_126|>", "<|LOC_127|>", "<|LOC_128|>", "<|LOC_129|>", "<|LOC_130|>", "<|LOC_131|>", "<|LOC_132|>", "<|LOC_133|>", "<|LOC_134|>", "<|LOC_135|>", "<|LOC_136|>", "<|LOC_137|>", "<|LOC_138|>", "<|LOC_139|>", "<|LOC_140|>", "<|LOC_141|>", "<|LOC_142|>", "<|LOC_143|>", "<|LOC_144|>", "<|LOC_145|>", "<|LOC_146|>", "<|LOC_147|>", "<|LOC_148|>", "<|LOC_149|>", "<|LOC_150|>", "<|LOC_151|>", "<|LOC_152|>", "<|LOC_153|>", "<|LOC_154|>", "<|LOC_155|>", "<|LOC_156|>", "<|LOC_157|>", "<|LOC_158|>", "<|LOC_159|>", "<|LOC_160|>", "<|LOC_161|>", "<|LOC_162|>", "<|LOC_163|>", "<|LOC_164|>", "<|LOC_165|>", "<|LOC_166|>", "<|LOC_167|>", "<|LOC_168|>", "<|LOC_169|>", "<|LOC_170|>", "<|LOC_171|>", "<|LOC_172|>", "<|LOC_173|>", "<|LOC_174|>", "<|LOC_175|>", "<|LOC_176|>", "<|LOC_177|>", "<|LOC_178|>", "<|LOC_179|>", "<|LOC_180|>", "<|LOC_181|>", "<|LOC_182|>", "<|LOC_183|>", "<|LOC_184|>", "<|LOC_185|>", "<|LOC_186|>", "<|LOC_187|>", "<|LOC_188|>", "<|LOC_189|>", "<|LOC_190|>", "<|LOC_191|>", "<|LOC_192|>", "<|LOC_193|>", "<|LOC_194|>", "<|LOC_195|>", "<|LOC_196|>", "<|LOC_197|>", "<|LOC_198|>", "<|LOC_199|>", "<|LOC_200|>", "<|LOC_201|>", "<|LOC_202|>", "<|LOC_203|>", "<|LOC_204|>", "<|LOC_205|>", "<|LOC_206|>", "<|LOC_207|>", "<|LOC_208|>", "<|LOC_209|>", "<|LOC_210|>", "<|LOC_211|>", "<|LOC_212|>", "<|LOC_213|>", "<|LOC_214|>", "<|LOC_215|>", "<|LOC_216|>", "<|LOC_217|>", "<|LOC_218|>", "<|LOC_219|>", "<|LOC_220|>", "<|LOC_221|>", "<|LOC_222|>", "<|LOC_223|>", "<|LOC_224|>", "<|LOC_225|>", "<|LOC_226|>", "<|LOC_227|>", "<|LOC_228|>", "<|LOC_229|>", "<|LOC_230|>", "<|LOC_231|>", "<|LOC_232|>", "<|LOC_233|>", "<|LOC_234|>", "<|LOC_235|>", "<|LOC_236|>", "<|LOC_237|>", "<|LOC_238|>", "<|LOC_239|>", "<|LOC_240|>", "<|LOC_241|>", "<|LOC_242|>", "<|LOC_243|>", "<|LOC_244|>", "<|LOC_245|>", "<|LOC_246|>", "<|LOC_247|>", "<|LOC_248|>", "<|LOC_249|>", "<|LOC_250|>", "<|LOC_251|>", "<|LOC_252|>", "<|LOC_253|>", "<|LOC_254|>", "<|LOC_255|>", "<|LOC_256|>", "<|LOC_257|>", "<|LOC_258|>", "<|LOC_259|>", "<|LOC_260|>", "<|LOC_261|>", "<|LOC_262|>", "<|LOC_263|>", "<|LOC_264|>", "<|LOC_265|>", "<|LOC_266|>", "<|LOC_267|>", "<|LOC_268|>", "<|LOC_269|>", "<|LOC_270|>", "<|LOC_271|>", "<|LOC_272|>", "<|LOC_273|>", "<|LOC_274|>", "<|LOC_275|>", "<|LOC_276|>", "<|LOC_277|>", "<|LOC_278|>", "<|LOC_279|>", "<|LOC_280|>", "<|LOC_281|>", "<|LOC_282|>", "<|LOC_283|>", "<|LOC_284|>", "<|LOC_285|>", "<|LOC_286|>", "<|LOC_287|>", "<|LOC_288|>", "<|LOC_289|>", "<|LOC_290|>", "<|LOC_291|>", "<|LOC_292|>", "<|LOC_293|>", "<|LOC_294|>", "<|LOC_295|>", "<|LOC_296|>", "<|LOC_297|>", "<|LOC_298|>", "<|LOC_299|>", "<|LOC_300|>", "<|LOC_301|>", "<|LOC_302|>", "<|LOC_303|>", "<|LOC_304|>", "<|LOC_305|>", "<|LOC_306|>", "<|LOC_307|>", "<|LOC_308|>", "<|LOC_309|>", "<|LOC_310|>", "<|LOC_311|>", "<|LOC_312|>", "<|LOC_313|>", "<|LOC_314|>", "<|LOC_315|>", "<|LOC_316|>", "<|LOC_317|>", "<|LOC_318|>", "<|LOC_319|>", "<|LOC_320|>", "<|LOC_321|>", "<|LOC_322|>", "<|LOC_323|>", "<|LOC_324|>", "<|LOC_325|>", "<|LOC_326|>", "<|LOC_327|>", "<|LOC_328|>", "<|LOC_329|>", "<|LOC_330|>", "<|LOC_331|>", "<|LOC_332|>", "<|LOC_333|>", "<|LOC_334|>", "<|LOC_335|>", "<|LOC_336|>", "<|LOC_337|>", "<|LOC_338|>", "<|LOC_339|>", "<|LOC_340|>", "<|LOC_341|>", "<|LOC_342|>", "<|LOC_343|>", "<|LOC_344|>", "<|LOC_345|>", "<|LOC_346|>", "<|LOC_347|>", "<|LOC_348|>", "<|LOC_349|>", "<|LOC_350|>", "<|LOC_351|>", "<|LOC_352|>", "<|LOC_353|>", "<|LOC_354|>", "<|LOC_355|>", "<|LOC_356|>", "<|LOC_357|>", "<|LOC_358|>", "<|LOC_359|>", "<|LOC_360|>", "<|LOC_361|>", "<|LOC_362|>", "<|LOC_363|>", "<|LOC_364|>", "<|LOC_365|>", "<|LOC_366|>", "<|LOC_367|>", "<|LOC_368|>", "<|LOC_369|>", "<|LOC_370|>", "<|LOC_371|>", "<|LOC_372|>", "<|LOC_373|>", "<|LOC_374|>", "<|LOC_375|>", "<|LOC_376|>", "<|LOC_377|>", "<|LOC_378|>", "<|LOC_379|>", "<|LOC_380|>", "<|LOC_381|>", "<|LOC_382|>", "<|LOC_383|>", "<|LOC_384|>", "<|LOC_385|>", "<|LOC_386|>", "<|LOC_387|>", "<|LOC_388|>", "<|LOC_389|>", "<|LOC_390|>", "<|LOC_391|>", "<|LOC_392|>", "<|LOC_393|>", "<|LOC_394|>", "<|LOC_395|>", "<|LOC_396|>", "<|LOC_397|>", "<|LOC_398|>", "<|LOC_399|>", "<|LOC_400|>", "<|LOC_401|>", "<|LOC_402|>", "<|LOC_403|>", "<|LOC_404|>", "<|LOC_405|>", "<|LOC_406|>", "<|LOC_407|>", "<|LOC_408|>", "<|LOC_409|>", "<|LOC_410|>", "<|LOC_411|>", "<|LOC_412|>", "<|LOC_413|>", "<|LOC_414|>", "<|LOC_415|>", "<|LOC_416|>", "<|LOC_417|>", "<|LOC_418|>", "<|LOC_419|>", "<|LOC_420|>", "<|LOC_421|>", "<|LOC_422|>", "<|LOC_423|>", "<|LOC_424|>", "<|LOC_425|>", "<|LOC_426|>", "<|LOC_427|>", "<|LOC_428|>", "<|LOC_429|>", "<|LOC_430|>", "<|LOC_431|>", "<|LOC_432|>", "<|LOC_433|>", "<|LOC_434|>", "<|LOC_435|>", "<|LOC_436|>", "<|LOC_437|>", "<|LOC_438|>", "<|LOC_439|>", "<|LOC_440|>", "<|LOC_441|>", "<|LOC_442|>", "<|LOC_443|>", "<|LOC_444|>", "<|LOC_445|>", "<|LOC_446|>", "<|LOC_447|>", "<|LOC_448|>", "<|LOC_449|>", "<|LOC_450|>", "<|LOC_451|>", "<|LOC_452|>", "<|LOC_453|>", "<|LOC_454|>", "<|LOC_455|>", "<|LOC_456|>", "<|LOC_457|>", "<|LOC_458|>", "<|LOC_459|>", "<|LOC_460|>", "<|LOC_461|>", "<|LOC_462|>", "<|LOC_463|>", "<|LOC_464|>", "<|LOC_465|>", "<|LOC_466|>", "<|LOC_467|>", "<|LOC_468|>", "<|LOC_469|>", "<|LOC_470|>", "<|LOC_471|>", "<|LOC_472|>", "<|LOC_473|>", "<|LOC_474|>", "<|LOC_475|>", "<|LOC_476|>", "<|LOC_477|>", "<|LOC_478|>", "<|LOC_479|>", "<|LOC_480|>", "<|LOC_481|>", "<|LOC_482|>", "<|LOC_483|>", "<|LOC_484|>", "<|LOC_485|>", "<|LOC_486|>", "<|LOC_487|>", "<|LOC_488|>", "<|LOC_489|>", "<|LOC_490|>", "<|LOC_491|>", "<|LOC_492|>", "<|LOC_493|>", "<|LOC_494|>", "<|LOC_495|>", "<|LOC_496|>", "<|LOC_497|>", "<|LOC_498|>", "<|LOC_499|>", "<|LOC_500|>", "<|LOC_501|>", "<|LOC_502|>", "<|LOC_503|>", "<|LOC_504|>", "<|LOC_505|>", "<|LOC_506|>", "<|LOC_507|>", "<|LOC_508|>", "<|LOC_509|>", "<|LOC_510|>", "<|LOC_511|>", "<|LOC_512|>", "<|LOC_513|>", "<|LOC_514|>", "<|LOC_515|>", "<|LOC_516|>", "<|LOC_517|>", "<|LOC_518|>", "<|LOC_519|>", "<|LOC_520|>", "<|LOC_521|>", "<|LOC_522|>", "<|LOC_523|>", "<|LOC_524|>", "<|LOC_525|>", "<|LOC_526|>", "<|LOC_527|>", "<|LOC_528|>", "<|LOC_529|>", "<|LOC_530|>", "<|LOC_531|>", "<|LOC_532|>", "<|LOC_533|>", "<|LOC_534|>", "<|LOC_535|>", "<|LOC_536|>", "<|LOC_537|>", "<|LOC_538|>", "<|LOC_539|>", "<|LOC_540|>", "<|LOC_541|>", "<|LOC_542|>", "<|LOC_543|>", "<|LOC_544|>", "<|LOC_545|>", "<|LOC_546|>", "<|LOC_547|>", "<|LOC_548|>", "<|LOC_549|>", "<|LOC_550|>", "<|LOC_551|>", "<|LOC_552|>", "<|LOC_553|>", "<|LOC_554|>", "<|LOC_555|>", "<|LOC_556|>", "<|LOC_557|>", "<|LOC_558|>", "<|LOC_559|>", "<|LOC_560|>", "<|LOC_561|>", "<|LOC_562|>", "<|LOC_563|>", "<|LOC_564|>", "<|LOC_565|>", "<|LOC_566|>", "<|LOC_567|>", "<|LOC_568|>", "<|LOC_569|>", "<|LOC_570|>", "<|LOC_571|>", "<|LOC_572|>", "<|LOC_573|>", "<|LOC_574|>", "<|LOC_575|>", "<|LOC_576|>", "<|LOC_577|>", "<|LOC_578|>", "<|LOC_579|>", "<|LOC_580|>", "<|LOC_581|>", "<|LOC_582|>", "<|LOC_583|>", "<|LOC_584|>", "<|LOC_585|>", "<|LOC_586|>", "<|LOC_587|>", "<|LOC_588|>", "<|LOC_589|>", "<|LOC_590|>", "<|LOC_591|>", "<|LOC_592|>", "<|LOC_593|>", "<|LOC_594|>", "<|LOC_595|>", "<|LOC_596|>", "<|LOC_597|>", "<|LOC_598|>", "<|LOC_599|>", "<|LOC_600|>", "<|LOC_601|>", "<|LOC_602|>", "<|LOC_603|>", "<|LOC_604|>", "<|LOC_605|>", "<|LOC_606|>", "<|LOC_607|>", "<|LOC_608|>", "<|LOC_609|>", "<|LOC_610|>", "<|LOC_611|>", "<|LOC_612|>", "<|LOC_613|>", "<|LOC_614|>", "<|LOC_615|>", "<|LOC_616|>", "<|LOC_617|>", "<|LOC_618|>", "<|LOC_619|>", "<|LOC_620|>", "<|LOC_621|>", "<|LOC_622|>", "<|LOC_623|>", "<|LOC_624|>", "<|LOC_625|>", "<|LOC_626|>", "<|LOC_627|>", "<|LOC_628|>", "<|LOC_629|>", "<|LOC_630|>", "<|LOC_631|>", "<|LOC_632|>", "<|LOC_633|>", "<|LOC_634|>", "<|LOC_635|>", "<|LOC_636|>", "<|LOC_637|>", "<|LOC_638|>", "<|LOC_639|>", "<|LOC_640|>", "<|LOC_641|>", "<|LOC_642|>", "<|LOC_643|>", "<|LOC_644|>", "<|LOC_645|>", "<|LOC_646|>", "<|LOC_647|>", "<|LOC_648|>", "<|LOC_649|>", "<|LOC_650|>", "<|LOC_651|>", "<|LOC_652|>", "<|LOC_653|>", "<|LOC_654|>", "<|LOC_655|>", "<|LOC_656|>", "<|LOC_657|>", "<|LOC_658|>", "<|LOC_659|>", "<|LOC_660|>", "<|LOC_661|>", "<|LOC_662|>", "<|LOC_663|>", "<|LOC_664|>", "<|LOC_665|>", "<|LOC_666|>", "<|LOC_667|>", "<|LOC_668|>", "<|LOC_669|>", "<|LOC_670|>", "<|LOC_671|>", "<|LOC_672|>", "<|LOC_673|>", "<|LOC_674|>", "<|LOC_675|>", "<|LOC_676|>", "<|LOC_677|>", "<|LOC_678|>", "<|LOC_679|>", "<|LOC_680|>", "<|LOC_681|>", "<|LOC_682|>", "<|LOC_683|>", "<|LOC_684|>", "<|LOC_685|>", "<|LOC_686|>", "<|LOC_687|>", "<|LOC_688|>", "<|LOC_689|>", "<|LOC_690|>", "<|LOC_691|>", "<|LOC_692|>", "<|LOC_693|>", "<|LOC_694|>", "<|LOC_695|>", "<|LOC_696|>", "<|LOC_697|>", "<|LOC_698|>", "<|LOC_699|>", "<|LOC_700|>", "<|LOC_701|>", "<|LOC_702|>", "<|LOC_703|>", "<|LOC_704|>", "<|LOC_705|>", "<|LOC_706|>", "<|LOC_707|>", "<|LOC_708|>", "<|LOC_709|>", "<|LOC_710|>", "<|LOC_711|>", "<|LOC_712|>", "<|LOC_713|>", "<|LOC_714|>", "<|LOC_715|>", "<|LOC_716|>", "<|LOC_717|>", "<|LOC_718|>", "<|LOC_719|>", "<|LOC_720|>", "<|LOC_721|>", "<|LOC_722|>", "<|LOC_723|>", "<|LOC_724|>", "<|LOC_725|>", "<|LOC_726|>", "<|LOC_727|>", "<|LOC_728|>", "<|LOC_729|>", "<|LOC_730|>", "<|LOC_731|>", "<|LOC_732|>", "<|LOC_733|>", "<|LOC_734|>", "<|LOC_735|>", "<|LOC_736|>", "<|LOC_737|>", "<|LOC_738|>", "<|LOC_739|>", "<|LOC_740|>", "<|LOC_741|>", "<|LOC_742|>", "<|LOC_743|>", "<|LOC_744|>", "<|LOC_745|>", "<|LOC_746|>", "<|LOC_747|>", "<|LOC_748|>", "<|LOC_749|>", "<|LOC_750|>", "<|LOC_751|>", "<|LOC_752|>", "<|LOC_753|>", "<|LOC_754|>", "<|LOC_755|>", "<|LOC_756|>", "<|LOC_757|>", "<|LOC_758|>", "<|LOC_759|>", "<|LOC_760|>", "<|LOC_761|>", "<|LOC_762|>", "<|LOC_763|>", "<|LOC_764|>", "<|LOC_765|>", "<|LOC_766|>", "<|LOC_767|>", "<|LOC_768|>", "<|LOC_769|>", "<|LOC_770|>", "<|LOC_771|>", "<|LOC_772|>", "<|LOC_773|>", "<|LOC_774|>", "<|LOC_775|>", "<|LOC_776|>", "<|LOC_777|>", "<|LOC_778|>", "<|LOC_779|>", "<|LOC_780|>", "<|LOC_781|>", "<|LOC_782|>", "<|LOC_783|>", "<|LOC_784|>", "<|LOC_785|>", "<|LOC_786|>", "<|LOC_787|>", "<|LOC_788|>", "<|LOC_789|>", "<|LOC_790|>", "<|LOC_791|>", "<|LOC_792|>", "<|LOC_793|>", "<|LOC_794|>", "<|LOC_795|>", "<|LOC_796|>", "<|LOC_797|>", "<|LOC_798|>", "<|LOC_799|>", "<|LOC_800|>", "<|LOC_801|>", "<|LOC_802|>", "<|LOC_803|>", "<|LOC_804|>", "<|LOC_805|>", "<|LOC_806|>", "<|LOC_807|>", "<|LOC_808|>", "<|LOC_809|>", "<|LOC_810|>", "<|LOC_811|>", "<|LOC_812|>", "<|LOC_813|>", "<|LOC_814|>", "<|LOC_815|>", "<|LOC_816|>", "<|LOC_817|>", "<|LOC_818|>", "<|LOC_819|>", "<|LOC_820|>", "<|LOC_821|>", "<|LOC_822|>", "<|LOC_823|>", "<|LOC_824|>", "<|LOC_825|>", "<|LOC_826|>", "<|LOC_827|>", "<|LOC_828|>", "<|LOC_829|>", "<|LOC_830|>", "<|LOC_831|>", "<|LOC_832|>", "<|LOC_833|>", "<|LOC_834|>", "<|LOC_835|>", "<|LOC_836|>", "<|LOC_837|>", "<|LOC_838|>", "<|LOC_839|>", "<|LOC_840|>", "<|LOC_841|>", "<|LOC_842|>", "<|LOC_843|>", "<|LOC_844|>", "<|LOC_845|>", "<|LOC_846|>", "<|LOC_847|>", "<|LOC_848|>", "<|LOC_849|>", "<|LOC_850|>", "<|LOC_851|>", "<|LOC_852|>", "<|LOC_853|>", "<|LOC_854|>", "<|LOC_855|>", "<|LOC_856|>", "<|LOC_857|>", "<|LOC_858|>", "<|LOC_859|>", "<|LOC_860|>", "<|LOC_861|>", "<|LOC_862|>", "<|LOC_863|>", "<|LOC_864|>", "<|LOC_865|>", "<|LOC_866|>", "<|LOC_867|>", "<|LOC_868|>", "<|LOC_869|>", "<|LOC_870|>", "<|LOC_871|>", "<|LOC_872|>", "<|LOC_873|>", "<|LOC_874|>", "<|LOC_875|>", "<|LOC_876|>", "<|LOC_877|>", "<|LOC_878|>", "<|LOC_879|>", "<|LOC_880|>", "<|LOC_881|>", "<|LOC_882|>", "<|LOC_883|>", "<|LOC_884|>", "<|LOC_885|>", "<|LOC_886|>", "<|LOC_887|>", "<|LOC_888|>", "<|LOC_889|>", "<|LOC_890|>", "<|LOC_891|>", "<|LOC_892|>", "<|LOC_893|>", "<|LOC_894|>", "<|LOC_895|>", "<|LOC_896|>", "<|LOC_897|>", "<|LOC_898|>", "<|LOC_899|>", "<|LOC_900|>", "<|LOC_901|>", "<|LOC_902|>", "<|LOC_903|>", "<|LOC_904|>", "<|LOC_905|>", "<|LOC_906|>", "<|LOC_907|>", "<|LOC_908|>", "<|LOC_909|>", "<|LOC_910|>", "<|LOC_911|>", "<|LOC_912|>", "<|LOC_913|>", "<|LOC_914|>", "<|LOC_915|>", "<|LOC_916|>", "<|LOC_917|>", "<|LOC_918|>", "<|LOC_919|>", "<|LOC_920|>", "<|LOC_921|>", "<|LOC_922|>", "<|LOC_923|>", "<|LOC_924|>", "<|LOC_925|>", "<|LOC_926|>", "<|LOC_927|>", "<|LOC_928|>", "<|LOC_929|>", "<|LOC_930|>", "<|LOC_931|>", "<|LOC_932|>", "<|LOC_933|>", "<|LOC_934|>", "<|LOC_935|>", "<|LOC_936|>", "<|LOC_937|>", "<|LOC_938|>", "<|LOC_939|>", "<|LOC_940|>", "<|LOC_941|>", "<|LOC_942|>", "<|LOC_943|>", "<|LOC_944|>", "<|LOC_945|>", "<|LOC_946|>", "<|LOC_947|>", "<|LOC_948|>", "<|LOC_949|>", "<|LOC_950|>", "<|LOC_951|>", "<|LOC_952|>", "<|LOC_953|>", "<|LOC_954|>", "<|LOC_955|>", "<|LOC_956|>", "<|LOC_957|>", "<|LOC_958|>", "<|LOC_959|>", "<|LOC_960|>", "<|LOC_961|>", "<|LOC_962|>", "<|LOC_963|>", "<|LOC_964|>", "<|LOC_965|>", "<|LOC_966|>", "<|LOC_967|>", "<|LOC_968|>", "<|LOC_969|>", "<|LOC_970|>", "<|LOC_971|>", "<|LOC_972|>", "<|LOC_973|>", "<|LOC_974|>", "<|LOC_975|>", "<|LOC_976|>", "<|LOC_977|>", "<|LOC_978|>", "<|LOC_979|>", "<|LOC_980|>", "<|LOC_981|>", "<|LOC_982|>", "<|LOC_983|>", "<|LOC_984|>", "<|LOC_985|>", "<|LOC_986|>", "<|LOC_987|>", "<|LOC_988|>", "<|LOC_989|>", "<|LOC_990|>", "<|LOC_991|>", "<|LOC_992|>", "<|LOC_993|>", "<|LOC_994|>", "<|LOC_995|>", "<|LOC_996|>", "<|LOC_997|>", "<|LOC_998|>", "<|LOC_999|>", "<|LOC_1000|>", "<|LOC_BEGIN|>", "<|LOC_END|>", "<|LOC_SEP|>", "<|CROP_COL_SEP|>", "<|CROP_ROW_SEP|>", "<|IMAGE_SEP|>"]}

tokenization_ernie4_5.py

@@ -0,0 +1,352 @@
+# Copyright (c) 2025 Baidu, Inc. All Rights Reserved.
+#
+# 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.
+"""Ernie4_5_Tokenizer"""
+
+import os
+from shutil import copyfile
+from typing import Dict, List, Optional, Tuple, Union
+import torch
+import numpy as np
+import sentencepiece as spm
+
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.tokenization_utils_base import (
+    PaddingStrategy,
+)
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+
+class Ernie4_5_Tokenizer(PreTrainedTokenizer):
+    """
+    Ernie4_5_Tokenizer
+        vocab_files_names (dict): Mapping vocabulary-related config name to actual filename.
+        model_input_names (List): Model input names expected by the tokenizer
+        padding_side (str): Padding side (where to add padding tokens)
+    """
+    vocab_files_names = {
+        "vocab_file": "tokenizer.model",
+    }
+    model_input_names = ["input_ids", "position_ids", "attention_mask", "labels"]
+    padding_side = "right"
+
+    def __init__(
+        self,
+        vocab_file,
+        bos_token="<s>",
+        cls_token="<cls>",
+        eos_token="</s>",
+        mask_token="<mask:0>",
+        pad_token="<pad>",
+        sep_token="<sep>",
+        unk_token="<unk>",
+        additional_special_tokens=None,
+        split_special_tokens=False,
+        tokenizer_alpha=None,
+        **kwargs
+    ):
+        """
+        Initialize the ERNIE tokenizer.
+
+        Args:
+            vocab_file (str): Path to the SentencePiece model file.
+            bos_token (str, optional): Beginning of sentence token. Defaults to "<s>".
+            cls_token (str, optional): Classification token. Defaults to "<cls>".
+            eos_token (str, optional): End of sentence token. Defaults to "</s>".
+            mask_token (str, optional): Mask token. Defaults to "<mask:0>".
+            pad_token (str, optional): Padding token. Defaults to "<pad>".
+            sep_token (str, optional): Separator token. Defaults to "<sep>".
+            unk_token (str, optional): Unknown token. Defaults to "<unk>".
+            additional_special_tokens (List[str], optional): Additional special tokens.
+                Defaults to ["<mask:1>", "<mask:7>"].
+            split_special_tokens (bool, optional): Whether to split special tokens. Defaults to False.
+            tokenizer_alpha (float, optional): Alpha parameter for SentencePiece sampling.
+            **kwargs: Additional keyword arguments passed to the parent class.
+        """
+
+        self.vocab_file = vocab_file
+        self.sp_model = spm.SentencePieceProcessor()
+        self.sp_model.Load(vocab_file)
+        self.tokenizer_alpha = tokenizer_alpha
+
+        if additional_special_tokens is None:
+            additional_special_tokens = ["<mask:1>", "<mask:7>"]
+        super().__init__(
+            bos_token=bos_token,
+            cls_token=cls_token,
+            eos_token=eos_token,
+            mask_token=mask_token,
+            pad_token=pad_token,
+            sep_token=sep_token,
+            unk_token=unk_token,
+            additional_special_tokens=additional_special_tokens,
+            split_special_tokens=split_special_tokens,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self):
+        """Returns the size of the vocabulary.
+
+        Returns:
+            int: The number of tokens in the vocabulary.
+        """
+        return self.sp_model.vocab_size()
+
+    def get_vocab(self):
+        """Get the vocabulary as a dictionary mapping tokens to their IDs.
+
+        Returns:
+            dict: A dictionary mapping tokens to their corresponding IDs.
+        """
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text):
+        """Tokenize text using SentencePiece.
+
+        Args:
+            text (str): The text to tokenize.
+
+        Returns:
+            list: A list of tokens.
+        """
+        if self.tokenizer_alpha is not None:
+            return self.sp_model.encode_as_pieces(
+                text,
+                enable_sampling=True,
+                nbest_size=-1,
+                alpha=self.tokenizer_alpha,
+            )
+        else:
+            return self.sp_model.encode_as_pieces(text)
+
+    def _convert_token_to_id(self, token):
+        """Convert a token (str) to an ID using the vocabulary.
+
+        Args:
+            token (str): The token to convert.
+
+        Returns:
+            int: The corresponding token ID.
+        """
+        return self.sp_model.piece_to_id(token)
+
+    def _convert_id_to_token(self, id):
+        """Convert an ID to a token (str) using the vocabulary.
+
+        Args:
+            id (int): The token ID to convert.
+
+        Returns:
+            str: The corresponding token.
+        """
+        if id >= self.vocab_size:
+            return self.unk_token
+        else:
+            return self.sp_model.id_to_piece(id)
+
+    def convert_tokens_to_string(self, tokens):
+        """Convert a sequence of tokens back to a single string.
+
+        Args:
+            tokens (List[str]): A list of tokens to convert.
+
+        Returns:
+            str: The reconstructed string.
+        """
+        current_sub_tokens = []
+        out_string = ""
+        prev_is_special = False
+        for token in tokens:
+            # make sure that special tokens are not decoded using sentencepiece model
+            if token in self.all_special_tokens:
+                if not prev_is_special:
+                    out_string += " "
+                out_string += self.sp_model.decode(current_sub_tokens) + token
+                prev_is_special = True
+                current_sub_tokens = []
+            else:
+                current_sub_tokens.append(token)
+                prev_is_special = False
+        out_string += self.sp_model.decode(current_sub_tokens)
+        return out_string
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        """Build model inputs by adding special tokens to sequences.
+
+        Args:
+            token_ids_0 (List[int]): List of token IDs for the first sequence.
+            token_ids_1 (List[int], optional): List of token IDs for the second sequence.
+
+        Returns:
+            List[int]: List of token IDs with special tokens added.
+        """
+        output = token_ids_0
+        last_cls_index = -1
+        last_sep_index = -1
+        if self.cls_token_id in output:
+            last_cls_index = len(output) - output[::-1].index(self.cls_token_id) - 1
+        if self.sep_token_id in output:
+            last_sep_index = len(output) - output[::-1].index(self.sep_token_id) - 1
+
+        if last_cls_index > last_sep_index:
+            next_token_id = self.sep_token_id
+        elif last_sep_index > last_cls_index:
+            next_token_id = self.cls_token_id
+        else:
+            output = [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
+            next_token_id = self.cls_token_id
+
+        output = [self.bos_token_id] + output
+        # Assume no markup in text if token_ids_1 is given.
+        if token_ids_1 is not None:
+            output = output + token_ids_1 + [next_token_id]
+        return output
+
+    def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
+        """Get a mask showing which tokens are special tokens.
+
+        Args:
+            token_ids_0 (List[int]): List of token IDs for the first sequence.
+            token_ids_1 (List[int], optional): List of token IDs for the second sequence.
+            already_has_special_tokens (bool): Whether the tokens already include special tokens.
+
+        Returns:
+            List[int]: A mask where 1 indicates special tokens and 0 indicates regular tokens.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(token_ids_0, token_ids_1, already_has_special_tokens=True)
+
+        # [bos_token, cls_token, tokens_0, sep_token]
+        if token_ids_1 is None:
+            return [1, 1] + ([0] * len(token_ids_0)) + [1]
+        # [bos_token, cls_token, tokens_0, sep_token, tokens_1, cls_token]
+        return [1, 1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+
+    def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """
+        Save the vocabulary and special tokens file to a directory.
+
+        Args:
+            save_directory (str): The directory in which to save the vocabulary.
+            filename_prefix (Optional[str]): Optional prefix for the saved filename.
+
+        Returns:
+            Tuple[str]: Paths to the files saved.
+
+        Raises:
+            ValueError: If the save_directory is not a valid directory.
+        """
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        out_vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"],
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, "wb") as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file,)
+
+    def _pad(
+        self,
+        encoded_inputs: Union[Dict],
+        max_length: Optional[int] = None,
+        padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
+        pad_to_multiple_of: Optional[int] = None,
+        padding_side: Optional[str] = None,
+        return_attention_mask: Optional[bool] = None,
+    ) -> dict:
+        """
+        Pad encoded inputs according to specified strategy.
+
+        Args:
+            encoded_inputs (Union[Dict]): Dictionary of encoded inputs.
+            max_length (Optional[int]): Maximum length to pad to.
+            padding_strategy (PaddingStrategy): Strategy for padding.
+            pad_to_multiple_of (Optional[int]): Pad to a multiple of this value.
+            return_attention_mask (Optional[bool]): Whether to return attention mask.
+
+        Returns:
+            dict: Dictionary with padded inputs and optional attention mask.
+
+        Raises:
+            ValueError: If attention_mask has unexpected type or invalid padding strategy.
+        """
+        if return_attention_mask is None:
+            return_attention_mask = "attention_mask" in self.model_input_names
+        if return_attention_mask:
+            required_input = encoded_inputs[self.model_input_names[0]]
+            if padding_strategy == PaddingStrategy.LONGEST:
+                max_length = len(required_input)
+            if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
+                max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
+            needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
+
+            if "attention_mask" in encoded_inputs and encoded_inputs["attention_mask"] is not None:
+                attention_mask = encoded_inputs.pop("attention_mask")
+                if isinstance(attention_mask, torch.Tensor):
+                    attention_mask = attention_mask.numpy()
+                elif isinstance(attention_mask, list):
+                    attention_mask = np.array(attention_mask)
+                elif not isinstance(attention_mask, np.ndarray):
+                    raise ValueError(f"Unexpected type {type(attention_mask)} of attention_mask, ")
+            else:
+                # Create default attention mask if none provided
+                attention_mask = np.tril(np.ones((len(required_input), len(required_input)), dtype=np.int64))
+                attention_mask = np.expand_dims(attention_mask, axis=0)
+
+            if needs_to_be_padded:
+                difference = max_length - len(required_input)
+                if self.padding_side == "right":
+                    if attention_mask.ndim == 1:
+                        pad_width = [(0, difference)]
+                    else:
+                        pad_width = [(0, 0), (0, difference), (0, difference)]
+                elif self.padding_side == "left":
+                    if attention_mask.ndim == 1:
+                        pad_width = [(difference, 0)]
+                    else:
+                        pad_width = [(0, 0), (difference, 0), (difference, 0)]
+                else:
+                    raise ValueError("Invalid padding strategy:" + str(self.padding_side))
+                attention_mask = np.pad(
+                    attention_mask,
+                    pad_width=pad_width,
+                    mode="constant",
+                    constant_values=0,
+                )
+
+        encoded_inputs = super()._pad(
+            encoded_inputs,
+            max_length,
+            padding_strategy=padding_strategy,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_attention_mask=False,
+        )
+        if return_attention_mask:
+            encoded_inputs["attention_mask"] = attention_mask.tolist()
+        return encoded_inputs
+

tokenizer.model

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

tokenizer_config.json

@@ -0,0 +1,22 @@
+{
+    "bos_token": "<s>",
+    "eos_token": "</s>",
+    "pad_token": "<unk>",
+    "unk_token": "<unk>",
+    "cls_token": "<|begin_of_sentence|>",
+    "sep_token": "<|end_of_sentence|>",
+    "mask_token": "<mask:1>",
+    "sys_start_token": "<mask:4>",
+    "sys_end_token": "<mask:5>",
+    "header_start_token": "<mask:6>",
+    "header_end_token": "<mask:7>",
+    "additional_special_tokens": null,
+    "chat_template": "{%- if not add_generation_prompt is defined -%}\n    {%- set add_generation_prompt = true -%}\n{%- endif -%}\n{%- if not cls_token is defined -%}\n    {%- set cls_token = \"<|begin_of_sentence|>\" -%}\n{%- endif -%}\n{%- if not sep_token is defined -%}\n    {%- set sep_token = \"<|end_of_sentence|>\" -%}\n{%- endif -%}\n{{- cls_token -}}\n{%- for message in messages -%}\n    {%- if message[\"role\"] == \"user\" -%}\n        {{- \"User: \" + message[\"content\"] + \"\n\" -}}\n    {%- elif message[\"role\"] == \"assistant\" -%}\n        {{- \"Assistant: \" + message[\"content\"] + sep_token -}}\n    {%- elif message[\"role\"] == \"system\" -%}\n        {{- message[\"content\"] + \"\n\" -}}\n    {%- endif -%}\n{%- endfor -%}\n{%- if add_generation_prompt -%}\n    {{- \"Assistant: \" -}}\n{%- endif -%}",
+    "tokenizer_class": "Ernie4_5_Tokenizer",
+    "auto_map": {
+        "AutoTokenizer": [
+            "tokenization_ernie4_5.Ernie4_5_Tokenizer",
+            "tokenization_ernie4_5.Ernie4_5_Tokenizer"
+        ]
+    }
+}