upload model

MODEL_LICENSE.md

@@ -0,0 +1,40 @@
+# LingoWhale-8B模型许可协议
+
+## 1. 定义
+- “发布方”：指发布源模型的LingoWhale-8B模型团队。
+- “源模型”：指根据本许可提供的LingoWhale-8B模型参数。
+- “使用方”：指根据本协议使用源模型的单位或个人。
+
+## 2. 许可内容
+根据本许可的条款和条件，发布方特此授予您非排他性、全球性、不可转让、不可再许可、可撤销、免版税的版权许可。
+
+上述版权声明和本许可声明应包含在此源模型的所有副本或重要部分中。
+
+## 3. 限制
+您不得出于任何军事或非法目的使用、复制、修改、合并、发布、分发、复制或创建此源模型的全部或部分衍生品。
+
+您不得利用此源模型从事任何危害国家安全和国家统一、危害社会公共利益、侵犯人身权益的行为。
+
+## 4. 免责声明
+此源模型“按原样”提供，不提供任何明示或暗示的保证，包括但不限于对适销性、特定用途的适用性和非侵权性的保证。在任何情况下，作者或版权持有人均不对任何索赔、损害或其他责任负责，无论是在合同诉讼、侵权行为还是其他方面，由源模型或源模型的使用或其他交易引起、由源模型引起或与之相关模型。
+
+## 5. 责任限制
+除适用法律禁止的范围外，在任何情况下且根据任何法律理论，无论是基于侵权行为、疏忽、合同、责任或其他原因，任何发布方均不对您承担任何直接、间接、特殊、偶然、示范性、 或间接损害，或任何其他商业损失，即使使用方已被告知此类损害的可能性。
+
+## 6. 争议解决
+本许可受中华人民共和国法律管辖并按其解释。 因本许可引起的或与本许可有关的任何争议，由发布方住所地人民法院管辖。
+
+请注意，许可证可能会更新到更全面的版本。 有关许可和版权的任何问题，请通过license@deeplang.ai与我们联系。
+
+## 7. 附则
+若您期望基于本协议的许可条件与限制，将此源模型或其衍生品用作商业用途，请您按照如下方式联系发布方，以进行登记并向发布方申请书面授权：
+
+1. 联系邮箱：license@deeplang.ai
+2. 需提交内容如下：
+
+| 选项 | 是否必须 | 说明 |
+| :---- | :----: | :---- |
+| 申请人姓名 | 是 | 请填写申请人真实姓名。 |
+| 申请企业名称 | 是 | 请填写完整的企业名称，自然人不能申请此源模型的商业许可。 |
+| 联系方式 | 是 | 请填写真实邮箱，用于接收授权书文件。 |
+| 使用目的和场景 | 是 | 请填写真实使用目的和场景。 |

README.md

