![](\"resources/en/hx_boxplot.png\")
\n\n![](\"resources/en/codegeex_training.png\")
\nLeft: the proportion of programming languages in CodeGeeX's training data. \n Right: the plot of training loss against the training steps of CodeGeeX.
\n\n**Code Corpus**: Our training data contains two parts. The first part is from open-sourced code datasets, [The Pile](https://pile.eleuther.ai/) and [CodeParrot](https://github.com/huggingface/transformers/tree/main/examples/research_projects/codeparrot). The Pile contains a subset of code corpus that collects public repositories with more than 100 stars from GitHub, from which we select codes in 23 popular programming languages. The second part is supplementary data directly scrapped from the public GitHub repositories that do not appear in previous datasets, including Python, Java and C++. To obtain data of potentially higher quality, repositories with at least one star and its size smaller than 10MB are chosen. A file is filtered out if it 1) has more than 100 characters per line on average, 2) is automatically generated, 3) has a ratio of alphabet less than 40%, or 4) is bigger than 100KB or smaller than 1KB. To help the model distinguish different languages, we add a language-specific prefix at the beginning of each segment in the form of ``[Comment sign] language: [LANG]``, e.g., ``# language: Python``. For tokenization, we use the same tokenizer as GPT-2 and process whitespaces as extra tokens, resulting in a vocabulary of 50,400 tokens. In total, the code corpus has 23 programming languages with 158.7B tokens.\n\n**Training**: We implement CodeGeeX in [Mindspore 1.7](https://www.mindspore.cn/) and train it on 1,536 Ascend 910 AI Processor (32GB). The model weights are under FP16 format, except that we use FP32 for layer-norm and softmax for higher precision and stability. The entire model consumes about 27GB of memory. To increase the training efficiency, we adopt an 8-way model parallel training together with 192-way data parallel training, with ZeRO-2 optimizer enabled. The micro-batch size is 16 and the global batch size reaches 3,072. Moreover, we adopt techniques to further boost the training efficiency including the element-wise operator fusion, fast gelu activation, matrix multiplication dimension optimization, etc. The entire training process takes nearly two months, spanning from April 18 to June 22, 2022, during which 850B tokens were passed for training, i.e., 5+ epochs.\n\n## HumanEval-X: A new benchmark for Multilingual Program Synthesis\nTo better evaluate the multilingual ability of code generation models, we propose a new benchmark HumanEval-X. While previous works evaluate multilingual program synthesis under semantic similarity (e.g., [CodeBLEU](https://arxiv.org/abs/2009.10297)) which is often misleading, HumanEval-X evaluates the functional correctness of the generated programs. HumanEval-X consists of 820 high-quality human-crafted data samples (each with test cases) in Python, C++, Java, JavaScript, and Go, and can be used for various tasks.\n\n![](\"resources/en/hx_tasks.png\")
\n\nAn illustration of tasks supported by HumanEval-X. Declarations, docstrings, and solutions are marked with red, green, and blue respectively. Code generation uses declaration and docstring as input, to generate solution. Code translation uses declaration in both languages and translate the solution in source language to the one in target language.
\n\nIn HumanEval-X, every sample in each language contains declaration, docstring, and solution, which can be combined in various ways to support different downstream tasks including generation, translation, summarization, etc. We currently focus on two tasks: **code generation** and **code translation**. For code generation, the model uses declaration and docstring as input to generate the solution. For code translation, the model uses declarations in both languages and the solution in the source language as input, to generate solutions in the target language. We remove the description during code translation to prevent the model from directly solving the problem. For both tasks, we use the unbiased pass@k metric proposed in [Codex](https://arxiv.org/abs/2107.03374): $\\text{pass}@k:= \\mathbb{E}[1-\\frac{\\tbinom{n-c}{k}}{\\tbinom{n}{k}}]$, with $n=200$ and $k\\in(1,10,100)$.\n\n### Multilingual Code Generation\n\n![](\"resources/en/hx_generattion_radar_horizon.png\")
\nLeft: the detailed pass@k (k=1,10,100) performance on code generation task for five languages in HumanEval-X. Right: the average performance of all languages of each model. CodeGeeX achieves the highest average performance compared with InCoder-6.7B, CodeGen-Multi-6B and CodeGen-Multi-16B.
\n\n\nWe compare CodeGeeX with two other open-sourced code generation models, [InCoder](https://github.com/dpfried/incoder) (from Meta) and [CodeGen](https://github.com/salesforce/CodeGen) (from Salesforce). Specifically, InCoder-6.7B, CodeGen-Multi-6B and CodeGen-Multi-16B are considered. CodeGeeX significantly outperforms models with smaller scales (by 7.5%~16.3%) and is competitive with CodeGen-Multi-16B with a larger scale (average performance 54.76% vs. 54.39%). CodeGeeX achieves the best average performance across languages.\n\n### Crosslingual Code Translation\n\n![](\"resources/en/hx_translation.png\"){kind=spacer}
\n\nResults on HumanEval-X code translation task. Best language-wise performance are bolded.
\n\nWe also evaluate the performance of translation across different programming languages. We test the zero-shot performance of CodeGeeX, as well as the fine-tuned CodeGeeX-13B-FT (fine-tuned using the training set of code translation tasks in [XLCoST](https://github.com/reddy-lab-code-research/XLCoST); Go is absent in the original set, we thus add a small set to it). The results indicate that models have a preference for languages, e.g., CodeGeeX is good at translating other languages to Python and C++, while CodeGen-Multi-16B is better at translating to JavaScript and Go; these could probably be due to the difference in language distribution in the training corpus. Among 20 translation pairs, we also observe that the performance of A-to-B and B-to-A are always negatively correlated, which might indicate that the current models are still not capable of learning all languages well. \n\n### How to use HumanEval-X and contribute to it?\n\nFor more details on how to use HumanEval-X, please see [usage](codegeex/benchmark/README.md). We highly welcome the community to contribute to HumanEval-X by adding more problems or extending it to other languages, please check out the [standard format](codegeex/benchmark/README.md#how-to-use-humaneval-x) of HumanEval-X and add a pull request. \n\nPlease kindly let us know if you have any comment or suggestion, via [codegeex@aminer.cn](mailto:codegeex@aminer.cn).\n\n\n![](\"resources/en/hx_examples.png\")
\n
\n\nExamples of Generation
\n\n\n
\nThis project is supported by the National Science Foundation for Distinguished Young Scholars (No. 61825602). \n\n### Lead Contributors\n\nQinkai Zheng ([Tsinghua KEG](http://keg.cs.tsinghua.edu.cn/glm-130b/)), Xiao Xia (Tsinghua KEG), Xu Zou (Tsinghua KEG)\n\n### Contributors\n\nTsinghua KEG---The Knowledge Engineering Group at Tsinghua: Aohan Zeng, Wendi Zheng, Lilong Xue\n\nZhilin Yang's Group at Tsinghua IIIS: Yifeng Liu, Yanru Chen, Yichen Xu (BUPT, work was done when visiting Tsinghua)\n\nPeng Cheng Laboratory: Qingyu Chen, Zhongqi Li, Gaojun Fan\n\nZhipu\\.AI: Yufei Xue, Shan Wang, Jiecai Shan, Haohan Jiang, Lu Liu, Xuan Xue, Peng Zhang\n\nAscend and Mindspore Team: Yifan Yao, Teng Su, Qihui Deng, Bin Zhou\n\n### Data Annotations\n\nRuijie Cheng (Tsinghua), Peinan Yu (Tsinghua), Jingyao Zhang (Zhipu\\.AI), Bowen Huang (Zhipu\\.AI), Shaoyu Wang (Zhipu\\.AI) \n \n### Advisors\n\n[Zhilin Yang](https://kimiyoung.github.io/) (Tsinghua IIIS), Yuxiao Dong (Tsinghua KEG), Wenguang Chen (Tsinghua PACMAN), Jie Tang (Tsinghua KEG)\n \n\n### Computation Sponsors\n\n[Peng Cheng Laboratory](https://www.pcl.ac.cn/index.html)\n\n[Zhipu.AI](https://www.zhipu.ai/)---an AI startup that aims to teach machines to think like humans\n\n### Project Leader \n\n[Jie Tang](http://keg.cs.tsinghua.edu.cn/jietang/) (Tsinghua KEG & BAAI)\n
\n\n## License\n\nOur code is licensed under the [Apache-2.0 license](LICENSE).\nOur model is licensed under the [license](MODEL_LICENSE).\n\n## Citation\n\nIf you find our work useful, please cite:\n\n```\n@inproceedings{zheng2023codegeex,\n title={CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X},\n author={Qinkai Zheng and Xiao Xia and Xu Zou and Yuxiao Dong and Shan Wang and Yufei Xue and Zihan Wang and Lei Shen and Andi Wang and Yang Li and Teng Su and Zhilin Yang and Jie Tang},\n booktitle={Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining},\n pages={5673--5684},\n year={2023}\n}\n```\n"
},
{
"path": "README_zh.md",
"content": "Acknowledgement
\n\nThis project is supported by the National Science Foundation for Distinguished Young Scholars (No. 61825602). \n\n### Lead Contributors\n\nQinkai Zheng ([Tsinghua KEG](http://keg.cs.tsinghua.edu.cn/glm-130b/)), Xiao Xia (Tsinghua KEG), Xu Zou (Tsinghua KEG)\n\n### Contributors\n\nTsinghua KEG---The Knowledge Engineering Group at Tsinghua: Aohan Zeng, Wendi Zheng, Lilong Xue\n\nZhilin Yang's Group at Tsinghua IIIS: Yifeng Liu, Yanru Chen, Yichen Xu (BUPT, work was done when visiting Tsinghua)\n\nPeng Cheng Laboratory: Qingyu Chen, Zhongqi Li, Gaojun Fan\n\nZhipu\\.AI: Yufei Xue, Shan Wang, Jiecai Shan, Haohan Jiang, Lu Liu, Xuan Xue, Peng Zhang\n\nAscend and Mindspore Team: Yifan Yao, Teng Su, Qihui Deng, Bin Zhou\n\n### Data Annotations\n\nRuijie Cheng (Tsinghua), Peinan Yu (Tsinghua), Jingyao Zhang (Zhipu\\.AI), Bowen Huang (Zhipu\\.AI), Shaoyu Wang (Zhipu\\.AI) \n \n### Advisors\n\n[Zhilin Yang](https://kimiyoung.github.io/) (Tsinghua IIIS), Yuxiao Dong (Tsinghua KEG), Wenguang Chen (Tsinghua PACMAN), Jie Tang (Tsinghua KEG)\n \n\n### Computation Sponsors\n\n[Peng Cheng Laboratory](https://www.pcl.ac.cn/index.html)\n\n[Zhipu.AI](https://www.zhipu.ai/)---an AI startup that aims to teach machines to think like humans\n\n### Project Leader \n\n[Jie Tang](http://keg.cs.tsinghua.edu.cn/jietang/) (Tsinghua KEG & BAAI)\n
![](\"resources/logo/codegeex_logo.png\"){kind=logo}
\n\n\n 🏠 主页 | 📖 博客 | 🪧 示例 | 🤖 模型下载 | 📄 论文 | 🌐 English\n
\n\n 🛠 VS Code, Jetbrains, Cloud Studio 插件 | 👋 欢迎加入 微信开发者交流群 \n
\n\n🌟 [CodeGeeX2](https://github.com/THUDM/CodeGeeX2) 已推出，更强，更快，更轻量。\n\n- [CodeGeeX: 多语言代码生成模型](#codegeex-多语言代码生成模型)\n - [新闻](#新闻)\n - [使用指南](#使用指南)\n - [安装](#安装)\n - [模型权重](#模型权重)\n - [用GPU进行推理](#用gpu进行推理)\n - [插件使用指南](#插件使用指南)\n - [CodeGeeX: 多语言代码生成模型](#codegeex-多语言代码生成模型-1)\n - [国产平台实现与训练](#国产平台实现与训练)\n - [HumanEval-X: 多语言代码生成基准](#humaneval-x-多语言代码生成基准)\n - [多语言代码生成](#多语言代码生成)\n - [跨语言代码翻译](#跨语言代码翻译)\n - [许可证](#许可证)\n - [引用](#引用)\n# CodeGeeX: 多语言代码生成模型\n\nCodeGeeX是一个具有130亿参数的多编程语言代码生成预训练模型。CodeGeeX采用华为MindSpore框架实现，在鹏城实验室“鹏城云脑II”中的192个节点（共1536个国产[昇腾910 AI处理器](https://e.huawei.com/cn/products/servers/ascend)）上训练而成。截至2022年6月22日，CodeGeeX历时两个月在20多种编程语言的代码语料库（>8500亿Token）上预训练得到。CodeGeeX有以下特点：\n* **高精度代码生成**：支持生成Python、C++、Java、JavaScript和Go等多种主流编程语言的代码，在HumanEval-X代码生成任务上取得47%~60%求解率，较其他开源基线模型有更佳的平均性能。[代码生成示例](https://models.aminer.cn/codegeex/zh-CN)\n* **跨语言代码翻译**：支持代码片段在不同编程语言间进行自动翻译转换，翻译结果正确率高，在HumanEval-X代码翻译任务上超越了其它基线模型。[代码翻译示例](https://models.aminer.cn/codegeex/zh-CN/codeTranslator)\n* **自动编程插件**：CodeGeeX插件现已上架VSCode插件市场（完全免费），用户可以通过其强大的少样本生成能力，自定义代码生成风格和能力，更好辅助代码编写。[插件下载](https://marketplace.visualstudio.com/items?itemName=aminer.codegeex)\n* **模型跨平台开源**: 所有代码和模型权重开源开放，用作研究用途。CodeGeeX同时支持昇腾和英伟达平台，可在单张昇腾910或英伟达V100/A100上实现推理。[申请模型权重](https://models.aminer.cn/codegeex/download/request)\n\n**全新多编程语言评测基准HumanEval-X**：HumanEval-X是第一个支持功能正确性评测的多语言、多任务的基准，包含820个人工编写的高质量代码生成题目、测试用例与参考答案，覆盖5种编程语言（Python、C++、Java、JavaScript、Go），支持代码生成与代码翻译能力的评测。[如何使用](codegeex/benchmark/README_zh.md)\n\n![](\"resources/zh/hx_boxplot_zh.png\")
\n\n在HumanEval-X代码生成任务上，与其它开源基线模型相比，CodeGeeX取得了最佳的平均性能。
\n\n## 新闻\n\n* 🌟 **2023-07-24**: [CodeGeeX2](https://github.com/THUDM/CodeGeeX2) 已推出，更强，更快，更轻量。支持超过100种语言，并具有多种新特性。\n\n* **2023-5-16**: CodeGeeX 论文已被 [KDD 2023, Long Beach](https://kdd.org/kdd2023/) 接收，并将在会上做相关报告。\n\n* **2023-03-30**: CodeGeeX 论文已发表在[arxiv](https://arxiv.org/abs/2303.17568)。\n\n* **2023-02-14**: CodeGeeX 现已支持 [Cloud Studio](https://cloudstudio.net/), 一款腾讯推出、十分好用的在线编辑器。单击此页面顶部的徽章可快速启动环境测试 CodeGeeX。\n\n* **2023-02-13**: 感谢 [OneFlow](https://github.com/Oneflow-Inc/oneflow) 加入了oneflow版推理支持，在FP16下比FasterTransformer还要快！更多优化细节请点击[这里](https://github.com/THUDM/CodeGeeX/pull/65).\n\n* **2023-02**: [CodeGeeX \"Coding With AI\"黑客松](https://dorahacks.io/hackathon/codegeex/)正在进行中，为CodeGeeX设计应用并赢取奖品（RTX 4090、DJI无人机等）！\n\n* **2022-12-31**: 我们在 [codegeex-fastertransformer](https://github.com/CodeGeeX/codegeex-fastertransformer) 中发布了 CodeGeeX 的 FasterTransformer 版本。INT8加速版本达到 <15ms/token 的平均速度。祝大家新年快乐！\n\n* **2022-12-13**: 我们开源了VS Code插件源码：[codegeex-vscode-extension](https://github.com/CodeGeeX/codegeex-vscode-extension)，参考 [QuickStart](https://github.com/CodeGeeX/codegeex-vscode-extension/blob/main/doc/quickstart.md) 开始开发吧！\n\n* **2022-12-11**: CodeGeeX for Jetbrains IDEs已上线，支持IntelliJ IDEA, PyCharm, GoLand, CLion等，[点击下载](https://plugins.jetbrains.com/plugin/20587-codegeex)。\n \n* **2022-12-04**: 我们开源了量化代码（需要更少的显存：27GB -> 15GB）以及模型并行代码（可以运行在多个显存至少8GB的GPUs上）。\n\n* **2022-09-30**: 我们开源了跨平台代码和模型权重，同时支持昇腾和英伟达平台。\n## 使用指南\n\nCodeGeeX最初使用Mindspore框架实现，并在昇腾910AI芯片上进行训练。为适配更多平台，我们将其转换到[Megatron-LM](https://github.com/NVIDIA/Megatron-LM)框架，支持Pytorch+GPU环境。\n### 安装\n\n需要Python 3.7+ / CUDA 11+ / PyTorch 1.10+ / DeepSpeed 0.6+，通过以下命令安装 ``codegeex``: \n```bash\ngit clone git@github.com:THUDM/CodeGeeX.git\ncd CodeGeeX\npip install -e .\n```\n\n### 模型权重\n\n通过[该链接](https://models.aminer.cn/codegeex/download/request)申请权重，您将收到一个包含临时下载链接文件```urls.txt```的邮件。推荐使用[aria2](https://aria2.github.io/)通过以下命令快速下载（请保证有足够的硬盘空间存放权重（～26GB））：\n```bash\naria2c -x 16 -s 16 -j 4 --continue=true -i urls.txt \n``` \n使用以下命令合并得到完整的权重：\n```bash\ncat codegeex_13b.tar.gz.* > codegeex_13b.tar.gz\ntar xvf codegeex_13b.tar.gz\n```\n\n### 用GPU进行推理\n\n尝试使用CodeGeeX模型生成第一个程序吧！首先，在配置文件``configs/codegeex_13b.sh``中写明存放权重的路径。其次，将提示（可以是任意描述或代码片段）写入文件``tests/test_prompt.txt``，运行以下脚本即可开始推理（需指定GPU序号）：\n```bash\n# On a single GPU (with more than 27GB RAM)\nbash ./scripts/test_inference.sh\n\n左侧：CodeGeeX训练数据中各编程语言占比。\n右侧：CodeGeeX训练损失函数随训练步数下降曲线。
\n\n**语料**：CodeGeeX的训练语料由两部分组成。第一部分是开源代码数据集，[The Pile](https://pile.eleuther.ai/) 与 [CodeParrot](https://github.com/huggingface/transformers/tree/main/examples/research_projects/codeparrot)。The Pile包含GitHub上拥有超过100颗星的一部分开源仓库，我们从中选取了23种编程语言的代码。第二部分是补充数据，直接从GitHub开源仓库中爬取Python、Java、C++代码；为了获取高质量数据，我们根据以下准则选取代码仓库：1)至少拥有1颗星；2)总大小<10MB；3)不在此前的开源代码数据集中。我们还去掉了符合下列任一条件的文件：1)平均每行长度大于100字符；2)由自动生成得到；3)含有的字母不足字母表内的40%；4)大于100KB或小于1KB。为了让模型区分不同语言，我们在每个样本的开头加上一个前缀，其形式为``[注释符] language: [语言]``，例如：``# language: Python``。我们使用与GPT-2相同的分词器，并将空格处理为特殊标识符，词表大小为50400。整个代码语料含有23种编程语言、总计1587亿个标识符（不含填充符）。\n\n### 国产平台实现与训练\n我们在[Mindspore 1.7](https://www.mindspore.cn/)框架上实现了CodeGeeX模型，并使用鹏城实验室的全国产计算平台上进行训练。具体来说，CodeGeeX使用了其一个计算集群中的1536个昇腾910 AI处理器（32GB）进行了两个月左右的训练（2022年4月18日至6月22日）。除了Layer-norm与Softmax使用FP32格式以获得更高的精度与稳定性，模型参数整体使用FP16格式，最终整个模型需要占用约27GB显存。为了增加训练效率，我们使用8路模型并行和192路数据并行的训练策略，微批大小为16、全局批大小为3072，并采用ZeRO-2优化器降低显存占用。\n\n在开发与训练过程中，我们和华为Mindspore团队合作，对MindSpore框架进行了部分优化，进而大幅度提升训练效率。比如，我们发现矩阵乘法的计算时间占比仅为22.9%，大量时间被用于各类其它算子，因此实现了一系列算子融合，包括单元素算子融合、层归一化算子融合、FastGelu与矩阵乘法融合、批量矩阵乘法与加法融合等；再比如我们还对矩阵乘法算子的维度实现自动搜索调优，使其搜索出效率最高的计算维度组合。这些优化为训练速度带来了显著提升，在同等GPU卡数规模下（128卡），昇腾910对CodeGeeX这一模型的训练效率从约为NVIDIA A100的16.7%提升至43%；在千卡规模下，昇腾910训练效率相比自身优化前提升近300%。使用优化后的软硬件训练时，CodeGeeX单日训练量可达到54.3B个标识符（含填充符），证明了国产深度学习平台与工具的快速迭代能力以及强大竞争力。\n\n## HumanEval-X: 多语言代码生成基准\n为了更好地评测代码生成模型的多语言生成能力，我们构建了一个新基准HumanEval-X。此前，多语言代码生成能力是基于语义相似度（比如[CodeBLEU](https://arxiv.org/abs/2009.10297)）衡量的，具有一定误导性；HumanEval-X则可用于衡量生成代码的功能正确性。HumanEval-X包含820个高质量手写样本，覆盖Python、C++、Java、JavaScript、Go，可用于多种任务。\n\n![](\"resources/zh/hx_tasks_zh.png\")
\n\nHumanEval-X支持的任务示例。声明、描述、解答分别用红、绿、蓝色标注。代码生成将声明与描述作为输入，输出解答。代码翻译将两种语言的声明与源语言的解答作为输入，输出目标语言的解答。
\n\nHumanEval-X中每个语言的样本，包含了声明、描述和解答，它们之间的组合可以支持不同的下游任务，包括生成、翻译、概括等。我们目前关注两个任务：**代码生成**与**代码翻译**。对于代码生成任务，模型将函数声明与文档字符串作为输入，输出函数实现；对于代码翻译任务，模型将两种语言的函数声明与源语言的实现作为输入，输出目标语言上的实现。我们在代码翻译任务中不将文档字符串输入模型，以避免模型直接通过描述生成答案。在两种任务下，我们都采用[Codex](https://arxiv.org/abs/2107.03374)所使用的无偏pass@k指标，判断生成代码的功能正确性: $\\text{pass}@k:= \\mathbb{E}[1-\\frac{\\tbinom{n-c}{k}}{\\tbinom{n}{k}}]$, $n=200$, $k\\in(1,10,100)$.\n\n### 多语言代码生成\n\n![](\"resources/zh/hx_generattion_radar_horizon_zh.png\")
\n\n左侧: HumanEval-X中五种语言具体的pass@k（k=1,10,100）性能。右侧: 模型在所有语言上的平均性能。CodeGeeX的平均表现优于InCoder-6.7B和CodeGen-Multi-6B/16B。
\n\n\n我们将CodeGeeX与另外两个开源代码生成模型进行比较，分别为Meta的[InCoder](https://github.com/dpfried/incoder)与Salesforce的[CodeGen](https://github.com/salesforce/CodeGen)，选取InCoder-6.7B、CodeGen-Multi-6B 与 CodeGen-Multi-16B。CodeGeeX能获得最佳的平均性能，显著超越了参数量更小的模型(7.5%~16.3%的提升)，与参数量更大的模型CodeGen-Multi-16B表现相当（平均性能 54.76% vs. 54.39%）。\n\n### 跨语言代码翻译\n\n![](\"resources/zh/hx_translation_zh.png\"){kind=spacer}
\n\nHumanEval-X上的代码翻译任务结果。加粗结果表示在每种语言pass@k上的最佳效果。
\n\n我们还评测了模型在多语言间代码翻译上的性能。对于CodeGeeX，我们评测了未经微调的CodeGeeX-13B与经过微调的CodeGeeX-13B-FT（使用[XLCoST](https://github.com/reddy-lab-code-research/XLCoST)中代码翻译任务的训练集与一部分Go语言数据微调）。如上表显示，模型对特定语言存在偏好，比如CodeGeeX擅长将其他语言翻译为Python与C++，而CodeGen-Multi-16B擅长翻译为JavaScript和Go，这可能是由于训练集中的语料占比存在差异。在20个翻译对中，我们还观察到两种语言互相翻译的表现常常是呈负相关的，这可能说明现有的模型还不足以学好所有的语言。\n\n\n\n\n\n\n
\n这一项目由国家自然科学基金杰出青年科学基金项目（No. 61825602）支持。\n \n\n#### 学生负责人\n\n郑勤锴（[清华大学知识工程实验室](https://keg.cs.tsinghua.edu.cn/glm-130b/zh/posts/glm-130b/)），夏箫（清华大学知识工程实验室），邹旭（清华大学知识工程实验室）\n \n#### 技术贡献\n\n清华大学知识工程实验室：曾奥涵，郑问迪，薛理龙\n\n清华大学交叉信息学院：刘益枫，陈彦儒，徐奕辰（北邮大四访问清华期间研究工作）\n\n鹏城实验室：陈庆玉，李忠琦，范高俊\n\n智谱AI：薛宇飞，王山，陕杰才，姜皓瀚，刘璐，薛旋，张鹏\n\n华为昇腾团队：姚逸璠，苏腾，邓启辉，周斌\n\n#### 数据标注\n程锐杰（清华大学），于沛楠（清华大学），张竞尧（智谱AI），黄铂文（智谱AI），王炤宇（智谱AI）\n \n#### 指导教师\n\n[杨植麟](https://kimiyoung.github.io/)（清华大学交叉信息学院），东昱晓（清华大学知识工程实验室），陈文光（清华大学PACMAN实验室），[唐杰](http://keg.cs.tsinghua.edu.cn/jietang/)（清华大学知识工程实验室）\n \n\n#### 计算资源支持\n\n[鹏城实验室](https://www.pcl.ac.cn/index.html)\n\n[智谱AI](https://www.zhipu.ai/)\n\n#### 项目总负责\n\n[唐杰](http://keg.cs.tsinghua.edu.cn/jietang/)（清华大学知识工程实验室 & 北京智源人工智能研究院）\n
\n\n如果遇到问题或有任何建议，欢迎通过邮件与我们联系[codegeex@aminer.cn](mailto:codegeex@aminer.cn).\n\n## 许可证\n\n代码使用[Apache-2.0许可证](LICENSE)\n模型使用[许可证](MODEL_LICENSE)\n\n## 引用\n\n如果觉得我们的工作有帮助，欢迎引用以下论文：\n\n```\n@inproceedings{zheng2023codegeex,\n title={CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X},\n author={Qinkai Zheng and Xiao Xia and Xu Zou and Yuxiao Dong and Shan Wang and Yufei Xue and Zihan Wang and Lei Shen and Andi Wang and Yang Li and Teng Su and Zhilin Yang and Jie Tang},\n booktitle={Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining},\n pages={5673--5684},\n year={2023}\n}\n```"
},
{
"path": "api/README_zh.md",
"content": "\n\n# 创建CodeGeeX API\n\n使用[天启 · API开放平台](https://tianqi.aminer.cn/open/)申请CodeGeeX API：\n\n致谢
\n\n这一项目由国家自然科学基金杰出青年科学基金项目（No. 61825602）支持。\n \n\n#### 学生负责人\n\n郑勤锴（[清华大学知识工程实验室](https://keg.cs.tsinghua.edu.cn/glm-130b/zh/posts/glm-130b/)），夏箫（清华大学知识工程实验室），邹旭（清华大学知识工程实验室）\n \n#### 技术贡献\n\n清华大学知识工程实验室：曾奥涵，郑问迪，薛理龙\n\n清华大学交叉信息学院：刘益枫，陈彦儒，徐奕辰（北邮大四访问清华期间研究工作）\n\n鹏城实验室：陈庆玉，李忠琦，范高俊\n\n智谱AI：薛宇飞，王山，陕杰才，姜皓瀚，刘璐，薛旋，张鹏\n\n华为昇腾团队：姚逸璠，苏腾，邓启辉，周斌\n\n#### 数据标注\n程锐杰（清华大学），于沛楠（清华大学），张竞尧（智谱AI），黄铂文（智谱AI），王炤宇（智谱AI）\n \n#### 指导教师\n\n[杨植麟](https://kimiyoung.github.io/)（清华大学交叉信息学院），东昱晓（清华大学知识工程实验室），陈文光（清华大学PACMAN实验室），[唐杰](http://keg.cs.tsinghua.edu.cn/jietang/)（清华大学知识工程实验室）\n \n\n#### 计算资源支持\n\n[鹏城实验室](https://www.pcl.ac.cn/index.html)\n\n[智谱AI](https://www.zhipu.ai/)\n\n#### 项目总负责\n\n[唐杰](http://keg.cs.tsinghua.edu.cn/jietang/)（清华大学知识工程实验室 & 北京智源人工智能研究院）\n
![](\"../resources/api/api_step_1.png\")
\n\n点击首页中的天启平台体验入口：\n![](\"../resources/api/api_step_2.png\")
\n点击API应用：\n![](\"../resources/api/api_step_3.png\")
\n输入任意名称，创建API应用。创建后会得到API Key/Secret，用于调用API：\n![](\"../resources/api/api_step_4.png\")
\n\n在API信息中，可以查看代码生成/代码翻译的请求地址和使用文档：\n![](\"../resources/api/api_step_5.png\")
\n\n根据文档中的描述使用API，Python版参考目录``api/codegeex-api-example-python``；JAVA版参考工程：``api/codegeex-api-example-java``\n"
},
{
"path": "api/codegeex-api-example-java/pom.xml",
"content": "\n![](\"../../resources/en/hx_tasks.png\")
\n\nAn illustration of tasks supported by HumanEval-X. Declarations, docstrings, and solutions are marked with red, green, and blue respectively. Code generation uses declaration and docstring as input, to generate solution. Code translation uses declaration in both languages and translate the solution in source language to the one in target language.
\n\nIn HumanEval-X, every sample in each language contains declaration, docstring, and solution, which can be combined in various ways to support different downstream tasks including generation, translation, summarization, etc. We currently focus on two tasks: **code generation** and **code translation**. For code generation, the model uses declaration and docstring as input to generate the solution. For code translation, the model uses declarations in both languages and the solution in the source language as input, to generate solutions in the target language. We remove the description during code translation to prevent the model from directly solving the problem. For both tasks, we use the unbiased pass@k metric proposed in [Codex](https://arxiv.org/abs/2107.03374): $\\text{pass}@k:= \\mathbb{E}[1-\\frac{\\tbinom{n-c}{k}}{\\tbinom{n}{k}}]$, with $n=200$ and $k\\in(1,10,100)$.\n\n## How to use HumanEval-X\n\nData are stored in ``codegeex/benchmark/humaneval-x/[LANG]/data/humaneval_[LANG].jsonl.gz``, using JSON list format. There are six keys:\n\n* ``task_id``: indicates the target language and ID of the problem. Language is one of [\"Python\", \"Java\", \"JavaScript\", \"CPP\", \"Go\"].\n* ``prompt``: the function declaration and docstring, used for code generation.\n* ``declaration``: only the function declaration, used for code translation. \n* ``canonical_solution``: human-crafted example solutions.\n* ``test``: hidden test samples, used for evaluation.\n* ``example_test``: public test samples (appeared in prompt), used for evaluation. \n\n### Evaluation Environment\n\nThe evaluation of the generated codes involves compiling and running in multiple programming languages. The versions of the programming language environments and packages we use are as follows:\n\n| Dependency | Version |\n| ---------- | -------- |\n| Python | 3.8.12 |\n| JDK | 18.0.2.1 |\n| Node.js | 16.14.0 |\n| js-md5 | 0.7.3 |\n| C++ | 11 |\n| g++ | 7.5.0 |\n| Boost | 1.71.0 |\n| OpenSSL | 3.0.0 |\n| go | 1.18.4 |\n\nIn order to save everyone the trouble of setting up the environments for these languages, we build a Docker image with the required environments and CodeGeeX installed.\n\nYou can directly pull the image from Docker Hub:\n\n```bash\ndocker pull rishubi/codegeex:latest\n```\n\nAlternatively, if you are familiar with Dockerfile, you can build the image from `codegeex/docker/Dockerfile` or configure the Dockerfile as you like it:\n\n```bash\ncd codegeex/docker\ndocker build [OPTIONS] .\n```\n\nAfter obtaining the image, you can build a container using the following command:\n\n```bash\ndocker run -it --gpus all --mount type=bind,source=\n\nHumanEval-X支持的任务示例。声明、描述、解答分别用红、绿、蓝色标注。代码生成将声明与描述作为输入，输出解答。代码翻译将两种语言的声明与源语言的解答作为输入，输出目标语言的解答。
\n\nHumanEval-X中每个语言的样本，包含了声明、描述和解答，它们之间的组合可以支持不同的下游任务，包括生成、翻译、概括等。我们目前关注两个任务：**代码生成**与**代码翻译**。对于代码生成任务，模型将函数声明与文档字符串作为输入，输出函数实现；对于代码翻译任务，模型将两种语言的函数声明与源语言的实现作为输入，输出目标语言上的实现。我们在代码翻译任务中不将文档字符串输入模型，以避免模型直接通过描述生成答案。在两种任务下，我们都采用[Codex](https://arxiv.org/abs/2107.03374)所使用的无偏pass@k指标：$\\text{pass}@k:= \\mathbb{E}[1-\\frac{\\tbinom{n-c}{k}}{\\tbinom{n}{k}}]$, $n=200$, $k\\in(1,10,100)$。\n\n## 如何使用HumanEval-X\n\n样本使用JSON列表格式存储在``codegeex/benchmark/humaneval-x/[LANG]/data/humaneval_[LANG].jsonl.gz``，每条样本包含6个部分：\n\n* ``task_id``: 题目的目标语言与ID。语言为[\"Python\", \"Java\", \"JavaScript\", \"CPP\", \"Go\"]中之一。\n* ``prompt``: 函数声明与描述，用于代码生成。\n* ``declaration``: 仅有函数声明，用于代码翻译。\n* ``canonical_solution``: 手写的示例解答。\n* ``test``: 隐藏测例，用于评测。\n* ``example_test``: 提示中出现的公开测例，用于评测。\n\n### 评测环境\n\n评测生成的代码需要使用多种语言编译、运行。我们使用的各编程语言依赖及所用包的版本如下：\n\n| 依赖 | 版本 |\n| ------- | -------- |\n| Python | 3.8.12 |\n| JDK | 18.0.2.1 |\n| Node.js | 16.14.0 |\n| js-md5 | 0.7.3 |\n| C++ | 11 |\n| g++ | 7.5.0 |\n| Boost | 1.71.0 |\n| OpenSSL | 3.0.0 |\n| go | 1.18.4 |\n\n为了省去使用者配置这些语言环境的麻烦，我们构建了一个Docker镜像，并在其中配置了所需要的环境。\n\n可以直接从Docker Hub拉取镜像：\n\n```bash\ndocker pull rishubi/codegeex:latest\n```\n\n如果您熟悉Dockerfile，也可以从`codegeex/docker/Dockerfile`构建镜像，或者修改之以定制自己的配置：\n\n```bash\ncd codegeex/docker\ndocker build [OPTIONS] .\n```\n\n获取镜像后，使用如下命令创建容器：\n\n```bash\ndocker run -it --gpus all --mount type=bind,source=\", \"langauge\": \"HTML\"}\n{\"code\": \"// Write a function that returns the sum of the numbers from 1 to n.\\n// For example, if n is 5, then the function should return 1 + 2 + 3 + 4 + 5.\\n// If n is 0, then the function should return 0.\\n// If n is less than 0, then the function should return -1.\\n/**\\n * @param {number} n\\n * @return {number}\\n */\\nfunction sum ($n) {\", \"langauge\": \"PHP\"}\n{\"code\": \"// Write a function that returns the sum of the numbers from 1 to n.\\n// For example, if n is 5, then the function should return 1 + 2 + 3 + 4 + 5.\\n\\nfunction sum(n) {\", \"langauge\": \"JavaScript\"}\n{\"code\": \"// Write a function that returns the sum of the numbers from 1 to n,\\n// but using a for loop instead of a while loop.\\n\\nfunction sumForLoop(n) {\", \"langauge\": \"TypeScript\"}\n{\"code\": \"// Write a function that returns the sum of the numbers from 1 to n,\\n// but using a for loop instead of a while loop.\\n\\nfunc sumN(n int) int {\", \"langauge\": \"Go\"}\n{\"code\": \"// Write a function that returns the sum of the numbers from 1 to n,\\n// but using a for loop instead of a while loop.\\n\\nfn sum_numbers(n: usize) -> usize {\", \"langauge\": \"Rust\"}\n{\"code\": \"-- Search all the records from the table CodeGeeX\\n-- Delete iterms with odd indices\", \"langauge\": \"SQL\"}\n{\"code\": \"// Write a function that returns the sum of the numbers from 1 to n.\\n// For example, if n is 5, then the function should return 1 + 2 + 3 + 4 + 5.\\n\\nfun sum(n: Int): Int {\", \"langauge\": \"Kotlin\"}\n{\"code\": \"! Write a function that returns the sum of the numbers from 1 to n.\\n! For example, if n is 5, then the function should return 1 + 2 + 3 + 4 + 5.\\n\\n! Use the following header:\\n! module sum_numbers\\n! end\\nmodule sum_numbers\", \"langauge\": \"Fortran\"}\n{\"code\": \"# Write a function that returns the sum of the numbers from 1 to n.\\n# For example, if n is 5, then the function should return 1 + 2 + 3 + 4 + 5.\\nsum_numbers <- function(n) {\", \"langauge\": \"R\"}\n"
},
{
"path": "deployment/server_gradio.py",
"content": "import json\nimport torch\nimport argparse\nimport gradio as gr\n\nimport codegeex\nfrom codegeex.torch import CodeGeeXModel\nfrom codegeex.tokenizer import CodeGeeXTokenizer\nfrom codegeex.quantization import quantize\nfrom codegeex.data.data_utils import LANGUAGE_TAG\nfrom codegeex.megatron.inference import set_random_seed\n\n\ndef model_provider(args):\n \"\"\"Build the model.\"\"\"\n\n model = CodeGeeXModel(\n args.hidden_size,\n args.num_layers,\n args.num_attention_heads,\n args.padded_vocab_size,\n args.max_position_embeddings\n )\n return model\n\n\ndef add_code_generation_args(parser):\n group = parser.add_argument_group(title=\"code generation\")\n group.add_argument(\n \"--num-layers\",\n type=int,\n default=39,\n )\n group.add_argument(\n \"--hidden-size\",\n type=int,\n default=5120,\n )\n group.add_argument(\n \"--num-attention-heads\",\n type=int,\n default=40,\n )\n group.add_argument(\n \"--padded-vocab-size\",\n type=int,\n default=52224,\n )\n group.add_argument(\n \"--max-position-embeddings\",\n type=int,\n default=2048,\n )\n group.add_argument(\n \"--tokenizer-path\",\n type=str,\n default=\"./tokenizer\",\n )\n group.add_argument(\n \"--example-path\",\n type=str,\n default=\"./\",\n )\n group.add_argument(\n \"--load\",\n type=str,\n )\n group.add_argument(\n \"--state-dict-path\",\n type=str,\n )\n group.add_argument(\n \"--micro-batch-size\",\n type=int,\n default=1,\n )\n group.add_argument(\n \"--quantize\",\n action=\"store_true\",\n )\n \n return parser\n\n\ndef main():\n parser = argparse.ArgumentParser()\n parser = add_code_generation_args(parser)\n args, _ = parser.parse_known_args()\n \n print(\"Loading tokenizer ...\")\n tokenizer = CodeGeeXTokenizer(\n tokenizer_path=args.tokenizer_path, \n mode=\"codegeex-13b\")\n\n print(\"Loading state dict ...\")\n state_dict = torch.load(args.load, map_location=\"cpu\")\n state_dict = state_dict[\"module\"]\n\n print(\"Building CodeGeeX model ...\")\n model = model_provider(args)\n model.load_state_dict(state_dict)\n model.eval()\n model.half()\n if args.quantize:\n model = quantize(model, weight_bit_width=8, backend=\"torch\")\n model.cuda()\n\n def predict(\n prompt, \n lang,\n seed, \n out_seq_length, \n temperature, \n top_k, \n top_p,\n ):\n set_random_seed(seed)\n if lang.lower() in LANGUAGE_TAG:\n prompt = LANGUAGE_TAG[lang.lower()] + \"\\n\" + prompt\n \n generated_code = codegeex.generate(\n model,\n tokenizer,\n prompt,\n out_seq_length=out_seq_length,\n seq_length=args.max_position_embeddings,\n top_k=top_k,\n top_p=top_p,\n temperature=temperature,\n micro_batch_size=args.micro_batch_size,\n backend=\"megatron\",\n verbose=True,\n )\n return prompt + generated_code\n\n examples = []\n with open(args.example_path, \"r\") as f:\n for line in f:\n examples.append(list(json.loads(line).values()))\n\n with gr.Blocks() as demo:\n gr.Markdown(\n \"\"\"\n ![](\"https://raw.githubusercontent.com/THUDM/CodeGeeX/main/resources/logo/codegeex_logo.png\")
\n \"\"\")\n gr.Markdown(\n \"\"\"\n \"\nMASTER_IP=$(cat $HOSTFILE | head -n 1)\ncat $HOSTFILE | awk '{print $1 \" slots=8\"}' > $SCRIPT_DIR/hostfile\necho \"MASTER_IP=$MASTER_IP\"\n\n# ====== Parameters ======\nDATA_PATH=\"\"\nCKPT_PATH=\"\"\nDS_CONFIG=ds_config.json\n# - 13b\nTP=1\nPP=1\nNLAYERS=39\nHIDDEN=5120\nNATTN_HEAD=40\nEMBED_VOCAB=52224\nGLOBAL_BATCH=560\nMICRO_BATCH=10\nNTRAIN_ITERS=100000\nEVAL_INT=10\nSAVE_INT=10\nTRIAL_TAG=\"13b-test\"\n# - trial\nTRIAL_NAME=\"pretrain-codegeex\"\n# - zero stage\nZERO_STAGE=2\n# - logging & output\nNOW=$(date +\"%Y%m%d_%H%M%S\")\nOUTPUT_DIR=\"-$TRIAL_NAME-$TRIAL_TAG\"\nTB_DIR=$OUTPUT_DIR/tb$NOW\nmkdir -p $OUTPUT_DIR\nmkdir -p $TB_DIR\n\n# Deepspeed config\ncat < $DS_CONFIG\n{\n \"train_batch_size\" : $GLOBAL_BATCH,\n \"train_micro_batch_size_per_gpu\": $MICRO_BATCH,\n \"steps_per_print\": 5,\n \"zero_optimization\": {\n \"stage\": $ZERO_STAGE,\n \"reduce_bucket_size\": 50000000,\n \"allgather_bucket_size\": 50000000,\n \"overlap_comm\": true,\n \"contiguous_gradients\": false\n },\n \"fp16\": {\n \"enabled\": true,\n \"loss_scale\": 0,\n \"loss_scale_window\": 500,\n \"hysteresis\": 2,\n \"min_loss_scale\": 1,\n \"initial_scale_power\": 12\n },\n \"wall_clock_breakdown\" : true\n}\nEOT\n\nds_args=\"\"\nds_args=\" --deepspeed ${ds_args}\"\nds_args=\" --no-pipeline-parallel ${ds_args}\"\nds_args=\" --deepspeed_config=$DS_CONFIG ${ds_args}\"\nds_args=\" --zero-stage=$ZERO_STAGE ${ds_args}\"\nds_args=\" --deepspeed-activation-checkpointing ${ds_args}\"\n\necho \"Launching deepspeed\"\ndeepspeed \\\n --hostfile hostfile \\\n --master_addr $MASTER_IP \\\n $MAIN_DIR/codegeex/megatron/tools/pretrain_codegeex.py \\\n --tensor-model-parallel-size $TP \\\n --pipeline-model-parallel-size $PP \\\n --no-pipeline-parallel \\\n --num-layers $NLAYERS \\\n --hidden-size $HIDDEN \\\n --make-vocab-size-divisible-by $EMBED_VOCAB \\\n --num-attention-heads $NATTN_HEAD \\\n --seq-length 512 \\\n --loss-scale 12 \\\n --max-position-embeddings 2048 \\\n --micro-batch-size $MICRO_BATCH \\\n --global-batch-size $GLOBAL_BATCH \\\n --train-iters $NTRAIN_ITERS \\\n --lr 1e-6 \\\n --min-lr 1e-7 \\\n --lr-decay-iters 100000 \\\n --lr-decay-style cosine \\\n --lr-warmup-iters 1000 \\\n --log-interval 1 \\\n --eval-iters 10 \\\n --eval-interval $EVAL_INT \\\n --data-path $DATA_PATH \\\n --vocab-file $MAIN_DIR/codegeex/tokenizer/vocab.json \\\n --merge-file $MAIN_DIR/codegeex/tokenizer/merges.txt \\\n --save-interval $SAVE_INT \\\n --save $OUTPUT_DIR \\\n --load $OUTPUT_DIR \\\n --load-state $CKPT_PATH \\\n --split 98,2,0 \\\n --clip-grad 1.0 \\\n --weight-decay 0.1 \\\n --adam-beta1 0.9 \\\n --adam-beta2 0.95 \\\n --fp16 \\\n --ln-fp16 \\\n --attention-softmax-in-fp32 \\\n --checkpoint-activations \\\n --override-lr-scheduler \\\n --tensorboard-dir $TB_DIR \\\n $ds_args |& tee ${OUTPUT_DIR}/$NOW.log"
},
{
"path": "scripts/gather_output.sh",
"content": "# This script is used to gather the distributed outputs of different ranks.\n\nOUTPUT_DIR=$1\nOUTPUT_PREFIX=$2\nIF_REMOVE_RANK_FILES=$3\n\necho \"$OUTPUT_DIR\"\necho \"$OUTPUT_PREFIX\"\n\nif [ -z \"$IF_REMOVE_RANK_FILES\" ]\nthen\n IF_REMOVE_RANK_FILES=0\nfi\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\n\n\nCMD=\"python $MAIN_DIR/codegeex/benchmark/gather_output.py \\\n --output_dir $OUTPUT_DIR \\\n --output_prefix $OUTPUT_PREFIX \\\n --if_remove_rank_files $IF_REMOVE_RANK_FILES\"\n\necho \"$CMD\"\neval \"$CMD\""
},
{
"path": "scripts/generate_humaneval_x.sh",
"content": "# This script is used to generate solutions of HumanEval-X.\n\nLANGUAGE=$1 # Target programming language, currently support one of [\"python\", \"java\", \"cpp\", \"js\", \"go\"]\nOUTPUT_PATH=$2 # Output path of the generated programs.\nHOSTLIST=$3 # Provide hostfile if generating distributedly\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\n# export CUDA settings\nexport CUDA_HOME=/usr/local/cuda-11.1/\n\n# import model configuration\nsource \"$MAIN_DIR/configs/codegeex_13b.sh\"\n\n# nccl options\nOPTIONS_NCCL=\"export NCCL_DEBUG=warn; export NCCL_IB_DISABLE=0; export NCCL_IB_GID_INDEX=3\"\nOPTIONS_PATH=\"export PATH=$PATH; export LD_LIBRARY_PATH=$LD_LIBRARY_PATH\"\nCWD=$(pwd)\n\n# set master ip for zmq server\nif [ -z \"$HOSTLIST\" ]; then\n ZMQ_ADDR=$(hostname -i)\n echo \"$ZMQ_ADDR\" > \"./hostfile\"\n HOSTLIST=\"./hostfile\"\nelse\n ZMQ_ADDR=$(cat $HOSTLIST | head -n 1)\nfi\necho \"master_ip: $ZMQ_ADDR\"\n\nNUM_SAMPLES=1\nMICRO_BSZ=1\nWORLD_SIZE=1\nTEMP=0.8\nTOPP=0.95\nSEED=42\nDATASET=humaneval\nTODAY=$(date +%y%m%d)\nCHANNEL_PORT=$(expr $RANDOM + 5000)\nMASTER_PORT=$(expr $RANDOM + 8000)\n\n# save log file\nLOG_DIR=$MAIN_DIR/log\nmkdir -p \"$LOG_DIR\"\nLOG_PATH=\"$LOG_DIR/$TODAY-generation.log\"\n\nif [ -z \"$LANGUAGE\" ]; then\n LANGUAGE=python\nfi\n\nif [ -z \"$INPUT_PATH\" ]; then\n INPUT_PATH=$MAIN_DIR/codegeex/benchmark/humaneval-x/$LANGUAGE/data/humaneval_$LANGUAGE.jsonl.gz\nfi\n\nif [ -z \"$OUTPUT_PATH\" ]; then\n OUTPUT_PATH=$MAIN_DIR/codegeex/benchmark/output/humaneval-x/codegeex/\n mkdir -p \"$OUTPUT_PATH\"\nfi\n\nJOB_ID=codegeex-ns$NUM_SAMPLES-t$TEMP-topp$TOPP-seed$SEED-$LANGUAGE\n\nRUN_CMD=\"python \\\n $MAIN_DIR/codegeex/benchmark/humaneval-x/generate_humaneval_x.py \\\n --hostfile $HOSTLIST \\\n --channel-ip $ZMQ_ADDR \\\n --channel-port $CHANNEL_PORT \\\n --master-port $MASTER_PORT \\\n --tokenizer-path $TOKENIZER_PATH \\\n --load-deepspeed \\\n --temperature $TEMP \\\n --top-p $TOPP \\\n --out-seq-length 1024 \\\n --micro-batch-size $MICRO_BSZ \\\n --samples-per-problem $NUM_SAMPLES \\\n --language-type $LANGUAGE \\\n --dataset $DATASET \\\n --input-path $INPUT_PATH \\\n --output-prefix $OUTPUT_PATH/$JOB_ID \\\n --gen-node-world-size $WORLD_SIZE \\\n --seed $SEED \\\n $MODEL_ARGS\"\n\nRUN_CMD=\"$OPTIONS_NCCL; $OPTIONS_PATH; $RUN_CMD\"\nRUN_CMD=\"cd $CWD; $RUN_CMD\"\n\nif (( WORLD_SIZE != 1 )); then\n RUN_CMD=\"pdsh -R ssh -w ^$HOSTLIST \\\"$RUN_CMD\\\"\"\nfi\n\necho \"$RUN_CMD\"\necho \"Writing log to $LOG_PATH\"\neval \"$RUN_CMD\" > \"$LOG_PATH\"\nbash $MAIN_DIR/scripts/gather_output.sh $OUTPUT_PATH $JOB_ID 1\n"
},
{
"path": "scripts/pretrain_codegeex.sh",
"content": "SCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\n\n# ====== Environment ======\n# - NCCL & IB\nexport NCCL_DEBUG=info\nexport NCCL_IB_DISABLE=0\nexport NCCL_IB_GID_INDEX=3\n\nHOSTFILE=\"\"\nMASTER_IP=$(cat $HOSTFILE | head -n 1)\ncat $HOSTFILE | awk '{print $1 \" slots=8\"}' > $SCRIPT_DIR/hostfile\necho \"MASTER_IP=$MASTER_IP\"\n\n# ====== Parameters ======\nDATA_PATH=\"\"\nCKPT_PATH=\"\"\nDS_CONFIG=ds_config.json\n# - 13b\nTP=1\nPP=1\nNLAYERS=39\nHIDDEN=5120\nNATTN_HEAD=40\nEMBED_VOCAB=52224\nGLOBAL_BATCH=560\nMICRO_BATCH=10\nNTRAIN_ITERS=100000\nEVAL_INT=10\nSAVE_INT=10\nTRIAL_TAG=\"13b-test\"\n# - trial\nTRIAL_NAME=\"pretrain-codegeex\"\n# - zero stage\nZERO_STAGE=2\n# - logging & output\nNOW=$(date +\"%Y%m%d_%H%M%S\")\nOUTPUT_DIR=\"-$TRIAL_NAME-$TRIAL_TAG\"\nTB_DIR=$OUTPUT_DIR/tb$NOW\nmkdir -p $OUTPUT_DIR\nmkdir -p $TB_DIR\n\n# Deepspeed config\ncat < $DS_CONFIG\n{\n \"train_batch_size\" : $GLOBAL_BATCH,\n \"train_micro_batch_size_per_gpu\": $MICRO_BATCH,\n \"steps_per_print\": 5,\n \"zero_optimization\": {\n \"stage\": $ZERO_STAGE,\n \"reduce_bucket_size\": 50000000,\n \"allgather_bucket_size\": 50000000,\n \"overlap_comm\": true,\n \"contiguous_gradients\": false\n },\n \"fp16\": {\n \"enabled\": true,\n \"loss_scale\": 0,\n \"loss_scale_window\": 500,\n \"hysteresis\": 2,\n \"min_loss_scale\": 1,\n \"initial_scale_power\": 12\n },\n \"wall_clock_breakdown\" : true\n}\nEOT\n\nds_args=\"\"\nds_args=\" --deepspeed ${ds_args}\"\nds_args=\" --no-pipeline-parallel ${ds_args}\"\nds_args=\" --deepspeed_config=$DS_CONFIG ${ds_args}\"\nds_args=\" --zero-stage=$ZERO_STAGE ${ds_args}\"\nds_args=\" --deepspeed-activation-checkpointing ${ds_args}\"\n\necho \"Launching deepspeed\"\ndeepspeed \\\n --hostfile hostfile \\\n --master_addr $MASTER_IP \\\n $MAIN_DIR/codegeex/megatron/tools/pretrain_codegeex.py \\\n --tensor-model-parallel-size $TP \\\n --pipeline-model-parallel-size $PP \\\n --no-pipeline-parallel \\\n --num-layers $NLAYERS \\\n --hidden-size $HIDDEN \\\n --make-vocab-size-divisible-by $EMBED_VOCAB \\\n --num-attention-heads $NATTN_HEAD \\\n --seq-length 512 \\\n --loss-scale 12 \\\n --max-position-embeddings 2048 \\\n --micro-batch-size $MICRO_BATCH \\\n --global-batch-size $GLOBAL_BATCH \\\n --train-iters $NTRAIN_ITERS \\\n --lr 2e-4 \\\n --min-lr 1e-7 \\\n --lr-decay-iters 100000 \\\n --lr-decay-style cosine \\\n --lr-warmup-iters 1500 \\\n --log-interval 1 \\\n --eval-iters 10 \\\n --eval-interval $EVAL_INT \\\n --data-path $DATA_PATH \\\n --vocab-file $MAIN_DIR/codegeex/tokenizer/vocab.json \\\n --merge-file $MAIN_DIR/codegeex/tokenizer/merges.txt \\\n --save-interval $SAVE_INT \\\n --save $OUTPUT_DIR \\\n --load $OUTPUT_DIR \\\n --load-state $CKPT_PATH \\\n --split 98,2,0 \\\n --clip-grad 1.0 \\\n --weight-decay 0.1 \\\n --adam-beta1 0.9 \\\n --adam-beta2 0.95 \\\n --fp16 \\\n --ln-fp16 \\\n --attention-softmax-in-fp32 \\\n --checkpoint-activations \\\n --override-lr-scheduler \\\n --tensorboard-dir $TB_DIR \\\n $ds_args |& tee ${OUTPUT_DIR}/$NOW.log\n\n"
},
{
"path": "scripts/process_pretrain_dataset.sh",
"content": "# Process dataset for CodeGeeX pretraining\n\nDATASET_PATH=$1\nOUTPUT_PATH=$2\nLANGUAGE=$3\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\nif [ -z \"$LANGUAGE\" ]; then\n LANGUAGE=python\nfi\n\nCMD=\"python $MAIN_DIR/codegeex/data/process_pretrain_dataset.py \\\n --dataset_path $DATASET_PATH \\\n --tokenizer_path $TOKENIZER_PATH \\\n --output_prefix $OUTPUT_PATH \\\n --language $LANGUAGE \\\n --mode pretrain \\\n --seq_len 2048\"\n\necho \"$CMD\"\neval \"$CMD\""
},
{
"path": "scripts/test_inference.sh",
"content": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\n# import model configuration\nsource \"$MAIN_DIR/configs/codegeex_13b.sh\"\n\n# export CUDA settings\nif [ -z \"$GPU\" ]; then\n GPU=0\nfi\n\nexport CUDA_HOME=/usr/local/cuda-11.1/\nexport CUDA_VISIBLE_DEVICES=$GPU\n\nif [ -z \"$PROMPT_FILE\" ]; then\n PROMPT_FILE=$MAIN_DIR/tests/test_prompt.txt\nfi\n\n# remove --greedy if using sampling\nCMD=\"python $MAIN_DIR/tests/test_inference.py \\\n --prompt-file $PROMPT_FILE \\\n --tokenizer-path $TOKENIZER_PATH \\\n --micro-batch-size 1 \\\n --out-seq-length 1024 \\\n --temperature 0.8 \\\n --top-p 0.95 \\\n --top-k 0 \\\n --greedy \\\n $MODEL_ARGS\"\n\necho \"$CMD\"\neval \"$CMD\"\n"
},
{
"path": "scripts/test_inference_oneflow.sh",
"content": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\n# import model configuration\nsource \"$MAIN_DIR/configs/codegeex_13b.sh\"\n\n# export CUDA settings\nif [ -z \"$GPU\" ]; then\n GPU=0\nfi\n\nexport CUDA_HOME=/usr/local/cuda-11.1/\nexport CUDA_VISIBLE_DEVICES=$GPU\n\nif [ -z \"$PROMPT_FILE\" ]; then\n PROMPT_FILE=$MAIN_DIR/tests/test_prompt.txt\nfi\n\n# remove --greedy if using sampling\nCMD=\"python $MAIN_DIR/tests/test_inference_oneflow.py \\\n --prompt-file $PROMPT_FILE \\\n --tokenizer-path $TOKENIZER_PATH \\\n --micro-batch-size 1 \\\n --out-seq-length 1024 \\\n --temperature 0.8 \\\n --top-p 0.95 \\\n --top-k 0 \\\n --greedy \\\n $MODEL_ARGS\"\n\necho \"$CMD\"\neval \"$CMD\"\n"
},
{
"path": "scripts/test_inference_oneflow_quantized.sh",
"content": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\n# import model configuration\nsource \"$MAIN_DIR/configs/codegeex_13b.sh\"\n\n# export CUDA settings\nif [ -z \"$GPU\" ]; then\n GPU=1\nfi\n\nexport CUDA_HOME=/usr/local/cuda-11.1/\nexport CUDA_VISIBLE_DEVICES=$GPU\n\nif [ -z \"$PROMPT_FILE\" ]; then\n PROMPT_FILE=$MAIN_DIR/tests/test_prompt.txt\nfi\n\n# remove --greedy if using sampling\nCMD=\"python $MAIN_DIR/tests/test_inference_oneflow.py \\\n --prompt-file $PROMPT_FILE \\\n --tokenizer-path $TOKENIZER_PATH \\\n --micro-batch-size 1 \\\n --out-seq-length 1024 \\\n --temperature 0.2 \\\n --top-p 0.95 \\\n --top-k 0 \\\n --quantize \\\n $MODEL_ARGS\"\n\necho \"$CMD\"\neval \"$CMD\"\n"
},
{
"path": "scripts/test_inference_paddle.sh",
"content": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\n# import model configuration\nsource \"$MAIN_DIR/configs/codegeex_13b_paddle.sh\"\n\n# export CUDA settings\nif [ -z \"$GPU\" ]; then\n GPU=0\nfi\n\nexport CUDA_HOME=/usr/local/cuda-11.1/\nexport CUDA_VISIBLE_DEVICES=$GPU\n\nif [ -z \"$PROMPT_FILE\" ]; then\n PROMPT_FILE=$MAIN_DIR/tests/test_prompt.txt\nfi\n\n# remove --greedy if using sampling\nCMD=\"python $MAIN_DIR/tests/test_inference_paddle.py \\\n --prompt-file $PROMPT_FILE \\\n --tokenizer-path $TOKENIZER_PATH \\\n --micro-batch-size 1 \\\n --out-seq-length 1024 \\\n --temperature 0.8 \\\n --top-p 0.95 \\\n --top-k 0 \\\n --greedy \\\n $MODEL_ARGS\"\n\necho \"$CMD\"\neval \"$CMD\"\n"
},
{
"path": "scripts/test_inference_parallel.sh",
"content": "# This script is used to test the inference of CodeGeeX.\n\nMP_SIZE=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\nif [ -z \"$MP_SIZE\" ]; then\n MP_SIZE=1\nfi\n\nif [ \"$MP_SIZE\" -eq 1 ]; then\n source \"$MAIN_DIR/configs/codegeex_13b.sh\"\n echo \"Load config from $MAIN_DIR/configs/codegeex_13b.sh\"\nelse\n source \"$MAIN_DIR/configs/codegeex_13b_parallel.sh\"\n echo \"Load config from $MAIN_DIR/configs/codegeex_13b_parallel.sh\"\nfi\n\n# export CUDA settings\nexport CUDA_HOME=/usr/local/cuda-11.1/\n# export CUDA_VISIBLE_DEVICES=0,1\n\nif [ -z \"$PROMPT_FILE\" ]; then\n PROMPT_FILE=$MAIN_DIR/tests/test_prompt.txt\nfi\n\n# remove --greedy if using sampling\nCMD=\"torchrun --nproc_per_node $MP_SIZE $MAIN_DIR/tests/test_inference_megatron.py \\\n --tensor-model-parallel-size $MP_SIZE \\\n --prompt-file $PROMPT_FILE \\\n --tokenizer-path $TOKENIZER_PATH \\\n --micro-batch-size 1 \\\n --out-seq-length 1024 \\\n --temperature 0.8 \\\n --top-p 0.95 \\\n --top-k 0 \\\n --greedy \\\n --use-cpu-initialization \\\n --ln-fp16 \\\n $MODEL_ARGS\"\n\necho \"$CMD\"\neval \"$CMD\"\n"
},
{
"path": "scripts/test_inference_quantized.sh",
"content": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\n# import model configuration\nsource \"$MAIN_DIR/configs/codegeex_13b.sh\"\n\n# export CUDA settings\nif [ -z \"$GPU\" ]; then\n GPU=0\nfi\n\nexport CUDA_HOME=/usr/local/cuda-11.1/\nexport CUDA_VISIBLE_DEVICES=$GPU\n\nif [ -z \"$PROMPT_FILE\" ]; then\n PROMPT_FILE=$MAIN_DIR/tests/test_prompt.txt\nfi\n\n# remove --greedy if using sampling\nCMD=\"python $MAIN_DIR/tests/test_inference.py \\\n --prompt-file $PROMPT_FILE \\\n --tokenizer-path $TOKENIZER_PATH \\\n --micro-batch-size 1 \\\n --out-seq-length 1024 \\\n --temperature 0.2 \\\n --top-p 0.95 \\\n --top-k 0 \\\n --quantize \\\n $MODEL_ARGS\"\n\necho \"$CMD\"\neval \"$CMD\"\n"
},
{
"path": "scripts/translate_humaneval_x.sh",
"content": "# This script is used to translate solutions of HumanEval-X.\n\nLANG_SRC_TYPE=$1 # Source programming language, currently support one of [\"python\", \"java\", \"cpp\", \"js\", \"go\"]\nLANG_TGT_TYPE=$2 # Target programming language, currently support one of [\"python\", \"java\", \"cpp\", \"js\", \"go\"]\nOUTPUT_PATH=$3 # Output path of the generated programs.\nHOSTLIST=$4 # Provide hostfile if generating distributedly\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\n# export CUDA settings\nexport CUDA_HOME=/usr/local/cuda-11.1/\n\n# import model configuration\nsource \"$MAIN_DIR/configs/codegeex_13b.sh\"\n\n# nccl options\nOPTIONS_NCCL=\"export NCCL_DEBUG=warn; export NCCL_IB_DISABLE=0; export NCCL_IB_GID_INDEX=3\"\nOPTIONS_PATH=\"export PATH=$PATH; export LD_LIBRARY_PATH=$LD_LIBRARY_PATH\"\nCWD=$(pwd)\n\n# set master ip for zmq server\nif [ -z \"$HOSTLIST\" ]; then\n ZMQ_ADDR=$(hostname -i)\n echo \"$ZMQ_ADDR\" > \"./hostfile\"\n HOSTLIST=\"./hostfile\"\nelse\n ZMQ_ADDR=$(cat $HOSTLIST | head -n 1)\nfi\necho \"master_ip: $ZMQ_ADDR\"\n\nNUM_SAMPLES=1\nMICRO_BSZ=1\nWORLD_SIZE=1\nTEMP=0.8\nTOPP=0.95\nSEED=42\nDATASET=humaneval\nTODAY=$(date +%y%m%d)\nCHANNEL_PORT=$(expr $RANDOM + 5000)\nMASTER_PORT=$(expr $RANDOM + 8000)\n\n# save log file\nLOG_DIR=$MAIN_DIR/log\nmkdir -p \"$LOG_DIR\"\nLOG_PATH=\"$LOG_DIR/$TODAY-translation.log\"\n\nif [ -z \"$LANG_SRC_TYPE\" ]\nthen\n LANG_SRC_TYPE=python\nfi\n\nif [ -z \"$LANG_TGT_TYPE\" ]\nthen\n LANG_TGT_TYPE=java\nfi\n\nif [ -z \"$INPUT_SRC_PATH\" ]\nthen\n INPUT_SRC_PATH=$MAIN_DIR/codegeex/benchmark/humaneval-x/$LANG_SRC_TYPE/data/humaneval_$LANG_SRC_TYPE.jsonl.gz\nfi\n\nif [ -z \"$INPUT_TGT_PATH\" ]\nthen\n INPUT_TGT_PATH=$MAIN_DIR/codegeex/benchmark/humaneval-x/$LANG_TGT_TYPE/data/humaneval_$LANG_TGT_TYPE.jsonl.gz\nfi\n\nif [ -z \"$OUTPUT_PATH\" ]; then\n OUTPUT_PATH=$MAIN_DIR/codegeex/benchmark/output/humaneval-x/codegeex/\n mkdir -p \"$OUTPUT_PATH\"\nfi\n\nJOB_ID=codegeex-ns$NUM_SAMPLES-t$TEMP-topp$TOPP-seed$SEED-$LANGUAGE\n\nRUN_CMD=\"python \\\n $MAIN_DIR/codegeex/benchmark/humaneval-x/translate_humaneval_x.py \\\n --hostfile $HOSTLIST \\\n --channel-ip $ZMQ_ADDR \\\n --channel-port $CHANNEL_PORT \\\n --master-port $MASTER_PORT \\\n --tokenizer-path $TOKENIZER_PATH \\\n --load-deepspeed \\\n --temperature $TEMP \\\n --top-p $TOPP \\\n --out-seq-length 1024 \\\n --micro-batch-size $MICRO_BSZ \\\n --samples-per-problem $NUM_SAMPLES \\\n --language-src-type $LANG_SRC_TYPE \\\n --language-tgt-type $LANG_TGT_TYPE \\\n --src-path $INPUT_SRC_PATH \\\n --tgt-path $INPUT_TGT_PATH \\\n --dataset $DATASET \\\n --output-prefix $OUTPUT_PATH/$JOB_ID \\\n --gen-node-world-size $WORLD_SIZE \\\n --seed $SEED \\\n $MODEL_ARGS\"\n\nRUN_CMD=\"$OPTIONS_NCCL; $OPTIONS_PATH; $RUN_CMD\"\nRUN_CMD=\"cd $CWD; $RUN_CMD\"\n\nif (( WORLD_SIZE != 1 )); then\n RUN_CMD=\"pdsh -R ssh -w ^$HOSTLIST \\\"$RUN_CMD\\\"\"\nfi\n\necho \"$RUN_CMD\"\necho \"Writing log to $LOG_PATH\"\neval \"$RUN_CMD\" > \"$LOG_PATH\"\nbash $MAIN_DIR/scripts/gather_output.sh $OUTPUT_PATH $JOB_ID 1\n"
},
{
"path": "setup.py",
"content": "from setuptools import setup, find_packages\n\nsetup(\n name=\"codegeex\",\n py_modules=[\"codegeex\"],\n version=\"1.0\",\n description=\"CodeGeeX: A Open Multilingual Code Generation Model.\",\n author=\"Qinkai Zheng\",\n packages=find_packages(),\n install_requires=[\n \"fire>=0.4.0\",\n \"ipython>=8.4.0\",\n \"numpy>=1.22.0\",\n \"pandas>=1.3.5\",\n \"pyzmq>=23.2.1\",\n \"regex>=2022.3.15\",\n \"setuptools>=58.0.4\",\n \"transformers>=4.22.0\",\n \"tokenizers>=0.11.0\",\n \"torch>=1.10.0\",\n \"tqdm>=4.63.0\",\n \"cpm_kernels\",\n \"deepspeed>0.6.1\",\n ],\n entry_points={}\n)\n"
},
{
"path": "tests/test_inference.py",
"content": "import time\nimport torch\nimport argparse\nimport numpy as np\n\nimport codegeex\nfrom codegeex.torch import CodeGeeXModel\nfrom codegeex.tokenizer import CodeGeeXTokenizer\nfrom codegeex.quantization import quantize\n\n\ndef model_provider(args):\n \"\"\"Build the model.\"\"\"\n\n model = CodeGeeXModel(\n args.hidden_size,\n args.num_layers,\n args.num_attention_heads,\n args.padded_vocab_size,\n args.max_position_embeddings\n )\n \n return model\n\n\ndef add_code_generation_args(parser):\n group = parser.add_argument_group(title=\"code generation\")\n group.add_argument(\n \"--num-layers\",\n type=int,\n default=39,\n )\n group.add_argument(\n \"--hidden-size\",\n type=int,\n default=5120,\n )\n group.add_argument(\n \"--num-attention-heads\",\n type=int,\n default=40,\n )\n group.add_argument(\n \"--padded-vocab-size\",\n type=int,\n default=52224,\n )\n group.add_argument(\n \"--max-position-embeddings\",\n type=int,\n default=2048,\n )\n group.add_argument(\n \"--temperature\",\n type=float,\n default=1.0,\n help=\"Sampling temperature.\",\n )\n group.add_argument(\n \"--greedy\",\n action=\"store_true\",\n default=False,\n help=\"Use greedy sampling.\",\n )\n group.add_argument(\n \"--top-p\",\n type=float,\n default=0.0,\n help=\"Top p sampling.\",\n )\n group.add_argument(\n \"--top-k\",\n type=int,\n default=0,\n help=\"Top k sampling.\",\n )\n group.add_argument(\n \"--out-seq-length\",\n type=int,\n default=2048,\n help=\"Size of the output generated text.\",\n )\n group.add_argument(\n \"--prompt-file\",\n type=str,\n default=\"./test_prompt.txt\",\n )\n group.add_argument(\n \"--tokenizer-path\",\n type=str,\n default=\"./tokenizer\",\n )\n group.add_argument(\n \"--load\",\n type=str,\n )\n group.add_argument(\n \"--state-dict-path\",\n type=str,\n )\n group.add_argument(\n \"--micro-batch-size\",\n type=int,\n default=1,\n )\n group.add_argument(\n \"--quantize\",\n action=\"store_true\",\n )\n group.add_argument(\n \"--interative\",\n action=\"store_true\",\n )\n \n return parser\n\n \ndef main():\n parser = argparse.ArgumentParser()\n parser = add_code_generation_args(parser)\n args, _ = parser.parse_known_args()\n \n print(\"Loading tokenizer ...\")\n tokenizer = CodeGeeXTokenizer(\n tokenizer_path=args.tokenizer_path, \n mode=\"codegeex-13b\")\n\n print(\"Loading state dict ...\")\n state_dict = torch.load(args.load, map_location=\"cpu\")\n state_dict = state_dict[\"module\"]\n\n print(\"Building CodeGeeX model ...\")\n model = model_provider(args)\n model.load_state_dict(state_dict)\n model.eval()\n model.half()\n if args.quantize:\n model = quantize(model, weight_bit_width=8, backend=\"torch\")\n model.cuda()\n torch.cuda.synchronize()\n \n with open(args.prompt_file, \"r\") as f:\n prompt = f.readlines()\n prompt = \"\".join(prompt)\n \n out_seq_lengths = [args.out_seq_length]\n for out_seq_length in out_seq_lengths: \n print(f\"Generating with out_seq_len {out_seq_length}...\")\n while True:\n print(\"\\nPlease Input Query (Ctrl-D to save multiple lines, 'stop' to exit) >>> \")\n prompts = []\n while True:\n try:\n line = input()\n except EOFError:\n break\n prompts.append(line)\n prompt = \"\\n\".join(prompts)\n prompt = prompt.strip()\n if not prompt:\n print('Query should not be empty!')\n continue\n if prompt == \"stop\":\n return \n try:\n t0 = time.perf_counter()\n generated_code = codegeex.generate(\n model,\n tokenizer,\n prompt,\n out_seq_length=out_seq_length,\n seq_length=args.max_position_embeddings,\n top_k=args.top_k,\n top_p=args.top_p,\n temperature=args.temperature,\n micro_batch_size=args.micro_batch_size,\n backend=\"megatron\",\n verbose=True,\n )\n t1 = time.perf_counter()\n print(\"Total generation time:\", t1 - t0)\n except (ValueError, FileNotFoundError) as e:\n print(e)\n continue\n \n print(\"Generation finished.\")\n\n\nif __name__ == \"__main__\":\n main()"
},
{
"path": "tests/test_inference_megatron.py",
"content": "import copy\nimport time\nimport torch\nimport numpy as np\n\nfrom codegeex.megatron import get_tokenizer, get_args, print_rank_0\nfrom codegeex.megatron.initialize import initialize_megatron\nfrom codegeex.megatron.model import CodeGeeXModel\nfrom codegeex.megatron.code_generation_utils import get_token_stream\nfrom codegeex.quantization import quantize\nfrom codegeex.megatron.training import get_model\nfrom codegeex.megatron.checkpointing import load_checkpoint\n\ntorch.set_printoptions(precision=8)\n\n\ndef set_random_seed(seed):\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n\n torch.backends.cudnn.deterministic = True\n torch.backends.cudnn.benchmark = False\n\n\ndef model_provider(pre_process=True, post_process=True):\n \"\"\"Build the model.\"\"\"\n \n print_rank_0(\"Building CodeGeeX model ...\")\n model = CodeGeeXModel(num_tokentypes=0,\n parallel_output=False)\n\n return model\n\n\ndef add_code_generation_args(parser):\n \"\"\"Code generation arguments.\"\"\"\n group = parser.add_argument_group(title=\"code generation\")\n\n group.add_argument(\n \"--temperature\",\n type=float,\n default=1.0,\n help=\"Sampling temperature.\",\n )\n group.add_argument(\n \"--greedy\",\n action=\"store_true\",\n default=False,\n help=\"Use greedy sampling.\",\n )\n group.add_argument(\n \"--top-p\",\n type=float,\n default=0.0,\n help=\"Top p sampling.\",\n )\n group.add_argument(\n \"--top-k\",\n type=int,\n default=0,\n help=\"Top k sampling.\",\n )\n group.add_argument(\n \"--out-seq-length\",\n type=int,\n default=2048,\n help=\"Size of the output generated text.\",\n )\n group.add_argument(\n \"--recompute\",\n action=\"store_true\",\n help=\"During generation recompute all attention \"\n \"instead of using previously computed keys/values.\",\n )\n group.add_argument(\n \"--ws-encoding-start-id\",\n type=int,\n default=10,\n help=\"Start id for whitespace encoding\",\n )\n group.add_argument(\n \"--ws-encoding-length\",\n type=int,\n default=10,\n help=\"Length of whitespace encoding\",\n )\n group.add_argument(\n \"--n-generation\",\n type=int,\n default=10,\n )\n group.add_argument(\n \"--eos-id\",\n type=int,\n default=50256,\n )\n group.add_argument(\n \"--prompt-file\",\n type=str,\n default=\"./test_prompt.txt\",\n )\n group.add_argument(\n \"--perf-file\",\n type=str,\n default=\"./perf_out.txt\",\n )\n group.add_argument(\n \"--perf-trace\",\n type=str,\n default=\"./perf_out.txt\",\n )\n group.add_argument(\n \"--use-torch-profile\",\n action=\"store_true\",\n )\n group.add_argument(\n \"--ln-fp32\",\n action=\"store_true\",\n )\n group.add_argument(\n '--bad-ids',\n nargs=\"*\",\n type=int,\n default=None,\n help='Identify the type of programming language to generate',\n )\n group.add_argument(\n \"--quantize\",\n action=\"store_true\",\n )\n\n return parser\n\n\ndef main():\n initialize_megatron(\n extra_args_provider=add_code_generation_args,\n args_defaults={\n 'no_load_rng': True,\n 'no_load_optim': True,\n }\n )\n\n args = get_args()\n set_random_seed(args.seed)\n\n print_rank_0(\"Loading tokenizer ...\")\n tokenizer = get_tokenizer()\n\n print_rank_0(\"Loading state dict ...\")\n\n model = get_model(model_provider)\n if args.load is not None:\n _ = load_checkpoint(model, None, None)\n \n assert len(model) == 1, \"Above condition should have caught this\"\n \n model = model[0]\n model.eval()\n if args.fp16 and args.ln_fp16:\n model.half()\n if args.quantize:\n model = quantize(model, weight_bit_width=8, backend=\"megatron\")\n\n with open(args.prompt_file, \"r\") as f:\n prompt = f.readlines()\n prompt = \"\".join(prompt)\n \n times = {}\n out_seq_lengths = [args.out_seq_length]\n micro_batch_size = args.micro_batch_size\n for out_seq_length in out_seq_lengths: \n print_rank_0(f\"Generating with out_seq_len {out_seq_length}...\")\n \n times[out_seq_length] = []\n for prompt in [prompt] * args.n_generation:\n t0 = time.perf_counter()\n tokens = tokenizer.tokenize(prompt)\n print_rank_0(tokens)\n print_rank_0(\"Current prompt:\")\n print_rank_0(prompt)\n n_token_prompt = len(tokens)\n print_rank_0(f\"N_token_prompt:{n_token_prompt}\")\n token_stream = get_token_stream(\n model,\n [copy.deepcopy(tokens) for _ in range(micro_batch_size)],\n micro_batch_size=micro_batch_size,\n topk=args.top_k,\n topp=args.top_p,\n temperature=args.temperature,\n )\n is_finished = [False for _ in range(micro_batch_size)]\n for i, generated in enumerate(token_stream):\n generated_tokens = generated[0]\n for j in range(micro_batch_size):\n if is_finished[j]:\n continue\n if generated_tokens[j].cpu().numpy()[-1] == tokenizer.eod or len(\n generated_tokens[j]) >= out_seq_length:\n is_finished[j] = True\n generated_tokens_ = generated_tokens[j].cpu().numpy().tolist()\n generated_code = tokenizer.detokenize(generated_tokens_[n_token_prompt:])\n t1 = time.perf_counter()\n print_rank_0(f\"Total generation time: {t1 - t0}, # Tokens: {len(generated_tokens_) - n_token_prompt}\")\n print_rank_0(f\"{(t1 - t0) / (len(generated_tokens_) - n_token_prompt)}s/token\")\n times[out_seq_length].append(t1 - t0)\n print_rank_0(\"================================= Generated code:\")\n print_rank_0(generated_code)\n t0 = time.perf_counter()\n \n if all(is_finished):\n break\n\n print_rank_0(times)\n for out_seq_length in times.keys():\n print_rank_0(f\"{out_seq_length}, {np.mean(times[out_seq_length])}\")\n \n print_rank_0(\"Generation finished.\")\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "tests/test_inference_oneflow.py",
"content": "\nimport os\nimport copy\nimport time\nimport oneflow as torch\nimport random\nimport argparse\nimport numpy as np\n\nfrom codegeex.oneflow.inference import get_token_stream\nfrom codegeex.oneflow import CodeGeeXModel\nfrom codegeex.tokenizer import CodeGeeXTokenizer\nfrom codegeex.quantization import quantize_oneflow\nos.environ[\"ONEFLOW_KERNEL_ENABLE_FUSED_LINEAR\"] = \"1\"\nos.environ[\"ONEFLOW_LINEAR_EMBEDDING_SKIP_INIT\"] = \"1\"\n\ndef model_provider(args):\n \"\"\"Build the model.\"\"\"\n\n model = CodeGeeXModel(\n args.hidden_size,\n args.num_layers,\n args.num_attention_heads,\n args.padded_vocab_size,\n args.max_position_embeddings\n )\n \n return model\n\n\ndef add_code_generation_args(parser):\n group = parser.add_argument_group(title=\"code generation\")\n group.add_argument(\n \"--num-layers\",\n type=int,\n default=39,\n )\n group.add_argument(\n \"--hidden-size\",\n type=int,\n default=5120,\n )\n group.add_argument(\n \"--num-attention-heads\",\n type=int,\n default=40,\n )\n group.add_argument(\n \"--padded-vocab-size\",\n type=int,\n default=52224,\n )\n group.add_argument(\n \"--max-position-embeddings\",\n type=int,\n default=2048,\n )\n group.add_argument(\n \"--temperature\",\n type=float,\n default=1.0,\n help=\"Sampling temperature.\",\n )\n group.add_argument(\n \"--greedy\",\n action=\"store_true\",\n default=False,\n help=\"Use greedy sampling.\",\n )\n group.add_argument(\n \"--top-p\",\n type=float,\n default=0.0,\n help=\"Top p sampling.\",\n )\n group.add_argument(\n \"--top-k\",\n type=int,\n default=0,\n help=\"Top k sampling.\",\n )\n group.add_argument(\n \"--out-seq-length\",\n type=int,\n default=2048,\n help=\"Size of the output generated text.\",\n )\n group.add_argument(\n \"--prompt-file\",\n type=str,\n default=\"./test_prompt.txt\",\n )\n group.add_argument(\n \"--tokenizer-path\",\n type=str,\n default=\"./tokenizer\",\n )\n group.add_argument(\n \"--load\",\n type=str,\n )\n group.add_argument(\n \"--state-dict-path\",\n type=str,\n )\n group.add_argument(\n \"--micro-batch-size\",\n type=int,\n default=1,\n )\n group.add_argument(\n \"--quantize\",\n action=\"store_true\",\n )\n \n return parser\n\n \ndef main():\n parser = argparse.ArgumentParser()\n parser = add_code_generation_args(parser)\n args, _ = parser.parse_known_args()\n \n print(\"Loading tokenizer ...\")\n tokenizer = CodeGeeXTokenizer(\n tokenizer_path=args.tokenizer_path, \n mode=\"codegeex-13b\")\n\n print(\"Loading state dict ...\")\n state_dict = torch.load(args.load, map_location=\"cpu\")\n state_dict = state_dict[\"module\"]\n\n print(\"Building CodeGeeX model ...\")\n model = model_provider(args)\n model.load_state_dict(state_dict)\n model.eval()\n model.half()\n if args.quantize:\n model = quantize_oneflow(model, weight_bit_width=8)\n model.cuda()\n torch.cuda.synchronize()\n with open(args.prompt_file, \"r\") as f:\n prompt = f.readlines()\n prompt = \"\".join(prompt)\n \n times = {}\n out_seq_lengths = [args.out_seq_length]\n micro_batch_size = args.micro_batch_size\n seq_length = args.max_position_embeddings\n for out_seq_length in out_seq_lengths: \n print(f\"Generating with out_seq_len {out_seq_length}...\")\n \n times[out_seq_length] = []\n for prompt in [prompt]:\n t0 = time.perf_counter()\n tokens = tokenizer.encode_code(prompt)\n print(tokens)\n print(\"Current prompt:\")\n print(prompt)\n n_token_prompt = len(tokens)\n print(\"N_token_prompt:\", n_token_prompt)\n token_stream = get_token_stream(\n model,\n tokenizer,\n seq_length,\n out_seq_length,\n [copy.deepcopy(tokens) for _ in range(micro_batch_size)],\n micro_batch_size=micro_batch_size,\n topk=args.top_k,\n topp=args.top_p,\n temperature=args.temperature,\n greedy=args.greedy,\n )\n is_finished = [False for _ in range(micro_batch_size)]\n for i, generated in enumerate(token_stream):\n generated_tokens = generated[0]\n for j in range(micro_batch_size):\n if is_finished[j]:\n continue\n generated_token_numpy = generated_tokens[j].numpy()\n if generated_token_numpy[-1] == tokenizer.eos_token_id or len(\n generated_tokens[j]) >= out_seq_length:\n is_finished[j] = True\n generated_tokens_ = generated_token_numpy.tolist()\n generated_code = tokenizer.decode_code(generated_tokens_[n_token_prompt:])\n generated_code = \"\".join(generated_code)\n t1 = time.perf_counter()\n print(\"Total generation time:\", t1 - t0, \"# Tokens:\", len(generated_tokens_) - n_token_prompt)\n print(f\"{(t1 - t0) / (len(generated_tokens_) - n_token_prompt)}s/token\")\n times[out_seq_length].append(t1 - t0)\n print(\"================================= Generated code:\")\n print(generated_code)\n \n if all(is_finished):\n break\n \n print(times)\n for out_seq_length in times.keys():\n print(out_seq_length, np.mean(times[out_seq_length]))\n \n print(\"Generation finished.\")\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "tests/test_inference_paddle.py",
"content": "\nimport os\nimport copy\nimport time\nimport paddle\nimport random\nimport argparse\nimport numpy as np\n\nfrom codegeex.paddle.inference import get_token_stream\nfrom codegeex.paddle import CodeGeeXModel\nfrom codegeex.tokenizer import CodeGeeXTokenizer\n\n\ndef model_provider(args):\n \"\"\"Build the model.\"\"\"\n\n old_dtype = paddle.get_default_dtype()\n paddle.set_default_dtype(\"float16\")\n model = CodeGeeXModel(\n args.hidden_size,\n args.num_layers,\n args.num_attention_heads,\n args.padded_vocab_size,\n args.max_position_embeddings\n )\n model.language_model.embedding.word_embeddings.to(dtype=\"float32\")\n model.language_model.embedding.position_embeddings.to(dtype=\"float32\")\n model.language_model.topQueryEmbedding.top_query_embeddings.to(dtype=\"float32\")\n for i in model.language_model.transformer.layers:\n i.input_layernorm.to(dtype=\"float32\")\n i.post_attention_layernorm.to(dtype=\"float32\")\n model.language_model.transformer.topQueryLayer.input_layernorm.to(dtype=\"float32\")\n model.language_model.transformer.topQueryLayer.post_attention_layernorm.to(dtype=\"float32\")\n model.language_model.transformer.final_layernorm.to(dtype=\"float32\")\n paddle.set_default_dtype(old_dtype)\n \n return model\n\n\ndef add_code_generation_args(parser):\n group = parser.add_argument_group(title=\"code generation\")\n group.add_argument(\n \"--num-layers\",\n type=int,\n default=39,\n )\n group.add_argument(\n \"--hidden-size\",\n type=int,\n default=5120,\n )\n group.add_argument(\n \"--num-attention-heads\",\n type=int,\n default=40,\n )\n group.add_argument(\n \"--padded-vocab-size\",\n type=int,\n default=52224,\n )\n group.add_argument(\n \"--max-position-embeddings\",\n type=int,\n default=2048,\n )\n group.add_argument(\n \"--temperature\",\n type=float,\n default=1.0,\n help=\"Sampling temperature.\",\n )\n group.add_argument(\n \"--greedy\",\n action=\"store_true\",\n default=False,\n help=\"Use greedy sampling.\",\n )\n group.add_argument(\n \"--top-p\",\n type=float,\n default=0.0,\n help=\"Top p sampling.\",\n )\n group.add_argument(\n \"--top-k\",\n type=int,\n default=0,\n help=\"Top k sampling.\",\n )\n group.add_argument(\n \"--out-seq-length\",\n type=int,\n default=2048,\n help=\"Size of the output generated text.\",\n )\n group.add_argument(\n \"--prompt-file\",\n type=str,\n default=\"./test_prompt.txt\",\n )\n group.add_argument(\n \"--tokenizer-path\",\n type=str,\n default=\"./tokenizer\",\n )\n group.add_argument(\n \"--load\",\n type=str,\n )\n group.add_argument(\n \"--state-dict-path\",\n type=str,\n )\n group.add_argument(\n \"--micro-batch-size\",\n type=int,\n default=1,\n )\n group.add_argument(\n \"--quantize\",\n action=\"store_true\",\n )\n \n return parser\n\n \ndef main():\n parser = argparse.ArgumentParser()\n parser = add_code_generation_args(parser)\n args, _ = parser.parse_known_args()\n \n print(\"Loading tokenizer ...\")\n tokenizer = CodeGeeXTokenizer(\n tokenizer_path=args.tokenizer_path, \n mode=\"codegeex-13b\")\n\n print(\"Loading state dict ...\")\n state_dict = paddle.load(args.load)\n state_dict = state_dict[\"module\"]\n\n print(\"Building CodeGeeX model ...\")\n model = model_provider(args)\n model.set_state_dict(state_dict)\n model.eval()\n model.to(dtype=\"float16\")\n if args.quantize:\n raise NotImplementedError(\"quantize\")\n \n with open(args.prompt_file, \"r\") as f:\n prompt = f.readlines()\n prompt = \"\".join(prompt)\n \n times = {}\n out_seq_lengths = [args.out_seq_length]\n micro_batch_size = args.micro_batch_size\n seq_length = args.max_position_embeddings\n for out_seq_length in out_seq_lengths: \n print(f\"Generating with out_seq_len {out_seq_length}...\")\n \n times[out_seq_length] = []\n for prompt in [prompt]:\n t0 = time.perf_counter()\n tokens = tokenizer.encode_code(prompt)\n print(tokens)\n print(\"Current prompt:\")\n print(prompt)\n n_token_prompt = len(tokens)\n print(\"N_token_prompt:\", n_token_prompt)\n token_stream = get_token_stream(\n model,\n tokenizer,\n seq_length,\n out_seq_length,\n [copy.deepcopy(tokens) for _ in range(micro_batch_size)],\n micro_batch_size=micro_batch_size,\n topk=args.top_k,\n topp=args.top_p,\n temperature=args.temperature,\n greedy=args.greedy,\n )\n is_finished = [False for _ in range(micro_batch_size)]\n for i, generated in enumerate(token_stream):\n generated_tokens = generated[0]\n for j in range(micro_batch_size):\n if is_finished[j]:\n continue\n if generated_tokens[j].cpu().numpy()[-1] == tokenizer.eos_token_id or len(\n generated_tokens[j]) >= out_seq_length:\n is_finished[j] = True\n generated_tokens_ = generated_tokens[j].cpu().numpy().tolist()\n generated_code = tokenizer.decode_code(generated_tokens_[n_token_prompt:])\n generated_code = \"\".join(generated_code)\n t1 = time.perf_counter()\n print(\"Total generation time:\", t1 - t0, \"# Tokens:\", len(generated_tokens_) - n_token_prompt)\n print(f\"{(t1 - t0) / (len(generated_tokens_) - n_token_prompt)}s/token\")\n times[out_seq_length].append(t1 - t0)\n print(\"================================= Generated code:\")\n print(generated_code)\n \n if all(is_finished):\n break\n \n print(times)\n for out_seq_length in times.keys():\n print(out_seq_length, np.mean(times[out_seq_length]))\n \n print(\"Generation finished.\")\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "tests/test_prompt.txt",
"content": "code translation\nJava:\npublic class Solution {\n public static boolean hasCloseElements(int[] nums, int threshold) {\n for (int i = 0; i < nums.length - 1; i++) {\n for (int j = i + 1; j < nums.length; j++) {\n if (Math.abs(nums[i] - nums[j]) < threshold) {\n return true;\n }\n }\n }\n return false;\n }\n}\nPython:\n"
},
{
"path": "vscode-extension/README.md",
"content": "\n\n🌐 中文\n\n\n\n\n\n\n\nWe introduce CodeGeeX, a large-scale multilingual code generation model with 13 billion parameters, pretrained on a large code corpus of more than 20 programming languages. With CodeGeeX, we can generate codes by only providing natural language descriptions, complete any code snippet, or translate codes to other programming languages, etc. CodeGeeX also provides customizable features (**Prompt Mode**) to help you configure your own programming assistant. Happy coding!\n\nPlease search \"codegeex\" in VS Code Marketplace to download it for free (minimum version requirement of VS Code is 1.68.0). For more information, please check out our [Homepage](https://models.aminer.cn/codegeex/) and [GitHub repo](https://github.com/THUDM/CodeGeeX).\n\nPlease kindly let us know if you encounter any problem or have any suggestion, via [codegeex@aminer.cn](mailto:codegeex@aminer.cn).\n\n- [Basic Usage](#basic-usage)\n- [Privacy](#privacy)\n- [Guidance](#guidance)\n - [Stealth mode](#stealth-mode)\n - [Interactive mode](#interactive-mode)\n - [Translation mode](#translation-mode)\n - [Prompt mode (Experimental)](#prompt-mode-experimental)\n\n## Basic Usage\nMake sure vscode version is >= 1.68.0. Install the extension and enable it globally. There are four modes of usage:\n\n- **Stealth mode**: Keep CodeGeeX activated, it will start generating codes when you stop writing (the icon at the bottom of VSCode starts spinning). When the generated code is shown in gray, just press ``Tab`` to insert the generated codes. \n- **Interactive mode**: Press ``Ctrl+Enter`` to activate the interactive mode, CodeGeeX will generate ``X`` candidates and show them in the right panel (``X`` can be modified in extension settings ``Candidate Num``). Then, select the best candidate by clicking on it.\n- **Translation mode**: Select code, and press ``Ctrl+Alt+T`` to activate the translation mode. Then, choose the language of the selected code. You will get the code translated into the same language as your current editor. Click on the ``use code`` button to insert the result. You can also configure in the settings whether to comment out the original code or to replace it.\n- **Prompt mode (experimental)**: Select codes to be used as input, then press ``Alt/Option+t`` to trigger the prompt mode. It will show a list of pre-defined prompt templates and choose one to generate codes with your input. This mode is fully customizable, you can add your own templates in the extension settings ``Prompt Templates``. \n\n## Privacy\nWe highly respect the privacy of your code. The code is only used as the input of CodeGeeX to assist your programming. At the first time of usage, we will ask if you agree to share the generated code only for research purpose (**disabled** by default).\n\n## Guidance\nPlease see the details and examples for how to use the three modes in CodeGeeX:\n### Stealth mode\nIn this mode, CodeGeeX will start generating codes when you stop writing (the icon at the bottom of VSCode starts spinning). When the generated code is shown in gray, just press ``Tab`` to insert the generated codes. You can also press ``Alt/Option+[`` or ``]`` to change between candidates. If you are not satisfied with the current suggestions, you can also press ``Alt/Option+N`` to get new suggestions. Change the number of candidates in the extension settings ``Candidate Num`` (more candidates will slow down the generation speed). **Note**: The generation always starts at the current position of your cursor, thus if you modify the code before the generation is finished, it will probably cause bugs. We keep working on making the generation faster.\n\n\n\n### Interactive mode\nIn this mode, press ``Ctrl+Enter`` to generate codes and visualize the candidates in another panel. Then, click on the best candidate to insert the generated codes to the current position of cursor. \n\n\n\n### Translation mode\nIn this mode, paste or tape some code from another language to the current file, select the code, and press ``Ctrl+Alt+T``. Then, choose the language of the selected code. Wait for a few seconds, you will get the code translated into the same language as your current editor. Click on the ``use code`` button to insert the result into the current position of your cursor. You can also configure in the settings whether to comment out the original code or to replace it.\n\n\n\n### Prompt mode (Experimental)\nIn this mode, you can add extra prompts to the input and implement some cool features, like code explanation, summarization, generation with specific coding style, and more. The principle behind is the few-shot ability of CodeGeeX. When you provide a few examples as extra prompts in the input, CodeGeeX will imitate what are done by these examples and generate codes accordingly. For example, you can give an example that explains each line of code. Select the code you want to explain, then press ``Alt/Option+t`` to trigger the prompt mode. It will show a list of pre-defined prompt templates and choose the ``explanation`` to generate codes with your input. Magically, the codes will be explained line by line.\n\n\n\nThe template of the above example looks like the following, which contains ``[Example code]``, ````, ``[Example code with explanation]`` and ``[Explanation head]``. ```` is where the selected code will be inserted. ```` means the first line of your input (which is used here to ensure the same function will be explained). When you use the prompt mode, CodeGeeX will combine your input with the template and use them all as the input to generate codes. \n\n```python\n# language: Python\n\ndef sum_squares(lst):\n sum = 0\n for i in range(len(lst)):\n if i % 3 == 0:\n lst[i] = lst[i]**2\n elif i % 4 == 0:\n lst[i] = lst[i]**3\n sum += lst[i]\n return sum\n\n\n\n# Explain the code line by line\ndef sum_squares(lst):\n # initialize sum\n sum = 0\n # loop through the list\n for i in range(len(lst)):\n # if the index is a multiple of 3\n if i % 3 == 0:\n # square the entry\n lst[i] = lst[i]**2\n # if the index is a multiple of 4\n elif i % 4 == 0:\n # cube the entry\n lst[i] = lst[i]**3\n # add the entry to the sum\n sum += lst[i]\n # return the sum\n return sum\n\n# Explain the code line by line\n\n```\n\nAnd here is another example for python docstring generation \n```python\ndef add_binary(a, b):\n '''\n Returns the sum of two decimal numbers in binary digits.\n\n Parameters:\n a (int): A decimal integer\n b (int): Another decimal integer\n\n Returns:\n binary_sum (str): Binary string of the sum of a and b\n '''\n binary_sum = bin(a+b)[2:]\n return binary_sum\n\n\n```\n\nThe templates are fully customizable, you can add your own templates in the extension settings ``Prompt Templates``. ``key`` is the name that you want to show in the list of templates, ``value`` is the path to the template file (``.txt``, ``.py``, ``.h``, etc). Try this feature and write your own templates, you can make the generated codes follow your coding style, generate with a specific function name, or add a specific comment, etc.\n"
},
{
"path": "vscode-extension/README_zh.md",
"content": "\n\n🌐 English\n\n\n\n\n\n\n\nCodeGeeX是一个具有130亿参数的多编程语言代码生成预训练模型，使用超过二十种编程语言训练得到。基于CodeGeeX开发的插件可以实现通过描述生成代码、补全代码、代码翻译等一系列功能。CodeGeeX同样提供可以定制的**提示模式（Prompt Mode）**，构建专属的编程助手。Happy Coding！\n\nVS Code插件市场搜索\"codegeex\"即可免费使用(需要VS Code版本不低于1.68.0)，更多关于CodeGeeX信息请见我们的[主页](https://models.aminer.cn/codegeex/) and [GitHub仓库](https://github.com/THUDM/CodeGeeX)。\n\n如使用过程中遇到问题或有任何改进意见，欢迎发送邮件到[codegeex@aminer.cn](mailto:codegeex@aminer.cn)反馈！\n\n- [基本用法](#基本用法)\n- [隐私声明](#隐私声明)\n- [使用指南](#使用指南)\n - [隐匿模式](#隐匿模式)\n - [交互模式](#交互模式)\n - [翻译模式](#翻译模式)\n - [提示模式（实验功能）](#提示模式实验功能)\n\n## 基本用法\n需要保证VS Code版本 >= 1.68.0。安装插件并全局激活CodeGeeX，有以下四种使用模式：\n\n- **隐匿模式**: 保持CodeGeeX处于激活状态，当您停止输入时，会从当前光标处开始生成（右下角CodeGeeX图标转圈表示正在生成）。 生成完毕之后会以灰色显示，按``Tab``即可插入生成结果。 \n- **交互模式**: 按``Ctrl+Enter``激活交互模式，CodeGeeX将生成``X``个候选，并显示在右侧窗口中（``X`` 数量可以在设置的``Candidate Num``中修改）。 点击候选代码上方的``use code``即可插入。\n- **翻译模式**: 选择代码，然后按下``Ctrl+Alt+T``激活翻译模式，CodeGeeX会把该代码翻译成匹配您当前编辑器语言的代码。点击翻译结果上方的``use code``插入。您还可以在设置中选择您希望插入的时候如何处理被翻译的代码，您可以选择注释它们或者覆盖它们。\n- **提示模式（实验功能）**: 选择需要作为输入的代码，按``Alt/Option+t``触发提示模式，会显示预定义模板列表，选择其中一个模板，即可将代码插入到模板中进行生成。 这个模式高度自定义，可以在设置中 ``Prompt Templates``修改或添加模板内容，为模型加入额外的提示。 \n\n## 隐私声明\n\n我们高度尊重用户代码的隐私，代码仅用来辅助编程。在您第一次使用时，我们会询问您是否同意将生成的代码用于研究用途，帮助CodeGeeX变得更好（该选项默认**关闭**）。\n## 使用指南\n\n以下是CodeGeeX几种模式的详细用法：\n\n### 隐匿模式\n\n在该模式中，CodeGeeX将在您停止输入时，从光标处开始生成（右下角CodeGeeX图标转圈表示正在生成）。生成完毕之后会以灰色显示，按``Tab``即可插入生成结果。 在生成多个候选的情况下，可以使用``Alt/Option+[`` 或 ``]``在几个候选间进行切换。如果你对现有建议不满意，可以使用``Alt/Option+N``去获得新的候选。可以在设置中改变``Candidate Num``（增加个数会导致生成速度相对变慢）。**注意**：生成总是从当前光标位置开始，如果您在生成结束前移动光标位置，可能会导致一些bugs。我们正在努力使生成速度变得更快以提升用户体验。\n\n\n\n### 交互模式\n\n在该模式中，按``Ctrl+Enter``激活交互模式，CodeGeeX将生成``X``个候选，并显示在右侧窗口中（``X`` 数量可以在设置的``Candidate Num``中修改）。 点击候选代码上方的``use code``即可插入结果到为当前光标位置。 \n\n\n\n### 翻译模式\n\n在当前的语言的文本编辑器中输入或者粘贴其他语言的代码，您用鼠标选择这些代码，然后按下``Ctrl+Alt+T``激活翻译模式，您根据提示选择该代码的语言，然后CodeGeeX会帮您把该代码翻译成匹配您当前编辑器语言的代码。点击翻译结果上方的``use code``即可插入。您还可以在设置中选择您希望插入的时候如何处理被翻译的代码，您可以选择注释它们或者覆盖它们。\n\n\n\n### 提示模式（实验功能）\n\n在该模式中，您可以在输入中添加额外的提示来实现一些有趣的功能，包括并不限于代码解释、概括、以特定风格生成等。该模式的原理是利用了CodeGeeX强大的少样本生成能力。当您在输入中提供一些例子时，CodeGeeX会模仿这些例子并实现相应的功能。比如，您可以自定义模板中提供一段逐行解释代码的例子。选择您想要解释的代码，按``Alt/Option+t``触发提示模式，选择您写好的模板（如``explanation``），CodeGeeX就会解释您输入的代码。以下我们会详细介绍如何制作模板。\n\n\n\n上述例子中的模板如下图所示，由``[示例代码]``, ````, ``[带解释的示例代码]`` and ``[输出函数头]`` 。````表示您选中的代码将会插入的位置。 ```` 这一句用来保证模型解释的是同一个函数。当使用提示模式时，CodeGeeX会将您选择的代码（插入到部分）和模板代码相结合，一起作为模型的输入。 \n\n```python\n# language: Python\n\ndef sum_squares(lst):\n sum = 0\n for i in range(len(lst)):\n if i % 3 == 0:\n lst[i] = lst[i]**2\n elif i % 4 == 0:\n lst[i] = lst[i]**3\n sum += lst[i]\n return sum\n\n\n\n# Explain the code line by line\ndef sum_squares(lst):\n # initialize sum\n sum = 0\n # loop through the list\n for i in range(len(lst)):\n # if the index is a multiple of 3\n if i % 3 == 0:\n # square the entry\n lst[i] = lst[i]**2\n # if the index is a multiple of 4\n elif i % 4 == 0:\n # cube the entry\n lst[i] = lst[i]**3\n # add the entry to the sum\n sum += lst[i]\n # return the sum\n return sum\n\n# Explain the code line by line\n\n```\n\n以下是另一个Python文档字符串生成的例子，CodeGeeX在您写新函数时会模仿该注释的格式：\n```python\ndef add_binary(a, b):\n '''\n Returns the sum of two decimal numbers in binary digits.\n\n Parameters:\n a (int): A decimal integer\n b (int): Another decimal integer\n\n Returns:\n binary_sum (str): Binary string of the sum of a and b\n '''\n binary_sum = bin(a+b)[2:]\n return binary_sum\n\n\n```\n\n模板文件是高度自定义化的，您可以将自定义模板添加到插件设置中的``Prompt Templates``中。 ``key``表示模板的名字， ``value``是模板文件的路径（可以是您电脑上的任一路径，``.txt``, ``.py``, ``.h``, 等格式文件均可）。通过该功能，您可以让CodeGeeX生成具有特定风格或功能的代码，快尝试定义自己的专属模板吧！"
}
]
\n \"\"\")\n gr.Markdown(\n \"\"\"\n \n 🏠 Homepage | 📖 Blog | 🪧 DEMO | 🛠 VS Code or Jetbrains Extensions | 💻 Source code | 🤖 Download Model\n
\n \"\"\")\n gr.Markdown(\n \"\"\"\n We introduce CodeGeeX, a large-scale multilingual code generation model with 13 billion parameters, pre-trained on a large code corpus of more than 20 programming languages. CodeGeeX supports 15+ programming languages for both code generation and translation. CodeGeeX is open source, please refer to our [GitHub](https://github.com/THUDM/CodeGeeX) for more details. This is a minimal-functional DEMO, for other DEMOs like code translation, please visit our [Homepage](https://codegeex.cn). We also offer free [VS Code](https://marketplace.visualstudio.com/items?itemName=aminer.codegeex) or [Jetbrains](https://plugins.jetbrains.com/plugin/20587-codegeex) extensions for full functionality. \n \"\"\")\n\n with gr.Row():\n with gr.Column():\n prompt = gr.Textbox(lines=13, placeholder='Please enter the description or select an example input below.',label='Input')\n with gr.Row():\n gen = gr.Button(\"Generate\")\n clr = gr.Button(\"Clear\")\n\n outputs = gr.Textbox(lines=15, label='Output')\n\n gr.Markdown(\n \"\"\"\n Generation Parameter\n \"\"\")\n with gr.Row():\n with gr.Column():\n lang = gr.Radio(\n choices=[\"C++\", \"C\", \"C#\", \"Python\", \"Java\", \"HTML\", \"PHP\", \"JavaScript\", \"TypeScript\", \"Go\",\n \"Rust\",\n \"SQL\", \"Kotlin\", \"R\", \"Fortran\"], value='lang', label='Programming Language',\n default=\"Python\")\n with gr.Column():\n seed = gr.Slider(maximum=10000, value=8888, step=1, label='Seed')\n with gr.Row():\n out_seq_length = gr.Slider(maximum=1024, value=128, minimum=1, step=1, label='Output Sequence Length')\n temperature = gr.Slider(maximum=1, value=0.2, minimum=0, label='Temperature')\n with gr.Row():\n top_k = gr.Slider(maximum=40, value=0, minimum=0, step=1, label='Top K')\n top_p = gr.Slider(maximum=1, value=1.0, minimum=0, label='Top P')\n\n inputs = [prompt, lang, seed, out_seq_length, temperature, top_k, top_p]\n gen.click(fn=predict, inputs=inputs, outputs=outputs)\n clr.click(fn=lambda value: gr.update(value=\"\"), inputs=clr, outputs=prompt)\n\n gr_examples = gr.Examples(examples=examples, inputs=[prompt, lang],\n label=\"Example Inputs (Click to insert an examplet it into the input box)\",\n examples_per_page=20)\n \n demo.launch(server_port=6007)\n\nif __name__ == '__main__':\n with torch.no_grad():\n main()" }, { "path": "requirements.txt", "content": "fire>=0.4.0\nipython>=8.4.0\nnumpy>=1.22.0\npandas>=1.3.5\npyzmq>=23.2.1\nregex>=2022.3.15\nsetuptools>=58.0.4\ntransformers>=4.22.0\ntorch>=1.10.0\ntqdm>=4.63.0\ncpm_kernels\ndeepspeed>0.6.1" }, { "path": "resources/zh/wechat.md", "content": "\n
\n\n
\n"
},
{
"path": "scripts/convert_ckpt_parallel.sh",
"content": "# This script is used to convert checkpoint model parallel partitions.\n\nLOAD_CKPT_PATH=$1 # Path to weights in .pt format.\nSAVE_CKPT_PATH=$2 # Path to save the output MP checkpoints.\nMP_SIZE=$3 # Model parallel size\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\nif [ -z \"$MP_SIZE\" ]; then\n MP_SIZE=1\nfi\n\n# export CUDA settings\nexport CUDA_HOME=/usr/local/cuda-11.1/\nexport CUDA_VISIBLE_DEVICES=0,1\n\n\nCMD=\"python $MAIN_DIR/codegeex/megatron/convert_ckpt_parallel.py \\\n --load-ckpt-path $LOAD_CKPT_PATH \\\n --save-ckpt-path $SAVE_CKPT_PATH \\\n --tokenizer-path $TOKENIZER_PATH \\\n --target-tensor-model-parallel-size $MP_SIZE \\\n --num-layers 39 \\\n --hidden-size 5120 \\\n --num-attention-heads 40 \\\n --max-position-embeddings 2048 \\\n --attention-softmax-in-fp32 \\\n --fp16 \\\n --micro-batch-size 1 \\\n --make-vocab-size-divisible-by 52224 \\\n --seq-length 2048\"\n\necho \"$CMD\"\neval \"$CMD\""
},
{
"path": "scripts/convert_mindspore_to_megatron.sh",
"content": "# This script is used to convert mindspore checkpoint to the megatron format.\n\nNPY_CKPT_PATH=$1 # Path to Mindspore exported weights in .npy format.\nSAVE_CKPT_PATH=$2 # Path to save the output .pt checkpoint.\nGPU=$3\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\nTOKENIZER_PATH=\"$MAIN_DIR/codegeex/tokenizer/\"\n\n# export CUDA settings\nif [ -z \"$GPU\" ]; then\n GPU=0\nfi\n\nexport CUDA_HOME=/usr/local/cuda-11.1/\nexport CUDA_VISIBLE_DEVICES=$GPU\n\n\nCMD=\"python $MAIN_DIR/codegeex/megatron/mindspore_to_megatron.py \\\n --npy-ckpt-path $NPY_CKPT_PATH \\\n --save-ckpt-path $SAVE_CKPT_PATH \\\n --tokenizer-path $TOKENIZER_PATH \\\n $MODEL_ARGS\"\n\necho \"$CMD\"\neval \"$CMD\""
},
{
"path": "scripts/evaluate_humaneval_x.py",
"content": "import argparse\nimport os\nfrom pathlib import Path\nfrom codegeex.benchmark.evaluate_humaneval_x import evaluate_functional_correctness\n#GLOBALS\nINPUT_FILE: str \nLANGUAGE: str \nN_WORKERS: int \nTIMEOUT: int \n\n\nparser = argparse.ArgumentParser(\"Debugging evaluate humaneval_x\")\n# Path to the .jsonl file that contains the generated codes.\nparser.add_argument(\"-s\",\"--samples\", type=str)\n\n# Target programming language, currently support one of [\"python\", \"java\", \"cpp\", \"js\", \"go\"]\nparser.add_argument(\"-l\",\"--language\", default=\"python\", type=str)\n\n# Number of parallel workers.\nparser.add_argument(\"-w\",\"--workers\", default=64, type=int)\n\n# Timeout in seconds.\nparser.add_argument(\"-t\",\"--timeout\", default=5, type=int)\n\nargs = parser.parse_args()\n\nINPUT_FILE = args.samples\nLANGUAGE = args.language \nN_WORKERS = args.workers \nTIMEOUT= args.timeout\n\n\n\nSCRIPT_PATH: str = Path(os.path.abspath(__file__))\nprint(SCRIPT_PATH)\nSCRIPT_DIR: str = os.path.dirname(SCRIPT_PATH)\nprint(SCRIPT_DIR)\nMAIN_DIR: str = os.path.dirname(SCRIPT_DIR)\nprint(MAIN_DIR)\n\nDATA_DIR=os.path.join(MAIN_DIR,\"codegeex/benchmark/humaneval-x/\" + LANGUAGE + \"/data/humaneval_\" + LANGUAGE + \".jsonl.gz\")\nprint(DATA_DIR)\n\nTMP_DIR=os.path.join(MAIN_DIR, \"/codegeex/benchmark/humaneval-x/\")\n\n\n#Debugging\nINPUT_FILE='/home/rog0d/Escritorio/CodeGeeX/generations/humaneval_rust_generations.jsonl.gz'\nLANGUAGE='rust'\nDATA_DIR=os.path.join(MAIN_DIR,\"codegeex/benchmark/humaneval-x/\" + LANGUAGE + \"/data/humaneval_\" + LANGUAGE + \".jsonl.gz\")\n\n\"\"\"\ninput_file: str = None,\n tmp_dir: str = \"./\",\n n_workers: int = 32,\n timeout: float = 5.0,\n problem_file: str = \"../data/humaneval_python.jsonl.gz\",\n out_dir: str = None,\n k: List[int] = [1, 10, 100],\n test_groundtruth: bool = False,\n example_test: bool = False,\n\n\"\"\"\n\nevaluate_functional_correctness(input_file=INPUT_FILE,\n n_workers=N_WORKERS,\n tmp_dir=TMP_DIR,\n problem_file=DATA_DIR,\n timeout=300.0)\n\n\n"
},
{
"path": "scripts/evaluate_humaneval_x.sh",
"content": "# This script is for evaluating the functional correctness of the generated codes of HumanEval-X.\n\nINPUT_FILE=$1 # Path to the .jsonl file that contains the generated codes.\nLANGUAGE=$2 # Target programming language, currently support one of [\"python\", \"java\", \"cpp\", \"js\", \"go\"]\nN_WORKERS=$3 # Number of parallel workers.\nTIMEOUT=$4 # Timeout in seconds.\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\n\necho \"$INPUT_FILE\"\n\nif [ -z \"$N_WORKERS\" ]\nthen\n N_WORKERS=64\nfi\n\nif [ -z \"$LANGUAGE\" ]\nthen\n LANGUAGE=python\nfi\n\nif [ -z \"$TIMEOUT\" ]\nthen\n TIMEOUT=5\nfi\n\nDATA_DIR=$MAIN_DIR/codegeex/benchmark/humaneval-x/$LANGUAGE/data/humaneval_$LANGUAGE.jsonl.gz\n\nif [ $LANGUAGE = go ]; then\n export PATH=$PATH:/usr/local/go/bin\nfi\n\nif [ $LANGUAGE = cpp ]; then\n export PATH=$PATH:/usr/bin/openssl\nfi\n\nCMD=\"python $MAIN_DIR/codegeex/benchmark/humaneval-x/evaluate_humaneval_x.py \\\n --input_file \"$INPUT_FILE\" \\\n --n_workers $N_WORKERS \\\n --tmp_dir $MAIN_DIR/codegeex/benchmark/humaneval-x/ \\\n --problem_file $DATA_DIR \\\n --timeout $TIMEOUT\"\n\necho \"$CMD\"\neval \"$CMD\""
},
{
"path": "scripts/finetune_codegeex.sh",
"content": "SCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\n\n# ====== Environment ======\n# - NCCL & IB\nexport NCCL_DEBUG=info\nexport NCCL_IB_DISABLE=0\nexport NCCL_IB_GID_INDEX=3\n\nHOSTFILE=\"\n\n扫码关注公众号加入「CodeGeeX交流群」
\nScan the QR code to join the \"CodeGeeX WeChat Group\"
\n