@@ -1,3 +1,289 @@
----
-license: apache-2.0
----
+<p align="left">
+    <a href="README_EN.md">English</a>&nbsp ｜ &nbsp中文
+</p>
+<br>
+
+<div align="center">
+<h1>
+LingoWhale-8B
+</h1>
+</div>
+
+<p align="center">
+🤗 <a href="https://huggingface.co/deeplang-ai/LingoWhale-8B" target="_blank">Hugging Face</a> • 🤖 <a href="https://www.modelscope.cn/models/DeepLang/LingoWhale-8B" target="_blank">ModelScope</a> • ⛵ <a href="https://wisemodel.cn/models/%E5%8C%97%E4%BA%AC%E6%B7%B1%E8%A8%80%E7%A7%91%E6%8A%80%E6%9C%89%E9%99%90%E8%B4%A3%E4%BB%BB%E5%85%AC%E5%8F%B8/LingoWhale-8B/" target="_blank">Wisemodel</a>
+</p>
+
+<div align="center">
+<strong>
+深言科技联合清华大学NLP实验室开源语鲸-8B模型 🎉
+</strong>
+</div>
+
+# 目录
+
+- [目录](#目录)
+- [模型介绍](#模型介绍)
+- [测评结果](#测评结果)
+- [生成样例](#生成样例)
+- [部署和推理](#部署和推理)
+- [微调方法](#微调方法)
+- [开源协议](#开源协议)
+
+# 模型介绍
+
+LingoWhale-8B是由深言科技推出的语鲸系列大模型中首个开源的中英双语大语言模型。
+
+LingoWhale-8B模型在数万亿token的高质量中英数据上进行预训练，具有强大的基础能力，在多个公开评测基准上均达到领先效果。在预训练阶段，模型使用8K的上下文长度进行训练，能够完成更长上下文的理解和生成任务。
+
+LingoWhale-8B模型对学术研究完全开放，使用方通过邮件申请并获得官方商用许可后，即可免费商用。
+
+在开源模型权重的同时，我们也提供了符合用户习惯的Huggingface推理接口以及LoRA等参数高效微调示例，便于开发者快速使用LingoWhale-8B模型。
+
+受模型参数量影响，大模型固有的幻觉问题、数学计算能力相对较弱等问题在LingoWhale-8B模型中仍然存在。请大家在使用前了解这些问题，评估可能存在的风险。后续版本的LingoWhale模型将会针对此类问题进行重点优化。
+
+
+# 测评结果
+
+我们在以下公开评测基准上进行了测试：
+
+- [C-Eval](https://arxiv.org/abs/2305.08322)是一个中文基础模型评估基准，包含了13948个多项选择题，涵盖了52个不同的学科和四个难度级别。它旨在评估中文语言模型的能力。我们使用该数据集的dev集作为few-shot的来源，在test集上进行了5-shot测试。
+- [MMLU](https://arxiv.org/abs/2009.03300)是一个英文基础模型评估基准，涵盖了基本数学、美国历史、计算机科学、法律等多个领域，共包含57个任务。它用于评估语言模型在不同领域任务上的表现。我们对模型进行了5-shot测试。
+- [CMMLU](https://arxiv.org/abs/2306.09212)是一个中文评估基准，涵盖了从基础学科到高级专业水平的67个主题。它用于评估中文语言模型在知识和推理能力方面的表现。我们使用该数据集的dev集作为few-shot的来源，在test集上进行了5-shot测试。
+- [Gaokao](https://arxiv.org/abs/2305.12474)是一个以中国高考题目为数据集的评估基准。它旨在提供测评中文语言模型在语言理解能力和逻辑推理能力方面的能力。我们只保留了其中的四选一的选择题，随机划分后对模型进行了5-shot测试。
+- [HumanEval](https://arxiv.org/abs/2107.03374)是一个包含上百个编程问题的英文评估基准。它用于评估语言模型在程序理解与生成能力方面的表现。我们采用了zero-shot计算Pass@1的方法对模型进行了测试。
+- [GSM8K](https://arxiv.org/abs/2110.14168)是一个由高质量、语言多样化的小学数学应用题组成的数据集。它要求根据给定的场景选择最合理的方案，用于评估语言模型在数学应用方面的能力。我们对模型进行了8-shot测试。
+- [BBH](https://arxiv.org/abs/2210.09261)是一个从204项Big-Bench评测基准任务中选择出的表现较差的任务单独形成的评估基准。它用于评估大型语言模型在具有挑战性的任务上的表现。我们对模型进行了3-shot测试。
+- [AGIEval](https://arxiv.org/abs/2304.06364)是一项考察基础模型类人能力的基准测试，专门用于评估基础模型在人类认知和问题解决相关任务中的能力。我们只保留了其中的四选一的选择题，随机划分后对模型进行了5-shot测试。
+
+
+这些评估基准提供了标准化的测试和度量，用于评估语言模型在不同任务和领域上的性能和能力。评测方法和测评结果如下表所示：
+
+| **Model**              |**C-Eval**| **MMLU** |**CMMLU** |**GAOKAO**|**HumanEval**|**GSM8K** | **BBH**  |**AGIEval**|
+|:-----------------------|:--------:|:--------:|:--------:|:--------:|:-----------:|:--------:|:--------:|:---------:|
+|                        |  5-shot  |  5-shot  |  5-shot  |  5-shot  |   0-shot    |  8-shot  |  3-shot  |  5-shot   |
+| **GPT-4**              |   68.4   |   83.9   |   70.3   |   66.2   |    69.5     |   90.0   |   75.1   |   63.3    |
+| **GPT-3.5 Turbo**      |   51.1   |   68.5   |   54.1   |   47.1   |    52.4     |   57.8   |   61.6   |   46.1    |
+| **LLaMA2-7B**          |   28.9   |   45.7   |   31.4   |   26.0   |    12.8     |   16.2   |   39.2   |   26.5    |
+| **ChatGLM2-6B**        |   50.2   |   45.9   |   49.0   |   49.4   |    9.2      |   28.9   |   31.7   |   45.3    |
+| **Baichuan2-7B-Base**  |   54.0   |   54.2   |   57.1   |   47.5   |    18.3     |   24.5   |   41.6   |   42.7    |
+| **Qwen-7B v1.1**       |   63.5   |   58.2   |   62.2   |     -    |    29.9     |   51.7   |   45.0   |     -     |
+| **LingoWhale-8B-base** |   63.6   |   60.2   |   62.8   |   50.3   |    32.9     |   55.0   |   47.5   |   43.8    |
+
+<span style="color:gray">对于以上所有对比模型，我们列出了其官方汇报结果并四舍五入保留一位小数。</span>
+
+# 生成样例
+
+> 生成以下样例的模型是使用开源SFT数据基于LingoWhale-8B模型进行微调后的对话模型
+
+<details><summary><b>数学计算</b></summary><p>
+
+**问题**：f(x)为偶函数， f(1) = 3, f(2) = 4, 那么f(1) + f(-2) ^2 等于多少？ 请一步一步说明
+
+**解答**：首先，根据题目中给出的条件，我们可以知道：
+
+- f(x) 是一个偶函数，即 f(-x) = f(x)。
+- f(1) = 3。
+- f(2) = 4。
+
+接下来，我们需要求出 f(1) + f(-2)^2 的值。
+由于 f(x) 是偶函数，所以 f(-2) = f(2)，因此 f(-2) = 4。
+将 f(1) 和 f(-2) 的值代入 f(1) + f(-2)^2 中，得到：
+f(1) + f(-2)^2 = 3 + (4)^2 = 3 + 16 = 19。
+因此，f(1) + f(-2)^2 的值为 19。
+
+</details>
+
+<details><summary><b>角色扮演</b></summary><p>
+
+**问题**：作为有丰富营销经验的品牌经理，请为我提供一条关于以下情境的品牌营销策略：
+情境：618期间，提高品牌销售额
+品牌定位：国产精品眉笔
+建议要求：提供具有创新性和实效性的品牌营销策略，帮助提升品牌的销售额， 并加入直播带货的方案。
+
+**解答**：在618期间，为了提高品牌销售额，我们可以采取以下品牌营销策略：
+
+1. 推出限时优惠活动：在618活动期间，我们可以在品牌官方网站上推出限时优惠活动，例如购买眉笔送赠品、打折促销等，以吸引消费者的注意力并促进销售。
+
+2. 利用社交媒体宣传：我们可以通过社交媒体平台（如微博、微信、抖音等）发布品牌宣传内容，包括产品介绍、使用技巧、用户评价等，以增加品牌曝光率和用户互动性。
+
+3. 与知名博主合作推广：我们可以邀请一些知名的美妆博主或网红来试用我们的产品，并在其社交媒体上发布推荐视频或图文，以吸引更多的目标受众。
+
+4. 直播带货：我们可以邀请一些知名的带货主播或明星来为我们的产品进行直播带货，通过他们的影响力来吸引更多的消费者关注和购买。
+
+5. 参与电商平台活动：我们可以积极参与各大电商平台的促销活动，例如满减、满赠等，以吸引更多的消费者购买我们的产品。
+
+综上所述，以上品牌营销策略可以帮助我们在618期间提高品牌销售额，同时加入直播带货的方案可以更好地吸引目标受众，提高销售效果。
+
+</details>
+
+<details><summary><b>诗歌创作</b></summary><p>
+
+**问题**：写一个七言诗，描述岳飞的诗。诗词希望是宋代诗的风格。
+
+**解答**：《岳飞》
+
+铁马金戈壮士魂，
+
+精忠报国志凌云。
+
+千古英雄谁可及，
+
+岳家军威震乾坤。
+
+</details>
+
+<details><summary><b>文化知识</b></summary><p>
+
+**问题**：我要去成都武侯祠参观，能帮我详细介绍下武侯的生平事迹吗？
+
+**解答**：诸葛亮（181年-234年），字孔明，号卧龙，徐州琅琊阳都（今山东临沂市沂南县）人，三国时期蜀汉丞相，杰出的政治家、军事家、文学家、书法家、发明家。早年随叔父诸葛玄到荆州，诸葛玄死后，诸葛亮就在隆中隐居。刘备三顾茅庐请出诸葛亮，联合东吴孙权于赤壁之战大败曹军，形成三国鼎足之势，又夺占荆益。刘备入蜀后，诸葛亮镇守荆州。刘备称帝后，诸葛亮被任命为丞相，主持朝政。后主刘禅继位，诸葛亮被封为武乡侯，领益州牧。勤勉谨慎，大小政事必亲自处理，赏罚严明；与东吴联盟，改善和西南各族的关系；实行屯田政策，加强战备。前后五次北伐中原，多以粮尽无功。终因积劳成疾，于234年病逝于五丈原（今陕西岐山县南），享年54岁。刘禅追封其为忠武侯，后世常以武侯尊称诸葛亮。东晋政权因其军事才能特追封他为武兴王。诸葛亮散文代表作有《出师表》、《诫子书》等。曾发明木牛流马、孔明灯等，并改造连弩，叫做诸葛连弩，可一弩十矢俱发。诸葛亮一生“鞠躬尽瘁、死而后已”，是中国传统文化中忠臣与智者的代表人物。
+
+</details>
+
+<details><summary><b>代码生成</b></summary><p>
+
+**问题**：你是一名编程开发者，能写一个正则帮忙判断一个字符串是否是合理的手机号码， 用Python实现。
+
+**解答**：
+```python
+import re
+
+def is_valid_phone(phone):
+    pattern = r'^1[3-9]\d{9}$'
+    return bool(re.match(pattern, phone))
+
+print(is_valid_phone('13800138000'))  # True
+print(is_valid_phone('12345678901'))  # False
+```
+
+</details>
+
+# 部署和推理
+
+## 安装依赖
+```shell
+pip install -r requirements.txt
+```
+
++ 如需使用flash attention需要安装[flash-attn](https://github.com/Dao-AILab/flash-attention)
++ 如需使用LoRA等少参数微调方法需要额外安装[peft](https://github.com/huggingface/peft)
+
+## Python 样例
+
+```python
+>>> from transformers import AutoModelForCausalLM, AutoTokenizer
+>>> tokenizer = AutoTokenizer.from_pretrained("deeplang-ai/LingoWhale-8B", trust_remote_code=True)
+>>> model = AutoModelForCausalLM.from_pretrained("deeplang-ai/LingoWhale-8B", device_map="auto", trust_remote_code=True)
+>>> inputs = tokenizer("陋室铭\n唐 刘禹锡\n", return_tensors="pt")
+>>> inputs = inputs.to("cuda:0")
+>>> pred = model.generate(**inputs, max_new_tokens=100, repetition_penalty=1.1)
+>>> print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))
+```
+
+# 微调方法
+模型微调样例代码以`transformers.Trainer`为基础，其中大部分参数和使用方法都可以参考Huggingface中[`Trainer`](https://huggingface.co/docs/transformers/v4.34.1/en/main_classes/trainer#trainer) 的教程和介绍。
+
+> 本章节旨在展示微调过程，并不对该微调配置下进行微调后的模型效果进行保证。
+
+## 单机训练
+下面是一个单机进行微调的例子，使用的数据为从[COIG](https://huggingface.co/datasets/BAAI/COIG)数据集中随机选取的10000条指令微调数据，可以使用自己的数据进行替换。
+
+```shell
+hostfile=""
+deepspeed --hostfile=$hostfile finetune/finetune.py  \
+    --report_to "none" \
+    --data_path "finetune/data/coig_10k.json" \
+    --model_name_or_path deeplang-ai/LingoWhale-8B \
+    --output_dir "output" \
+    --model_max_length 2048 \
+    --num_train_epochs 4 \
+    --per_device_train_batch_size 4 \
+    --gradient_accumulation_steps 1 \
+    --save_strategy epoch \
+    --learning_rate 2e-5 \
+    --lr_scheduler_type constant \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.98 \
+    --adam_epsilon 1e-8 \
+    --max_grad_norm 1.0 \
+    --weight_decay 1e-4 \
+    --warmup_ratio 0.0 \
+    --logging_steps 1 \
+    --gradient_checkpointing True \
+    --deepspeed finetune/ds_config.json \
+    --bf16 True \
+    --tf32 True
+```
+
+若要替换为自己的数据，可以使用如下格式的json文件。
+```json
+[
+    {
+        "id": 0,
+        "conversations": [
+            {
+                "from": "human",
+                "value": "请问什么是“模式年龄”？"
+            },
+            {
+                "from": "model",
+                "value": "模式年龄是指利用放射性衰变规律假定地质样品形成时的初始同位素组成计算得到的年龄。"
+            },
+            ...
+        ]
+    },
+    ...
+]
+```
+
+## 多机训练
+
+多机器训练需要编辑如下格式的`hostfile`文件。其中，每一行表示一个机器，`ip_address-X`为各个机器对应的ip地址，`slots`内容表示机器可用GPU数量。内容格式如下：
+
+```
+ip_address-1 slots=8
+ip_address-2 slots=8
+ip_address-3 slots=8
+ip_address-4 slots=8
+...
+```
+
+同时指定hostfile参数为`hostfile`文件路径，然后运行如下命令即可启动多机训练。
+
+```shell
+hostfile="/path/to/hostfile"
+deepspeed --hostfile=$hostfile finetune/finetune.py  \
+    --report_to "none" \
+    --data_path "finetune/data/coig_10k.json" \
+    --model_name_or_path deeplang-ai/LingoWhale-8B \
+    --output_dir "output" \
+    --model_max_length 2048 \
+    --num_train_epochs 4 \
+    --per_device_train_batch_size 4 \
+    --gradient_accumulation_steps 1 \
+    --save_strategy epoch \
+    --learning_rate 2e-5 \
+    --lr_scheduler_type constant \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.98 \
+    --adam_epsilon 1e-8 \
+    --max_grad_norm 1.0 \
+    --weight_decay 1e-4 \
+    --warmup_ratio 0.0 \
+    --logging_steps 1 \
+    --gradient_checkpointing True \
+    --deepspeed finetune/ds_config.json \
+    --bf16 True \
+    --tf32 True
+```
+## 少参数微调
+通过使用[peft](https://github.com/huggingface/peft)，可以轻松调LoRA, Prefix-Tuning等少参数微调的方法。目前在代码中集合了LoRA的训练方法，可以通过加入`--use_lora True`启动。
+
+使用LoRA训练的checkpoint可以通过下面的代码读取和调用：
+```python
+from peft import AutoPeftModelForCausalLM
+model = AutoPeftModelForCausalLM.from_pretrained("output", trust_remote_code=True)
+```
+
+# 开源协议
+社区使用LingoWhale-8B模型需要遵循[Apache 2.0](http://www.apache.org/licenses/LICENSE-2.0)和[《LingoWhale-8B模型许可协议》](MODEL_LICENSE.md)。若您期望将此源模型或其衍��品用作商业用途，请参考[《LingoWhale-8B模型许可协议》](MODEL_LICENSE.md)。

README_EN.md

@@ -0,0 +1,295 @@
+<p align="left">
+    English</a>&nbsp ｜ &nbsp<a href="README.md">中文</a>
+</p>
+<br>
+
+<div align="center">
+<h1>
+LingoWhale-8B
+</h1>
+</div>
+
+<p align="center">
+🤗 <a href="https://huggingface.co/deeplang-ai/LingoWhale-8B" target="_blank">Hugging Face</a> • 🤖 <a href="https://www.modelscope.cn/models/DeepLang/LingoWhale-8B" target="_blank">ModelScope</a> • ⛵ <a href="https://wisemodel.cn/models/%E5%8C%97%E4%BA%AC%E6%B7%B1%E8%A8%80%E7%A7%91%E6%8A%80%E6%9C%89%E9%99%90%E8%B4%A3%E4%BB%BB%E5%85%AC%E5%8F%B8/LingoWhale-8B/" target="_blank">Wisemodel</a>
+</p>
+
+<div align="center">
+<strong>
+LingoWhale-8B model open-sourced by DeepLangAI in collaboration with THUNLP Lab 🎉
+</strong>
+</div>
+
+# Table of Contents
+
+- [Introduction](#introduction)
+- [Evaluation](#evaluation)
+- [Generated Examples](#generated-examples)
+- [Deployment and Inference](#deployment-and-inference)
+- [Fine-tuning](#fine-tuning)
+- [Open Source License](#open-source-license)
+
+# Introduction
+
+LingoWhale-8B is the first open-source model in the LingoWhale series introduced by DeepLangAI. It's a bilingual (Chinese-English) large language model.
+
+LingoWhale-8B has been pre-trained on a large volume of high-quality bilingual data and exhibits powerful capabilities as a foundation model. It has achieved leading results on multiple public benchmarks. During its pre-training phase, the model was trained with a context window of 8K, allowing it to comprehend and generate longer sequences.
+
+LingoWhale-8B is fully open for academic research. Users can apply for commercial use by email, and once granted official commercial permission, they can use it for free.
+
+Along with open-sourcing the model weights, we also provide a Huggingface inference interface and parameter efficient fine-tuning examples like LoRA, making it easier for developers to use the LingoWhale-8B model.
+
+Due to the scale of model parameters, intrinsic issues of large language models like hallucination and relatively weak mathematical computation capabilities persist in LingoWhale-8B. Please understand these issues and evaluate the possible risks before using the model. Future versions of the LingoWhale model will focus on optimizing these areas.
+
+# Evaluation
+
+We tested on the following public evaluation benchmarks:
+
+- [C-Eval](https://arxiv.org/abs/2305.08322) is a Chinese foundation model evaluation benchmark consisting of 13,948 multiple-choice questions, covering 52 different subjects and four difficulty levels. It aims to assess the capability of Chinese language models. We used the dev set of this dataset as a few-shot source and conducted a 5-shot test on the test set.
+  
+- [MMLU](https://arxiv.org/abs/2009.03300) is an English foundation model evaluation benchmark that spans various domains like basic mathematics, American history, computer science, law, among others, with a total of 57 tasks. It evaluates language models' performance on different domain tasks. We performed a 5-shot test on this benchmark.
+  
+- [CMMLU](https://arxiv.org/abs/2306.09212) is a Chinese evaluation benchmark that encompasses 67 topics ranging from basic subjects to advanced professional levels. It evaluates Chinese language models' performance in knowledge and reasoning capabilities. We used the dev set of this dataset as a few-shot source and conducted a 5-shot test on the test set.
+  
+- [Gaokao](https://arxiv.org/abs/2305.12474) is an evaluation benchmark based on the dataset of Chinese college entrance examination questions. It aims to provide an assessment of Chinese language models in terms of language comprehension and logical reasoning capabilities. We retained only the four-option multiple-choice questions from it and conducted a 5-shot test after random partitioning.
+  
+- [HumanEval](https://arxiv.org/abs/2107.03374) is an English evaluation benchmark containing over one hundred coding problems. It assesses language models' abilities in code comprehension and generation. We adopted a zero-shot setting and the Pass@1 metric for testing the model.
+  
+- [GSM8K](https://arxiv.org/abs/2110.14168) is a dataset composed of high-quality elementary school math application problems. It requires the models to select the most appropriate solution based on the provided scenario and evaluates the models' capabilities in mathematical application. We conducted an 8-shot test on this benchmark.
+  
+- [BBH](https://arxiv.org/abs/2210.09261) is an evaluation benchmark formed from a selection of challenging tasks out of 204 Big-Bench benchmark tasks. We performed a 3-shot test on this benchmark.
+  
+- [AGIEval](https://arxiv.org/abs/2304.06364) is a benchmark to examine foundation models' human-like capabilities, specifically assessing foundational models' abilities in human cognition and problem-solving tasks. We retained only the four-option multiple-choice questions from it and conducted a 5-shot test after random partitioning.
+  
+These evaluation benchmarks provide standardized tests and metrics to assess language models' performance and capabilities across various tasks and domains. The evaluation results are shown in the table below:
+
+| **Model**              |**C-Eval**| **MMLU** |**CMMLU** |**GAOKAO**|**HumanEval**|**GSM8K** | **BBH**  |**AGIEval**|
+|:-----------------------|:--------:|:--------:|:--------:|:--------:|:-----------:|:--------:|:--------:|:---------:|
+|                        |  5-shot  |  5-shot  |  5-shot  |  5-shot  |   0-shot    |  8-shot  |  3-shot  |  5-shot   |
+| **GPT-4**              |   68.4   |   83.9   |   70.3   |   66.2   |    69.5     |   90.0   |   75.1   |   63.3    |
+| **GPT-3.5 Turbo**      |   51.1   |   68.5   |   54.1   |   47.1   |    52.4     |   57.8   |   61.6   |   46.1    |
+| **LLaMA2-7B**          |   28.9   |   45.7   |   31.4   |   26.0   |    12.8     |   16.2   |   39.2   |   26.5    |
+| **ChatGLM2-6B**        |   50.2   |   45.9   |   49.0   |   49.4   |    9.2      |   28.9   |   31.7   |   45.3    |
+| **Baichuan2-7B-Base**  |   54.0   |   54.2   |   57.1   |   47.5   |    18.3     |   24.5   |   41.6   |   42.7    |
+| **Qwen-7B v1.1**       |   63.5   |   58.2   |   62.2   |     -    |    29.9     |   51.7   |   45.0   |     -     |
+| **LingoWhale-8B-base** |   63.6   |   60.2   |   62.8   |   50.3   |    32.9     |   55.0   |   47.5   |   43.8    |
+
+<span style="color:gray">For all of the above comparison models, we list their official reports and round them to one decimal place.</span>
+
+# Generated Examples
+
+The following examples are generated by a Chat model fine-tuned based on the LingoWhale-8B model using open source SFT data.
+
+<details><summary><b>Mathematical calculations</b></summary><p>
+
+**Query**：f(x)为偶函数， f(1) = 3, f(2) = 4, 那么f(1) + f(-2) ^2 等于多少？ 请一步一步说明
+
+**Answer**：首先，根据题目中给出的条件，我们可以知道：
+
+- f(x) 是一个偶函数，即 f(-x) = f(x)。
+- f(1) = 3。
+- f(2) = 4。
+
+接下来，我们需要求出 f(1) + f(-2)^2 的值。
+由于 f(x) 是偶函数，所以 f(-2) = f(2)，因此 f(-2) = 4。
+将 f(1) 和 f(-2) 的值代入 f(1) + f(-2)^2 中，得到：
+f(1) + f(-2)^2 = 3 + (4)^2 = 3 + 16 = 19。
+因此，f(1) + f(-2)^2 的值为 19。
+
+</details>
+
+<details><summary><b>Role playing</b></summary><p>
+
+**Query**：作为有丰富营销经验的品牌经理，请为我提供一条关于以下情境的品牌营销策略：
+情境：618期间，提高品牌销售额
+品牌定位：国产精品眉笔
+建议要求：提供具有创新性和实效性的品牌营销策略，帮助提升品牌的销售额， 并加入直播带货的方案。
+
+**Answer**：在618期间，为了提高品牌销售额，我们可以采取以下品牌营销策略：
+
+1. 推出限时优惠活动：在618活动期间，我们可以在品牌官方网站上推出限时优惠活动，例如购买眉笔送赠品、打折促销等，以吸引消费者的注意力并促进销售。
+
+2. 利用社交媒体宣传：我们可以通过社交媒体平台（如微博、微信、抖音等）发布品牌宣传内容，包括产品介绍、使用技巧、用户评价等，以增加品牌曝光率和用户互动性。
+
+3. 与知名博主合作推广：我们可以邀请一些知名的美妆博主或网红来试用我们的产品，并在其社交媒体上发布推荐视频或图文，以吸引更多的目标受众。
+
+4. 直播带货：我们可以邀请一些知名的带货主播或明星来为我们的产品进行直播带货，通过他们的影响力来吸引更多的消费者关注和购买。
+
+5. 参与电商平台活动：我们可以积极参与各大电商平台的促销活动，例如满减、满赠等，以吸引更多的消费者购买我们的产品。
+
+综上所述，以上品牌营销策略可以帮助我们在618期间提高品牌销售额，同时加入直播带货的方案可以更好地吸引目标受众，提高销售效果。
+
+</details>
+
+<details><summary><b>Poetry Writing</b></summary><p>
+
+**Query**：写一个七言诗，描述岳飞的诗。诗词希望是宋代诗的风格。
+
+**Answer**：《岳飞》
+
+铁马金戈壮士魂，
+
+精忠报国志凌云。
+
+千古英雄谁可及，
+
+岳家军威震乾坤。
+
+</details>
+
+<details><summary><b>Cultural knowledge</b></summary><p>
+
+**Query**：我要去成都武侯祠参观，能帮我详细介绍下武侯的生平事迹吗？
+
+**Answer**：诸葛亮（181年-234年），字孔明，号卧龙，徐州琅琊阳都（今山东临沂市沂南县）人，三国时期蜀汉丞相，杰出的政治家、军事家、文学家、书法家、发明家。早年随叔父诸葛玄到荆州，诸葛玄死后，诸葛亮就在隆中隐居。刘备三顾茅庐请出诸葛亮，联合东吴孙权于赤壁之战大败曹军，形成三国鼎足之势，又夺占荆益。刘备入蜀后，诸葛亮镇守荆州。刘备称帝后，诸葛亮被任命为丞相，主持朝政。后主刘禅继位，诸葛亮被封为武乡侯，领益州牧。勤勉谨慎，大小政事必亲自处理，赏罚严明；与东吴联盟，改善和西南各族��关系；实行屯田政策，加强战备。前后五次北伐中原，多以粮尽无功。终因积劳成疾，于234年病逝于五丈原（今陕西岐山县南），享年54岁。刘禅追封其为忠武侯，后世常以武侯尊称诸葛亮。东晋政权因其军事才能特追封他为武兴王。诸葛亮散文代表作有《出师表》、《诫子书》等。曾发明木牛流马、孔明灯等，并改造连弩，叫做诸葛连弩，可一弩十矢俱发。诸葛亮一生“鞠躬尽瘁、死而后已”，是中国传统文化中忠臣与智者的代表人物。
+
+</details>
+
+<details><summary><b>Code generation</b></summary><p>
+
+**Query**：你是一名编程开发者，能写一个正则帮忙判断一个字符串是否是合理的手机号码， 用Python实现。
+
+**Answer**：
+```python
+import re
+
+def is_valid_phone(phone):
+    pattern = r'^1[3-9]\d{9}$'
+    return bool(re.match(pattern, phone))
+
+print(is_valid_phone('13800138000'))  # True
+print(is_valid_phone('12345678901'))  # False
+```
+
+</details>
+
+# Deployment and Inference
+
+## Install Dependencies
+```shell
+pip install -r requirements.txt
+```
+
++ To use flash attention, you need to install [flash-attn](https://github.com/Dao-AILab/flash-attention)
+
++ For using LoRA or other parameter efficient fine-tuning methods, please install [peft](https://github.com/huggingface/peft)
+
+## Python Example
+
+```python
+>>> from transformers import AutoModelForCausalLM, AutoTokenizer
+>>> tokenizer = AutoTokenizer.from_pretrained("deeplang-ai/LingoWhale-8B", trust_remote_code=True)
+>>> model = AutoModelForCausalLM.from_pretrained("deeplang-ai/LingoWhale-8B", device_map="auto", trust_remote_code=True)
+>>> inputs = tokenizer("陋室铭\n唐 刘禹锡\n", return_tensors="pt")
+>>> inputs = inputs.to("cuda:0")
+>>> pred = model.generate(**inputs, max_new_tokens=100, repetition_penalty=1.1)
+>>> print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))
+```
+
+# Fine-tuning
+The fine-tuning example is based on `transformers.Trainer`. For a more deteiled guide on the arguments usage, please refer to Huggingface [`Trainer`](https://huggingface.co/docs/transformers/v4.34.1/en/main_classes/trainer#trainer) tutorial.
+
+The aim of this section is to showcase the fine-tuning process. No guarantees are made about the model performance under this fine-tuning configuration.
+
+## Single-Machine Training
+Below is an example of fine-tuning on a single machine. The data used is 10,000 randomly selected instruction fine-tuning data from the COIG dataset. You can replace it with your own data.
+
+```shell
+hostfile=""
+deepspeed --hostfile=$hostfile finetune/finetune.py  \
+    --report_to "none" \
+    --data_path "finetune/data/coig_10k.json" \
+    --model_name_or_path deeplang-ai/LingoWhale-8B \
+    --output_dir "output" \
+    --model_max_length 2048 \
+    --num_train_epochs 4 \
+    --per_device_train_batch_size 16 \
+    --gradient_accumulation_steps 1 \
+    --save_strategy epoch \
+    --learning_rate 2e-5 \
+    --lr_scheduler_type constant \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.98 \
+    --adam_epsilon 1e-8 \
+    --max_grad_norm 1.0 \
+    --weight_decay 1e-4 \
+    --warmup_ratio 0.0 \
+    --logging_steps 1 \
+    --gradient_checkpointing True \
+    --deepspeed finetune/ds_config.json \
+    --bf16 True \
+    --tf32 True
+```
+
+To use you own data, please convert it to the json format below
+```json
+[
+    {
+        "id": 0,
+        "conversations": [
+            {
+                "from": "human",
+                "value": "请问什么是“模式年龄”？"
+            },
+            {
+                "from": "model",
+                "value": "模式年龄是指利用放射性衰变规律假定地质样品形成时的初始同位素组成计算得到的年龄。"
+            },
+            ...
+        ]
+    },
+    ...
+]
+```
+
+## Multi-Machine Training
+
+For multi-machine training, you need to create a `hostfile` in the following format. Each line represents a machine. `ip_address-X` refers to the IP address of each machine, and the `slots` content indicates the number of available GPUs on the machine. The content format is as follows:
+
+```
+ip_address-1 slots=8
+ip_address-2 slots=8
+ip_address-3 slots=8
+ip_address-4 slots=8
+...
+```
+
+Next, specify the `hostfile` argument using the path of the `hostfile`, and run the following command to start multi-machine training.
+
+```shell
+hostfile="/path/to/hostfile"
+deepspeed --hostfile=$hostfile finetune/finetune.py  \
+    --report_to "none" \
+    --data_path "finetune/data/coig_10k.json" \
+    --model_name_or_path deeplang-ai/LingoWhale-8B \
+    --output_dir "output" \
+    --model_max_length 2048 \
+    --num_train_epochs 4 \
+    --per_device_train_batch_size 16 \
+    --gradient_accumulation_steps 1 \
+    --save_strategy epoch \
+    --learning_rate 2e-5 \
+    --lr_scheduler_type constant \
+    --adam_beta1 0.9 \
+    --adam_beta2 0.98 \
+    --adam_epsilon 1e-8 \
+    --max_grad_norm 1.0 \
+    --weight_decay 1e-4 \
+    --warmup_ratio 0.0 \
+    --logging_steps 1 \
+    --gradient_checkpointing True \
+    --deepspeed finetune/ds_config.json \
+    --bf16 True \
+    --tf32 True
+```
+
+## Parameter-Efficient Fine-Tuning
+By using [peft](https://github.com/huggingface/peft), you can easily apply parameter-efficient fine-tuning methods like LoRA, Prefix-Tuning, etc. The training method for LoRA is currently integrated into the code, which can be activated by adding `--use_lora True`.
+
+LoRA checkpoints can be loaded using the following code:
+```python
+from peft import AutoPeftModelForCausalLM
+model = AutoPeftModelForCausalLM.from_pretrained("output", trust_remote_code=True)
+```
+
+# Open Source License
+The community use of the LingoWhale-8B model must adhere to the [Apache 2.0](http://www.apache.org/licenses/LICENSE-2.0) and the [LingoWhale-8B Model License Agreement](MODEL_LICENSE.md). If you wish to use this source model or its derivatives for commercial purposes, please refer to [LingoWhale-8B Model License Agreement](MODEL_LICENSE.md).

config.json

@@ -0,0 +1,33 @@
+{
+  "architectures": [
+    "LingoWhaleForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_lingowhale.LingoWhaleConfig",
+    "AutoModelForCausalLM": "modeling_lingowhale.LingoWhaleForCausalLM"
+  },
+  "tokenizer_class": "LingoWhaleTokenizer",
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 11008,
+  "max_position_embeddings": 8192,
+  "model_max_length": 8192,
+  "model_type": "lingowhale",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 36,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "_from_model_config": true,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.31.0",
+  "use_cache": true,
+  "emb_dropout_prob": 0.0,
+  "attn_dropout_prob": 0.0,
+  "vocab_size": 96000,
+  "rope_theta": 10000.0,
+  "use_flash_attention": true
+}

configuration_lingowhale.py

@@ -0,0 +1,74 @@
+# Copyright 2023 DeepLang AI. All Rights Reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class LingoWhaleConfig(PretrainedConfig):
+    model_type = "lingowhale"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=96000,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=36,
+        num_attention_heads=32,
+        hidden_act="silu",
+        max_position_embeddings=8192,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        emb_dropout_prob=0.0,
+        attn_dropout_prob=0.0,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        use_flash_attention=True,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.emb_dropout_prob = emb_dropout_prob
+        self.attn_dropout_prob = attn_dropout_prob
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.use_flash_attention = use_flash_attention
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,14 @@
+{
+  "pad_token_id": 0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "user_token_id": 3,
+  "assistant_token_id": 4,
+  "max_new_tokens": 2048,
+  "temperature": 0.9,
+  "top_k": 1,
+  "top_p": 1,
+  "repetition_penalty": 1.1,
+  "do_sample": false,
+  "transformers_version": "4.31.0"
+}

modeling_lingowhale.py

@@ -0,0 +1,941 @@
+# Copyright 2023 DeepLang AI. All Rights Reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import os
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from torch.nn import functional as F
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutputWithPast,
+                                           CausalLMOutputWithPast)
+from transformers.utils import logging
+
+from .configuration_lingowhale import LingoWhaleConfig
+
+logger = logging.get_logger(__name__)
+
+try:
+    from einops import rearrange
+except ImportError:
+    rearrange = None
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_unpadded_func
+except ImportError:
+    try:
+        from flash_attn.flash_attn_interface import \
+            flash_attn_varlen_func as flash_attn_unpadded_func
+    except ImportError:
+        flash_attn_unpadded_func = None
+
+
+# Copied from transformers.models.bart.modeling_bart._make_causal_mask
+def _make_causal_mask(
+    input_ids_shape: torch.Size,
+    dtype: torch.dtype,
+    device: torch.device,
+    past_key_values_length: int = 0,
+):
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full(
+        (tgt_len, tgt_len),
+        torch.tensor(torch.finfo(dtype).min, device=device),
+        device=device,
+    )
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat(
+            [
+                torch.zeros(tgt_len,
+                            past_key_values_length,
+                            dtype=dtype,
+                            device=device),
+                mask,
+            ],
+            dim=-1,
+        )
+    return mask[None, None, :, :].expand(bsz, 1, tgt_len,
+                                         tgt_len + past_key_values_length)
+
+
+# Copied from transformers.models.bart.modeling_bart._expand_mask
+def _expand_mask(mask: torch.Tensor,
+                 dtype: torch.dtype,
+                 tgt_len: Optional[int] = None):
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+
+    bsz, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len,
+                                                  src_len).to(dtype)
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(inverted_mask.to(torch.bool),
+                                     torch.finfo(dtype).min)
+
+
+class LingoWhaleRMSNorm(torch.nn.Module):
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class LingoWhaleRotaryEmbedding(torch.nn.Module):
+
+    def __init__(self,
+                 dim,
+                 max_position_embeddings=2048,
+                 base=10000,
+                 device=None):
+        super().__init__()
+        self.inv_freq = 1.0 / (base**(
+            torch.arange(0, dim, 2).float().to(device) / dim))
+        self.max_seq_len_cached = max_position_embeddings
+        t = torch.arange(self.max_seq_len_cached,
+                         device=self.inv_freq.device,
+                         dtype=torch.float32)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32)
+        self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
+        if seq_len > self.max_seq_len_cached:
+            self.max_seq_len_cached = seq_len
+            t = torch.arange(
+                self.max_seq_len_cached,
+                device=self.inv_freq.device,
+                dtype=torch.float32,
+            )
+            freqs = torch.outer(t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32).to(
+                x.device)
+            self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32).to(
+                x.device)
+        elif self.cos_cached.device != x.device:
+            self.cos_cached = self.cos_cached.to(x.device)
+            self.sin_cached = self.sin_cached.to(x.device)
+        return (
+            self.cos_cached[:, :, :seq_len, ...],
+            self.sin_cached[:, :, :seq_len, ...],
+        )
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., :x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos_, sin_, position_ids):
+    cos = cos_.squeeze(1).squeeze(0)  # [seq_len, dim]
+    sin = sin_.squeeze(1).squeeze(0)  # [seq_len, dim]
+    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
+    k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
+    return q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+
+class LingoWhaleMLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_and_up_proj = nn.Linear(self.hidden_size,
+                                          self.intermediate_size * 2,
+                                          bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size,
+                                   self.hidden_size,
+                                   bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        gate_and_up = self.gate_and_up_proj(x)
+        [gate, up] = torch.chunk(gate_and_up, 2, dim=-1)
+
+        acted = self.act_fn(gate)
+        tmp = acted * up
+
+        result = self.down_proj(tmp)
+
+        return result
+
+
+class LingoWhaleAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: LingoWhaleConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.dropout_p = config.attn_dropout_prob
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads}).")
+        self.qkv_proj = nn.Linear(self.hidden_size,
+                                  3 * self.hidden_size,
+                                  bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim,
+                                self.hidden_size,
+                                bias=False)
+        self.attention_dropout = torch.nn.Dropout(self.dropout_p)
+        self._init_rope()
+
+    def attention_mask_func(self, attention_scores, attention_mask):
+        attention_scores.masked_fill_(attention_mask, -10000.0)
+        return attention_scores
+
+    def forward_torch_softmax(self, input, mask):
+        input = input.float()
+        mask_output = (self.attention_mask_func(input, mask)
+                       if mask is not None else input)
+        probs = torch.nn.Softmax(dim=-1)(mask_output)
+
+        probs = probs.bfloat16()
+
+        return probs
+
+    def _self_attention(self, query_layer, key_layer, value_layer,
+                        attention_mask):
+        output_size = (
+            query_layer.size(1),
+            query_layer.size(2),
+            query_layer.size(0),
+            key_layer.size(0),
+        )
+
+        # [sq, b, np, hn] -> [sq, b * np, hn]
+        query_layer = query_layer.reshape(output_size[2],
+                                          output_size[0] * output_size[1], -1)
+
+        # [sk, b, np, hn] -> [sk, b * np, hn]
+        key_layer = key_layer.reshape(output_size[3],
+                                      output_size[0] * output_size[1], -1)
+
+        matmul_input_buffer = torch.randn(
+            (output_size[0] * output_size[1], output_size[2], output_size[3]),
+            dtype=query_layer.dtype,
+            device=query_layer.device,
+        )
+        norm_factor = math.sqrt(key_layer.shape[-1])
+        # Raw attention scores. [b * np, sq, sk]
+        matmul_result = torch.baddbmm(
+            matmul_input_buffer,
+            query_layer.transpose(0, 1),  # [b * np, sq, hn]
+            key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
+            beta=0.0,
+            alpha=(1.0 / norm_factor),
+        )
+
+        # change view to [b, np, sq, sk]
+        attention_scores = matmul_result.view(*output_size)
+
+        # attention scores and attention mask [b, np, sq, sk]
+        attention_probs = self.forward_torch_softmax(attention_scores,
+                                                     attention_mask)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+
+        attention_probs = self.attention_dropout(attention_probs)
+
+        # =========================
+        # Context layer. [sq, b, hp]
+        # =========================
+
+        # value_layer -> context layer.
+        # [sk, b, np, hn] --> [b, np, sq, hn]
+
+        # context layer shape: [b, np, sq, hn]
+        output_size = (
+            value_layer.size(1),
+            value_layer.size(2),
+            query_layer.size(0),
+            value_layer.size(3),
+        )
+
+        # change view [sk, b * np, hn]
+        value_layer = value_layer.reshape(value_layer.size(0),
+                                          output_size[0] * output_size[1], -1)
+
+        # change view [b * np, sq, sk]
+        attention_probs = attention_probs.view(output_size[0] * output_size[1],
+                                               output_size[2], -1)
+
+        # matmul: [b * np, sq, hn]
+        context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
+
+        # change view [b, np, sq, hn]
+        context_layer = context_layer.view(*output_size)
+
+        # [b, np, sq, hn] --> [sq, b, np, hn]
+        context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
+
+        # [sq, b, np, hn] --> [sq, b, hp]
+        new_context_layer_shape = context_layer.size()[:-2] + (
+            self.hidden_size, )
+
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        return context_layer
+
+    def _self_attention_flash(self, q, k, v):
+        batch_size, seqlen_q = q.shape[0], q.shape[1]
+        seqlen_k = k.shape[1]
+
+        q, k, v = [rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v]]
+        cu_seqlens_q = torch.arange(
+            0,
+            (batch_size + 1) * seqlen_q,
+            step=seqlen_q,
+            dtype=torch.int32,
+            device=q.device,
+        )
+
+        if self.training:
+            # during training q,k,v always have same seqlen
+            assert seqlen_k == seqlen_q
+
+            is_causal = True
+            cu_seqlens_k = cu_seqlens_q
+            dropout_p = self.dropout_p
+        else:
+            # turn off FA causal mask after first inference autoregressive iteration
+            # only on first autoregressive step q,k,v have same seqlen
+            is_causal = seqlen_q == seqlen_k
+            cu_seqlens_k = torch.arange(
+                0,
+                (batch_size + 1) * seqlen_k,
+                step=seqlen_k,
+                dtype=torch.int32,
+                device=q.device,
+            )
+            dropout_p = 0
+
+        output = flash_attn_unpadded_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqlen_q,
+            seqlen_k,
+            dropout_p,
+            causal=is_causal,
+        )
+
+        output = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
+        return output
+
+    def _init_rope(self):
+        self.rotary_emb = LingoWhaleRotaryEmbedding(
+            self.head_dim,
+            max_position_embeddings=self.max_position_embeddings,
+            base=self.rope_theta,
+        )
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return (tensor.view(bsz, seq_len, self.num_heads,
+                            self.head_dim).transpose(1, 2).contiguous())
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
+               Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        proj = self.qkv_proj(hidden_states)
+        proj = (proj.unflatten(-1,
+                               (3, self.hidden_size)).unsqueeze(0).transpose(
+                                   0, -2).squeeze(-2))
+
+        query_states = (proj[0].view(bsz, q_len, self.num_heads,
+                                     self.head_dim).transpose(1, 2))
+        key_states = (proj[1].view(bsz, q_len, self.num_heads,
+                                   self.head_dim).transpose(1, 2))
+        value_states = (proj[2].view(bsz, q_len, self.num_heads,
+                                     self.head_dim).transpose(1, 2))
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin, position_ids)
+
+        # [bsz, nh, t, hd]
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        query_states = query_states.transpose(1, 2).transpose(0, 1)
+        value_states = value_states.transpose(1, 2).transpose(0, 1)
+        key_states = key_states.transpose(1, 2).transpose(0, 1)
+        attention_mask = attention_mask < -0.5
+
+        if self.config.use_flash_attention and flash_attn_unpadded_func is not None:
+            assert (
+                rearrange is not None
+            ), "Please install einops first, e.g., with pip install einops"
+            q, k, v = [
+                rearrange(x, "s b ... -> b s ...").contiguous()
+                for x in (query_states, key_states, value_states)
+            ]
+            attn_output = self._self_attention_flash(q, k, v)
+            attn_output = rearrange(attn_output,
+                                    "b s h d -> s b (h d)").contiguous()
+        else:
+            attn_output = self._self_attention(query_states, key_states,
+                                               value_states, attention_mask)
+        attn_output = attn_output.transpose(0, 1)
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LingoWhaleDecoderLayer(nn.Module):
+
+    def __init__(self, config: LingoWhaleConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = LingoWhaleAttention(config=config)
+        self.mlp = LingoWhaleMLP(config)
+        self.input_layernorm = LingoWhaleRMSNorm(config.hidden_size,
+                                                 eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LingoWhaleRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor,
+                                                 torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        outputs = (hidden_states, )
+
+        if output_attentions:
+            outputs += (self_attn_weights, )
+
+        if use_cache:
+            outputs += (present_key_value, )
+
+        return outputs
+
+
+class LingoWhalePreTrainedModel(PreTrainedModel):
+    config_class = LingoWhaleConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LingoWhaleDecoderLayer"]
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, LingoWhaleModel):
+            module.gradient_checkpointing = value
+
+
+class LingoWhaleModel(LingoWhalePreTrainedModel):
+
+    def __init__(self, config: LingoWhaleConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size,
+                                         self.padding_idx)
+        self.layers = nn.ModuleList([
+            LingoWhaleDecoderLayer(config)
+            for _ in range(config.num_hidden_layers)
+        ])
+        self.norm = LingoWhaleRMSNorm(config.hidden_size,
+                                      eps=config.rms_norm_eps)
+        self.drop = nn.Dropout(config.emb_dropout_prob)
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
+    def _prepare_decoder_attention_mask(self, attention_mask, input_shape,
+                                        inputs_embeds, past_key_values_length):
+        # create causal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(attention_mask,
+                                              inputs_embeds.dtype,
+                                              tgt_len=input_shape[-1]).to(
+                                                  inputs_embeds.device)
+            combined_attention_mask = (expanded_attn_mask
+                                       if combined_attention_mask is None else
+                                       expanded_attn_mask +
+                                       combined_attention_mask)
+
+        return combined_attention_mask
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = (output_attentions if output_attentions is not None
+                             else self.config.output_attentions)
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = (return_dict if return_dict is not None else
+                       self.config.use_return_dict)
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError(
+                "You have to specify either decoder_input_ids or decoder_inputs_embeds"
+            )
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length,
+                seq_length + past_key_values_length,
+                dtype=torch.long,
+                device=device,
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        # embed positions
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past),
+                dtype=torch.bool,
+                device=inputs_embeds.device,
+            )
+        attention_mask = self._prepare_decoder_attention_mask(
+            attention_mask,
+            (batch_size, seq_length),
+            inputs_embeds,
+            past_key_values_length,
+        )
+
+        hidden_states = inputs_embeds
+        hidden_states = self.drop(hidden_states)
+        hidden_states = hidden_states.to(dtype=torch.bfloat16)
+
+        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
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states, )
+
+            past_key_value = (past_key_values[idx]
+                              if past_key_values is not None else None)
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, output_attentions, None)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    None,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (
+                    layer_outputs[2 if output_attentions else 1], )
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1], )
+
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states, )
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(
+                v for v in
+                [hidden_states, next_cache, all_hidden_states, all_self_attns]
+                if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+class LingoWhaleForCausalLM(LingoWhalePreTrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LingoWhaleModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = torch.nn.Linear(config.hidden_size,
+                                       config.vocab_size,
+                                       bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
+        *model_args,
+        config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
+        cache_dir: Optional[Union[str, os.PathLike]] = None,
+        ignore_mismatched_sizes: bool = False,
+        force_download: bool = False,
+        local_files_only: bool = False,
+        token: Optional[Union[str, bool]] = None,
+        revision: str = "main",
+        use_safetensors: bool = None,
+        **kwargs,
+    ):
+        # Load config if we don't provide a configuration
+        if not isinstance(config, PretrainedConfig):
+            config_path = (config if config is not None else
+                           pretrained_model_name_or_path)
+            config, model_kwargs = cls.config_class.from_pretrained(
+                config_path,
+                cache_dir=cache_dir,
+                return_unused_kwargs=True,
+                force_download=force_download,
+                resume_download=False,
+                proxies=None,
+                local_files_only=local_files_only,
+                token=token,
+                revision=revision,
+                subfolder="",
+                _from_auto=False,
+                _from_pipeline=None,
+                **kwargs,
+            )
+        else:
+            model_kwargs = kwargs
+        if "torch_dtype" not in kwargs:
+            kwargs["torch_dtype"] = config.torch_dtype
+        return super(LingoWhaleForCausalLM, cls).from_pretrained(
+            pretrained_model_name_or_path,
+            *model_args,
+            config=config,
+            cache_dir=cache_dir,
+            ignore_mismatched_sizes=ignore_mismatched_sizes,
+            force_download=force_download,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            use_safetensors=use_safetensors,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = (output_attentions if output_attentions is not None
+                             else self.config.output_attentions)
+        output_hidden_states = (output_hidden_states
+                                if output_hidden_states is not None else
+                                self.config.output_hidden_states)
+        return_dict = (return_dict if return_dict is not None else
+                       self.config.use_return_dict)
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            softmax_normalizer = shift_logits.max(-1).values**2
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits, ) + outputs[1:]
+            return (loss, ) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        **kwargs,
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update({
+            "position_ids": position_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "attention_mask": attention_mask,
+        })
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (tuple(
+                past_state.index_select(0, beam_idx)
+                for past_state in layer_past), )
+        return reordered_past

tokenization_lingowhale.py

@@ -0,0 +1,264 @@
+# Copyright 2023 DeepLang AI. All Rights Reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple
+
+import sentencepiece as spm
+from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {},
+    "tokenizer_file": {},
+}
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {}
+
+
+class LingoWhaleTokenizer(PreTrainedTokenizer):
+    """
+    Construct a LingoWhale tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        unk_token="<!!UNK!!>",
+        bos_token="<!!BOS!!>",
+        eos_token="<!!EOS!!>",
+        pad_token=None,
+        sp_model_kwargs: Optional[Dict[str, Any]] = None,
+        add_bos_token=True,
+        add_eos_token=False,
+        clean_up_tokenization_spaces=False,
+        **kwargs,
+    ):
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+        bos_token = (AddedToken(bos_token, lstrip=False, rstrip=False)
+                     if isinstance(bos_token, str) else bos_token)
+        eos_token = (AddedToken(eos_token, lstrip=False, rstrip=False)
+                     if isinstance(eos_token, str) else eos_token)
+        unk_token = (AddedToken(unk_token, lstrip=False, rstrip=False)
+                     if isinstance(unk_token, str) else unk_token)
+        pad_token = (AddedToken(pad_token, lstrip=False, rstrip=False)
+                     if isinstance(pad_token, str) else pad_token)
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            add_bos_token=add_bos_token,
+            add_eos_token=add_eos_token,
+            sp_model_kwargs=self.sp_model_kwargs,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+        self.vocab_file = vocab_file
+        self.add_bos_token = add_bos_token
+        self.add_eos_token = add_eos_token
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(vocab_file)
+
+    @property
+    def unk_token_length(self):
+        return len(self.sp_model.encode(str(self.unk_token)))
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["sp_model"] = None
+        return state
+
+    def __setstate__(self, d):
+        self.__dict__ = d
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(self.vocab_file)
+
+    @property
+    def vocab_size(self):
+        """Returns vocab size"""
+        return self.sp_model.get_piece_size()
+
+    def get_vocab(self):
+        """Returns vocab as a dict"""
+        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):
+        """Returns a tokenized string."""
+        return self.sp_model.encode(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.sp_model.piece_to_id(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        token = self.sp_model.IdToPiece(index)
+        return token
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        current_sub_tokens = []
+        out_string = ""
+        prev_is_special = False
+        for i, token in enumerate(tokens):
+            # make sure that special tokens are not decoded using sentencepiece model
+            if token in self.all_special_tokens:
+                if not prev_is_special and i != 0:
+                    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 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.
+
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        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 "") +
+            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 build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = bos_token_id + token_ids_0 + eos_token_id
+
+        if token_ids_1 is not None:
+            output = output + bos_token_id + token_ids_1 + eos_token_id
+
+        return output
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0,
+                token_ids_1=token_ids_1,
+                already_has_special_tokens=True,
+            )
+
+        bos_token_id = [1] if self.add_bos_token else []
+        eos_token_id = [1] if self.add_eos_token else []
+
+        if token_ids_1 is None:
+            return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
+        return (bos_token_id + ([0] * len(token_ids_0)) + eos_token_id +
+                bos_token_id + ([0] * len(token_ids_1)) + eos_token_id)
+
+    def create_token_type_ids_from_sequences(
+            self,
+            token_ids_0: List[int],
+            token_ids_1: Optional[List[int]] = None) -> List[int]:
+        """
+        Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
+        sequence pair mask has the following format:
+
+        ```
+        0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+        | first sequence    | second sequence |
+        ```
+
+        if token_ids_1 is None, only returns the first portion of the mask (0s).
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of ids.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
+        """
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)
+
+        if token_ids_1 is not None:
+            output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)
+
+        return output

tokenizer_config.json

@@ -0,0 +1,47 @@
+{
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_lingowhale.LingoWhaleTokenizer",
+      null
+    ]
+  },
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "use_fast": false,
+  "clean_up_tokenization_spaces": false,
+  "model_max_length": 8192,
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LingoWhaleTokenizer",
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<!!BOS!!>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": true
+  },
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<!!EOS!!>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": true
+  },
+  "pad_token": {
+    "__type": "AddedToken",
+    "content": "<!!UNK!!>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": true
+  },
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<!!UNK!!>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": true
+  }
+}