Showing preview only (3,090K chars total). Download the full file or copy to clipboard to get everything.
Repository: zai-org/CodeGeeX
Branch: main
Commit: 2838420b7b44
Files: 189
Total size: 2.9 MB
Directory structure:
gitextract_96pnh25s/
├── .gitmodules
├── LICENSE
├── MODEL_LICENSE
├── README.md
├── README_zh.md
├── api/
│ ├── README_zh.md
│ ├── codegeex-api-example-java/
│ │ ├── pom.xml
│ │ └── src/
│ │ └── main/
│ │ └── java/
│ │ └── cn/
│ │ └── aminer/
│ │ └── codegeex/
│ │ └── example/
│ │ ├── CodeGenerationExample.java
│ │ └── pojo/
│ │ └── Payload.java
│ └── codegeex-api-example-python/
│ └── generation_example.py
├── codegeex/
│ ├── __init__.py
│ ├── benchmark/
│ │ ├── README.md
│ │ ├── README_zh.md
│ │ ├── __init__.py
│ │ ├── evaluate_humaneval_x.py
│ │ ├── execution.py
│ │ ├── gather_output.py
│ │ ├── humaneval-x/
│ │ │ ├── __init__.py
│ │ │ ├── cpp/
│ │ │ │ └── evaluation/
│ │ │ │ └── test.cpp
│ │ │ ├── evaluate_humaneval_x.py
│ │ │ ├── generate_humaneval_x.py
│ │ │ ├── go/
│ │ │ │ └── evaluation/
│ │ │ │ ├── go.mod
│ │ │ │ └── go.sum
│ │ │ └── translate_humaneval_x.py
│ │ ├── inspect_result.py
│ │ ├── metric.py
│ │ ├── rust/
│ │ │ └── Cargo.toml
│ │ └── utils.py
│ ├── data/
│ │ ├── __init__.py
│ │ ├── data_utils.py
│ │ ├── process_pretrain_dataset.py
│ │ ├── processor.py
│ │ └── types.py
│ ├── docker/
│ │ └── Dockerfile
│ ├── kernels/
│ │ ├── __init__.py
│ │ └── quantization.fatbin
│ ├── megatron/
│ │ ├── __init__.py
│ │ ├── arguments.py
│ │ ├── checkpointing.py
│ │ ├── code_generation_utils.py
│ │ ├── convert_ckpt_parallel.py
│ │ ├── data/
│ │ │ ├── __init__.py
│ │ │ ├── blendable_dataset.py
│ │ │ ├── data_samplers.py
│ │ │ ├── dataset_utils.py
│ │ │ ├── helpers.cpp
│ │ │ ├── indexed_dataset.py
│ │ │ └── prompt_dataset.py
│ │ ├── enums.py
│ │ ├── global_vars.py
│ │ ├── inference.py
│ │ ├── initialize.py
│ │ ├── learning_rates.py
│ │ ├── memory.py
│ │ ├── merge_ckpt_parallel.py
│ │ ├── microbatches.py
│ │ ├── mindspore_to_megatron.py
│ │ ├── model/
│ │ │ ├── __init__.py
│ │ │ ├── codegeex_model.py
│ │ │ ├── distributed.py
│ │ │ ├── language_model.py
│ │ │ ├── module.py
│ │ │ ├── transformer.py
│ │ │ └── utils.py
│ │ ├── mpu/
│ │ │ ├── __init__.py
│ │ │ ├── cross_entropy.py
│ │ │ ├── data.py
│ │ │ ├── initialize.py
│ │ │ ├── layers.py
│ │ │ ├── mappings.py
│ │ │ ├── random.py
│ │ │ └── utils.py
│ │ ├── optimizer/
│ │ │ ├── __init__.py
│ │ │ ├── clip_grads.py
│ │ │ ├── grad_scaler.py
│ │ │ └── optimizer.py
│ │ ├── p2p_communication.py
│ │ ├── schedules.py
│ │ ├── tokenizer/
│ │ │ ├── __init__.py
│ │ │ ├── gpt2_tokenization.py
│ │ │ └── tokenizer.py
│ │ ├── tools/
│ │ │ ├── collect_env.py
│ │ │ ├── finetune_codegeex.py
│ │ │ └── pretrain_codegeex.py
│ │ ├── training.py
│ │ └── utils.py
│ ├── mindspore/
│ │ ├── configs/
│ │ │ ├── 13B.sh
│ │ │ ├── 13B_128p_save_1p.sh
│ │ │ ├── 13B_128p_save_8p_ckpt.sh
│ │ │ ├── 13B_1p_to_torch.sh
│ │ │ ├── 13B_finetune.sh
│ │ │ ├── 13B_generate.sh
│ │ │ ├── 13B_generate_1p.sh
│ │ │ ├── 13B_generate_1p_values.sh
│ │ │ ├── 13B_generate_finetune.sh
│ │ │ ├── 13B_generate_humaneval.sh
│ │ │ └── 13B_generate_values.sh
│ │ ├── convertion_1p.py
│ │ ├── finetune.py
│ │ ├── generation.py
│ │ ├── generation_1p.py
│ │ ├── generation_batch.py
│ │ ├── generation_finetune.py
│ │ ├── generation_humaneval.py
│ │ ├── generation_values.py
│ │ ├── generation_values_1p.py
│ │ ├── save_1p_ckpt_from_8p_ckpt.py
│ │ ├── save_8p_ckpt.py
│ │ ├── scripts/
│ │ │ ├── custom_tune_bank_new/
│ │ │ │ └── Ascend910ProA/
│ │ │ │ ├── cube/
│ │ │ │ │ └── repository_ascend910ProA_matmul.json
│ │ │ │ └── vector/
│ │ │ │ └── Ascend910ProA_AiCore_32_v001_20220509_200939_588817.json
│ │ │ ├── layer_norm.py
│ │ │ ├── layer_norm_x_backprop_v2.py
│ │ │ ├── ma-pre-start.sh
│ │ │ ├── run_modelarts.py
│ │ │ ├── run_modelarts_gen_finetune.py
│ │ │ └── run_modelarts_gen_humaneval_x.py
│ │ ├── src/
│ │ │ ├── __init__.py
│ │ │ ├── adam.py
│ │ │ ├── callbacks.py
│ │ │ ├── code_tokenizer.py
│ │ │ ├── dataset.py
│ │ │ ├── dataset_finetune.py
│ │ │ ├── generate.py
│ │ │ ├── generate_finetune.py
│ │ │ ├── generate_greedy.py
│ │ │ ├── generate_humaneval.py
│ │ │ ├── metrics.py
│ │ │ ├── pangu_alpha.py
│ │ │ ├── pangu_alpha_config.py
│ │ │ ├── pangu_alpha_fp16_predict.py
│ │ │ ├── pangu_alpha_wrapcell.py
│ │ │ ├── pangu_alpha_wrapcell_finetune.py
│ │ │ ├── preprocess.py
│ │ │ ├── sat_dataset.py
│ │ │ ├── tokenization_jieba.py
│ │ │ └── utils.py
│ │ └── train.py
│ ├── oneflow/
│ │ ├── __init__.py
│ │ ├── codegeex_model.py
│ │ └── inference.py
│ ├── paddle/
│ │ ├── __init__.py
│ │ ├── codegeex_model.py
│ │ ├── inference.py
│ │ └── pt_to_pdparams.py
│ ├── quantization/
│ │ ├── __init__.py
│ │ ├── quantize.py
│ │ └── quantize_oneflow.py
│ ├── tokenizer/
│ │ ├── __init__.py
│ │ ├── added_tokens.json
│ │ ├── merges.txt
│ │ ├── special_tokens_map.json
│ │ ├── tokenizer.py
│ │ ├── tokenizer_config.json
│ │ └── vocab.json
│ └── torch/
│ ├── __init__.py
│ ├── codegeex_model.py
│ ├── get_ckpt_qkv.py
│ └── inference.py
├── configs/
│ ├── codegeex_13b.sh
│ ├── codegeex_13b_paddle.sh
│ └── codegeex_13b_parallel.sh
├── deployment/
│ ├── example_inputs.jsonl
│ └── server_gradio.py
├── requirements.txt
├── resources/
│ └── zh/
│ └── wechat.md
├── scripts/
│ ├── convert_ckpt_parallel.sh
│ ├── convert_mindspore_to_megatron.sh
│ ├── evaluate_humaneval_x.py
│ ├── evaluate_humaneval_x.sh
│ ├── finetune_codegeex.sh
│ ├── gather_output.sh
│ ├── generate_humaneval_x.sh
│ ├── pretrain_codegeex.sh
│ ├── process_pretrain_dataset.sh
│ ├── test_inference.sh
│ ├── test_inference_oneflow.sh
│ ├── test_inference_oneflow_quantized.sh
│ ├── test_inference_paddle.sh
│ ├── test_inference_parallel.sh
│ ├── test_inference_quantized.sh
│ └── translate_humaneval_x.sh
├── setup.py
├── tests/
│ ├── test_inference.py
│ ├── test_inference_megatron.py
│ ├── test_inference_oneflow.py
│ ├── test_inference_paddle.py
│ └── test_prompt.txt
└── vscode-extension/
├── README.md
└── README_zh.md
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitmodules
================================================
[submodule "vscode-extension/codegeex-vscode-extension"]
path = vscode-extension/codegeex-vscode-extension
url = git@github.com:CodeGeeX/codegeex-vscode-extension.git
================================================
FILE: LICENSE
================================================
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
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.
================================================
FILE: MODEL_LICENSE
================================================
The CodeGeeX License
1. Definitions
“Licensor” means the CodeGeeX Model Team that distributes its Software.
“Software” means the CodeGeeX model parameters made available under this license.
2. License Grant
Subject to the terms and conditions of this License, the Licensor hereby grants to you a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable, royalty-free copyright license to use the Software solely for your non-commercial research purposes.
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
3. Restriction
You will not use, copy, modify, merge, publish, distribute, reproduce, or create derivative works of the Software, in whole or in part, for any commercial, military, or illegal purposes.
You will not use the Software for any act that may undermine China's national security and national unity, harm the public interest of society, or infringe upon the rights and interests of human beings.
4. Disclaimer
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
5. Limitation of Liability
EXCEPT TO THE EXTENT PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER BASED IN TORT, NEGLIGENCE, CONTRACT, LIABILITY, OR OTHERWISE WILL ANY LICENSOR BE LIABLE TO YOU FOR ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES, OR ANY OTHER COMMERCIAL LOSSES, EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
6. Dispute Resolution
This license shall be governed and construed in accordance with the laws of People’s Republic of China. Any dispute arising from or in connection with this License shall be submitted to Haidian District People's Court in Beijing.
Note that the license is subject to update to a more comprehensive version. For any questions related to the license and copyright, please contact us at report@aminer.cn.
================================================
FILE: README.md
================================================
<img src="resources/logo/codegeex_logo.png">
<p align="center">
🏠 <a href="https://codegeex.cn" target="_blank">Homepage</a> | 📖 <a href="https://models.aminer.cn/codegeex/blog/" target="_blank">Blog</a> | 🪧 <a href="https://models.aminer.cn/codegeex/playground" target="_blank">DEMO</a> | 🤖 <a href="https://codegeex.cn/download/request" target="_blank">Download Model</a> | 📄 <a href="https://arxiv.org/abs/2303.17568" target="_blank">Paper</a> | 🌐 <a href="README_zh.md" target="_blank">中文</a>
</p>
<p align="center">
🛠 <a href="https://marketplace.visualstudio.com/items?itemName=aminer.codegeex" target="_blank">VS Code</a>, <a href="https://plugins.jetbrains.com/plugin/20587-codegeex" target="_blank">Jetbrains</a>, <a href="https://plugins.jetbrains.com/plugin/20587-codegeex" target="_blank">Cloud Studio</a> supported | 👋 Join our <a href="https://discord.gg/8gjHdkmAN6" target="_blank">Discord</a>, <a href="https://join.slack.com/t/codegeexworkspace/shared_invite/zt-1s118ffrp-mpKKhQD0tKBmzNZVCyEZLw" target="_blank">Slack</a>, <a href="https://t.me/+IipIayJ32B1jOTg1" target="_blank">Telegram</a>, <a href="resources/zh/wechat.md"target="_blank">WeChat</a>
</p>
🌟 The newest [CodeGeeX4](https://github.com/THUDM/CodeGeeX4) has been released. | 最新一代 [CodeGeeX4](https://github.com/THUDM/CodeGeeX4) 模型已经正式开源。
- [CodeGeeX: A Multilingual Code Generation Model](#codegeex-a-multilingual-code-generation-model)
- [News](#news)
- [Getting Started](#getting-started)
- [Installation](#installation)
- [Model Weights](#model-weights)
- [Inference on GPUs](#inference-on-gpus)
- [VS Code and Jetbrains Extension Guidance](#vs-code-and-jetbrains-extension-guidance)
- [CodeGeeX: Architecture, Code Corpus, and Implementation](#codegeex-architecture-code-corpus-and-implementation)
- [HumanEval-X: A new benchmark for Multilingual Program Synthesis](#humaneval-x-a-new-benchmark-for-multilingual-program-synthesis)
- [Multilingual Code Generation](#multilingual-code-generation)
- [Crosslingual Code Translation](#crosslingual-code-translation)
- [How to use HumanEval-X and contribute to it?](#how-to-use-humaneval-x-and-contribute-to-it)
- [License](#license)
- [Citation](#citation)
# CodeGeeX: A Multilingual Code Generation Model
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. As of **June 22**, 2022, CodeGeeX has been trained on more than 850 billion tokens on a cluster of 1,536 [Ascend 910 AI Processors](https://e.huawei.com/en/products/servers/ascend). CodeGeeX has several unique features:
* **Multilingual Code Generation**: CodeGeeX has good performance for generating executable programs in several mainstream programming languages, including Python, C++, Java, JavaScript, Go, etc. [DEMO](https://models.aminer.cn/codegeex)
* **Crosslingual Code Translation**: CodeGeeX supports the translation of code snippets between different languages. Simply by one click, CodeGeeX can transform a program into any expected language with a high accuracy. [DEMO](https://models.aminer.cn/codegeex/codeTranslator)
* **Customizable Programming Assistant**: CodeGeeX is available in the VS Code extension marketplace **for free**. It supports code completion, explanation, summarization and more, which empower users with a better coding experience. [VS Code Extension](https://marketplace.visualstudio.com/items?itemName=aminer.codegeex)
* **Open-Source and Cross-Platform**: All codes and model weights are publicly available for research purposes. CodeGeeX supports both Ascend and NVIDIA platforms. It supports inference in a single Ascend 910, NVIDIA V100 or A100. [Apply Model Weights](https://models.aminer.cn/codegeex/download/request)
**HumanEval-X for Realistic Multilingual Benchmarking.** To help standardize the evaluation of multilingual code generation and translation, we develop and release the **HumanEval-X** Benchmark. HumanEval-X is a new multilingual benchmark that contains **820 human-crafted** coding problems in **5** programming languages (Python, C++, Java, JavaScript, and Go), each of these problems is associated with tests and solutions. [Usage](codegeex/benchmark/README.md) [🤗 Available in HuggingFace](https://huggingface.co/datasets/THUDM/humaneval-x)
<img src="resources/en/hx_boxplot.png">
<p align="center"><i>CodeGeeX achieves the highest average performance compared with other open-sourced multilingual baselines.</i> </p>
## News
* 🌟 **2023-07-24**: [CodeGeeX2](https://github.com/THUDM/CodeGeeX2) has been released, more powerful, faster, and lightweight. Support 100+ languages and many new features.
* **2023-5-16**: CodeGeeX paper has been accepted by [KDD 2023, Long Beach](https://kdd.org/kdd2023/) and will be represented during the conference.
* **2023-03-30**: CodeGeeX paper is now available at [arxiv](https://arxiv.org/abs/2303.17568).
* **2023-02-14**: CodeGeeX now supports [Cloud Studio](https://cloudstudio.net/), a fantastic web IDE from Tencent. Click on the badge on top of this page to quickly launch an environment to test CodeGeeX.
* **2023-02-13**: Thanks a lot to [OneFlow](https://github.com/Oneflow-Inc/oneflow) team for adding oneflow backend for CodeGeeX's inference (Even faster than FasterTransformer under FP16!). Check more details [here](https://github.com/THUDM/CodeGeeX/pull/65).
* **2023-02**: We are hosting [CodeGeeX "Coding With AI" Hackathon](https://dorahacks.io/hackathon/codegeex/), design cool applications based on CodeGeeX and win prizes (RTX 4090, DJI drone, etc)!
* **2022-12-31**: We release the FasterTransformer version of CodeGeeX in [codegeex-fastertransformer](https://github.com/CodeGeeX/codegeex-fastertransformer). The INT8 accelerated version reaches an a verage speed of <15ms/token. Happy new year to everyone!
* **2022-12-13**: We release the source code of CodeGeeX VS Code extension in [codegeex-vscode-extension](https://github.com/CodeGeeX/codegeex-vscode-extension). Follow [QuickStart](https://github.com/CodeGeeX/codegeex-vscode-extension/blob/main/doc/quickstart.md) to start development.
* **2022-12-11**: CodeGeeX is now available for Jetbrains IDEs (IntelliJ IDEA, PyCharm, GoLand, CLion, etc), download it [here](https://plugins.jetbrains.com/plugin/20587-codegeex).
* **2022-12-04**: We release source code of quantization (requires less GPU RAM: 27GB -> 15GB) and model parallelism (possible to run on multiple GPUs with <8G RAM).
* **2022-09-30**: We release the cross-platform source code and models weights for both Ascend and NVIDIA platforms.
## Getting Started
CodeGeeX is initially implemented in Mindspore and trained Ascend 910 AI Processors. We provide a torch-compatible version based on [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) to facilitate usage on GPU platforms.
### Installation
Python 3.7+ / CUDA 11+ / PyTorch 1.10+ / DeepSpeed 0.6+ are required. Install ``codegeex`` package via:
```bash
git clone git@github.com:THUDM/CodeGeeX.git
cd CodeGeeX
pip install -e .
```
Or use [CodeGeeX docker](https://hub.docker.com/r/codegeex/codegeex) to quickly set up the environment (with [nvidia-docker](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker) installed):
```bash
docker pull codegeex/codegeex:latest
# To enable GPU support, clarify device ids with --device
docker run --gpus '"device=0,1"' -it --ipc=host --name=codegeex codegeex/codegeex
```
### Model Weights
Apply and download model weights through this [link](https://models.aminer.cn/codegeex/download/request). You'll receive by mail ```urls.txt``` that contains temporary download links. We recommend you to use [aria2](https://aria2.github.io/) to download it via the following command (Please make sure you have enough disk space to download the checkpoint (~26GB)):
```bash
aria2c -x 16 -s 16 -j 4 --continue=true -i urls.txt
```
Run the following command to get the full model weights:
```bash
cat codegeex_13b.tar.gz.* > codegeex_13b.tar.gz
tar xvf codegeex_13b.tar.gz
```
### Inference on GPUs
Have a try on generating the first program with CodeGeeX. First, specify the path of the model weights in ``configs/codegeex_13b.sh``. Second, write the prompt (natural language description or code snippet) into a file, e.g., ``tests/test_prompt.txt``, then run the following script:
```bash
# On a single GPU (with more than 27GB RAM)
bash ./scripts/test_inference.sh <GPU_ID> ./tests/test_prompt.txt
# With quantization (with more than 15GB RAM)
bash ./scripts/test_inference_quantized.sh <GPU_ID> ./tests/test_prompt.txt
# On multiple GPUs (with more than 6GB RAM, need to first convert ckpt to MP_SIZE partitions)
bash ./scripts/convert_ckpt_parallel.sh <LOAD_CKPT_PATH> <SAVE_CKPT_PATH> <MP_SIZE>
bash ./scripts/test_inference_parallel.sh <MP_SIZE> ./tests/test_prompt.txt
```
### VS Code and Jetbrains Extension Guidance
Based on CodeGeeX, we also develop free extentions for VS Code and Jetbrains IDEs, and more in the future.
For VS Code, search "codegeex" in Marketplace or install it [here](https://marketplace.visualstudio.com/items?itemName=aminer.codegeex). Detailed instructions can be found in
[VS Code Extension Guidance](vscode-extension/README.md). For developers, we have also released the source code in [codegeex-vscode-extension](https://github.com/CodeGeeX/codegeex-vscode-extension), please follow [QuickStart](https://github.com/CodeGeeX/codegeex-vscode-extension/blob/main/doc/quickstart.md) to start development.
For Jetbrains IDEs, search "codegeex" in Plugins or install it [here](https://plugins.jetbrains.com/plugin/20587-codegeex).
Make sure your IDE version is 2021.1 or later. CodeGeeX now supports IntelliJ IDEA, PyCharm, GoLand, CLion, Android Studio, AppCode, Aqua, DataSpell, DataGrip, Rider, RubyMine, and WebStorm.
## CodeGeeX: Architecture, Code Corpus, and Implementation
**Architecture**: CodeGeeX is a large-scale pre-trained programming language model based on transformers. It is a left-to-right autoregressive decoder, which takes code and natural language as input and predicts the probability of the next token. CodeGeeX contains 40 transformer layers with a hidden size of 5,120 for self-attention blocks and 20,480 for feed-forward layers, making its size reach 13 billion parameters. It supports a maximum sequence length of 2,048.
<img src="resources/en/codegeex_training.png">
<p align="center"><i><b>Left:</b> the proportion of programming languages in CodeGeeX's training data.
<b>Right:</b> the plot of training loss against the training steps of CodeGeeX.</i></p>
**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.
**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.
## HumanEval-X: A new benchmark for Multilingual Program Synthesis
To 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.
<img src="resources/en/hx_tasks.png">
<p align="center"><i>An illustration of tasks supported by <b>HumanEval-X</b>. Declarations, docstrings, and solutions are marked with red, green, and blue respectively. <b>Code generation</b> uses declaration and docstring as input, to generate solution. <b>Code translation</b> uses declaration in both languages and translate the solution in source language to the one in target language.</i></p>
In 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)$.
### Multilingual Code Generation
<img src="resources/en/hx_generattion_radar_horizon.png">
<p align="center"><i><b>Left</b>: the detailed pass@k (k=1,10,100) performance on code generation task for five languages in HumanEval-X. <b>Right</b>: 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.</i></p>
We 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.
### Crosslingual Code Translation
<img src="resources/en/hx_translation.png">
<p align="center"><i>Results on HumanEval-X <b>code translation</b> task. Best language-wise performance are <b>bolded</b>.</i></p>
We 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.
### How to use HumanEval-X and contribute to it?
For 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.
Please kindly let us know if you have any comment or suggestion, via [codegeex@aminer.cn](mailto:codegeex@aminer.cn).
<details>
<summary><b>Examples of Generation</b></summary>
<img src="resources/en/hx_examples.png">
</details>
<details>
<summary><b>Acknowledgement</b></summary>
<br/>
This project is supported by the National Science Foundation for Distinguished Young Scholars (No. 61825602).
### Lead Contributors
Qinkai Zheng ([Tsinghua KEG](http://keg.cs.tsinghua.edu.cn/glm-130b/)), Xiao Xia (Tsinghua KEG), Xu Zou (Tsinghua KEG)
### Contributors
Tsinghua KEG---The Knowledge Engineering Group at Tsinghua: Aohan Zeng, Wendi Zheng, Lilong Xue
Zhilin Yang's Group at Tsinghua IIIS: Yifeng Liu, Yanru Chen, Yichen Xu (BUPT, work was done when visiting Tsinghua)
Peng Cheng Laboratory: Qingyu Chen, Zhongqi Li, Gaojun Fan
Zhipu\.AI: Yufei Xue, Shan Wang, Jiecai Shan, Haohan Jiang, Lu Liu, Xuan Xue, Peng Zhang
Ascend and Mindspore Team: Yifan Yao, Teng Su, Qihui Deng, Bin Zhou
### Data Annotations
Ruijie Cheng (Tsinghua), Peinan Yu (Tsinghua), Jingyao Zhang (Zhipu\.AI), Bowen Huang (Zhipu\.AI), Shaoyu Wang (Zhipu\.AI)
### Advisors
[Zhilin Yang](https://kimiyoung.github.io/) (Tsinghua IIIS), Yuxiao Dong (Tsinghua KEG), Wenguang Chen (Tsinghua PACMAN), Jie Tang (Tsinghua KEG)
### Computation Sponsors
[Peng Cheng Laboratory](https://www.pcl.ac.cn/index.html)
[Zhipu.AI](https://www.zhipu.ai/)---an AI startup that aims to teach machines to think like humans
### Project Leader
[Jie Tang](http://keg.cs.tsinghua.edu.cn/jietang/) (Tsinghua KEG & BAAI)
</details>
## License
Our code is licensed under the [Apache-2.0 license](LICENSE).
Our model is licensed under the [license](MODEL_LICENSE).
## Citation
If you find our work useful, please cite:
```
@inproceedings{zheng2023codegeex,
title={CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X},
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},
booktitle={Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining},
pages={5673--5684},
year={2023}
}
```
================================================
FILE: README_zh.md
================================================
<img src="resources/logo/codegeex_logo.png">
<p align="center">
🏠 <a href="https://models.aminer.cn/codegeex/zh-CN" target="_blank">主页</a> | 📖 <a href="https://models.aminer.cn/codegeex/blog/index_zh.html" target="_blank">博客</a> | 🪧 <a href="https://models.aminer.cn/codegeex/zh-CN/playground" target="_blank">示例</a> | 🤖 <a href="https://models.aminer.cn/codegeex/download/request" target="_blank">模型下载</a> | 📄 <a href="https://arxiv.org/abs/2303.17568" target="_blank">论文</a> | 🌐 <a href="https://github.com/THUDM/CodeGeeX/blob/main/README.md" target="_blank">English</a>
</p>
<p align="center">
🛠 <a href="https://marketplace.visualstudio.com/items?itemName=aminer.codegeex" target="_blank">VS Code</a>, <a href="https://plugins.jetbrains.com/plugin/20587-codegeex" target="_blank">Jetbrains</a>, <a href="https://plugins.jetbrains.com/plugin/20587-codegeex" target="_blank">Cloud Studio</a> 插件 | 👋 欢迎加入 <a href="resources/zh/wechat.md"target="_blank">微信开发者交流群</a>
</p>
🌟 [CodeGeeX2](https://github.com/THUDM/CodeGeeX2) 已推出，更强，更快，更轻量。
- [CodeGeeX: 多语言代码生成模型](#codegeex-多语言代码生成模型)
- [新闻](#新闻)
- [使用指南](#使用指南)
- [安装](#安装)
- [模型权重](#模型权重)
- [用GPU进行推理](#用gpu进行推理)
- [插件使用指南](#插件使用指南)
- [CodeGeeX: 多语言代码生成模型](#codegeex-多语言代码生成模型-1)
- [国产平台实现与训练](#国产平台实现与训练)
- [HumanEval-X: 多语言代码生成基准](#humaneval-x-多语言代码生成基准)
- [多语言代码生成](#多语言代码生成)
- [跨语言代码翻译](#跨语言代码翻译)
- [许可证](#许可证)
- [引用](#引用)
# CodeGeeX: 多语言代码生成模型
CodeGeeX是一个具有130亿参数的多编程语言代码生成预训练模型。CodeGeeX采用华为MindSpore框架实现，在鹏城实验室“鹏城云脑II”中的192个节点（共1536个国产[昇腾910 AI处理器](https://e.huawei.com/cn/products/servers/ascend)）上训练而成。截至2022年6月22日，CodeGeeX历时两个月在20多种编程语言的代码语料库（>8500亿Token）上预训练得到。CodeGeeX有以下特点：
* **高精度代码生成**：支持生成Python、C++、Java、JavaScript和Go等多种主流编程语言的代码，在HumanEval-X代码生成任务上取得47%~60%求解率，较其他开源基线模型有更佳的平均性能。[代码生成示例](https://models.aminer.cn/codegeex/zh-CN)
* **跨语言代码翻译**：支持代码片段在不同编程语言间进行自动翻译转换，翻译结果正确率高，在HumanEval-X代码翻译任务上超越了其它基线模型。[代码翻译示例](https://models.aminer.cn/codegeex/zh-CN/codeTranslator)
* **自动编程插件**：CodeGeeX插件现已上架VSCode插件市场（完全免费），用户可以通过其强大的少样本生成能力，自定义代码生成风格和能力，更好辅助代码编写。[插件下载](https://marketplace.visualstudio.com/items?itemName=aminer.codegeex)
* **模型跨平台开源**: 所有代码和模型权重开源开放，用作研究用途。CodeGeeX同时支持昇腾和英伟达平台，可在单张昇腾910或英伟达V100/A100上实现推理。[申请模型权重](https://models.aminer.cn/codegeex/download/request)
**全新多编程语言评测基准HumanEval-X**：HumanEval-X是第一个支持功能正确性评测的多语言、多任务的基准，包含820个人工编写的高质量代码生成题目、测试用例与参考答案，覆盖5种编程语言（Python、C++、Java、JavaScript、Go），支持代码生成与代码翻译能力的评测。[如何使用](codegeex/benchmark/README_zh.md)
<img src="resources/zh/hx_boxplot_zh.png">
<p align="center"><i>在HumanEval-X代码生成任务上，与其它开源基线模型相比，CodeGeeX取得了最佳的平均性能。</i> </p>
## 新闻
* 🌟 **2023-07-24**: [CodeGeeX2](https://github.com/THUDM/CodeGeeX2) 已推出，更强，更快，更轻量。支持超过100种语言，并具有多种新特性。
* **2023-5-16**: CodeGeeX 论文已被 [KDD 2023, Long Beach](https://kdd.org/kdd2023/) 接收，并将在会上做相关报告。
* **2023-03-30**: CodeGeeX 论文已发表在[arxiv](https://arxiv.org/abs/2303.17568)。
* **2023-02-14**: CodeGeeX 现已支持 [Cloud Studio](https://cloudstudio.net/), 一款腾讯推出、十分好用的在线编辑器。单击此页面顶部的徽章可快速启动环境测试 CodeGeeX。
* **2023-02-13**: 感谢 [OneFlow](https://github.com/Oneflow-Inc/oneflow) 加入了oneflow版推理支持，在FP16下比FasterTransformer还要快！更多优化细节请点击[这里](https://github.com/THUDM/CodeGeeX/pull/65).
* **2023-02**: [CodeGeeX "Coding With AI"黑客松](https://dorahacks.io/hackathon/codegeex/)正在进行中，为CodeGeeX设计应用并赢取奖品（RTX 4090、DJI无人机等）！
* **2022-12-31**: 我们在 [codegeex-fastertransformer](https://github.com/CodeGeeX/codegeex-fastertransformer) 中发布了 CodeGeeX 的 FasterTransformer 版本。INT8加速版本达到 <15ms/token 的平均速度。祝大家新年快乐！
* **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) 开始开发吧！
* **2022-12-11**: CodeGeeX for Jetbrains IDEs已上线，支持IntelliJ IDEA, PyCharm, GoLand, CLion等，[点击下载](https://plugins.jetbrains.com/plugin/20587-codegeex)。
* **2022-12-04**: 我们开源了量化代码（需要更少的显存：27GB -> 15GB）以及模型并行代码（可以运行在多个显存至少8GB的GPUs上）。
* **2022-09-30**: 我们开源了跨平台代码和模型权重，同时支持昇腾和英伟达平台。
## 使用指南
CodeGeeX最初使用Mindspore框架实现，并在昇腾910AI芯片上进行训练。为适配更多平台，我们将其转换到[Megatron-LM](https://github.com/NVIDIA/Megatron-LM)框架，支持Pytorch+GPU环境。
### 安装
需要Python 3.7+ / CUDA 11+ / PyTorch 1.10+ / DeepSpeed 0.6+，通过以下命令安装 ``codegeex``:
```bash
git clone git@github.com:THUDM/CodeGeeX.git
cd CodeGeeX
pip install -e .
```
### 模型权重
通过[该链接](https://models.aminer.cn/codegeex/download/request)申请权重，您将收到一个包含临时下载链接文件```urls.txt```的邮件。推荐使用[aria2](https://aria2.github.io/)通过以下命令快速下载（请保证有足够的硬盘空间存放权重（～26GB））：
```bash
aria2c -x 16 -s 16 -j 4 --continue=true -i urls.txt
```
使用以下命令合并得到完整的权重：
```bash
cat codegeex_13b.tar.gz.* > codegeex_13b.tar.gz
tar xvf codegeex_13b.tar.gz
```
### 用GPU进行推理
尝试使用CodeGeeX模型生成第一个程序吧！首先，在配置文件``configs/codegeex_13b.sh``中写明存放权重的路径。其次，将提示（可以是任意描述或代码片段）写入文件``tests/test_prompt.txt``，运行以下脚本即可开始推理（需指定GPU序号）：
```bash
# On a single GPU (with more than 27GB RAM)
bash ./scripts/test_inference.sh <GPU_ID> ./tests/test_prompt.txt
# With quantization (with more than 15GB RAM)
bash ./scripts/test_inference_quantized.sh <GPU_ID> ./tests/test_prompt.txt
# On multiple GPUs (with more than 6GB RAM, need to first convert ckpt to MP_SIZE partitions)
bash ./scripts/convert_ckpt_parallel.sh <LOAD_CKPT_PATH> <SAVE_CKPT_PATH> <MP_SIZE>
bash ./scripts/test_inference_parallel.sh <MP_SIZE> ./tests/test_prompt.txt
```
### 插件使用指南
基于CodeGeeX，我们开发了免费的插件，支持 VS Code 与 Jetbrains IDEs，未来会支持更多平台。
VS Code版本，在应用市场搜索“codegeex”或通过[该链接](https://marketplace.visualstudio.com/items?itemName=aminer.codegeex)安装。详细的使用指南在[CodeGeeX VS Code插件使用指南](vscode-extension/README_zh.md)。我们也开源了VS Code插件源码：[codegeex-vscode-extension](https://github.com/CodeGeeX/codegeex-vscode-extension)，参考[QuickStart](https://github.com/CodeGeeX/codegeex-vscode-extension/blob/main/doc/quickstart_zh.md) 开始开发吧！
Jetbrains版本，在Plugins市场搜索“codegeex”或通过[该链接](https://plugins.jetbrains.com/plugin/20587-codegeex)安装。
请确保IDE版本在2021.1或更高。CodeGeeX目前支持 IntelliJ IDEA, PyCharm, GoLand, CLion, Android Studio, AppCode, Aqua, DataSpell, DataGrip, Rider, RubyMine, WebStorm。
## CodeGeeX: 多语言代码生成模型
**架构**：CodeGeeX是一个基于transformers的大规模预训练编程语言模型。它是一个从左到右生成的自回归解码器，将代码或自然语言标识符（token）作为输入，预测下一个标识符的概率分布。CodeGeeX含有40个transformer层，每层自注意力块的隐藏层维数为5120，前馈层维数为20480，总参数量为130亿。模型支持的最大序列长度为2048。
<img src="resources/en/codegeex_training.png">
<p align="center"><i><b>左侧：</b>CodeGeeX训练数据中各编程语言占比。
<b>右侧：</b>CodeGeeX训练损失函数随训练步数下降曲线。</i></p>
**语料**：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亿个标识符（不含填充符）。
### 国产平台实现与训练
我们在[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优化器降低显存占用。
在开发与训练过程中，我们和华为Mindspore团队合作，对MindSpore框架进行了部分优化，进而大幅度提升训练效率。比如，我们发现矩阵乘法的计算时间占比仅为22.9%，大量时间被用于各类其它算子，因此实现了一系列算子融合，包括单元素算子融合、层归一化算子融合、FastGelu与矩阵乘法融合、批量矩阵乘法与加法融合等；再比如我们还对矩阵乘法算子的维度实现自动搜索调优，使其搜索出效率最高的计算维度组合。这些优化为训练速度带来了显著提升，在同等GPU卡数规模下（128卡），昇腾910对CodeGeeX这一模型的训练效率从约为NVIDIA A100的16.7%提升至43%；在千卡规模下，昇腾910训练效率相比自身优化前提升近300%。使用优化后的软硬件训练时，CodeGeeX单日训练量可达到54.3B个标识符（含填充符），证明了国产深度学习平台与工具的快速迭代能力以及强大竞争力。
## HumanEval-X: 多语言代码生成基准
为了更好地评测代码生成模型的多语言生成能力，我们构建了一个新基准HumanEval-X。此前，多语言代码生成能力是基于语义相似度（比如[CodeBLEU](https://arxiv.org/abs/2009.10297)）衡量的，具有一定误导性；HumanEval-X则可用于衡量生成代码的功能正确性。HumanEval-X包含820个高质量手写样本，覆盖Python、C++、Java、JavaScript、Go，可用于多种任务。
<img src="resources/zh/hx_tasks_zh.png">
<p align="center"><i><b>HumanEval-X</b>支持的任务示例。<font style='background-color:#F8CECC'>声明</font>、<font style='background-color:#D5E8D4'>描述</font>、<font style='background-color:#DAE8FC'>解答</font>分别用红、绿、蓝色标注。<i>代码生成</i>将声明与描述作为输入，输出解答。<i>代码翻译</i>将两种语言的声明与源语言的解答作为输入，输出目标语言的解答。</i></p>
HumanEval-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)$.
### 多语言代码生成
<img src="resources/zh/hx_generattion_radar_horizon_zh.png">
<p align="center"><i><b>左侧</b>: HumanEval-X中五种语言具体的pass@k（k=1,10,100）性能。<b>右侧</b>: 模型在所有语言上的平均性能。CodeGeeX的平均表现优于InCoder-6.7B和CodeGen-Multi-6B/16B。</i></p>
我们将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%）。
### 跨语言代码翻译
<img src="resources/zh/hx_translation_zh.png">
<p align="center"><i>HumanEval-X上的<b>代码翻译</b>任务结果。<b>加粗</b>结果表示在每种语言pass@k上的最佳效果。</i></p>
我们还评测了模型在多语言间代码翻译上的性能。对于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个翻译对中，我们还观察到两种语言互相翻译的表现常常是呈负相关的，这可能说明现有的模型还不足以学好所有的语言。
<details>
<summary><b>在线生成与翻译DEMO</b></summary>
<img src="resources/en/hx_examples.png">
我们为上述两个任务开发了DEMO：<a href="https://models.aminer.cn/codegeex/zh-CN/playground" target="_blank">代码生成</a>和<a href="https://models.aminer.cn/codegeex/zh-CN/codeTranslator" target="_blank">代码翻译</a>，欢迎点击体验！
</details>
<details>
<summary><b>致谢</b></summary>
<br/>
这一项目由国家自然科学基金杰出青年科学基金项目（No. 61825602）支持。
#### 学生负责人
郑勤锴（[清华大学知识工程实验室](https://keg.cs.tsinghua.edu.cn/glm-130b/zh/posts/glm-130b/)），夏箫（清华大学知识工程实验室），邹旭（清华大学知识工程实验室）
#### 技术贡献
清华大学知识工程实验室：曾奥涵，郑问迪，薛理龙
清华大学交叉信息学院：刘益枫，陈彦儒，徐奕辰（北邮大四访问清华期间研究工作）
鹏城实验室：陈庆玉，李忠琦，范高俊
智谱AI：薛宇飞，王山，陕杰才，姜皓瀚，刘璐，薛旋，张鹏
华为昇腾团队：姚逸璠，苏腾，邓启辉，周斌
#### 数据标注
程锐杰（清华大学），于沛楠（清华大学），张竞尧（智谱AI），黄铂文（智谱AI），王炤宇（智谱AI）
#### 指导教师
[杨植麟](https://kimiyoung.github.io/)（清华大学交叉信息学院），东昱晓（清华大学知识工程实验室），陈文光（清华大学PACMAN实验室），[唐杰](http://keg.cs.tsinghua.edu.cn/jietang/)（清华大学知识工程实验室）
#### 计算资源支持
[鹏城实验室](https://www.pcl.ac.cn/index.html)
[智谱AI](https://www.zhipu.ai/)
#### 项目总负责
[唐杰](http://keg.cs.tsinghua.edu.cn/jietang/)（清华大学知识工程实验室 & 北京智源人工智能研究院）
</details>
如果遇到问题或有任何建议，欢迎通过邮件与我们联系[codegeex@aminer.cn](mailto:codegeex@aminer.cn).
## 许可证
代码使用[Apache-2.0许可证](LICENSE)
模型使用[许可证](MODEL_LICENSE)
## 引用
如果觉得我们的工作有帮助，欢迎引用以下论文：
```
@inproceedings{zheng2023codegeex,
title={CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X},
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},
booktitle={Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining},
pages={5673--5684},
year={2023}
}
```
================================================
FILE: api/README_zh.md
================================================
# 创建CodeGeeX API
使用[天启 · API开放平台](https://tianqi.aminer.cn/open/)申请CodeGeeX API：
<img src="../resources/api/api_step_1.png">
点击首页中的天启平台体验入口：
<img src="../resources/api/api_step_2.png">
点击API应用：
<img src="../resources/api/api_step_3.png">
输入任意名称，创建API应用。创建后会得到API Key/Secret，用于调用API：
<img src="../resources/api/api_step_4.png">
在API信息中，可以查看代码生成/代码翻译的请求地址和使用文档：
<img src="../resources/api/api_step_5.png">
根据文档中的描述使用API，Python版参考目录``api/codegeex-api-example-python``；JAVA版参考工程：``api/codegeex-api-example-java``
================================================
FILE: api/codegeex-api-example-java/pom.xml
================================================
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<groupId>cn.aminer</groupId>
<artifactId>codegeex-api-example-java</artifactId>
<version>1.0-SNAPSHOT</version>
<!--
如果没有下面的这段encoding配置，会导致编译的时候输出WARNING信息：
[WARNING] Using platform encoding (UTF-8 actually) to copy filtered resources, i.e. build is platform dependent!
-->
<properties>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
</properties>
<build>
<plugins>
<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-compiler-plugin</artifactId>
<version>3.8.1</version>
<configuration>
<!-- 一般而言，target与source是保持一致的，但是，有时候为了让程序能在其他版本的jdk中运行(对于低版本目标jdk，源代码中不能使用
低版本jdk中不支持的语法)，会存在target不同于source的情况 -->
<source>1.8</source><!-- 源代码使用的JDK版本 -->
<target>1.8</target><!-- 需要生成的目标class文件的编译版本 -->
<encoding>UTF-8</encoding><!-- 字符集编码 -->
</configuration>
</plugin>
<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-assembly-plugin</artifactId>
<version>3.3.0</version>
<configuration>
<descriptorRefs>
<descriptorRef>jar-with-dependencies</descriptorRef>
</descriptorRefs>
</configuration>
<!-- 有下面这段executions，才能打出包含所有dependency的fat jar -->
<executions>
<execution>
<phase>package</phase><!-- 指定在打包节点执行jar包合并操作 -->
<goals>
<goal>single</goal><!-- 该模块只运行一次 -->
</goals>
</execution>
</executions>
</plugin>
</plugins>
</build>
<dependencies>
<dependency>
<groupId>com.fasterxml.jackson.module</groupId>
<artifactId>jackson-module-parameter-names</artifactId>
<version>2.6.6</version>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.datatype</groupId>
<artifactId>jackson-datatype-jdk8</artifactId>
<version>2.6.6</version>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.datatype</groupId>
<artifactId>jackson-datatype-jsr310</artifactId>
<version>2.6.6</version>
</dependency>
<dependency>
<groupId>com.squareup.okhttp3</groupId>
<artifactId>okhttp</artifactId>
</dependency>
<dependency>
<groupId>org.slf4j</groupId>
<artifactId>slf4j-log4j12</artifactId>
</dependency>
<dependency>
<groupId>log4j</groupId>
<artifactId>log4j</artifactId>
</dependency>
<dependency>
<groupId>org.projectlombok</groupId>
<artifactId>lombok</artifactId>
<scope>provided</scope>
</dependency>
</dependencies>
<dependencyManagement>
<dependencies>
<dependency>
<groupId>com.fasterxml.jackson.module</groupId>
<artifactId>jackson-module-parameter-names</artifactId>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.datatype</groupId>
<artifactId>jackson-datatype-jdk8</artifactId>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.datatype</groupId>
<artifactId>jackson-datatype-jsr310</artifactId>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-databind</artifactId>
</dependency>
<dependency>
<groupId>com.squareup.okhttp3</groupId>
<artifactId>okhttp</artifactId>
<version>4.10.0</version>
</dependency>
<dependency>
<groupId>log4j</groupId>
<artifactId>log4j</artifactId>
<version>1.2.17</version>
</dependency>
<dependency>
<groupId>org.slf4j</groupId>
<artifactId>slf4j-log4j12</artifactId>
<version>1.7.5</version>
</dependency>
<dependency>
<groupId>org.projectlombok</groupId>
<artifactId>lombok</artifactId>
<version>1.18.20</version>
<scope>provided</scope>
</dependency>
</dependencies>
</dependencyManagement>
<repositories>
<repository>
<!-- Maven 自带的中央仓库使用的id为central，如果其他的仓库声明也是用该id，就会覆盖中央仓库的配置 -->
<id>central</id>
<name>ALiYun</name>
<url>http://maven.aliyun.com/nexus/content/groups/public</url>
</repository>
</repositories>
</project>
================================================
FILE: api/codegeex-api-example-java/src/main/java/cn/aminer/codegeex/example/CodeGenerationExample.java
================================================
package cn.aminer.codegeex.example;
import cn.aminer.codegeex.example.pojo.Payload;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.*;
import java.io.IOException;
/**
* 调用 CodeGeeX API 生成代码的例子。
*
* @author Darran Zhang @ codelast.com
* @version 2023-01-20
*/
public class CodeGenerationExample {
public static final String API_KEY = "your_api_key"; // 在"天启开放平台"上申请到的API Key
public static final String API_SECRET = "your_api_secret"; // 在"天启开放平台"上申请到的API Secret
public static final int NUMBER = 3; // 生成几个候选
public static final String LANGUAGE = "Java"; // 编程语言
public static final String REQUEST_URL = "https://tianqi.aminer.cn/api/v2/multilingual_code_generate"; // 请求地址
public static void main(String[] args) throws Exception {
CodeGenerationExample example = new CodeGenerationExample();
String prompt = "// use OkHttpClient library to write a function to perform http post request\n\n" +
"public class HttpPost {\n" +
" public static void main(String[] args) {\n";
example.generateCode(prompt);
}
/**
* 生成代码。
*
* @param prompt 待补全的代码
*/
public void generateCode(String prompt) throws Exception {
ObjectMapper objectMapper = new ObjectMapper();
Payload payload = new Payload().setApiKey(API_KEY).setApiSecret(API_SECRET).setPrompt(prompt).setNumber(NUMBER)
.setLanguage(LANGUAGE);
String response = performHttpPost(REQUEST_URL, objectMapper.writeValueAsString(payload));
System.out.println(response);
}
/**
* 发起 HTTP POST 请求。
*
* @param url 请求的URL
* @param payload 请求的JSON数据
* @return 请求返回的内容，若出错则返回 null。
*/
public String performHttpPost(String url, String payload) {
HttpUrl.Builder builder = null;
try {
HttpUrl httpUrl = HttpUrl.parse(url);
if (httpUrl != null) {
builder = httpUrl.newBuilder();
}
} catch (IllegalArgumentException e) {
System.out.println("failed to create HttpUrl.Builder from url " + url + ":" + e);
}
if (builder == null) {
return null;
}
OkHttpClient client = new OkHttpClient();
RequestBody requestBody = RequestBody.create(payload, MediaType.parse("application/json; charset=utf-8"));
Request request = new Request.Builder()
.url(builder.build())
.post(requestBody)
.build();
try {
Response response = client.newCall(request).execute();
ResponseBody body = response.body();
if (body == null) {
System.out.println("null response body");
return null;
}
return body.string();
} catch (IOException e) {
System.out.println("failed to send POST request: " + e);
}
return null;
}
}
================================================
FILE: api/codegeex-api-example-java/src/main/java/cn/aminer/codegeex/example/pojo/Payload.java
================================================
package cn.aminer.codegeex.example.pojo;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
import lombok.Data;
import lombok.experimental.Accessors;
/**
* 发送到 CodeGeex API 的请求中包含的JSON payload对象。
*
* @author Darran Zhang @ codelast.com
* @version 2023-01-20
*/
@JsonIgnoreProperties(ignoreUnknown = true)
@Data
@Accessors(chain = true)
public class Payload {
@JsonProperty("apikey")
String apiKey; // 在"天启开放平台"上申请到的API Key
@JsonProperty("apisecret")
String apiSecret; // 在"天启开放平台"上申请到的API Secret
String prompt; // 待补全的代码
@JsonProperty("n")
int number; // 生成几个候选
@JsonProperty("lang")
String language; // 编程语言
}
================================================
FILE: api/codegeex-api-example-python/generation_example.py
================================================
# encoding:utf-8
import json
import requests
'''
Code Generation
'''
API_KEY = "" # Get from Tianqi console. 从控制台获取
API_SECRET = "" # Get from Tianqi console. 从控制台获取
PROMPT = "from typing import List\n\ndef has_close_elements(numbers: List[float], threshold: float) -> bool:\n " \
"\"\"\" Check if in given list of numbers, are any two numbers closer to each other than\n given " \
"threshold.\n >>> has_close_elements([1.0, 2.0, 3.0], 0.5)\n False\n >>> has_close_elements(" \
"[1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)\n True\n \"\"\"\n"
NUMBER = 3
LANG = "Python"
request_url = "https://tianqi.aminer.cn/api/v2/"
api = 'multilingual_code_generate'
# Request is in json format. 指定请求参数格式为json
headers = {'Content-Type': 'application/json'}
request_url = request_url + api
data = {
"apikey": API_KEY,
"apisecret": API_SECRET,
"prompt": PROMPT,
"n": NUMBER,
"lang": LANG
}
def main():
response = requests.post(request_url, headers=headers, data=json.dumps(data))
if response:
print(response.json())
if __name__ == '__main__':
main()
================================================
FILE: codegeex/__init__.py
================================================
import copy
from typing import *
from codegeex.tokenizer import CodeGeeXTokenizer
from codegeex.torch.inference import get_token_stream
def get_model(
backend: str = "megatron",
quantized: bool = False,
):
pass
def generate(
model,
tokenizer: CodeGeeXTokenizer,
prompt: str,
out_seq_length: int,
seq_length: int = 2048,
top_k: int = 0,
top_p: float = 1.0,
temperature: float = 1.0,
micro_batch_size: int = 1,
backend: str = "megatron",
greedy: bool = False,
verbose: bool = False,
):
tokens = tokenizer.encode_code(prompt)
n_token_prompt = len(tokens)
if verbose:
print(f"Current prompt:\n{prompt}")
print("N_token_prompt:", n_token_prompt)
generated_codes = []
if backend == "megatron":
token_stream = get_token_stream(
model,
tokenizer,
seq_length,
out_seq_length,
[copy.deepcopy(tokens) for _ in range(micro_batch_size)],
micro_batch_size=micro_batch_size,
topk=top_k,
topp=top_p,
temperature=temperature,
greedy=greedy,
)
is_finished = [False for _ in range(micro_batch_size)]
for i, generated in enumerate(token_stream):
generated_tokens = generated[0]
for j in range(micro_batch_size):
if is_finished[j]:
continue
if generated_tokens[j].cpu().numpy()[-1] == tokenizer.eos_token_id or len(generated_tokens[j]) >= out_seq_length:
is_finished[j] = True
generated_tokens_ = generated_tokens[j].cpu().numpy().tolist()
generated_code = tokenizer.decode_code(generated_tokens_[n_token_prompt:])
generated_code = "".join(generated_code)
generated_codes.append(generated_code)
if verbose:
print(f"\nGenerated code {i}:\n{generated_code}")
if all(is_finished):
break
return generated_codes
================================================
FILE: codegeex/benchmark/README.md
================================================
# HumanEval-X: A new benchmark for Multilingual Program Synthesis
🌐 <a href="README_zh.md" target="_blank">中文</a>
HumanEval-X is a new benchmark for better evaluating the multilingual ability of code generation models. 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.
<img src="../../resources/en/hx_tasks.png">
<p align="center"><i>An illustration of tasks supported by <b>HumanEval-X</b>. Declarations, docstrings, and solutions are marked with red, green, and blue respectively. <b>Code generation</b> uses declaration and docstring as input, to generate solution. <b>Code translation</b> uses declaration in both languages and translate the solution in source language to the one in target language.</i></p>
In 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)$.
## How to use HumanEval-X
Data are stored in ``codegeex/benchmark/humaneval-x/[LANG]/data/humaneval_[LANG].jsonl.gz``, using JSON list format. There are six keys:
* ``task_id``: indicates the target language and ID of the problem. Language is one of ["Python", "Java", "JavaScript", "CPP", "Go"].
* ``prompt``: the function declaration and docstring, used for code generation.
* ``declaration``: only the function declaration, used for code translation.
* ``canonical_solution``: human-crafted example solutions.
* ``test``: hidden test samples, used for evaluation.
* ``example_test``: public test samples (appeared in prompt), used for evaluation.
### Evaluation Environment
The 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:
| Dependency | Version |
| ---------- | -------- |
| Python | 3.8.12 |
| JDK | 18.0.2.1 |
| Node.js | 16.14.0 |
| js-md5 | 0.7.3 |
| C++ | 11 |
| g++ | 7.5.0 |
| Boost | 1.71.0 |
| OpenSSL | 3.0.0 |
| go | 1.18.4 |
In 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.
You can directly pull the image from Docker Hub:
```bash
docker pull rishubi/codegeex:latest
```
Alternatively, if you are familiar with Dockerfile, you can build the image from `codegeex/docker/Dockerfile` or configure the Dockerfile as you like it:
```bash
cd codegeex/docker
docker build [OPTIONS] .
```
After obtaining the image, you can build a container using the following command:
```bash
docker run -it --gpus all --mount type=bind,source=<LOCAL PATH>,target=<PATH IN CONTAINER> [OPTIONS] <IMAGE NAME:TAG>
```
### Evaluation
We recommend evaluating in [the provided image](#evaluation-environment). To evaluate the generated samples, save generated codes in the following JSON list format:
```
{"task_id": "../..", "generation: "..."}
{"task_id": "../..", "generation: "..."}
...
```
and evaluate them using the following script under the root directory of the repository (<font color='red'>please execute with caution, the generated codes might have unexpected behaviours though with very low possibility. See the warnings in [execution.py](execution.py) and uncomment the execution lines at your own risk</font>):
```bash
bash scripts/evaluate_humaneval_x.sh <RESULT_FILE> <LANG> <N_WORKERS>
```
================================================
FILE: codegeex/benchmark/README_zh.md
================================================
# HumanEval-X: 多语言代码生成基准
🌐 <a href="README.md" target="_blank">English</a>
为了更好地评测代码生成模型的多语言生成能力，我们构建了一个新基准HumanEval-X。此前，多语言代码生成能力是基于语义相似度（比如[CodeBLEU](https://arxiv.org/abs/2009.10297)）衡量的，具有一定误导性；HumanEval-X则可用于衡量生成代码的功能正确性。HumanEval-X包含820个高质量手写样本，覆盖Python、C++、Java、JavaScript、Go，可用于多种任务。
<img src="../../resources/en/hx_tasks.png">
<p align="center"><i><b>HumanEval-X</b>支持的任务示例。<font style='background-color:#F8CECC'>声明</font>、<font style='background-color:#D5E8D4'>描述</font>、<font style='background-color:#DAE8FC'>解答</font>分别用红、绿、蓝色标注。<i>代码生成</i>将声明与描述作为输入，输出解答。<i>代码翻译</i>将两种语言的声明与源语言的解答作为输入，输出目标语言的解答。</i></p>
HumanEval-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)$。
## 如何使用HumanEval-X
样本使用JSON列表格式存储在``codegeex/benchmark/humaneval-x/[LANG]/data/humaneval_[LANG].jsonl.gz``，每条样本包含6个部分：
* ``task_id``: 题目的目标语言与ID。语言为["Python", "Java", "JavaScript", "CPP", "Go"]中之一。
* ``prompt``: 函数声明与描述，用于代码生成。
* ``declaration``: 仅有函数声明，用于代码翻译。
* ``canonical_solution``: 手写的示例解答。
* ``test``: 隐藏测例，用于评测。
* ``example_test``: 提示中出现的公开测例，用于评测。
### 评测环境
评测生成的代码需要使用多种语言编译、运行。我们使用的各编程语言依赖及所用包的版本如下：
| 依赖 | 版本 |
| ------- | -------- |
| Python | 3.8.12 |
| JDK | 18.0.2.1 |
| Node.js | 16.14.0 |
| js-md5 | 0.7.3 |
| C++ | 11 |
| g++ | 7.5.0 |
| Boost | 1.71.0 |
| OpenSSL | 3.0.0 |
| go | 1.18.4 |
为了省去使用者配置这些语言环境的麻烦，我们构建了一个Docker镜像，并在其中配置了所需要的环境。
可以直接从Docker Hub拉取镜像：
```bash
docker pull rishubi/codegeex:latest
```
如果您熟悉Dockerfile，也可以从`codegeex/docker/Dockerfile`构建镜像，或者修改之以定制自己的配置：
```bash
cd codegeex/docker
docker build [OPTIONS] .
```
获取镜像后，使用如下命令创建容器：
```bash
docker run -it --gpus all --mount type=bind,source=<LOCAL PATH>,target=<PATH IN CONTAINER> [OPTIONS] <IMAGE NAME:TAG>
```
### 评测
我们推荐使用给定的[评测环境](#评测环境)进行评测。在评测前，将生成的代码以如下JSON列表形式存储：
```
{"task_id": "../..", "generation: "..."}
{"task_id": "../..", "generation: "..."}
...
```
并在本仓库的根目录下使用如下指令（<font color='red'>请谨慎执行，生成的代码可能有极低概率产生意外行为。在[execution.py](execution.py)中查看警告并取消执行代码的注释，风险自负</font>）：
```bash
bash scripts/evaluate_humaneval_x.sh <RESULT_FILE> <LANG> <N_WORKERS>
```
================================================
FILE: codegeex/benchmark/__init__.py
================================================
================================================
FILE: codegeex/benchmark/evaluate_humaneval_x.py
================================================
import os
import sys
import fire
import json
import gzip
import regex
import numpy as np
from typing import *
from tqdm.auto import tqdm
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from codegeex.benchmark.utils import read_dataset, IMPORT_HELPER
from codegeex.benchmark.metric import estimate_pass_at_k
from codegeex.benchmark.execution import check_correctness
LANGUAGE_NAME = {
"cpp" : "CPP",
"go" : "Go",
"java" : "Java",
"js" : "JavaScript",
"python": "Python",
}
def process_humaneval_test(sample, problems, example_test=False):
task_id = sample["task_id"]
language = task_id.split("/")[0].lower()
prompt = sample["prompt"]
if example_test and "example_test" in problems[task_id] and problems[task_id]["example_test"] != "":
test = problems[task_id]["example_test"]
else:
test = problems[task_id]["test"]
code = sample["generation"]
# Pre-process for different languages
if language == "python":
code_ = []
for line in code.split("\n"):
if (len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t'):
break
code_.append(line)
code = "\n".join(code_)
test_setup = "\n".join(IMPORT_HELPER["python"]) + "\n"
test_string = test_setup + prompt + code + "\n" + test + "\n"
elif language == "cpp":
test_set_up = ""
for s in IMPORT_HELPER["cpp"]:
if s not in prompt:
test_set_up += s + "\n"
test_string = test_set_up + "\n" + prompt + code + "\n" + test
elif language == "java":
test_string = prompt + code + "\n" + test
elif language == "js" or language == "javascript":
test_string = prompt + code + "\n" + test
elif language == "go":
import_string = problems[task_id]["import"]
prompt = prompt.replace(import_string, "")
if example_test and "example_test" in problems[task_id]:
test = problems[task_id]["example_test"]
else:
test = problems[task_id]["test"]
test_setup = problems[task_id]["test_setup"]
other_pkgs = []
for pkg in IMPORT_HELPER["go"]:
if pkg not in test_setup:
p = pkg.split("/")[-1]
if p + "." in code:
other_pkgs.append(f"\"{pkg}\"")
if other_pkgs:
import_other_pkgs = "import (\n" + " ".join([p + "\n" for p in other_pkgs]) + ")"
test_string = test_setup + "\n" + import_other_pkgs + "\n" + prompt + code + "\n" + test
else:
test_string = test_setup + "\n" + prompt + code + "\n" + test
elif language == "rust":
main = "\nfn main(){ \n } \n"
declaration = problems[task_id]["declaration"]
test_string = main + declaration + prompt + code + test
return test_string
def stream_jsonl_all(filename: str) -> Iterable[Dict]:
results = []
if filename.endswith(".gz"):
fp = gzip.open(open(filename, "rb"), "rt")
else:
fp = open(filename, "r")
for line in fp:
if any(not x.isspace() for x in line):
results.append(json.loads(line))
fp.close()
return results
def evaluate_functional_correctness(
input_file: str = None,
tmp_dir: str = "./",
n_workers: int = 32,
timeout: float = 500.0,
problem_file: str = "../data/humaneval_python.jsonl.gz",
out_dir: str = None,
k: List[int] = [1, 10, 100],
test_groundtruth: bool = False,
example_test: bool = False,
):
if example_test:
print("Example test...")
problems = read_dataset(problem_file,
dataset_type="humaneval")
sample_jsonl = stream_jsonl_all(input_file)
if example_test:
suffix = "_example_test.jsonl"
else:
suffix = "_results.jsonl"
if out_dir is not None:
if not os.path.exists(out_dir):
os.makedirs(out_dir)
out_file = os.path.join(out_dir, input_file.split('/')[-1].replace(".jsonl", suffix))
else:
out_file = os.path.join(input_file.replace(".jsonl", suffix))
if "/codegeex/benchmark/humaneval-x/" in input_file:
test_groundtruth = True
if "-to-" in input_file:
translation_mode = True
else:
translation_mode = False
with ThreadPoolExecutor(max_workers=n_workers) as executor:
futures = []
completion_id = Counter()
n_samples = 0
results = defaultdict(list)
if test_groundtruth:
print("Testing ground truth...")
for sample in tqdm(problems.values()):
task_id = sample["task_id"]
lang = task_id.split("/")[0].lower()
if lang == "javascript":
lang = "js"
tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation")
sample["generation"] = sample["canonical_solution"]
sample["test_code"] = process_humaneval_test(sample, problems, example_test)
if sample["test_code"] is None:
continue
args = (task_id, sample, lang, timeout, tmp_dir_, completion_id[task_id])
future = executor.submit(check_correctness, *args)
futures.append(future)
completion_id[task_id] += 1
n_samples += 1
else:
print("Reading samples...")
for sample in tqdm(sample_jsonl):
task_id = sample["task_id"]
lang = task_id.split("/")[0].lower()
if translation_mode:
task_id = sample["task_id"].split("/")[-1]
lang = regex.findall("-to-.*-", input_file)[0].split("-to-")[-1].rstrip("-")
for l in LANGUAGE_NAME:
if l in lang:
lang = l
break
task_id = f"{LANGUAGE_NAME[lang]}/{task_id}"
if lang == "javascript":
lang = "js"
tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation")
sample["task_id"] = task_id
sample["test_code"] = process_humaneval_test(sample, problems, example_test)
if sample["test_code"] is None:
continue
if "completion_id" in sample:
completion_id_ = sample["completion_id"]
else:
completion_id_ = completion_id[task_id]
args = (task_id, sample, lang, timeout, tmp_dir_, completion_id_)
future = executor.submit(check_correctness, *args)
futures.append(future)
completion_id[task_id] += 1
n_samples += 1
print(completion_id)
if len(completion_id) == len(problems):
evaluate_pass_at_k = True
else:
evaluate_pass_at_k = False
print("Running test suites...")
for future in tqdm(as_completed(futures), total=len(futures)):
result = future.result()
results[result["task_id"]].append((result["completion_id"], result))
# Calculate pass@k.
total, correct = [], []
for result in results.values():
passed = [r[1]["passed"] for r in result]
total.append(len(passed))
correct.append(sum(passed))
total = np.array(total)
correct = np.array(correct)
if evaluate_pass_at_k:
ks = k
pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean()
for k in ks if (total >= k).all()}
print(pass_at_k)
else:
print("Total:", np.sum(total))
print("Correct:", np.sum(correct))
print("Writing to: ", out_file)
if out_file.endswith(".gz"):
fp = gzip.GzipFile(fileobj=open(out_file, "wb"), mode="wb")
for res in results.values():
for r in res:
fp.write((json.dumps(r[1]) + "\n").encode("utf-8"))
else:
fp = open(out_file, 'w')
for res in results.values():
for r in res:
fp.write(json.dumps(r[1]) + "\n")
fp.close()
print("Evaluation finished.")
def main():
fire.Fire(evaluate_functional_correctness)
if __name__ == "__main__":
sys.exit(main())
================================================
FILE: codegeex/benchmark/execution.py
================================================
import contextlib
import faulthandler
import io
import multiprocessing
import os
import platform
import signal
import random
import subprocess
import tempfile
import gzip
import json
from typing import *
def dicts_to_jsonl(data_list: list, filename: str, compress: bool = True) -> None:
"""
Method saves list of dicts into jsonl file.
:param data: (list) list of dicts to be stored,
:param filename: (str) path to the output file. If suffix .jsonl is not given then methods appends
.jsonl suffix into the file.
:param compress: (bool) should file be compressed into a gzip archive?
"""
sjsonl = '.jsonl'
sgz = '.gz'
# Check filename
if not filename.endswith(sjsonl):
filename = filename + sjsonl
# Save data
if compress:
filename = filename + sgz
with gzip.open(filename, 'w') as compressed:
for ddict in data_list:
jout = json.dumps(ddict) + '\n'
jout = jout.encode('utf-8')
compressed.write(jout)
else:
with open(filename, 'w') as out:
for ddict in data_list:
jout = json.dumps(ddict) + '\n'
out.write(jout)
def check_correctness(
task_id: str,
sample: dict,
language_type: str,
timeout: float = 3.0,
tmp_dir: str = None,
completion_id: Optional[int] = None,
) -> Dict:
"""
Evaluates the functional correctness of a completion by running the test
suite provided in the problem.
"""
def unsafe_execute(tmp_dir):
random_id = random.uniform(1, 1000)
if "python" in language_type.lower():
with create_tempdir():
# These system calls are needed when cleaning up tempdir.
import os
import shutil
rmtree = shutil.rmtree
rmdir = os.rmdir
chdir = os.chdir
# Disable functionalities that can make destructive changes to the test.
reliability_guard()
try:
exec_globals = {}
with swallow_io():
with time_limit(timeout):
# WARNING
# This program exists to execute untrusted model-generated code. Although
# it is highly unlikely that model-generated code will do something overtly
# malicious in response to this test suite, model-generated code may act
# destructively due to a lack of model capability or alignment.
# Users are strongly encouraged to sandbox this evaluation suite so that it
# does not perform destructive actions on their host or network.
# Once you have read this disclaimer and taken appropriate precautions,
# uncomment the following line and proceed at your own risk:
exec(sample["test_code"], exec_globals)
result.append("passed")
except TimeoutException:
result.append("timed out")
except AssertionError as e:
result.append(f"failed: AssertionError")
except BaseException as e:
result.append(f"failed: {e}")
# Needed for cleaning up.
shutil.rmtree = rmtree
os.rmdir = rmdir
os.chdir = chdir
elif "go" in language_type.lower():
assert tmp_dir is not None, "Go should be evaluated in a dir where necessary module files installed."
import os
import shutil
if "tmp" not in tmp_dir:
tmp_dir = os.path.join(tmp_dir, "tmp")
tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
os.chdir(tmp_dir)
open(f"main_test.go", 'w').write(sample["test_code"])
try:
exec_result = None
with time_limit(timeout):
# WARNING
# This program exists to execute untrusted model-generated code. Although
# it is highly unlikely that model-generated code will do something overtly
# malicious in response to this test suite, model-generated code may act
# destructively due to a lack of model capability or alignment.
# Users are strongly encouraged to sandbox this evaluation suite so that it
# does not perform destructive actions on their host or network.
# Once you have read this disclaimer and taken appropriate precautions,
# uncomment the following line and proceed at your own risk:
exec_result = subprocess.run(["go", "test", f"-timeout={timeout}s", "main_test.go"], timeout=timeout, capture_output=True)
if exec_result.returncode == 0:
result.append("passed")
else:
if exec_result.stderr:
try:
err = exec_result.stderr.decode()
except:
err = exec_result.stderr
else:
try:
err = exec_result.stdout.decode()
except:
err = exec_result.stdout
result.append(f"failed: {err}")
except TimeoutException:
result.append("timed out")
shutil.rmtree(tmp_dir)
elif "js" in language_type.lower():
import os
import shutil
if "tmp" not in tmp_dir:
tmp_dir = os.path.join(tmp_dir, "tmp")
tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
os.chdir(tmp_dir)
open(f"test.js", 'w').write(sample["test_code"])
try:
exec_result = None
with time_limit(timeout):
# WARNING
# This program exists to execute untrusted model-generated code. Although
# it is highly unlikely that model-generated code will do something overtly
# malicious in response to this test suite, model-generated code may act
# destructively due to a lack of model capability or alignment.
# Users are strongly encouraged to sandbox this evaluation suite so that it
# does not perform destructive actions on their host or network.
# Once you have read this disclaimer and taken appropriate precautions,
# uncomment the following line and proceed at your own risk:
exec_result = subprocess.run(["node", "test.js"], timeout=timeout, capture_output=True)
if exec_result.stderr.decode():
err = exec_result.stderr.decode()
result.append(f"failed: {err}")
elif exec_result.stdout.decode():
err = exec_result.stdout.decode()
result.append(f"failed: {err}")
else:
result.append("passed")
except TimeoutException:
result.append("timed out")
shutil.rmtree(tmp_dir)
elif "cpp" in language_type.lower():
import os
import shutil
if "tmp" not in tmp_dir:
tmp_dir = os.path.join(tmp_dir, "tmp")
tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
os.chdir(tmp_dir)
open(f"test.cpp", 'w').write(sample["test_code"])
if "162" in task_id:
compilation_result = subprocess.run(["/usr/bin/g++", "-std=c++11", "test.cpp", "-lcrypto", "-lssl"],
timeout=timeout,
capture_output=True)
else:
compilation_result = subprocess.run(["/usr/bin/g++", "-std=c++11", "test.cpp"], timeout=timeout,
capture_output=True)
if compilation_result.returncode != 0:
if compilation_result.stderr:
err = compilation_result.stderr.decode()
else:
err = compilation_result.stdout.decode()
result.append(f"failed: compilation error: {err}")
else:
try:
exec_result = None
with time_limit(timeout):
# WARNING
# This program exists to execute untrusted model-generated code. Although
# it is highly unlikely that model-generated code will do something overtly
# malicious in response to this test suite, model-generated code may act
# destructively due to a lack of model capability or alignment.
# Users are strongly encouraged to sandbox this evaluation suite so that it
# does not perform destructive actions on their host or network.
# Once you have read this disclaimer and taken appropriate precautions,
# uncomment the following line and proceed at your own risk:
exec_result = subprocess.run(["./a.out"], timeout=timeout, capture_output=True)
if exec_result.returncode == 0:
result.append("passed")
else:
if exec_result.stderr:
try:
err = exec_result.stderr.decode()
except:
err = exec_result.stderr
else:
try:
err = exec_result.stdout.decode()
except:
err = exec_result.stdout
result.append(f"failed: {err}")
except TimeoutException:
result.append("timed out")
shutil.rmtree(tmp_dir)
elif "rust" in language_type.lower():
import os
WD: str = os.path.dirname(os.path.abspath(__file__))
RUST_DIR: str = os.path.join(WD, "rust")
RUST_SRC: str = os.path.join(RUST_DIR, "src")
RUST_BIN: str = os.path.join(RUST_SRC, "bin")
RUST_TMP_DIR: str = os.path.join(RUST_DIR, "tmp")
RUST_LOGS: str = os.path.join(RUST_TMP_DIR, "logs")
RUST_EXT: str = ".rs"
# Create mandatory tmp directories
os.makedirs(RUST_TMP_DIR, exist_ok=True)
os.makedirs(RUST_LOGS, exist_ok=True)
os.makedirs(RUST_SRC, exist_ok=True)
os.makedirs(RUST_BIN, exist_ok=True)
with tempfile.NamedTemporaryFile(dir = RUST_BIN, delete=False) as f:
#temporal file name
file_prefix = sample["task_id"].lower().replace("/", "_")
file_name:str = file_prefix +RUST_EXT
os.rename(f.name, os.path.join(RUST_BIN, file_name))
# Sample to pure Rust function
rust_code: str = sample["test_code"]
# dump the rust source code in the target temporal file
f.write(rust_code.encode('utf-8'))
# Proceed towards Rust binaries compilation. Therefore move to Rust module root dir.
os.chdir(RUST_DIR)
# Two possible outcomes
# Pass OR Fail compilation
log_filename: str = file_prefix + ".jsonl"
log_path: str = os.path.join(RUST_LOGS, log_filename)
cargo_check: str = "cargo check --bin " + file_prefix + " --message-format json >> " + log_path
# Compilation build status
returned_val_compilation: int
# Overwrite file content
if os.path.exists(log_path):
if(file_size := os.path.getsize(log_path)) >= 0:
os.remove(log_path)
returned_val_compilation = os.system(cargo_check)
else:
returned_val_compilation = os.system(cargo_check)
# 0 means success
if returned_val_compilation == 0:
#Execution pipeline
cargo_test: str = "cargo test --bin " +file_prefix+ " --message-format json >> " + log_path
returned_val_execution = os.system(cargo_test)
if returned_val_execution == 0:
result.append("passed")
else:
result.append(f"failed: execution error")
else:
result.append(f"failed: compilation error")
elif "java" in language_type.lower():
assert tmp_dir is not None, "Java should be evaluated in a temporary dir."
import os
import shutil
if "tmp" not in tmp_dir:
tmp_dir = os.path.join(tmp_dir, "tmp")
tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}")
if not os.path.exists(tmp_dir):
os.makedirs(tmp_dir)
os.chdir(tmp_dir)
open(os.path.join(tmp_dir, "Main.java"), 'w').write(sample["test_code"])
res = "failed: unknown error"
compile_returncode = -1
for _ in range(5):
try:
compilation_result = subprocess.run(['javac', os.path.join(tmp_dir, "Main.java")], timeout=5,
capture_output=True)
compile_returncode = compilation_result.returncode
break
except subprocess.TimeoutExpired as e:
continue
if compile_returncode != 0:
res = "failed: compilation error"
else:
exec_result = None
try:
# WARNING
# This program exists to execute untrusted model-generated code. Although
# it is highly unlikely that model-generated code will do something overtly
# malicious in response to this test suite, model-generated code may act
# destructively due to a lack of model capability or alignment.
# Users are strongly encouraged to sandbox this evaluation suite so that it
# does not perform destructive actions on their host or network.
# Once you have read this disclaimer and taken appropriate precautions,
# uncomment the following line and proceed at your own risk:
# exec_result = subprocess.run([f'java', '-cp', tmp_dir, 'Main'], timeout=timeout, capture_output=True)
if exec_result.returncode == 0:
res = "passed"
elif exec_result.returncode == 1:
if "AssertionError" in exec_result.stderr.decode('unicode-escape'):
res = "failed: wrong answer"
else:
res = f"failed: {exec_result.stderr.decode()}"
except subprocess.TimeoutExpired as e:
res = "time out"
except BaseException as e:
res = f"failed: {e}"
result.append(res)
shutil.rmtree(tmp_dir)
manager = multiprocessing.Manager()
result = manager.list()
p = multiprocessing.Process(target=unsafe_execute, args=(tmp_dir,))
p.start()
p.join(timeout=timeout + 1)
if p.is_alive():
p.kill()
if not result:
result.append("timed out")
return {
"task_id" : task_id,
"completion_id": completion_id,
"test_code" : sample["test_code"],
"prompt" : sample["prompt"],
"generation" : sample["generation"],
"result" : result[0],
"passed" : result[0] == "passed",
"finish" : -1 if "finish" not in sample else sample["finish"],
"file" : "" if "file" not in sample else sample["file"],
"output" : [] if "output" not in sample else sample["output"],
}
# Copyright (c) OpenAI (https://openai.com)
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ============================================================================
@contextlib.contextmanager
def time_limit(seconds: float):
def signal_handler(signum, frame):
raise TimeoutException("Timed out!")
signal.setitimer(signal.ITIMER_REAL, seconds)
signal.signal(signal.SIGALRM, signal_handler)
try:
yield
finally:
signal.setitimer(signal.ITIMER_REAL, 0)
@contextlib.contextmanager
def swallow_io():
stream = WriteOnlyStringIO()
with contextlib.redirect_stdout(stream):
with contextlib.redirect_stderr(stream):
with redirect_stdin(stream):
yield
@contextlib.contextmanager
def create_tempdir():
with tempfile.TemporaryDirectory() as dirname:
with chdir(dirname):
yield dirname
class TimeoutException(Exception):
pass
class WriteOnlyStringIO(io.StringIO):
""" StringIO that throws an exception when it's read from """
def read(self, *args, **kwargs):
raise IOError
def readline(self, *args, **kwargs):
raise IOError
def readlines(self, *args, **kwargs):
raise IOError
def readable(self, *args, **kwargs):
""" Returns True if the IO object can be read. """
return False
class redirect_stdin(contextlib._RedirectStream): # type: ignore
_stream = 'stdin'
@contextlib.contextmanager
def chdir(root):
if root == ".":
yield
return
cwd = os.getcwd()
os.chdir(root)
try:
yield
except BaseException as exc:
raise exc
finally:
os.chdir(cwd)
def reliability_guard(maximum_memory_bytes: Optional[int] = None):
"""
This disables various destructive functions and prevents the generated code
from interfering with the test (e.g. fork bomb, killing other processes,
removing filesystem files, etc.)
WARNING
This function is NOT a security sandbox. Untrusted code, including, model-
generated code, should not be blindly executed outside of one. See the
Codex paper for more information about OpenAI's code sandbox, and proceed
with caution.
"""
if maximum_memory_bytes is not None:
import resource
resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes))
resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes))
if not platform.uname().system == 'Darwin':
resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))
faulthandler.disable()
import builtins
builtins.exit = None
builtins.quit = None
import os
os.environ['OMP_NUM_THREADS'] = '1'
os.kill = None
os.system = None
os.putenv = None
os.remove = None
os.removedirs = None
os.rmdir = None
os.fchdir = None
os.setuid = None
os.fork = None
os.forkpty = None
os.killpg = None
os.rename = None
os.renames = None
os.truncate = None
os.replace = None
os.unlink = None
os.fchmod = None
os.fchown = None
os.chmod = None
os.chown = None
os.chroot = None
os.fchdir = None
os.lchflags = None
os.lchmod = None
os.lchown = None
os.getcwd = None
os.chdir = None
import shutil
shutil.rmtree = None
shutil.move = None
shutil.chown = None
import subprocess
subprocess.Popen = None # type: ignore
__builtins__['help'] = None
import sys
sys.modules['ipdb'] = None
sys.modules['joblib'] = None
sys.modules['resource'] = None
sys.modules['psutil'] = None
sys.modules['tkinter'] = None
================================================
FILE: codegeex/benchmark/gather_output.py
================================================
import os
import sys
import fire
import glob
def gather_output(
output_dir: str = "./output",
output_prefix: str = None,
if_remove_rank_files: int = 0,
):
if output_prefix is None:
output_list = glob.glob(output_dir + "/*")
else:
output_list = glob.glob(os.path.join(output_dir, output_prefix + "*"))
for output_file in output_list:
if "rank0" in output_file:
output_prefix_ = output_file.split("_rank0.jsonl")[0]
rank_files = glob.glob(output_prefix_ + "_rank*")
with open(output_prefix_ + ".jsonl", "w") as f_out:
for rank_file in rank_files:
with open(rank_file, "r") as f_in:
for line in f_in:
f_out.write(line)
if if_remove_rank_files:
os.remove(rank_file)
print(f"Removing {rank_file}...")
if output_prefix is None:
output_list = glob.glob(output_dir + "/*")
else:
output_list = glob.glob(os.path.join(output_dir, output_prefix + "*"))
for output_file in output_list:
if "rank" in output_file or "_unfinished" in output_file or "all" in output_file or "_result" in output_file:
continue
if "_finished" not in output_file:
continue
output_prefix_ = output_file.split("_finished.jsonl")[0]
files = [output_file, output_prefix_ + "_unfinished.jsonl"]
with open(output_prefix_ + "_all.jsonl", "w") as f_out:
for f in files:
with open(f, "r") as f_in:
for line in f_in:
f_out.write(line)
print("Gathering finished. Saved in {}".format(output_prefix_ + "_all.jsonl"))
def main():
fire.Fire(gather_output)
if __name__ == "__main__":
sys.exit(main())
================================================
FILE: codegeex/benchmark/humaneval-x/__init__.py
================================================
================================================
FILE: codegeex/benchmark/humaneval-x/cpp/evaluation/test.cpp
================================================
/*
Input to this function is a string containing multiple groups of nested parentheses. Your goal is to
separate those group into separate strings and return the vector of those.
Separate groups are balanced (each open brace is properly closed) and not nested within each other
Ignore any spaces in the input string.
>>> separate_paren_groups("( ) (( )) (( )( ))")
{"()", "(())", "(()())"}
*/
#include<stdio.h>
#include<vector>
#include<string>
using namespace std;
vector<string> separate_paren_groups(string paren_string){
vector<string> all_parens;
string current_paren;
int level=0;
char chr;
int i;
for (i=0;i<paren_string.length();i++)
{
chr=paren_string[i];
if (chr=='(')
{
level+=1;
current_paren+=chr;
}
if (chr==')')
{
level-=1;
current_paren+=chr;
if (level==0){
all_parens.push_back(current_paren);
current_paren="";
}
}
}
return all_parens;
}
#undef NDEBUG
#include<assert.h>
bool issame(vector<string> a,vector<string>b){
if (a.size()!=b.size()) return false;
for (int i=0;i<a.size();i++)
{
if (a[i]!=b[i]) return false;
}
return true;
}
int main(){
assert (issame(separate_paren_groups("(()()) ((())) () ((())()())"),{"(()())", "((()))", "()", "((())()())"}));
assert (issame(separate_paren_groups("() (()) ((())) (((())))"), {"()", "(())", "((()))", "(((())))" }));
assert (issame(separate_paren_groups("(()(())((())))") ,{ "(()(())((())))" }));
assert (issame(separate_paren_groups("( ) (( )) (( )( ))") ,{"()", "(())", "(()())"}));
}
================================================
FILE: codegeex/benchmark/humaneval-x/evaluate_humaneval_x.py
================================================
import os
import sys
import fire
import json
import gzip
import regex
import numpy as np
from typing import *
from tqdm.auto import tqdm
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from codegeex.benchmark.utils import read_dataset, IMPORT_HELPER
from codegeex.benchmark.metric import estimate_pass_at_k
from codegeex.benchmark.execution import check_correctness
LANGUAGE_NAME = {
"cpp" : "CPP",
"go" : "Go",
"java" : "Java",
"js" : "JavaScript",
"python": "Python",
}
def process_humaneval_test(sample, problems, example_test=False):
task_id = sample["task_id"]
language = task_id.split("/")[0].lower()
prompt = sample["prompt"]
if example_test and "example_test" in problems[task_id] and problems[task_id]["example_test"] != "":
test = problems[task_id]["example_test"]
else:
test = problems[task_id]["test"]
code = sample["generation"]
# Pre-process for different languages
if language == "python":
code_ = []
for line in code.split("\n"):
if (len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t'):
break
code_.append(line)
code = "\n".join(code_)
test_setup = "\n".join(IMPORT_HELPER["python"]) + "\n"
test_string = test_setup + prompt + code + "\n" + test + "\n"
elif language == "cpp":
test_set_up = ""
for s in IMPORT_HELPER["cpp"]:
if s not in prompt:
test_set_up += s + "\n"
test_string = test_set_up + "\n" + prompt + code + "\n" + test
elif language == "java":
test_string = prompt + code + "\n" + test
elif language == "js" or language == "javascript":
test_string = prompt + code + "\n" + test
elif language == "go":
import_string = problems[task_id]["import"]
prompt = prompt.replace(import_string, "")
if example_test and "example_test" in problems[task_id]:
test = problems[task_id]["example_test"]
else:
test = problems[task_id]["test"]
test_setup = problems[task_id]["test_setup"]
other_pkgs = []
for pkg in IMPORT_HELPER["go"]:
if pkg not in test_setup:
p = pkg.split("/")[-1]
if p + "." in code:
other_pkgs.append(f"\"{pkg}\"")
if other_pkgs:
import_other_pkgs = "import (\n" + " ".join([p + "\n" for p in other_pkgs]) + ")"
test_string = test_setup + "\n" + import_other_pkgs + "\n" + prompt + code + "\n" + test
else:
test_string = test_setup + "\n" + prompt + code + "\n" + test
elif language == "rust":
main = "\nfn main(){ \n } \n"
declaration = problems[task_id]["declaration"]
test_string = main + declaration + prompt + code + test
return test_string
def stream_jsonl_all(filename: str) -> Iterable[Dict]:
results = []
if filename.endswith(".gz"):
fp = gzip.open(open(filename, "rb"), "rt")
else:
fp = open(filename, "r")
for line in fp:
if any(not x.isspace() for x in line):
results.append(json.loads(line))
fp.close()
return results
def evaluate_functional_correctness(
input_file: str = None,
tmp_dir: str = "./",
n_workers: int = 32,
timeout: float = 500.0,
problem_file: str = "../data/humaneval_python.jsonl.gz",
out_dir: str = None,
k: List[int] = [1, 10, 100],
test_groundtruth: bool = False,
example_test: bool = False,
):
if example_test:
print("Example test...")
problems = read_dataset(problem_file,
dataset_type="humaneval")
sample_jsonl = stream_jsonl_all(input_file)
if example_test:
suffix = "_example_test.jsonl"
else:
suffix = "_results.jsonl"
if out_dir is not None:
if not os.path.exists(out_dir):
os.makedirs(out_dir)
out_file = os.path.join(out_dir, input_file.split('/')[-1].replace(".jsonl", suffix))
else:
out_file = os.path.join(input_file.replace(".jsonl", suffix))
if "/codegeex/benchmark/humaneval-x/" in input_file:
test_groundtruth = True
if "-to-" in input_file:
translation_mode = True
else:
translation_mode = False
with ThreadPoolExecutor(max_workers=n_workers) as executor:
futures = []
completion_id = Counter()
n_samples = 0
results = defaultdict(list)
if test_groundtruth:
print("Testing ground truth...")
for sample in tqdm(problems.values()):
task_id = sample["task_id"]
lang = task_id.split("/")[0].lower()
if lang == "javascript":
lang = "js"
tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation")
sample["generation"] = sample["canonical_solution"]
sample["test_code"] = process_humaneval_test(sample, problems, example_test)
if sample["test_code"] is None:
continue
args = (task_id, sample, lang, timeout, tmp_dir_, completion_id[task_id])
future = executor.submit(check_correctness, *args)
futures.append(future)
completion_id[task_id] += 1
n_samples += 1
else:
print("Reading samples...")
for sample in tqdm(sample_jsonl):
task_id = sample["task_id"]
lang = task_id.split("/")[0].lower()
if translation_mode:
task_id = sample["task_id"].split("/")[-1]
lang = regex.findall("-to-.*-", input_file)[0].split("-to-")[-1].rstrip("-")
for l in LANGUAGE_NAME:
if l in lang:
lang = l
break
task_id = f"{LANGUAGE_NAME[lang]}/{task_id}"
if lang == "javascript":
lang = "js"
tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation")
sample["task_id"] = task_id
sample["test_code"] = process_humaneval_test(sample, problems, example_test)
if sample["test_code"] is None:
continue
if "completion_id" in sample:
completion_id_ = sample["completion_id"]
else:
completion_id_ = completion_id[task_id]
args = (task_id, sample, lang, timeout, tmp_dir_, completion_id_)
future = executor.submit(check_correctness, *args)
futures.append(future)
completion_id[task_id] += 1
n_samples += 1
print(completion_id)
if len(completion_id) == len(problems):
evaluate_pass_at_k = True
else:
evaluate_pass_at_k = False
print("Running test suites...")
for future in tqdm(as_completed(futures), total=len(futures)):
result = future.result()
results[result["task_id"]].append((result["completion_id"], result))
# Calculate pass@k.
total, correct = [], []
for result in results.values():
passed = [r[1]["passed"] for r in result]
total.append(len(passed))
correct.append(sum(passed))
total = np.array(total)
correct = np.array(correct)
if evaluate_pass_at_k:
ks = k
pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean()
for k in ks if (total >= k).all()}
print(pass_at_k)
else:
print("Total:", np.sum(total))
print("Correct:", np.sum(correct))
print("Writing to: ", out_file)
if out_file.endswith(".gz"):
fp = gzip.GzipFile(fileobj=open(out_file, "wb"), mode="wb")
for res in results.values():
for r in res:
fp.write((json.dumps(r[1]) + "\n").encode("utf-8"))
else:
fp = open(out_file, 'w')
for res in results.values():
for r in res:
fp.write(json.dumps(r[1]) + "\n")
fp.close()
print("Evaluation finished.")
def main():
fire.Fire(evaluate_functional_correctness)
if __name__ == "__main__":
sys.exit(main())
================================================
FILE: codegeex/benchmark/humaneval-x/generate_humaneval_x.py
================================================
import argparse
import logging
import os
import random
import socket
import time
from typing import *
import torch
import zmq
from codegeex.benchmark.utils import read_dataset, process_extra_prompt
from codegeex.megatron import get_args
from codegeex.megatron.inference import run_generation_distributed, model_provider
from codegeex.megatron.initialize import initialize_megatron
from codegeex.quantization import quantize
logging.getLogger("torch").setLevel(logging.WARNING)
def add_code_generation_args(parser):
"""Code generation arguments."""
group = parser.add_argument_group(title="code generation")
group.add_argument(
"--hostfile",
type=str,
default="./hostfile",
)
group.add_argument(
"--channel-ip",
type=str,
default=None,
help="IP for ZeroMQ channel",
)
group.add_argument(
"--channel-port",
type=int,
default=5555,
help="Port for ZeroMQ channel",
)
group.add_argument(
"--master-port",
type=int,
default=6666,
help="Port for torch distributed",
)
group.add_argument(
"--temperature",
type=float,
default=1.0,
help="Sampling temperature.",
)
group.add_argument(
"--greedy",
action="store_true",
default=False,
help="Use greedy sampling.",
)
group.add_argument(
"--top-p",
type=float,
default=0.0,
help="Top p sampling.",
)
group.add_argument(
"--top-k",
type=int,
default=0,
help="Top k sampling.",
)
group.add_argument(
"--out-seq-length",
type=int,
default=1024,
help="Size of the output generated text.",
)
group.add_argument(
"--input-path",
type=str,
default="./benchmark/humaneval/HumanEval.jsonl",
help="Get input path",
)
group.add_argument(
"--num-samples",
type=int,
default=0,
help="Number of samples to generate",
)
group.add_argument(
"--recompute",
action="store_true",
help="During generation recompute all attention "
"instead of using previously computed keys/values.",
)
group.add_argument(
"--load-deepspeed",
action="store_true",
help="Load DeepSpeed checkpoint",
)
group.add_argument(
"--ws-encoding-start-id",
type=int,
default=None,
help="Start id for whitespace encoding",
)
group.add_argument(
"--ws-encoding-length",
type=int,
default=None,
help="Length of whitespace encoding",
)
group.add_argument(
"--dataset",
type=str,
default="humaneval",
)
group.add_argument(
"--samples-per-problem",
type=int,
default=200,
help="Number of samples to generate for each problem",
)
group.add_argument(
"--output-prefix",
type=str,
default="./output/humaneval",
help="Prefix for output files",
)
group.add_argument(
"--gen-node-rank",
type=int,
default=None,
)
group.add_argument(
"--gen-node-world-size",
type=int,
default=None,
)
group.add_argument(
"--gen-world-size",
type=int,
default=1,
help="Number of machines to use for generation",
)
group.add_argument(
"--gen-rank",
type=int,
default=0,
help="Machine rank for human eval generation",
)
group.add_argument(
"--extra-prompt",
type=str,
default=None,
help="Extra prompt to use for human eval generation",
)
group.add_argument(
"--verbose-interval",
type=int,
default=100,
)
group.add_argument(
"--problem-split",
type=str,
default="test",
)
group.add_argument(
"--prompt-type",
type=str,
default="notag",
)
group.add_argument(
"--num-devices-per-node",
type=int,
default=None,
)
group.add_argument(
"--return-scores",
action="store_true",
)
group.add_argument(
'--language-type',
default=None,
help='Identify the type of programming language to generate',
)
group.add_argument(
'--bad-ids',
nargs="*",
type=int,
default=None,
help='Specify bad ids that will not be used',
)
group.add_argument(
"--quantize",
action="store_true",
)
return parser
def main(node_rank: int, local_rank: int, master_port: int, num_devices: int):
"""Main program."""
os.environ["WORLD_SIZE"] = str(num_devices)
os.environ["RANK"] = str(local_rank)
os.environ["MASTER_ADDR"] = "0.0.0.0"
os.environ["MASTER_PORT"] = f"{master_port}"
initialize_megatron(
extra_args_provider=add_code_generation_args,
args_defaults={
"tokenizer_type": "GPT2BPETokenizer",
"no_load_rng" : True,
"no_load_optim" : True,
},
)
# set_random_seed(node_rank * num_devices + local_rank)
args = get_args()
if args.num_layers_per_virtual_pipeline_stage is not None:
print("Interleaved pipeline schedule is not yet supported for text generation.")
exit()
world_size = args.gen_node_world_size * num_devices
args.gen_rank = num_devices * node_rank + local_rank
args.gen_world_size = world_size
print(f"Generating on rank {args.gen_rank} of {args.gen_world_size}")
# Set up model and load checkpoint.
state_dict = torch.load(args.load, map_location="cpu")
state_dict = state_dict["module"]
print("Building CodeGeeX model ...")
model = model_provider()
model.load_state_dict(state_dict)
model.eval()
if args.fp16 and args.ln_fp16:
model.half()
if args.quantize:
model = quantize(model, weight_bit_width=8, backend="megatron")
model.cuda()
# Generate samples.
run_generation_distributed(model)
print(f"(gen_rank={args.gen_rank}, rank={local_rank}) finished, waiting ...")
torch.distributed.barrier()
def server():
print(f"[ server ] starting ...", flush=True)
parser = argparse.ArgumentParser()
parser.add_argument(
"--channel-ip",
type=str,
default=None,
help="IP for ZeroMQ channel",
)
parser.add_argument(
"--channel-port",
type=int,
default=5555,
help="Port for ZeroMQ channel",
)
parser.add_argument(
"--master-port",
type=int,
default=6666,
help="Port for torch distributed",
)
parser.add_argument(
"--samples-per-problem",
type=int,
default=200,
help="Number of samples to generate for each problem",
)
parser.add_argument(
"--gen-node-world-size",
type=int,
default=1,
help="Number of machines to use for generation",
)
parser.add_argument(
"--input-path",
type=str,
default="",
help="Get input path",
)
parser.add_argument(
"--problem-split",
type=str,
default="test",
)
parser.add_argument(
"--micro-batch-size",
type=int,
default=1,
)
parser.add_argument(
'--language-type',
default=None,
help='Identify the type of programming language to generate',
)
args = parser.parse_known_args()[0]
entries = read_dataset(args.input_path, dataset_type="humaneval")
assert args.samples_per_problem % args.micro_batch_size == 0, "samples_per_problem should be divisible by micro_batch_size"
for entry in entries.values():
entry["prompt"] = process_extra_prompt(entry["prompt"], args.language_type)
res = []
for entry in entries.values():
res.extend([entry] * (args.samples_per_problem // args.micro_batch_size))
random.shuffle(res)
all_entries = res
# setup zeromq channel
print(f"[ server ] starting up on port {args.channel_port}", flush=True)
context = zmq.Context()
print(f"[ server ] creating socket", flush=True)
socket = context.socket(zmq.REP)
print(f"[ server ] binding to port {args.channel_port}", flush=True)
socket.bind(f"tcp://*:{args.channel_port}")
print(
f"[ server ] loaded {len(entries)} entries, generated {len(all_entries)} samples",
flush=True,
)
remaining_entries = all_entries.copy()
running_workers = args.gen_node_world_size * torch.cuda.device_count()
num_finished = 0
print(f"[ server ] listening for requests ...", flush=True)
start_time = time.perf_counter()
while True:
# Wait for next request from client
msg = socket.recv_json()
rank = msg["rank"]
action = msg["action"]
if action == "pull":
if len(remaining_entries) == 0:
print(f"[ server ] Shutting down worker {rank}", flush=True)
socket.send_json({"task_id": None})
running_workers -= 1
if running_workers == 0:
print(f"[ server ] All workers finished", flush=True)
break
else:
entry = remaining_entries.pop()
time_elapsed = time.perf_counter() - start_time
print(f"[ server ] Sending entry {entry['task_id']} to worker {rank}", flush=True)
remaining = (
len(remaining_entries)
/ (len(all_entries) - len(remaining_entries))
* time_elapsed
)
time_per_sampple = 0.0 if num_finished == 0 else time_elapsed / num_finished / args.micro_batch_size
print(
f"[ server ] total {len(all_entries)}, assigned {len(all_entries) - len(remaining_entries)}, "
f"finished {num_finished}, "
f"elapsed {time_elapsed:.4f}",
f"speed {time_per_sampple:.4f}s/sample",
f"remaining {remaining:.4f}",
flush=True,
)
socket.send_json({"task_id": entry})
else:
if action == "success":
print(f"[ server ] {msg['task_id']} is finished", flush=True)
socket.send_json({"pong": 1})
else:
print(f"[ server ] {msg['task_id']} is not finished", flush=True)
remaining_entries.append(msg['task_id'])
socket.send_json({"pong": 1})
break
num_finished += 1
if __name__ == "__main__":
torch.multiprocessing.set_start_method("spawn")
parser = argparse.ArgumentParser()
parser.add_argument(
"--hostfile",
type=str,
default="./hostfile",
)
parser.add_argument(
"--master-port",
type=int,
default=5666,
)
parser.add_argument(
"--seed",
type=int,
default=42,
)
args = parser.parse_known_args()[0]
print("start method: " + torch.multiprocessing.get_start_method())
processes = []
num_devices = torch.cuda.device_count()
hosts = open(args.hostfile, "r").readlines()
hosts = [host.strip() for host in hosts]
master_port = args.master_port
node_rank = None
for i in range(len(hosts)):
if hosts[i] == socket.gethostbyname(socket.gethostname()):
node_rank = i
break
assert (
node_rank is not None
), f"Could not find hostname ({socket.gethostbyname(socket.gethostname())}) in hostlist"
# launch server
if socket.gethostbyname(socket.gethostname()) == hosts[0]:
server_process = torch.multiprocessing.Process(target=server)
print(f"Launching server ...")
server_process.start()
processes.append(server_process)
for i in range(num_devices):
local_rank = i
print(f"launching local rank {i}")
p = torch.multiprocessing.Process(
target=main,
args=(node_rank, local_rank, master_port, num_devices),
)
p.start()
processes.append(p)
for p in processes:
p.join()
================================================
FILE: codegeex/benchmark/humaneval-x/go/evaluation/go.mod
================================================
module humanEval
go 1.18
require (
github.com/go-openapi/inflect v0.19.0
github.com/stretchr/testify v1.8.0
)
require (
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
================================================
FILE: codegeex/benchmark/humaneval-x/go/evaluation/go.sum
================================================
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/go-openapi/inflect v0.19.0 h1:9jCH9scKIbHeV9m12SmPilScz6krDxKRasNNSNPXu/4=
github.com/go-openapi/inflect v0.19.0/go.mod h1:lHpZVlpIQqLyKwJ4N+YSc9hchQy/i12fJykb83CRBH4=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
================================================
FILE: codegeex/benchmark/humaneval-x/translate_humaneval_x.py
================================================
import os
import zmq
import time
import torch
import random
import socket
import logging
import argparse
from typing import *
from codegeex.benchmark.utils import read_translation_dataset
from codegeex.megatron import get_args
from codegeex.megatron.inference import run_generation_distributed, model_provider
from codegeex.megatron.initialize import initialize_megatron
logging.getLogger("torch").setLevel(logging.WARNING)
def add_code_generate_args(parser):
"""Code generation arguments."""
group = parser.add_argument_group(title="code generation")
group.add_argument(
"--hostfile",
type=str,
default="./hostfile",
)
group.add_argument(
"--channel-ip",
type=str,
default=None,
help="IP for ZeroMQ channel",
)
group.add_argument(
"--channel-port",
type=int,
default=5555,
help="Port for ZeroMQ channel",
)
group.add_argument(
"--master-port",
type=int,
default=5666,
)
group.add_argument(
"--temperature",
type=float,
default=1.0,
help="Sampling temperature.",
)
group.add_argument(
"--greedy",
action="store_true",
default=False,
help="Use greedy sampling.",
)
group.add_argument(
"--top-p",
type=float,
default=0.0,
help="Top p sampling.",
)
group.add_argument(
"--top-k",
type=int,
default=0,
help="Top k sampling.",
)
group.add_argument(
"--out-seq-length",
type=int,
default=1024,
help="Size of the output generated text.",
)
group.add_argument(
"--input-path",
type=str,
default="./benchmark/humaneval/HumanEval.jsonl",
help="Get input path",
)
group.add_argument(
"--num-samples",
type=int,
default=0,
help="Number of samples to generate",
)
group.add_argument(
"--recompute",
action="store_true",
help="During generation recompute all attention "
"instead of using previously computed keys/values.",
)
group.add_argument(
"--load-deepspeed",
action="store_true",
help="Load DeepSpeed checkpoint",
)
group.add_argument(
"--ws-encoding-start-id",
type=int,
default=None,
help="Start id for whitespace encoding",
)
group.add_argument(
"--ws-encoding-length",
type=int,
default=None,
help="Length of whitespace encoding",
)
group.add_argument(
"--dataset",
type=str,
default="humaneval",
)
group.add_argument(
"--samples-per-problem",
type=int,
default=200,
help="Number of samples to generate for each problem",
)
group.add_argument(
"--output-prefix",
type=str,
default="./output/humaneval",
help="Prefix for output files",
)
group.add_argument(
"--gen-node-rank",
type=int,
default=None,
)
group.add_argument(
"--gen-node-world-size",
type=int,
default=None,
)
group.add_argument(
"--gen-world-size",
type=int,
default=1,
help="Number of machines to use for generation",
)
group.add_argument(
"--gen-rank",
type=int,
default=0,
help="Machine rank for human eval generation",
)
group.add_argument(
"--extra-prompt",
type=str,
default=None,
help="Extra prompt to use for human eval generation",
)
group.add_argument(
"--verbose-interval",
type=int,
default=100,
)
group.add_argument(
"--problem-split",
type=str,
default="test",
)
group.add_argument(
"--prompt-type",
type=str,
default="notag",
)
group.add_argument(
"--num-devices-per-node",
type=int,
default=None,
)
group.add_argument(
"--return-scores",
action="store_true",
)
group.add_argument(
"--free-guidance",
action="store_true",
)
group.add_argument(
"--guide-temp",
type=float,
default=1.5,
)
group.add_argument(
"--attention-upweight",
type=float,
default=None,
)
group.add_argument(
'--bad-ids',
nargs="*",
type=int,
default=None,
help='Identify the type of programming language to generate',
)
group.add_argument(
"--src-path",
type=str,
default="",
help="Get source path",
)
group.add_argument(
"--tgt-path",
type=str,
default="",
help="Get target path",
)
group.add_argument(
'--language-src-type',
type=str,
default=None,
help='Identify the type of programming language',
)
group.add_argument(
'--language-tgt-type',
type=str,
default=None,
help='Identify the type of programming language to translate',
)
return parser
def main(node_rank: int, local_rank: int, master_port: int, num_devices: int):
"""Main program."""
os.environ["WORLD_SIZE"] = str(num_devices)
os.environ["RANK"] = str(local_rank)
os.environ["MASTER_ADDR"] = "0.0.0.0"
os.environ["MASTER_PORT"] = f"{master_port}"
initialize_megatron(
extra_args_provider=add_code_generate_args,
args_defaults={
"tokenizer_type": "GPT2BPETokenizer",
"no_load_rng" : True,
"no_load_optim" : True,
},
)
# set_random_seed(node_rank * num_devices + local_rank)
args = get_args()
if args.num_layers_per_virtual_pipeline_stage is not None:
print("Interleaved pipeline schedule is not yet supported for text generation.")
exit()
world_size = args.gen_node_world_size * num_devices
args.gen_rank = num_devices * node_rank + local_rank
args.gen_world_size = world_size
print(f"Generating on rank {args.gen_rank} of {args.gen_world_size}")
# Set up model and load checkpoint.
state_dict = torch.load(args.load, map_location="cpu")
state_dict = state_dict["module"]
print("Building CodeGeeX model ...")
model = model_provider()
model.load_state_dict(state_dict)
model.eval()
if args.fp16 and args.ln_fp16:
model.half()
model.cuda()
# Generate samples.
run_generation_distributed(model)
print(f"(gen_rank={args.gen_rank}, rank={local_rank}) finished, waiting ...")
torch.distributed.barrier()
def server():
print(f"[ server ] starting ...", flush=True)
parser = argparse.ArgumentParser()
parser.add_argument(
"--channel-ip",
type=str,
default=None,
help="IP for ZeroMQ channel",
)
parser.add_argument(
"--channel-port",
type=int,
default=5555,
help="Port for ZeroMQ channel",
)
parser.add_argument(
"--master-port",
type=int,
default=6666,
help="Port for torch distributed",
)
parser.add_argument(
"--samples-per-problem",
type=int,
default=200,
help="Number of samples to generate for each problem",
)
parser.add_argument(
"--gen-node-world-size",
type=int,
default=1,
help="Number of machines to use for generation",
)
parser.add_argument(
"--src-path",
type=str,
default="",
help="Get source path",
)
parser.add_argument(
"--tgt-path",
type=str,
default="",
help="Get target path",
)
parser.add_argument(
"--problem-split",
type=str,
default="test",
)
parser.add_argument(
"--micro-batch-size",
type=int,
default=1,
)
parser.add_argument(
'--language-src-type',
type=str,
default=None,
help='Identify the type of programming language',
)
parser.add_argument(
'--language-tgt-type',
type=str,
default=None,
help='Identify the type of programming language to translate',
)
args = parser.parse_known_args()[0]
entries = read_translation_dataset(args.src_path,
args.tgt_path,
lang_src=args.language_src_type,
lang_tgt=args.language_tgt_type,
dataset_type="humaneval")
assert args.samples_per_problem % args.micro_batch_size == 0, "samples_per_problem should be divisible by micro_batch_size"
res = []
for entry in entries.values():
res.extend([entry] * (args.samples_per_problem // args.micro_batch_size))
random.shuffle(res)
all_entries = res
# setup zeromq channel
print(f"[ server ] starting up on port {args.channel_port}", flush=True)
context = zmq.Context()
print(f"[ server ] creating socket", flush=True)
socket = context.socket(zmq.REP)
print(f"[ server ] binding to port {args.channel_port}", flush=True)
socket.bind(f"tcp://*:{args.channel_port}")
print(
f"[ server ] loaded {len(entries)} entries, generated {len(all_entries)} samples",
flush=True,
)
remaining_entries = all_entries.copy()
running_workers = args.gen_node_world_size * torch.cuda.device_count()
num_finished = 0
print(f"[ server ] listening for requests ...", flush=True)
start_time = time.perf_counter()
while True:
# Wait for next request from client
msg = socket.recv_json()
rank = msg["rank"]
action = msg["action"]
if action == "pull":
if len(remaining_entries) == 0:
print(f"[ server ] Shutting down worker {rank}", flush=True)
socket.send_json({"task_id": None})
running_workers -= 1
if running_workers == 0:
print(f"[ server ] All workers finished", flush=True)
break
else:
entry = remaining_entries.pop()
time_elapsed = time.perf_counter() - start_time
print(f"[ server ] Sending entry {entry['task_id']} to worker {rank}", flush=True)
remaining = (
len(remaining_entries)
/ (len(all_entries) - len(remaining_entries))
* time_elapsed
)
time_per_sampple = 0.0 if num_finished == 0 else time_elapsed / num_finished / args.micro_batch_size
print(
f"[ server ] total {len(all_entries)}, assigned {len(all_entries) - len(remaining_entries)}, "
f"finished {num_finished}, "
f"elapsed {time_elapsed:.4f}",
f"speed {time_per_sampple:.4f}s/sample",
f"remaining {remaining:.4f}",
flush=True,
)
socket.send_json({"task_id": entry})
else:
if action == "success":
print(f"[ server ] {msg['task_id']} is finished", flush=True)
socket.send_json({"pong": 1})
else:
print(f"[ server ] {msg['task_id']} is not finished", flush=True)
remaining_entries.append(msg['task_id'])
socket.send_json({"pong": 1})
break
num_finished += 1
if __name__ == "__main__":
torch.multiprocessing.set_start_method("spawn")
parser = argparse.ArgumentParser()
parser.add_argument(
"--hostfile",
type=str,
default="./hostfile",
)
parser.add_argument(
"--master-port",
type=int,
default=5666,
)
parser.add_argument(
"--seed",
type=int,
default=42,
)
args = parser.parse_known_args()[0]
print("start method: " + torch.multiprocessing.get_start_method())
processes = []
num_devices = torch.cuda.device_count()
hosts = open(args.hostfile, "r").readlines()
hosts = [host.strip() for host in hosts]
master_port = args.master_port
node_rank = None
for i in range(len(hosts)):
if hosts[i] == socket.gethostbyname(socket.gethostname()):
node_rank = i
break
assert (
node_rank is not None
), f"Could not find hostname ({socket.gethostbyname(socket.gethostname())}) in hostlist"
# launch server
if socket.gethostbyname(socket.gethostname()) == hosts[0]:
server_process = torch.multiprocessing.Process(target=server)
print(f"Launching server ...")
server_process.start()
processes.append(server_process)
for i in range(num_devices):
local_rank = i
print(f"launching local rank {i}")
p = torch.multiprocessing.Process(
target=main,
args=(node_rank, local_rank, master_port, num_devices),
)
p.start()
processes.append(p)
for p in processes:
p.join()
================================================
FILE: codegeex/benchmark/inspect_result.py
================================================
import os
import sys
import fire
import json
import glob
import numpy as np
import pandas as pd
from collections import defaultdict
from codegeex.benchmark.metric import estimate_pass_at_k
error_types = {
"python" : [
"accepted",
"assertion error",
"undefined error",
"runtime error",
"syntax error",
"timeout error",
"type error",
],
"java" : [
"accepted",
"compilation error",
"assertion error",
"timeout error",
"index error",
"class cast error",
"stack overflow",
"null pointer error",
"unsupported operation",
"number format error",
"no such element",
"illegal argument",
"out of memory",
"arithmetic error",
"others",
],
"cpp" : [
"accepted",
"compilation error",
"assertion error",
"range error",
"invalid argument",
"pointer error",
"out of memory",
"package error",
"others",
],
"javascript": [
"accepted",
"assertion error",
"undefined error",
"runtime error",
"syntax error",
"timeout error",
"range error",
"type error",
],
"go" : [
"accepted",
"assertion error",
"undefined error",
"runtime error",
"syntax error",
"timeout error",
"type error",
"notused error",
],
}
def inspect_result(
input_dir: str = None,
input_file: str = None,
output_dir: str = None,
pass_at_k_outpath: str = None,
):
if input_dir is not None:
input_files = glob.glob(input_dir + "/*_results.jsonl")
else:
input_files = [input_file]
if output_dir is not None:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
else:
output_dir = "/"
pass_at_k_outs = []
for input_file in input_files:
result_stats = defaultdict(dict)
df = None
incompleted = False
with open(input_file, "r") as f:
for i, line in enumerate(f):
obj = json.loads(line)
task_id = obj["task_id"]
language_type = task_id.split("/")[0].lower()
if language_type not in error_types:
if language_type == "humaneval":
language_type = "python"
elif language_type == "C++":
language_type = "cpp"
else:
incompleted = True
break
if task_id not in result_stats.keys():
default_stats = {}
default_stats["task_id"] = task_id
default_stats["n_sample"] = 0
for k in error_types[language_type]:
default_stats[k] = 0
result_stats[task_id] = default_stats.copy()
if df is None:
df = pd.DataFrame(columns=error_types[language_type])
result_stats[task_id]["n_sample"] += 1
if "failed" in obj["result"]:
error = obj["result"]
if language_type == "python":
if "assertionerror" in error.lower():
result_stats[task_id]["assertion error"] += 1
elif "syntax" in error.lower() or "indent" in error.lower() or "literal" in error.lower():
result_stats[task_id]["syntax error"] += 1
elif "not defined" in error.lower():
result_stats[task_id]["undefined error"] += 1
elif "timeout" in error.lower():
result_stats[task_id]["timeout error"] += 1
elif "type" in error.lower():
result_stats[task_id]["type error"] += 1
else:
result_stats[task_id]["runtime error"] += 1
elif language_type == "java":
if "wrong answer" in error:
result_stats[task_id]["assertion error"] += 1
elif "compilation error" in error:
result_stats[task_id]["compilation error"] += 1
elif "time out" in error:
result_stats[task_id]["timeout error"] += 1
elif "IndexOutOfBounds" in error:
result_stats[task_id]["index error"] += 1
elif "UnsupportedOperation" in error:
result_stats[task_id]["unsupported operation"] += 1
elif "ClassCast" in error:
result_stats[task_id]["class cast error"] += 1
elif "NullPointer" in error:
result_stats[task_id]["null pointer error"] += 1
elif "NumberFormat" in error:
result_stats[task_id]["number format error"] += 1
elif "NoSuchElement" in error:
result_stats[task_id]["no such element"] += 1
elif "StackOverflow" in error:
result_stats[task_id]["stack overflow"] += 1
elif "Arithmetic" in error:
result_stats[task_id]["arithmetic error"] += 1
elif "OutOfMemory" in error:
result_stats[task_id]["out of memory"] += 1
elif "IllegalArgument" in error:
result_stats[task_id]["illegal argument"] += 1
else:
result_stats[task_id]["others"] += 1
elif language_type == "cpp":
if "compilation error" in error.lower():
result_stats[task_id]["compilation error"] += 1
elif "int main(): assertion" in error.lower():
result_stats[task_id]["assertion error"] += 1
elif "out_of_range" in error.lower():
result_stats[task_id]['range error'] += 1
elif "corrupted top size" in error.lower():
result_stats[task_id]['range error'] += 1
elif "length_error" in error.lower():
result_stats[task_id]['range error'] += 1
elif "invalid_argument" in error.lower():
result_stats[task_id]['invalid argument'] += 1
elif "invalid pointer" in error.lower():
result_stats[task_id]['pointer error'] += 1
elif "double free" in error.lower():
result_stats[task_id]['pointer error'] += 1
elif "free()" in error.lower():
result_stats[task_id]['pointer error'] += 1
elif "logic_error" in error.lower():
result_stats[task_id]['pointer error'] += 1
elif "sysmalloc: assertion" in error.lower():
result_stats[task_id]['pointer error'] += 1
elif "stack smashing" in error.lower():
result_stats[task_id]['out of memory'] += 1
elif "bad_alloc" in error.lower():
result_stats[task_id]['out of memory'] += 1
elif "terminate called after throwing an instance of" in error.lower():
result_stats[task_id]['package error'] += 1
else:
result_stats[task_id]["others"] += 1
elif language_type == "javascript":
if "Assertion failed" in error:
result_stats[task_id]["assertion error"] += 1
elif "SyntaxError" in error:
result_stats[task_id]["syntax error"] += 1
elif "ReferenceError" in error:
result_stats[task_id]["undefined error"] += 1
elif "timed out" in error:
result_stats[task_id]["timeout error"] += 1
elif "TypeError" in error:
result_stats[task_id]["type error"] += 1
elif "RangeError" in error:
result_stats[task_id]["range error"] += 1
else:
result_stats[task_id]["runtime error"] += 1
elif language_type == "go":
if "Error: \tNot equal:" in error:
result_stats[task_id]["assertion error"] += 1
elif "undefined" in error:
result_stats[task_id]["undefined error"] += 1
elif "expected" in error and "found" in error:
result_stats[task_id]["syntax error"] += 1
elif "illegal" in error:
result_stats[task_id]["syntax error"] += 1
elif "unexpected" in error:
result_stats[task_id]["syntax error"] += 1
elif "FAIL" in error:
result_stats[task_id]["runtime error"] += 1
elif "timed out" in error:
result_stats[task_id]["timeout error"] += 1
elif "not used" in error:
result_stats[task_id]['notused error'] += 1
elif "type" in error:
result_stats[task_id]['type error'] += 1
else:
incompleted = True
break
else:
if obj["passed"]:
result_stats[task_id]["accepted"] += 1
if incompleted:
print(f"Language not supported, aborted. {input_file}")
else:
try:
total, correct = [], []
for k, res in result_stats.items():
total.append(res["n_sample"])
correct.append(res["accepted"])
df_res = pd.DataFrame(res, index=[int(k.split("/")[-1])])
df = pd.concat([df, df_res], axis=0)
total = np.array(total)
correct = np.array(correct)
ks = [1, 10, 100, 1000]
pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean()
for k in ks if (total >= k).all()}
print(pass_at_k)
pass_at_k["file"] = input_file
pass_at_k["n"] = res["n_sample"]
pass_at_k_outs.append(pass_at_k)
output_prefix = input_file.split("/")[-1].split(".jsonl")[0]
output_file = os.path.join(output_dir, output_prefix + "_stats.xlsx")
df = df.sort_index(ascending=True)
df.to_excel(output_file)
print(f"Stats saved in {output_file}")
except Exception as e:
print(e)
print(f"Data incompleted, aborted. {input_file}")
if pass_at_k_outpath is not None:
jsonl_path = os.path.join(output_dir, pass_at_k_outpath)
with open(jsonl_path, "w") as f_out:
for p in pass_at_k_outs:
f_out.write(json.dumps(p) + "\n")
print(f"Pass at k saved in {jsonl_path}")
def main():
fire.Fire(inspect_result)
if __name__ == "__main__":
sys.exit(main())
================================================
FILE: codegeex/benchmark/metric.py
================================================
# Copyright (c) OpenAI (https://openai.com)
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ============================================================================
import itertools
import numpy as np
from typing import *
def estimate_pass_at_k(
num_samples: Union[int, List[int], np.ndarray],
num_correct: Union[List[int], np.ndarray],
k: int
) -> np.ndarray:
"""
Estimates pass@k of each problem and returns them in an array.
"""
def estimator(n: int, c: int, k: int) -> float:
"""
Calculates 1 - comb(n - c, k) / comb(n, k).
"""
if n - c < k:
return 1.0
return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))
if isinstance(num_samples, int):
num_samples_it = itertools.repeat(num_samples, len(num_correct))
else:
assert len(num_samples) == len(num_correct)
num_samples_it = iter(num_samples)
return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)])
================================================
FILE: codegeex/benchmark/rust/Cargo.toml
================================================
[package]
name = "rust"
version = "0.1.0"
edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
rand = "0.4"
regex = "1"
md5 = "0.7.0"
================================================
FILE: codegeex/benchmark/utils.py
================================================
import os
from typing import *
from codegeex.data.data_utils import stream_jsonl, LANGUAGE_TAG
IMPORT_HELPER = {
"python": [
"import math",
"import re",
"import sys",
"import copy",
"import datetime",
"import itertools",
"import collections",
"import heapq",
"import statistics",
"import functools",
"import hashlib",
"import numpy",
"import numpy as np",
"import string",
"from typing import *",
"from collections import *",
],
"go" : [
"math",
"strings",
"fmt",
"strconv",
"time",
"bytes",
"regexp",
"sort",
"math/rand",
"crypto/md5",
],
"cpp" : [
"#include<stdlib.h>",
"#include<algorithm>",
"#include<math.h>",
"#include<stdio.h>",
"#include<vector>",
"#include<string>",
"#include<climits>",
"#include<cstring>",
"#include<iostream>",
],
}
def read_dataset(
data_file: str = None,
dataset_type: str = "humaneval",
num_shot=None,
) -> Dict:
if num_shot is not None:
print(f"{num_shot}-shot setting...")
if "humaneval" in dataset_type.lower():
if data_file is None:
current_path = os.path.dirname(os.path.abspath(__file__))
data_file = os.path.join(current_path, "..", "humaneval-x", "python", "data", "humaneval_python.jsonl.gz")
dataset = {task["task_id"]: task for task in stream_jsonl(data_file)}
else:
raise f"Dataset: {dataset_type} not supported."
return dataset
def read_translation_dataset(
data_file_src: str = None,
data_file_tgt: str = None,
lang_src: str = None,
lang_tgt: str = None,
dataset_type: str = "humaneval",
) -> Dict:
if "humaneval" in dataset_type.lower():
dataset_src = {task["task_id"]: task for task in stream_jsonl(data_file_src)}
dataset_tgt = {task["task_id"].split("/")[-1]: task for task in stream_jsonl(data_file_tgt)}
for k, sample in dataset_src.items():
prompt = "code translation\n"
if lang_src == "cpp":
prompt += "C++:\n"
elif lang_src == "js":
prompt += "JavaScript:\n"
else:
prompt += f"{lang_src}:\n".capitalize()
prompt += dataset_src[k]["declaration"] + "\n" + dataset_src[k]["canonical_solution"].rstrip() + "\n"
if lang_tgt == "cpp":
prompt += "C++:\n"
elif lang_tgt == "js":
prompt += "JavaScript:\n"
else:
prompt += f"{lang_tgt}:\n".capitalize()
prompt += dataset_tgt[k.split("/")[-1]]["declaration"]
dataset_src[k]["prompt"] = prompt
else:
raise f"Dataset: {dataset_type} not supported."
return dataset_src
def process_extra_prompt(prompt: str, language_type: str = None) -> str:
"""
Processes the extra prompt.
"""
language = language_type.lower()
if language in LANGUAGE_TAG:
extra_prompt = LANGUAGE_TAG[language] + "\n"
else:
extra_prompt = ""
return extra_prompt + prompt
def is_code_generation_finished(
code: str,
language_type: str = None,
dataset: str = None,
):
"""
Checks whether the generated code is finished.
"""
if language_type is None or dataset is None:
return False
if "humaneval" in dataset.lower():
if language_type.lower() == "python":
for line in code.split("\n"):
if len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t':
return True
end_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint"]
for w in end_words:
if w in code:
return True
elif language_type.lower() == "java":
if code.count("{") + 1 == code.count("}"):
return True
elif language_type.lower() == "go":
if code.count("{") + 1 == code.count("}"):
return True
elif language_type.lower() == "js":
if code.count("{") + 1 == code.count("}"):
return True
elif language_type.lower() == "cpp":
if code.count("{") + 1 == code.count("}"):
return True
return False
def cleanup_code(
code: str,
language_type: str = None,
dataset: str = None,
):
"""
Cleans up the generated code.
"""
if language_type is None or dataset is None:
return code
if "humaneval" in dataset.lower():
if language_type.lower() == "python":
end_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint", "\nassert"]
for w in end_words:
if w in code:
code = code[:code.rfind(w)]
elif language_type.lower() == "java":
main_pos = code.find("public static void main")
if main_pos != -1:
code = code[:main_pos] + '}'
if '}' in code:
code = code[:code.rfind('}')] + '}'
if code.count('{') + 1 == code.count('}'):
code += "\n}"
elif language_type.lower() == "go":
end_words = ["\n//", "\nfunc main("]
for w in end_words:
if w in code:
code = code[:code.rfind(w)]
if '}' in code:
code = code[:code.rfind('}')] + '}'
elif language_type.lower() == "cpp":
if '}' in code:
code = code[:code.rfind('}')] + '}'
elif language_type.lower() == "js":
if '}' in code:
code = code[:code.rfind('}')] + '}'
return code
================================================
FILE: codegeex/data/__init__.py
================================================
================================================
FILE: codegeex/data/data_utils.py
================================================
import os
import gzip
import json
from typing import *
LANGUAGE_TAG = {
"c" : "// language: C",
"c++" : "// language: C++",
"cpp" : "// language: C++",
"c#" : "// language: C#",
"csharp" : "// language: C#",
"css" : "/* language: CSS */",
"cuda" : "// language: Cuda",
"dart" : "// language: Dart",
"lua" : "// language: Lua",
"objectivec" : "// language: Objective-C",
"objective-c" : "// language: Objective-C",
"objective-c++": "// language: Objective-C++",
"python" : "# language: Python",
"perl" : "# language: Perl",
"prolog" : f"% language: Prolog",
"swift" : "// language: swift",
"lisp" : "; language: Lisp",
"java" : "// language: Java",
"scala" : "// language: Scala",
"tex" : f"% language: TeX",
"vue" : "<!--language: Vue-->",
"markdown" : "<!--language: Markdown-->",
"html" : "<!--language: HTML-->",
"php" : "// language: PHP",
"js" : "// language: JavaScript",
"javascript" : "// language: JavaScript",
"typescript" : "// language: TypeScript",
"go" : "// language: Go",
"shell" : "# language: Shell",
"rust" : "// language: Rust",
"sql" : "-- language: SQL",
"kotlin" : "// language: Kotlin",
"vb" : "' language: Visual Basic",
"ruby" : "# language: Ruby",
"pascal" : "// language: Pascal",
"r" : "# language: R",
"fortran" : "!language: Fortran",
"lean" : "-- language: Lean",
"matlab" : f"% language: Matlab",
"delphi" : "{language: Delphi}",
"scheme" : "; language: Scheme",
"basic" : "' language: Basic",
"assembly" : "; language: Assembly",
"groovy" : "// language: Groovy",
"abap" : "* language: Abap",
"gdscript" : "# language: GDScript",
"haskell" : "-- language: Haskell",
"julia" : "# language: Julia",
"elixir" : "# language: Elixir",
"excel" : "' language: Excel",
"clojure" : "; language: Clojure",
"actionscript" : "// language: ActionScript",
"solidity" : "// language: Solidity",
"powershell" : "# language: PowerShell",
"erlang" : f"% language: Erlang",
"cobol" : "// language: Cobol",
}
def stream_jsonl(filename: str) -> Iterable[Dict]:
"""
Parses each jsonl line and yields it as a dictionary
"""
if filename.endswith(".gz"):
with open(filename, "rb") as gzfp:
with gzip.open(gzfp, "rt") as fp:
for line in fp:
if any(not x.isspace() for x in line):
yield json.loads(line)
else:
with open(filename, "r") as fp:
for line in fp:
if any(not x.isspace() for x in line):
yield json.loads(line)
def write_jsonl(filename: str, data: Iterable[Dict], append: bool = False):
"""
Writes an iterable of dictionaries to jsonl
"""
if append:
mode = "ab"
else:
mode = "wb"
filename = os.path.expanduser(filename)
if filename.endswith(".gz"):
with open(filename, mode) as fp:
with gzip.GzipFile(fileobj=fp, mode="wb") as gzfp:
for x in data:
gzfp.write((json.dumps(x) + "\n").encode("utf-8"))
else:
with open(filename, mode) as fp:
for x in data:
fp.write((json.dumps(x) + "\n").encode("utf-8"))
def sliding_window(
prompt_tokens: list,
code_tokens: list,
seq_len: int,
sliding_stride: int,
minimum_code_len: int = 1,
) -> Iterable[Tuple[list, list]]:
"""
Generate a series of (prompt, code) pairs by sliding the window over the code.
"""
prompt_len = len(prompt_tokens)
code_len = len(code_tokens)
total_len = prompt_len + code_len
start_idx = max(0, prompt_len - seq_len + minimum_code_len) # at least `minimum_code_len` code token should be in the window
end_idx = max(0, total_len - seq_len)
start_idx = min(start_idx, end_idx)
for i in range(start_idx, end_idx + 1, sliding_stride):
current_prompt = prompt_tokens[i:i + seq_len]
current_code = code_tokens[max(i - prompt_len, 0):i - prompt_len + seq_len]
yield current_prompt, current_code
if (end_idx - start_idx) % sliding_stride != 0:
current_prompt = prompt_tokens[end_idx:end_idx + seq_len]
current_code = code_tokens[max(end_idx - prompt_len, 0):end_idx - prompt_len + seq_len]
yield current_prompt, current_code
================================================
FILE: codegeex/data/process_pretrain_dataset.py
================================================
import os
import glob
import fire
import torch
import multiprocessing
from typing import *
from tqdm.auto import tqdm
from time import perf_counter
from black import format_str, FileMode
from codegeex.data.types import PromptDataset, PromptSample
from codegeex.data.processor import PromptDatasetProcessor
from codegeex.data.data_utils import stream_jsonl, LANGUAGE_TAG
from codegeex.megatron.data.indexed_dataset import make_mmap_builder
from codegeex.tokenizer import CodeGeeXTokenizer
def try_format_code(code: str):
# Auto-correct to PEP8 format (Change tab to 4-whitespaces;
# add whitespace around some special symbols;
# reformat line length < 100, etc.)
try:
res = format_str(code, mode=FileMode(line_length=200))
except Exception as e:
res = code
print(e)
print("Wrong python format: {}".format(code))
return res
def load_pretrain_dataset(dataset_path: Union[str, List[str]]) -> Dict:
if type(dataset_path) is str:
dataset_path = [dataset_path]
for p in dataset_path:
if not os.path.isdir(p):
if p.endswith(".gz") or p.endswith(".jsonl"):
print(f"loading from {p}")
yield from stream_jsonl(p)
else:
p_list = glob.glob(p + "/*")
for p_ in p_list:
if p_.endswith(".gz") or p_.endswith(".jsonl"):
print(f"loading from {p_}")
yield from stream_jsonl(p_)
def process_sample(
sample: Dict,
language: str=None,
mode: str="pretrain",
) -> Iterable[PromptSample]:
if mode == "pretrain":
prompt = ""
else:
prompt = sample["prompt"]
try:
if language is not None and language in LANGUAGE_TAG.keys():
code = LANGUAGE_TAG[language] + "\n" + sample["code"]
else:
code = sample["code"]
except Exception as e:
print(e)
print("The key 'code' is missing in data. Aborted")
exit(0)
yield PromptSample(prompt, code)
def generate_prompt_samples(
dataset: Iterable[Dict],
language: str = None,
mode: str = "pretrain",
) -> PromptDataset:
for sample in dataset:
yield from process_sample(sample, language, mode)
def main(
tokenizer_path: str,
dataset_path: Union[str, List[str]],
output_prefix: str,
language: str = None,
mode: str = "pretrain",
discard_overlong: bool = False,
sliding_stride: int = 200,
num_workers: int = 32,
seq_len: int = 2048,
):
DATA_KEYS = ["input_ids", "attention_mask", "labels"]
# create output dir
os.makedirs(os.path.dirname(output_prefix), exist_ok=True)
tokenizer = CodeGeeXTokenizer(tokenizer_path=tokenizer_path)
pad_token_id = tokenizer.eos_token_id
dataset = load_pretrain_dataset(dataset_path)
prompt_dataset = generate_prompt_samples(dataset, language=language, mode=mode)
if num_workers == 0:
num_workers = multiprocessing.cpu_count()
pool = multiprocessing.Pool(num_workers)
output_bin_files = {}
output_idx_files = {}
builders = {}
for key in DATA_KEYS:
output_bin_files[key] = "{}_{}.bin".format(output_prefix, key)
output_idx_files[key] = "{}_{}.idx".format(output_prefix, key)
builders[key] = make_mmap_builder(
output_bin_files[key],
vocab_size=None, # magic number, should change it
)
# NOTE that we use seq_len + 1 instead of seq_len, since the input tokens will be shifted by one.
processor = PromptDatasetProcessor(
tokenize=tokenizer.encode_code,
pad_token=pad_token_id,
max_seq_len=seq_len + 1,
discard_overlong=discard_overlong,
sliding_stride=sliding_stride,
eod_token=pad_token_id)
processor.start_time = perf_counter()
doc_iter = pool.imap_unordered(processor.process_sample_strict,
prompt_dataset,
chunksize=20)
for doc_idx, docs in tqdm(enumerate(doc_iter, start=1)):
processor.doc_processed += 1
for doc in docs:
processor.doc_generated += 1
for key in DATA_KEYS:
builders[key].add_item(torch.IntTensor(doc[key]))
for key in DATA_KEYS:
builders[key].finalize(output_idx_files[key])
if __name__ == "__main__":
fire.Fire(main)
================================================
FILE: codegeex/data/processor.py
================================================
from typing import *
from time import perf_counter
from codegeex.data.data_utils import sliding_window
from codegeex.data.types import PromptSample, LabelSample
class PromptDatasetProcessor(object):
def __init__(
self,
tokenize: Callable,
pad_token: int,
keep_order: bool = False,
max_seq_len: int = 2048,
sliding_stride: int = 200,
discard_overlong: bool = True,
eod_token: int = None,
preprocess: Callable = None,
):
super(PromptDatasetProcessor, self).__init__()
self._keep_order = keep_order
self._max_seq_len = max_seq_len
self._sliding_stride = sliding_stride
self._tokenize = tokenize
self._pad_token = pad_token
self._discard_overlong = discard_overlong
self._eod_token = eod_token
self._preprocess = preprocess
self.doc_processed = 0
self.doc_generated = 0
self.start_time = 0
def pad_seq(self, prompt_tokens: List[int], code_tokens: List[int], extra: dict = None) -> Dict[str, List[int]]:
total_length = len(prompt_tokens) + len(code_tokens)
assert total_length <= self._max_seq_len, f"padding sequence: {total_length} > {self._max_seq_len}"
pad_len = self._max_seq_len - total_length
input_ids = prompt_tokens + code_tokens + [self._pad_token] * pad_len
attention_mask = [1] * len(prompt_tokens) + [1] * len(code_tokens) + [0] * pad_len
labels = [-100] * len(prompt_tokens) + code_tokens + [-100] * pad_len
return {
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": labels,
}
def process_sample(self, sample: PromptSample) -> Iterable[Dict[str, List[int]]]:
"""
Process a sample.
"""
prompt_tokens = self._tokenize(sample.prompt)
code_tokens = self._tokenize(sample.code)
if self._eod_token is not None:
code_tokens.append(self._eod_token)
if len(prompt_tokens) + len(code_tokens) > self._max_seq_len:
if self._discard_overlong:
return
for p, t in sliding_window(prompt_tokens, code_tokens, self._max_seq_len, self._sliding_stride, self._sliding_stride):
yield self.pad_seq(p, t)
else:
yield self.pad_seq(prompt_tokens, code_tokens, extra=sample.extra)
def process_sample_strict(self, sample: PromptSample) -> List[Dict[str, List[int]]]:
"""
Instead of processing lazily, we turn the iterable into a list.
"""
if sample is None:
return None
return list(self.process_sample(sample))
def process_sample_(self, sample) -> List[Dict[str, List[int]]]:
prompt_sample = self._preprocess(sample)
return self.process_sample_strict(prompt_sample)
def report(self):
duration = perf_counter() - self.start_time
process_speed = self.doc_processed * 1.0 / duration
gen_speed = self.doc_generated * 1.0 / duration
print(f">>> processed: {self.doc_processed} in {duration:.2f}s, speed: {process_speed:.2f} docs/s")
print(f"... generated: {self.doc_generated} in {duration:.2f}s, speed: {gen_speed:.2f} docs/s")
class LabelDatasetProcessor(object):
def __init__(
self,
tokenize: Callable,
pad_token: int,
keep_order: bool = False,
max_seq_len: int = 2048,
sliding_stride: int = 200,
discard_overlong: bool = True,
eod_token: int = None,
preprocess: Callable = None,
):
super(LabelDatasetProcessor, self).__init__()
self._keep_order = keep_order
self._max_seq_len = max_seq_len
self._sliding_stride = sliding_stride
self._tokenize = tokenize
self._pad_token = pad_token
self._discard_overlong = discard_overlong
self._eod_token = eod_token
self._preprocess = preprocess
self.doc_processed = 0
self.doc_generated = 0
self.start_time = 0
def pad_seq(self, prompt_tokens: List[int], label: int, extra: dict = None) -> Dict[str, List[int]]:
total_length = len(prompt_tokens)
assert total_length <= self._max_seq_len, f"padding sequence: {total_length} > {self._max_seq_len}"
pad_len = self._max_seq_len - total_length
input_ids = prompt_tokens + [self._pad_token] * pad_len
attention_mask = [1] * len(prompt_tokens) + [0] * pad_len
label = [label]
return {
"input_ids": input_ids,
"attention_mask": attention_mask,
"length": [len(prompt_tokens)],
"labels": label
}
def process_sample(self, sample: LabelSample) -> Iterable[Dict[str, List[int]]]:
"""
Process a sample.
"""
prompt_tokens = self._tokenize(sample.prompt)
label = sample.label
if len(prompt_tokens) > self._max_seq_len:
if self._discard_overlong:
return
prompt_tokens=prompt_tokens[-self._max_seq_len:]
yield self.pad_seq(prompt_tokens, label, extra=sample.extra)
def process_sample_strict(self, sample: LabelSample) -> List[Dict[str, List[int]]]:
"""
Instead of processing lazily, we turn the iterable into a list.
"""
if sample is None:
return None
return list(self.process_sample(sample))
def process_sample_(self, sample) -> List[Dict[str, List[int]]]:
prompt_sample = self._preprocess(sample)
return self.process_sample_strict(prompt_sample)
def report(self):
duration = perf_counter() - self.start_time
process_speed = self.doc_processed * 1.0 / duration
gen_speed = self.doc_generated * 1.0 / duration
print(f">>> processed: {self.doc_processed} in {duration:.2f}s, speed: {process_speed:.2f} docs/s")
print(f"... generated: {self.doc_generated} in {duration:.2f}s, speed: {gen_speed:.2f} docs/s")
================================================
FILE: codegeex/data/types.py
================================================
from typing import *
from dataclasses import dataclass
@dataclass
class PromptSample:
prompt: str
code: str
extra: dict = None
PromptDataset = Iterable[PromptSample]
@dataclass
class LabelSample:
prompt: str
label: int
extra: dict = None
LabelDataset = Iterable[LabelSample]
================================================
FILE: codegeex/docker/Dockerfile
================================================
FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime
RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak \
&& sed -i "s@http://.*archive.ubuntu.com@https://mirrors.tuna.tsinghua.edu.cn@g" /etc/apt/sources.list \
&& sed -i "s@http://.*security.ubuntu.com@https://mirrors.tuna.tsinghua.edu.cn@g" /etc/apt/sources.list \
&& pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple \
&& apt-get update && apt-get install -y curl npm git nano \
&& pip install fire zmq transformers tokenizers \
&& mkdir /workspace/download
RUN curl -o /workspace/download/go.tar.gz -SL https://go.dev/dl/go1.18.4.linux-amd64.tar.gz \
&& tar -zxf /workspace/download/go.tar.gz -C /usr/local && rm /workspace/download/go.tar.gz
ENV PATH=/bin:/usr/local/go/bin:$PATH
RUN curl -o /workspace/download/node.tar.gz -SL https://nodejs.org/download/release/v16.14.0/node-v16.14.0-linux-x64.tar.gz \
&& mkdir -p /usr/local/lib/nodejs && tar -zxf /workspace/download/node.tar.gz -C /usr/local/lib/nodejs && mv /usr/local/lib/nodejs/node-v16.14.0-linux-x64 /usr/local/lib/nodejs/node \
&& rm /workspace/download/node.tar.gz && npm install -g js-md5@0.7.3
ENV PATH=/usr/local/lib/nodejs/node/bin:$PATH
ENV NODE_PATH=/usr/local/lib/node_modules
RUN apt-get install -y build-essential && apt-get install -y g++ \
&& curl -o /workspace/download/boost_1_71_0.tar.gz -SL https://boostorg.jfrog.io/artifactory/main/release/1.71.0/source/boost_1_71_0.tar.gz \
&& tar -zxf /workspace/download/boost_1_71_0.tar.gz && rm /workspace/download/boost_1_71_0.tar.gz && cd boost_1_71_0 \
&& ./bootstrap.sh --prefix=/usr/ && ./b2 && ./b2 install \
&& cd .. && rm -r boost_1_71_0
RUN curl -o /workspace/download/openssl.tar.gz -SL https://www.openssl.org/source/old/3.0/openssl-3.0.0.tar.gz \
&& tar -zxf /workspace/download/openssl.tar.gz && cd openssl-3.0.0 && ./Configure && make && make install \
&& rm /workspace/download/openssl.tar.gz && rm -r /workspace/openssl-3.0.0
ENV PATH=/usr/bin/openssl:$PATH
RUN curl -o /workspace/download/jdk.tar.gz -SL https://download.oracle.com/java/18/latest/jdk-18_linux-x64_bin.tar.gz \
&& mkdir /usr/java && tar -zxf /workspace/download/jdk.tar.gz -C /usr/java && rm /workspace/download/jdk.tar.gz \
&& java_path=`ls /usr/java/${path}` && echo "export JAVA_HOME=/usr/java/${java_path}" >> ~/.profile
RUN git clone https://github.com/THUDM/CodeGeeX.git && pip install -e /workspace/CodeGeeX \
&& tar -zxf /workspace/CodeGeeX/codegeex/benchmark/humaneval-x/go/evaluation/vendor.tar.gz -C /workspace/CodeGeeX/codegeex/benchmark/humaneval-x/go/evaluation
SHELL ["sh", "-lc"]
RUN update-alternatives --install /usr/bin/java java $JAVA_HOME/bin/java 20000 \
&& update-alternatives --install /usr/bin/javac javac $JAVA_HOME/bin/javac 20000 \
&& rm -r /workspace/download \
&& env
================================================
FILE: codegeex/kernels/__init__.py
================================================
import pkg_resources
import torch
import ctypes
from typing import List
from cpm_kernels.kernels.base import LazyKernelCModule, KernelFunction, round_up
RESOURCE_PACKAGE_NAME = __name__
class Kernel:
def __init__(self, filename: str, function_names: List[str]):
filename = filename + ".fatbin"
if not pkg_resources.resource_exists(RESOURCE_PACKAGE_NAME, filename):
raise RuntimeError("File `%s` not found in `%s`" % (filename, RESOURCE_PACKAGE_NAME))
self.filename = filename
self.code = pkg_resources.resource_string(RESOURCE_PACKAGE_NAME, filename)
self._function_names = function_names
self._cmodule = LazyKernelCModule(self.code)
for name in self._function_names:
setattr(self, name, KernelFunction(self._cmodule, name))
kernels = Kernel(
"quantization",
[
"int4WeightCompression",
"int4WeightExtractionFloat",
"int4WeightExtractionHalf",
"int8WeightExtractionFloat",
"int8WeightExtractionHalf",
],
)
def compress_int4_weight(weight: torch.Tensor): # (n, m)
with torch.cuda.device(weight.device):
n, m = weight.size(0), weight.size(1)
assert m % 2 == 0
m = m // 2
out = torch.empty(n, m, dtype=torch.int8, device="cuda")
stream = torch.cuda.current_stream()
gridDim = (n, 1, 1)
blockDim = (min(round_up(m, 32), 1024), 1, 1)
kernels.int4WeightCompression(
gridDim,
blockDim,
0,
stream,
[ctypes.c_void_p(weight.data_ptr()), ctypes.c_void_p(out.data_ptr()), ctypes.c_int32(n), ctypes.c_int32(m)],
)
return out
def extract_weight_to_half(weight: torch.Tensor, scale_list: torch.Tensor, source_bit_width: int):
if source_bit_width == 8:
func = kernels.int8WeightExtractionHalf
elif source_bit_width == 4:
func = kernels.int4WeightExtractionHalf
else:
assert False, "Unsupported bit-width"
with torch.cuda.device(weight.device):
n, m = weight.size(0), weight.size(1)
out = torch.empty(n, m * (8 // source_bit_width), dtype=torch.half, device="cuda")
stream = torch.cuda.current_stream()
gridDim = (n, 1, 1)
blockDim = (min(round_up(m, 32), 1024), 1, 1)
func(
gridDim,
blockDim,
0,
stream,
[
ctypes.c_void_p(weight.data_ptr()),
ctypes.c_void_p(scale_list.data_ptr()),
ctypes.c_void_p(out.data_ptr()),
ctypes.c_int32(n),
ctypes.c_int32(m),
],
)
return out
if __name__ == "__main__":
weight = torch.randn(4, 32).to(torch.int8).cuda()
scale = torch.ones(weight.size(0)).to(torch.half).cuda()
print(weight)
b = compress_int4_weight(weight)
print(b)
a = extract_weight_to_half(b, scale, source_bit_width=4)
print(a)
================================================
FILE: codegeex/megatron/__init__.py
================================================
# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
from .global_vars import get_args
from .global_vars import get_current_global_batch_size
from .global_vars import get_num_microbatches
from .global_vars import update_num_microbatches
from .global_vars import get_tokenizer
from .global_vars import get_tensorboard_writer
from .global_vars import get_adlr_autoresume
from .global_vars import get_timers
from .initialize import initialize_megatron
def print_rank_0(message):
"""If distributed is initialized, print only on rank 0."""
if torch.distributed.is_initialized():
if torch.distributed.get_rank() == 0:
print(message, flush=True)
else:
print(message, flush=True)
def is_last_rank():
return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1)
def print_rank_last(message):
"""If distributed is initialized, print only on last rank."""
if torch.distributed.is_initialized():
if is_last_rank():
print(message, flush=True)
else:
print(message, flush=True)
================================================
FILE: codegeex/megatron/arguments.py
================================================
# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Megatron arguments."""
import argparse
import os
import torch
import deepspeed
def parse_args(extra_args_provider=None, defaults={}, ignore_unknown_args=False):
"""Parse all arguments."""
parser = argparse.ArgumentParser(
description="Megatron-LM Arguments", allow_abbrev=False
)
# Standard arguments.
parser = _add_network_size_args(parser)
parser = _add_regularization_args(parser)
parser = _add_training_args(parser)
parser = _add_initialization_args(parser)
parser = _add_learning_rate_args(parser)
parser = _add_checkpointing_args(parser)
parser = _add_mixed_precision_args(parser)
parser = _add_distributed_args(parser)
parser = _add_validation_args(parser)
parser = _add_data_args(parser)
parser = _add_autoresume_args(parser)
parser = _add_biencoder_args(parser)
parser = _add_vit_args(parser)
parser = _add_logging_args(parser)
parser = _add_zero_args(parser)
parser = _add_memoryopt_args(parser)
parser = _add_activation_checkpoint_args(parser)
parser = _add_inference_args(parser)
# Custom arguments.
if extra_args_provider is not None:
parser = extra_args_provider(parser)
parser = deepspeed.add_config_arguments(parser)
# Parse.
if ignore_unknown_args:
args, _ = parser.parse_known_args()
else:
args = parser.parse_args()
# helper argument to set deepspeed pipeline parallel or not
args.ds_pipeline_enabled = not args.no_pipeline_parallel
# Distributed args.
args.rank = int(os.getenv("RANK", "0"))
args.world_size = int(os.getenv("WORLD_SIZE", "1"))
# Tensor model parallel size.
args.tensor_model_parallel_size = min(
args.tensor_model_parallel_size, args.world_size
)
assert (
args.world_size % args.tensor_model_parallel_size == 0
), "world size" " ({}) is not divisible by tensor model parallel size ({})".format(
args.world_size, args.tensor_model_parallel_size
)
# Pipeline model parallel size.
args.pipeline_model_parallel_size = min(
args.pipeline_model_parallel_size,
(args.world_size // args.tensor_model_parallel_size),
)
# Checks.
if args.no_pipeline_parallel:
assert (
args.pipeline_model_parallel_size == 1
), "pipeline_model_parallel_size must be 1 if pipeline parallel is disabled"
model_parallel_size = (
args.pipeline_model_parallel_size * args.tensor_model_parallel_size
)
assert args.world_size % model_parallel_size == 0, (
"world size is not"
" divisible by tensor parallel size ({}) times pipeline parallel "
"size ({})".format(
args.world_size,
args.tensor_model_parallel_size,
args.pipeline_model_parallel_size,
)
)
args.data_parallel_size = args.world_size // model_parallel_size
if args.rank == 0:
print(
"using world size: {}, data-parallel-size: {}, "
"tensor-model-parallel size: {}, "
"pipeline-model-parallel size: {} ".format(
args.world_size,
args.data_parallel_size,
args.tensor_model_parallel_size,
args.pipeline_model_parallel_size,
),
flush=True,
)
# Deprecated arguments
assert args.batch_size is None, (
"--batch-size argument is no longer " "valid, use --micro-batch-size instead"
)
del args.batch_size
assert args.warmup is None, (
"--warmup argument is no longer valid, use " "--lr-warmup-fraction instead"
)
del args.warmup
assert args.model_parallel_size is None, (
"--model-parallel-size is no "
"longer valid, use --tensor-model-parallel-size instead"
)
del args.model_parallel_size
# Set input defaults.
for key in defaults:
# For default to be valid, it should not be provided in the
# arguments that are passed to the program. We check this by
# ensuring the arg is set to None.
if getattr(args, key) is not None:
if args.force_default:
print(
"WARNING: overriding arguments for {key}:{v2} \
with default {key}:{v}".format(
key=key, v=defaults[key], v2=getattr(args, key)
),
flush=True,
)
setattr(args, key, defaults[key])
else:
if args.rank == 0:
print(
"WARNING: overriding default arguments for {key}:{v} \
with {key}:{v2}".format(
key=key, v=defaults[key], v2=getattr(args, key)
),
flush=True,
)
else:
setattr(args, key, defaults[key])
# Batch size.
assert args.micro_batch_size is not None
assert args.micro_batch_size > 0
if args.global_batch_size is None:
args.global_batch_size = args.micro_batch_size * args.data_parallel_size
if args.rank == 0:
print(
"setting global batch size to {}".format(args.global_batch_size),
flush=True,
)
assert args.global_batch_size > 0
if args.num_layers_per_virtual_pipeline_stage is not None:
assert args.pipeline_model_parallel_size > 2, (
"pipeline-model-parallel size should be greater than 2 with "
"interleaved schedule"
)
assert args.num_layers % args.num_layers_per_virtual_pipeline_stage == 0, (
"number of layers is not divisible by number of layers per virtual "
"pipeline stage"
)
args.virtual_pipeline_model_parallel_size = (
args.num_layers // args.pipeline_model_parallel_size
) // args.num_layers_per_virtual_pipeline_stage
else:
args.virtual_pipeline_model_parallel_size = None
# Parameters dtype.
args.params_dtype = torch.float
if args.fp16:
assert not args.bf16
args.params_dtype = torch.half
if args.bf16:
assert not args.fp16
args.params_dtype = torch.bfloat16
# bfloat16 requires gradient accumulation and all-reduce to
# be done in fp32.
if not args.accumulate_allreduce_grads_in_fp32:
args.accumulate_allreduce_grads_in_fp32 = True
if args.rank == 0:
print(
"accumulate and all-reduce gradients in fp32 for "
"bfloat16 data type.",
flush=True,
)
if args.rank == 0:
print("using {} for parameters ...".format(args.params_dtype), flush=True)
# If we do accumulation and all-reduces in fp32, we need to have
# local DDP and we should set the use-contiguous-buffers-in-ddp.
if args.accumulate_allreduce_grads_in_fp32:
assert args.DDP_impl == "local"
args.use_contiguous_buffers_in_ddp = True
if args.dataloader_type is None:
args.dataloader_type = "single"
# Consumed tokens.
args.consumed_train_samples = 0
args.consumed_valid_samples = 0
args.consumed_train_tokens = 0
# Iteration-based training.
if args.train_iters:
# If we use iteration-based training, make sure the
# sample-based options are off.
assert args.train_samples is None, "expected iteration-based training"
assert (
args.lr_decay_samples is None
), "expected iteration-based learning rate decay"
assert (
args.lr_warmup_samples == 0
), "expected iteration-based learning rate warmup"
assert (
args.rampup_batch_size is None
), "expected no batch-size rampup for iteration-based training"
if args.lr_warmup_fraction is not None:
assert (
args.lr_warmup_iters == 0
), "can only specify one of lr-warmup-fraction and lr-warmup-iters"
# Sample-based training.
if args.train_samples:
# If we use sample-based training, make sure the
# iteration-based options are off.
assert args.train_iters is None, "expected sample-based training"
assert args.lr_decay_iters is None, "expected sample-based learning rate decay"
assert args.lr_warmup_iters == 0, "expected sample-based learnig rate warmup"
if args.lr_warmup_fraction is not None:
assert args.lr_warmup_samples == 0, (
"can only specify one of lr-warmup-fraction " "and lr-warmup-samples"
)
# Check required arguments.
required_args = [
"num_layers",
"hidden_size",
"num_attention_heads",
"max_position_embeddings",
]
for req_arg in required_args:
_check_arg_is_not_none(args, req_arg)
# args.learned_position_embeddings = args.learned_position_embeddings > 0
# Checks.
if args.ffn_hidden_size is None:
args.ffn_hidden_size = 4 * args.hidden_size
if args.kv_channels is None:
assert args.hidden_size % args.num_attention_heads == 0
args.kv_channels = args.hidden_size // args.num_attention_heads
if args.seq_length is not None:
assert args.encoder_seq_length is None
args.encoder_seq_length = args.seq_length
else:
assert args.encoder_seq_length is not None
args.seq_length = args.encoder_seq_length
if args.seq_length is not None:
assert args.max_position_embeddings >= args.seq_length
if args.decoder_seq_length is not None:
assert args.max_position_embeddings >= args.decoder_seq_length
if args.lr is not None:
assert args.min_lr <= args.lr
if args.save is not None:
assert args.save_interval is not None
# Mixed precision checks.
if args.fp16_lm_cross_entropy:
assert args.fp16, "lm cross entropy in fp16 only support in fp16 mode."
if args.fp32_residual_connection:
assert (
args.fp16 or args.bf16
), "residual connection in fp32 only supported when using fp16 or bf16."
# Activation checkpointing.
if args.distribute_checkpointed_activations:
assert args.checkpoint_activations, (
"for distribute-checkpointed-activations to work you "
"need to enable checkpoint-activations"
)
_print_args(args)
return args
def _print_args(args):
"""Print arguments."""
if args.rank == 0:
print("------------------------ arguments ------------------------", flush=True)
str_list = []
for arg in vars(args):
dots = "." * (48 - len(arg))
str_list.append(" {} {} {}".format(arg, dots, getattr(args, arg)))
for arg in sorted(str_list, key=lambda x: x.lower()):
print(arg, flush=True)
print("-------------------- end of arguments ---------------------", flush=True)
def _check_arg_is_not_none(args, arg):
assert getattr(args, arg) is not None, "{} argument is None".format(arg)
def _add_network_size_args(parser):
group = parser.add_argument_group(title="network size")
group.add_argument(
"--num-layers",
type=int,
default=None,
help="Number of transformer layers.",
)
group.add_argument(
"--hidden-size",
type=int,
default=None,
help="Transformer hidden size.",
)
group.add_argument(
"--reward-growth",
type=str,
default="constant",
choices=["constant", "linear", "quadratic"],
help="Reward growth function.",
)
group.add_argument(
"--ffn-hidden-size",
type=int,
default=None,
help="Transformer Feed-Forward Network hidden size. "
"This is set to 4*hidden-size if not provided",
)
group.add_argument(
"--num-attention-heads",
type=int,
default=None,
help="Number of transformer attention heads.",
)
group.add_argument(
"--kv-channels",
type=int,
default=None,
help="Projection weights dimension in multi-head "
"attention. This is set to "
" args.hidden_size // args.num_attention_heads "
"if not provided.",
)
group.add_argument(
"--scale-embeddings",
action="store_true",
help="Scale embeddings by sqrt(d_model).",
)
group.add_argument(
"--max-position-embeddings",
type=int,
default=None,
help="Maximum number of position embeddings to use. "
"This is the size of position embedding.",
)
group.add_argument(
"--no-learned-position-embeddings",
action="store_true",
help="Do not learn position embeddings. ",
)
group.add_argument(
"--make-vocab-size-divisible-by",
type=int,
default=128,
help="Pad the vocab size to be divisible by this value."
"This is added for computational efficieny reasons.",
)
group.add_argument(
"--layernorm-epsilon", type=float, default=1e-5, help="Layer norm epsilon."
)
group.add_argument(
"--apply-residual-connection-post-layernorm",
action="store_true",
help="If set, use original BERT residula connection " "ordering.",
)
group.add_argument('--scaled-upper-triang-masked-softmax-fusion',
action='store_true',
help='Enable fusion of query_key_value_scaling '
'time (upper diagonal) masking, softmax.')
group.add_argument(
"--openai-gelu",
action="store_true",
help="Use OpenAIs GeLU implementation. This option"
"should not be used unless for backward compatibility"
"reasons.",
)
group.add_argument(
"--onnx-safe",
type=bool,
required=False,
help="Use workarounds for known problems with " "Torch ONNX exporter",
)
group.add_argument(
"--bert-no-binary-head",
action="store_false",
help="Disable BERT binary head.",
dest="bert_binary_head",
)
group.add_argument(
"--compress",
action="store_true",
)
return parser
def _add_logging_args(parser):
group = parser.add_argument_group(title="logging")
group.add_argument(
"--log-params-norm",
action="store_true",
help="If set, calculate and log parameters norm.",
)
group.add_argument(
"--log-num-zeros-in-grad",
action="store_true",
help="If set, calculate and log the number of zeros in gradient.",
)
group.add_argument(
"--tensorboard-log-interval",
type=int,
default=1,
help="Report to tensorboard interval.",
)
group.add_argument(
"--tensorboard-queue-size",
type=int,
default=1000,
help="Size of the tensorboard queue for pending events "
"and summaries before one of the ‘add’ calls forces a "
"flush to disk.",
)
group.add_argument(
"--log-timers-to-tensorboard",
action="store_true",
help="If set, write timers to tensorboard.",
)
group.add_argument(
"--log-batch-size-to-tensorboard",
action="store_true",
help="If set, write batch-size to tensorboard.",
)
group.add_argument(
"--no-log-learnig-rate-to-tensorboard",
action="store_false",
help="Disable learning rate logging to tensorboard.",
dest="log_learning_rate_to_tensorboard",
)
group.add_argument(
"--no-log-loss-scale-to-tensorboard",
action="store_false",
help="Disable loss-scale logging to tensorboard.",
dest="log_loss_scale_to_tensorboard",
)
group.add_argument(
"--log-validation-ppl-to-tensorboard",
action="store_true",
help="If set, write validation perplexity to " "tensorboard.",
)
group.add_argument(
"--wandb-logging",
action="store_true",
help="If set, log training progress to wandb.",
)
group.add_argument(
"--wandb-log-interval",
type=int,
default=1,
help="Log to wandb every N steps.",
)
return parser
def _add_regularization_args(parser):
group = parser.add_argument_group(title="regularization")
group.add_argument(
"--attention-dropout",
type=float,
default=0.1,
help="Post attention dropout probability.",
)
group.add_argument(
"--hidden-dropout",
type=float,
default=0.1,
help="Dropout probability for hidden state transformer.",
)
group.add_argument(
"--weight-decay",
type=float,
default=0.01,
help="Weight decay coefficient for L2 regularization.",
)
group.add_argument(
"--tempering",
type=float,
default=None,
help="Tempering coefficient for the model.",
)
gitextract_96pnh25s/
├── .gitmodules
├── LICENSE
├── MODEL_LICENSE
├── README.md
├── README_zh.md
├── api/
│ ├── README_zh.md
│ ├── codegeex-api-example-java/
│ │ ├── pom.xml
│ │ └── src/
│ │ └── main/
│ │ └── java/
│ │ └── cn/
│ │ └── aminer/
│ │ └── codegeex/
│ │ └── example/
│ │ ├── CodeGenerationExample.java
│ │ └── pojo/
│ │ └── Payload.java
│ └── codegeex-api-example-python/
│ └── generation_example.py
├── codegeex/
│ ├── __init__.py
│ ├── benchmark/
│ │ ├── README.md
│ │ ├── README_zh.md
│ │ ├── __init__.py
│ │ ├── evaluate_humaneval_x.py
│ │ ├── execution.py
│ │ ├── gather_output.py
│ │ ├── humaneval-x/
│ │ │ ├── __init__.py
│ │ │ ├── cpp/
│ │ │ │ └── evaluation/
│ │ │ │ └── test.cpp
│ │ │ ├── evaluate_humaneval_x.py
│ │ │ ├── generate_humaneval_x.py
│ │ │ ├── go/
│ │ │ │ └── evaluation/
│ │ │ │ ├── go.mod
│ │ │ │ └── go.sum
│ │ │ └── translate_humaneval_x.py
│ │ ├── inspect_result.py
│ │ ├── metric.py
│ │ ├── rust/
│ │ │ └── Cargo.toml
│ │ └── utils.py
│ ├── data/
│ │ ├── __init__.py
│ │ ├── data_utils.py
│ │ ├── process_pretrain_dataset.py
│ │ ├── processor.py
│ │ └── types.py
│ ├── docker/
│ │ └── Dockerfile
│ ├── kernels/
│ │ ├── __init__.py
│ │ └── quantization.fatbin
│ ├── megatron/
│ │ ├── __init__.py
│ │ ├── arguments.py
│ │ ├── checkpointing.py
│ │ ├── code_generation_utils.py
│ │ ├── convert_ckpt_parallel.py
│ │ ├── data/
│ │ │ ├── __init__.py
│ │ │ ├── blendable_dataset.py
│ │ │ ├── data_samplers.py
│ │ │ ├── dataset_utils.py
│ │ │ ├── helpers.cpp
│ │ │ ├── indexed_dataset.py
│ │ │ └── prompt_dataset.py
│ │ ├── enums.py
│ │ ├── global_vars.py
│ │ ├── inference.py
│ │ ├── initialize.py
│ │ ├── learning_rates.py
│ │ ├── memory.py
│ │ ├── merge_ckpt_parallel.py
│ │ ├── microbatches.py
│ │ ├── mindspore_to_megatron.py
│ │ ├── model/
│ │ │ ├── __init__.py
│ │ │ ├── codegeex_model.py
│ │ │ ├── distributed.py
│ │ │ ├── language_model.py
│ │ │ ├── module.py
│ │ │ ├── transformer.py
│ │ │ └── utils.py
│ │ ├── mpu/
│ │ │ ├── __init__.py
│ │ │ ├── cross_entropy.py
│ │ │ ├── data.py
│ │ │ ├── initialize.py
│ │ │ ├── layers.py
│ │ │ ├── mappings.py
│ │ │ ├── random.py
│ │ │ └── utils.py
│ │ ├── optimizer/
│ │ │ ├── __init__.py
│ │ │ ├── clip_grads.py
│ │ │ ├── grad_scaler.py
│ │ │ └── optimizer.py
│ │ ├── p2p_communication.py
│ │ ├── schedules.py
│ │ ├── tokenizer/
│ │ │ ├── __init__.py
│ │ │ ├── gpt2_tokenization.py
│ │ │ └── tokenizer.py
│ │ ├── tools/
│ │ │ ├── collect_env.py
│ │ │ ├── finetune_codegeex.py
│ │ │ └── pretrain_codegeex.py
│ │ ├── training.py
│ │ └── utils.py
│ ├── mindspore/
│ │ ├── configs/
│ │ │ ├── 13B.sh
│ │ │ ├── 13B_128p_save_1p.sh
│ │ │ ├── 13B_128p_save_8p_ckpt.sh
│ │ │ ├── 13B_1p_to_torch.sh
│ │ │ ├── 13B_finetune.sh
│ │ │ ├── 13B_generate.sh
│ │ │ ├── 13B_generate_1p.sh
│ │ │ ├── 13B_generate_1p_values.sh
│ │ │ ├── 13B_generate_finetune.sh
│ │ │ ├── 13B_generate_humaneval.sh
│ │ │ └── 13B_generate_values.sh
│ │ ├── convertion_1p.py
│ │ ├── finetune.py
│ │ ├── generation.py
│ │ ├── generation_1p.py
│ │ ├── generation_batch.py
│ │ ├── generation_finetune.py
│ │ ├── generation_humaneval.py
│ │ ├── generation_values.py
│ │ ├── generation_values_1p.py
│ │ ├── save_1p_ckpt_from_8p_ckpt.py
│ │ ├── save_8p_ckpt.py
│ │ ├── scripts/
│ │ │ ├── custom_tune_bank_new/
│ │ │ │ └── Ascend910ProA/
│ │ │ │ ├── cube/
│ │ │ │ │ └── repository_ascend910ProA_matmul.json
│ │ │ │ └── vector/
│ │ │ │ └── Ascend910ProA_AiCore_32_v001_20220509_200939_588817.json
│ │ │ ├── layer_norm.py
│ │ │ ├── layer_norm_x_backprop_v2.py
│ │ │ ├── ma-pre-start.sh
│ │ │ ├── run_modelarts.py
│ │ │ ├── run_modelarts_gen_finetune.py
│ │ │ └── run_modelarts_gen_humaneval_x.py
│ │ ├── src/
│ │ │ ├── __init__.py
│ │ │ ├── adam.py
│ │ │ ├── callbacks.py
│ │ │ ├── code_tokenizer.py
│ │ │ ├── dataset.py
│ │ │ ├── dataset_finetune.py
│ │ │ ├── generate.py
│ │ │ ├── generate_finetune.py
│ │ │ ├── generate_greedy.py
│ │ │ ├── generate_humaneval.py
│ │ │ ├── metrics.py
│ │ │ ├── pangu_alpha.py
│ │ │ ├── pangu_alpha_config.py
│ │ │ ├── pangu_alpha_fp16_predict.py
│ │ │ ├── pangu_alpha_wrapcell.py
│ │ │ ├── pangu_alpha_wrapcell_finetune.py
│ │ │ ├── preprocess.py
│ │ │ ├── sat_dataset.py
│ │ │ ├── tokenization_jieba.py
│ │ │ └── utils.py
│ │ └── train.py
│ ├── oneflow/
│ │ ├── __init__.py
│ │ ├── codegeex_model.py
│ │ └── inference.py
│ ├── paddle/
│ │ ├── __init__.py
│ │ ├── codegeex_model.py
│ │ ├── inference.py
│ │ └── pt_to_pdparams.py
│ ├── quantization/
│ │ ├── __init__.py
│ │ ├── quantize.py
│ │ └── quantize_oneflow.py
│ ├── tokenizer/
│ │ ├── __init__.py
│ │ ├── added_tokens.json
│ │ ├── merges.txt
│ │ ├── special_tokens_map.json
│ │ ├── tokenizer.py
│ │ ├── tokenizer_config.json
│ │ └── vocab.json
│ └── torch/
│ ├── __init__.py
│ ├── codegeex_model.py
│ ├── get_ckpt_qkv.py
│ └── inference.py
├── configs/
│ ├── codegeex_13b.sh
│ ├── codegeex_13b_paddle.sh
│ └── codegeex_13b_parallel.sh
├── deployment/
│ ├── example_inputs.jsonl
│ └── server_gradio.py
├── requirements.txt
├── resources/
│ └── zh/
│ └── wechat.md
├── scripts/
│ ├── convert_ckpt_parallel.sh
│ ├── convert_mindspore_to_megatron.sh
│ ├── evaluate_humaneval_x.py
│ ├── evaluate_humaneval_x.sh
│ ├── finetune_codegeex.sh
│ ├── gather_output.sh
│ ├── generate_humaneval_x.sh
│ ├── pretrain_codegeex.sh
│ ├── process_pretrain_dataset.sh
│ ├── test_inference.sh
│ ├── test_inference_oneflow.sh
│ ├── test_inference_oneflow_quantized.sh
│ ├── test_inference_paddle.sh
│ ├── test_inference_parallel.sh
│ ├── test_inference_quantized.sh
│ └── translate_humaneval_x.sh
├── setup.py
├── tests/
│ ├── test_inference.py
│ ├── test_inference_megatron.py
│ ├── test_inference_oneflow.py
│ ├── test_inference_paddle.py
│ └── test_prompt.txt
└── vscode-extension/
├── README.md
└── README_zh.md
SYMBOL INDEX (1224 symbols across 114 files)
FILE: api/codegeex-api-example-java/src/main/java/cn/aminer/codegeex/example/CodeGenerationExample.java
class CodeGenerationExample (line 15) | public class CodeGenerationExample {
method main (line 22) | public static void main(String[] args) throws Exception {
method generateCode (line 35) | public void generateCode(String prompt) throws Exception {
method performHttpPost (line 50) | public String performHttpPost(String url, String payload) {
FILE: api/codegeex-api-example-java/src/main/java/cn/aminer/codegeex/example/pojo/Payload.java
class Payload (line 14) | @JsonIgnoreProperties(ignoreUnknown = true)
FILE: api/codegeex-api-example-python/generation_example.py
function main (line 33) | def main():
FILE: codegeex/__init__.py
function get_model (line 8) | def get_model(
function generate (line 15) | def generate(
FILE: codegeex/benchmark/evaluate_humaneval_x.py
function process_humaneval_test (line 27) | def process_humaneval_test(sample, problems, example_test=False):
function stream_jsonl_all (line 85) | def stream_jsonl_all(filename: str) -> Iterable[Dict]:
function evaluate_functional_correctness (line 99) | def evaluate_functional_correctness(
function main (line 234) | def main():
FILE: codegeex/benchmark/execution.py
function dicts_to_jsonl (line 15) | def dicts_to_jsonl(data_list: list, filename: str, compress: bool = True...
function check_correctness (line 44) | def check_correctness(
function time_limit (line 409) | def time_limit(seconds: float):
function swallow_io (line 422) | def swallow_io():
function create_tempdir (line 431) | def create_tempdir():
class TimeoutException (line 437) | class TimeoutException(Exception):
class WriteOnlyStringIO (line 441) | class WriteOnlyStringIO(io.StringIO):
method read (line 444) | def read(self, *args, **kwargs):
method readline (line 447) | def readline(self, *args, **kwargs):
method readlines (line 450) | def readlines(self, *args, **kwargs):
method readable (line 453) | def readable(self, *args, **kwargs):
class redirect_stdin (line 458) | class redirect_stdin(contextlib._RedirectStream): # type: ignore
function chdir (line 463) | def chdir(root):
function reliability_guard (line 477) | def reliability_guard(maximum_memory_bytes: Optional[int] = None):
FILE: codegeex/benchmark/gather_output.py
function gather_output (line 7) | def gather_output(
function main (line 51) | def main():
FILE: codegeex/benchmark/humaneval-x/cpp/evaluation/test.cpp
function separate_paren_groups (line 13) | vector<string> separate_paren_groups(string paren_string){
function issame (line 43) | bool issame(vector<string> a,vector<string>b){
function main (line 51) | int main(){
FILE: codegeex/benchmark/humaneval-x/evaluate_humaneval_x.py
function process_humaneval_test (line 27) | def process_humaneval_test(sample, problems, example_test=False):
function stream_jsonl_all (line 85) | def stream_jsonl_all(filename: str) -> Iterable[Dict]:
function evaluate_functional_correctness (line 99) | def evaluate_functional_correctness(
function main (line 234) | def main():
FILE: codegeex/benchmark/humaneval-x/generate_humaneval_x.py
function add_code_generation_args (line 21) | def add_code_generation_args(parser):
function main (line 202) | def main(node_rank: int, local_rank: int, master_port: int, num_devices:...
function server (line 250) | def server():
FILE: codegeex/benchmark/humaneval-x/translate_humaneval_x.py
function add_code_generate_args (line 19) | def add_code_generate_args(parser):
function main (line 228) | def main(node_rank: int, local_rank: int, master_port: int, num_devices:...
function server (line 274) | def server():
FILE: codegeex/benchmark/inspect_result.py
function inspect_result (line 73) | def inspect_result(
function main (line 281) | def main():
FILE: codegeex/benchmark/metric.py
function estimate_pass_at_k (line 27) | def estimate_pass_at_k(
FILE: codegeex/benchmark/utils.py
function read_dataset (line 51) | def read_dataset(
function read_translation_dataset (line 69) | def read_translation_dataset(
function process_extra_prompt (line 102) | def process_extra_prompt(prompt: str, language_type: str = None) -> str:
function is_code_generation_finished (line 115) | def is_code_generation_finished(
function cleanup_code (line 151) | def cleanup_code(
FILE: codegeex/data/data_utils.py
function stream_jsonl (line 67) | def stream_jsonl(filename: str) -> Iterable[Dict]:
function write_jsonl (line 84) | def write_jsonl(filename: str, data: Iterable[Dict], append: bool = False):
function sliding_window (line 104) | def sliding_window(
FILE: codegeex/data/process_pretrain_dataset.py
function try_format_code (line 19) | def try_format_code(code: str):
function load_pretrain_dataset (line 32) | def load_pretrain_dataset(dataset_path: Union[str, List[str]]) -> Dict:
function process_sample (line 49) | def process_sample(
function generate_prompt_samples (line 72) | def generate_prompt_samples(
function main (line 81) | def main(
FILE: codegeex/data/processor.py
class PromptDatasetProcessor (line 8) | class PromptDatasetProcessor(object):
method __init__ (line 9) | def __init__(
method pad_seq (line 34) | def pad_seq(self, prompt_tokens: List[int], code_tokens: List[int], ex...
method process_sample (line 48) | def process_sample(self, sample: PromptSample) -> Iterable[Dict[str, L...
method process_sample_strict (line 66) | def process_sample_strict(self, sample: PromptSample) -> List[Dict[str...
method process_sample_ (line 75) | def process_sample_(self, sample) -> List[Dict[str, List[int]]]:
method report (line 79) | def report(self):
class LabelDatasetProcessor (line 88) | class LabelDatasetProcessor(object):
method __init__ (line 89) | def __init__(
method pad_seq (line 114) | def pad_seq(self, prompt_tokens: List[int], label: int, extra: dict = ...
method process_sample (line 128) | def process_sample(self, sample: LabelSample) -> Iterable[Dict[str, Li...
method process_sample_strict (line 143) | def process_sample_strict(self, sample: LabelSample) -> List[Dict[str,...
method process_sample_ (line 152) | def process_sample_(self, sample) -> List[Dict[str, List[int]]]:
method report (line 156) | def report(self):
FILE: codegeex/data/types.py
class PromptSample (line 6) | class PromptSample:
class LabelSample (line 15) | class LabelSample:
FILE: codegeex/kernels/__init__.py
class Kernel (line 11) | class Kernel:
method __init__ (line 12) | def __init__(self, filename: str, function_names: List[str]):
function compress_int4_weight (line 37) | def compress_int4_weight(weight: torch.Tensor): # (n, m)
function extract_weight_to_half (line 58) | def extract_weight_to_half(weight: torch.Tensor, scale_list: torch.Tenso...
FILE: codegeex/megatron/__init__.py
function print_rank_0 (line 28) | def print_rank_0(message):
function is_last_rank (line 37) | def is_last_rank():
function print_rank_last (line 41) | def print_rank_last(message):
FILE: codegeex/megatron/arguments.py
function parse_args (line 25) | def parse_args(extra_args_provider=None, defaults={}, ignore_unknown_arg...
function _print_args (line 302) | def _print_args(args):
function _check_arg_is_not_none (line 315) | def _check_arg_is_not_none(args, arg):
function _add_network_size_args (line 319) | def _add_network_size_args(parser):
function _add_logging_args (line 426) | def _add_logging_args(parser):
function _add_regularization_args (line 495) | def _add_regularization_args(parser):
function _add_training_args (line 589) | def _add_training_args(parser):
function _add_initialization_args (line 748) | def _add_initialization_args(parser):
function _add_inference_args (line 773) | def _add_inference_args(parser):
function _add_learning_rate_args (line 806) | def _add_learning_rate_args(parser):
function _add_checkpointing_args (line 898) | def _add_checkpointing_args(parser):
function _add_mixed_precision_args (line 973) | def _add_mixed_precision_args(parser):
function _add_distributed_args (line 1040) | def _add_distributed_args(parser):
function _add_validation_args (line 1128) | def _add_validation_args(parser):
function _add_data_args (line 1152) | def _add_data_args(parser):
function _add_autoresume_args (line 1298) | def _add_autoresume_args(parser):
function _add_biencoder_args (line 1316) | def _add_biencoder_args(parser):
function _add_vit_args (line 1425) | def _add_vit_args(parser):
function _add_zero_args (line 1453) | def _add_zero_args(parser):
function _add_memoryopt_args (line 1485) | def _add_memoryopt_args(parser):
function _add_activation_checkpoint_args (line 1520) | def _add_activation_checkpoint_args(parser):
FILE: codegeex/megatron/checkpointing.py
function set_checkpoint_version (line 31) | def set_checkpoint_version(value):
function get_checkpoint_version (line 38) | def get_checkpoint_version():
function check_checkpoint_args (line 43) | def check_checkpoint_args(checkpoint_args):
function ensure_directory_exists (line 75) | def ensure_directory_exists(filename):
function get_checkpoint_name (line 82) | def get_checkpoint_name(checkpoints_path, iteration, release=False):
function get_checkpoint_tracker_filename (line 104) | def get_checkpoint_tracker_filename(checkpoints_path):
function save_checkpoint (line 110) | def save_checkpoint(iteration, model, optimizer, lr_scheduler):
function _transpose_first_dim (line 205) | def _transpose_first_dim(t, num_splits, num_splits_first, model):
function fix_query_key_value_ordering (line 249) | def fix_query_key_value_ordering(model, checkpoint_version):
function load_deepspeed_state (line 282) | def load_deepspeed_state(model):
function load_checkpoint (line 301) | def load_checkpoint(
function load_biencoder_checkpoint (line 508) | def load_biencoder_checkpoint(
FILE: codegeex/megatron/code_generation_utils.py
function get_batch (line 35) | def get_batch(context_tokens, micro_batch_size=None):
function get_batch_ (line 56) | def get_batch_(context_tokens):
function top_k_logits (line 75) | def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")):
function generate_samples_input_from_file (line 105) | def generate_samples_input_from_file(model):
function generate_samples_eval (line 225) | def generate_samples_eval(model, context, max_gen_length, eos_token_id):
function generate_samples_interactive_code_contest (line 252) | def generate_samples_interactive_code_contest(model, print_frequency=10):
function generate_samples_interactive (line 377) | def generate_samples_interactive(model, print_frequency=24):
function generate_samples_unconditional (line 487) | def generate_samples_unconditional(model):
function generate_and_write_samples_unconditional (line 528) | def generate_and_write_samples_unconditional(model):
function pad_batch (line 540) | def pad_batch(batch, pad_id, args):
function topk_sampling (line 550) | def topk_sampling(logits: torch.FloatTensor, num_samples: int):
function nuclear_sampling (line 566) | def nuclear_sampling(logits: torch.FloatTensor, temperature: float, top_...
function sample_topk_tokens (line 579) | def sample_topk_tokens(model,
function nuclear_sample_tokens (line 600) | def nuclear_sample_tokens(model,
class Beam (line 621) | class Beam:
method __repr__ (line 625) | def __repr__(self):
method get_code (line 628) | def get_code(self):
function expand_beams (line 632) | def expand_beams(beams: List[Beam], num_beams: int, model) -> List[Beam]:
function beam_search (line 661) | def beam_search(model, context_tokens, num_beams: int):
class Handle (line 727) | class Handle:
method __repr__ (line 731) | def __repr__(self):
method is_finished (line 734) | def is_finished(self):
method derived (line 737) | def derived(self, new_token: int, log_prob: float):
function expand_handles (line 742) | def expand_handles(handles: List[Handle], temperature: float, top_p: flo...
function generate_nuclear_sampling (line 771) | def generate_nuclear_sampling(model, context_tokens, num_samples: int, t...
function forward_step (line 801) | def forward_step(
function get_token_stream (line 841) | def get_token_stream(
function switch (line 900) | def switch(val1, val2, boolean):
function sample_sequence_batch (line 905) | def sample_sequence_batch(
FILE: codegeex/megatron/convert_ckpt_parallel.py
function get_change_ckpt_args (line 8) | def get_change_ckpt_args(parser):
function get_element_from_dict_by_path (line 33) | def get_element_from_dict_by_path(d, path):
function main (line 48) | def main():
FILE: codegeex/megatron/data/blendable_dataset.py
class BlendableDataset (line 25) | class BlendableDataset(torch.utils.data.Dataset):
method __init__ (line 26) | def __init__(self, datasets, weights):
method __len__ (line 63) | def __len__(self):
method __getitem__ (line 66) | def __getitem__(self, idx):
FILE: codegeex/megatron/data/data_samplers.py
function build_pretraining_data_loader (line 24) | def build_pretraining_data_loader(dataset, consumed_samples):
class MegatronPretrainingSampler (line 62) | class MegatronPretrainingSampler:
method __init__ (line 63) | def __init__(
method __len__ (line 99) | def __len__(self):
method get_start_end_idx (line 102) | def get_start_end_idx(self):
method __iter__ (line 107) | def __iter__(self):
class MegatronPretrainingRandomSampler (line 123) | class MegatronPretrainingRandomSampler:
method __init__ (line 124) | def __init__(
method __len__ (line 157) | def __len__(self):
method __iter__ (line 160) | def __iter__(self):
FILE: codegeex/megatron/data/dataset_utils.py
function get_datasets_weights_and_num_samples (line 33) | def get_datasets_weights_and_num_samples(data_prefix, train_valid_test_n...
function compile_helper (line 66) | def compile_helper():
function get_a_and_b_segments (line 81) | def get_a_and_b_segments(sample, np_rng):
function truncate_segments (line 113) | def truncate_segments(tokens_a, tokens_b, len_a, len_b, max_num_tokens, ...
function create_tokens_and_tokentypes (line 133) | def create_tokens_and_tokentypes(tokens_a, tokens_b, cls_id, sep_id):
function is_start_piece (line 163) | def is_start_piece(piece):
function create_masked_lm_predictions (line 172) | def create_masked_lm_predictions(
function pad_and_convert_to_numpy (line 396) | def pad_and_convert_to_numpy(
function get_indexed_dataset_ (line 429) | def get_indexed_dataset_(data_prefix, data_impl, skip_warmup):
function get_train_valid_test_split_ (line 450) | def get_train_valid_test_split_(splits_string, size):
function get_samples_mapping (line 477) | def get_samples_mapping(
FILE: codegeex/megatron/data/helpers.cpp
function build_blending_indices (line 36) | void build_blending_indices(py::array_t<uint8_t>& dataset_index,
function build_sample_idx (line 99) | py::array build_sample_idx(const py::array_t<int32_t>& sizes_,
function get_target_sample_len (line 188) | inline int32_t get_target_sample_len(const int32_t short_seq_ratio,
function build_mapping_impl (line 204) | py::array build_mapping_impl(const py::array_t<int64_t>& docs_,
function build_mapping (line 440) | py::array build_mapping(const py::array_t<int64_t>& docs_,
function build_blocks_mapping_impl (line 470) | py::array build_blocks_mapping_impl(const py::array_t<int64_t>& docs_,
function build_blocks_mapping (line 687) | py::array build_blocks_mapping(const py::array_t<int64_t>& docs_,
function PYBIND11_MODULE (line 712) | PYBIND11_MODULE(helpers, m) {
FILE: codegeex/megatron/data/indexed_dataset.py
function __best_fitting_dtype (line 26) | def __best_fitting_dtype(vocab_size=None):
function get_available_dataset_impl (line 33) | def get_available_dataset_impl():
function infer_dataset_impl (line 37) | def infer_dataset_impl(path):
function make_builder (line 55) | def make_builder(out_file, impl, vocab_size=None):
function make_dataset (line 64) | def make_dataset(path, impl, skip_warmup=False):
function dataset_exists (line 83) | def dataset_exists(path, impl):
function read_longs (line 90) | def read_longs(f, n):
function write_longs (line 96) | def write_longs(f, a):
function __best_fitting_dtype (line 112) | def __best_fitting_dtype(vocab_size=None):
function make_mmap_builder (line 119) | def make_mmap_builder(out_file, vocab_size=None):
function code (line 125) | def code(dtype):
function index_file_path (line 132) | def index_file_path(prefix_path):
function data_file_path (line 136) | def data_file_path(prefix_path):
function create_doc_idx (line 140) | def create_doc_idx(sizes):
class IndexedDataset (line 148) | class IndexedDataset(torch.utils.data.Dataset):
method __init__ (line 153) | def __init__(self, path):
method read_index (line 159) | def read_index(self, path):
method read_data (line 177) | def read_data(self, path):
method check_index (line 180) | def check_index(self, i):
method __del__ (line 184) | def __del__(self):
method __getitem__ (line 189) | def __getitem__(self, idx):
method __len__ (line 213) | def __len__(self):
method num_tokens (line 216) | def num_tokens(self, index):
method size (line 219) | def size(self, index):
method exists (line 223) | def exists(path):
method supports_prefetch (line 229) | def supports_prefetch(self):
class IndexedCachedDataset (line 233) | class IndexedCachedDataset(IndexedDataset):
method __init__ (line 234) | def __init__(self, path):
method supports_prefetch (line 240) | def supports_prefetch(self):
method prefetch (line 243) | def prefetch(self, indices):
method __getitem__ (line 268) | def __getitem__(self, idx):
class IndexedDatasetBuilder (line 285) | class IndexedDatasetBuilder(object):
method __init__ (line 296) | def __init__(self, out_file, dtype=np.int32):
method add_item (line 305) | def add_item(self, tensor):
method end_document (line 312) | def end_document(self):
method merge_file_ (line 315) | def merge_file_(self, another_file):
method finalize (line 335) | def finalize(self, index_file):
function _warmup_mmap_file (line 350) | def _warmup_mmap_file(path):
class MMapIndexedDataset (line 356) | class MMapIndexedDataset(torch.utils.data.Dataset):
class Index (line 357) | class Index(object):
method writer (line 361) | def writer(cls, path, dtype):
method __init__ (line 406) | def __init__(self, path, skip_warmup=False):
method __del__ (line 442) | def __del__(self):
method dtype (line 447) | def dtype(self):
method sizes (line 451) | def sizes(self):
method doc_idx (line 455) | def doc_idx(self):
method __getitem__ (line 459) | def __getitem__(self, i):
method __len__ (line 462) | def __len__(self):
method __init__ (line 465) | def __init__(self, path, skip_warmup=False):
method __getstate__ (line 474) | def __getstate__(self):
method __setstate__ (line 477) | def __setstate__(self, state):
method _do_init (line 480) | def _do_init(self, path, skip_warmup):
method __del__ (line 494) | def __del__(self):
method __len__ (line 499) | def __len__(self):
method __getitem__ (line 503) | def __getitem__(self, idx):
method get (line 524) | def get(self, idx, offset=0, length=None):
method sizes (line 540) | def sizes(self):
method doc_idx (line 544) | def doc_idx(self):
method get_doc_idx (line 547) | def get_doc_idx(self):
method set_doc_idx (line 550) | def set_doc_idx(self, doc_idx_):
method supports_prefetch (line 554) | def supports_prefetch(self):
method exists (line 558) | def exists(path):
class MMapIndexedDatasetBuilder (line 564) | class MMapIndexedDatasetBuilder(object):
method __init__ (line 565) | def __init__(self, out_file, dtype=np.int64):
method add_item (line 571) | def add_item(self, tensor):
method end_document (line 576) | def end_document(self):
method merge_file_ (line 579) | def merge_file_(self, another_file):
method finalize (line 591) | def finalize(self, index_file):
FILE: codegeex/megatron/data/prompt_dataset.py
function build_train_valid_test_datasets (line 31) | def build_train_valid_test_datasets(
function _build_train_valid_test_datasets (line 96) | def _build_train_valid_test_datasets(
function get_indexed_dataset_ (line 164) | def get_indexed_dataset_(data_prefix, data_impl, skip_warmup):
class PromptDataset (line 179) | class PromptDataset(torch.utils.data.Dataset):
method __init__ (line 180) | def __init__(
method __len__ (line 229) | def __len__(self):
method __getitem__ (line 234) | def __getitem__(self, idx):
function _build_index_mappings (line 253) | def _build_index_mappings(
function _num_epochs (line 314) | def _num_epochs(samples_per_epoch, num_samples):
function _build_doc_idx (line 319) | def _build_doc_idx(documents, num_epochs, np_rng, separate_last_epoch):
FILE: codegeex/megatron/enums.py
class LayerType (line 19) | class LayerType(enum.Enum):
class AttnType (line 24) | class AttnType(enum.Enum):
class AttnMaskType (line 29) | class AttnMaskType(enum.Enum):
FILE: codegeex/megatron/global_vars.py
function get_args (line 36) | def get_args():
function get_num_microbatches (line 42) | def get_num_microbatches():
function get_current_global_batch_size (line 46) | def get_current_global_batch_size():
function update_num_microbatches (line 50) | def update_num_microbatches(consumed_samples, consistency_check=True):
function get_tokenizer (line 54) | def get_tokenizer():
function get_tensorboard_writer (line 60) | def get_tensorboard_writer():
function get_adlr_autoresume (line 66) | def get_adlr_autoresume():
function get_timers (line 72) | def get_timers():
function set_global_variables (line 78) | def set_global_variables(
function _parse_args (line 95) | def _parse_args(extra_args_provider=None, defaults={}, ignore_unknown_ar...
function _build_num_microbatches_calculator (line 107) | def _build_num_microbatches_calculator(args):
function _build_tokenizer (line 117) | def _build_tokenizer(args):
function rebuild_tokenizer (line 125) | def rebuild_tokenizer(args):
function _set_tensorboard_writer (line 131) | def _set_tensorboard_writer(args):
function _set_adlr_autoresume (line 157) | def _set_adlr_autoresume(args):
function _set_timers (line 175) | def _set_timers():
function _ensure_var_is_initialized (line 182) | def _ensure_var_is_initialized(var, name):
function _ensure_var_is_not_initialized (line 187) | def _ensure_var_is_not_initialized(var, name):
class _Timer (line 192) | class _Timer:
method __init__ (line 195) | def __init__(self, name):
method start (line 201) | def start(self):
method stop (line 208) | def stop(self):
method reset (line 215) | def reset(self):
method elapsed (line 220) | def elapsed(self, reset=True):
class Timers (line 237) | class Timers:
method __init__ (line 240) | def __init__(self):
method __call__ (line 243) | def __call__(self, name):
method write (line 248) | def write(self, names, writer, iteration, normalizer=1.0, reset=False):
method log (line 258) | def log(self, names, normalizer=1.0, reset=True):
FILE: codegeex/megatron/inference.py
function model_provider (line 18) | def model_provider():
function set_random_seed (line 27) | def set_random_seed(seed):
function run_generation_distributed (line 35) | def run_generation_distributed(model):
FILE: codegeex/megatron/initialize.py
function initialize_megatron (line 44) | def initialize_megatron(
function _compile_dependencies (line 105) | def _compile_dependencies():
function setup_deepspeed_random_and_activation_checkpointing (line 171) | def setup_deepspeed_random_and_activation_checkpointing(args):
function _initialize_distributed (line 209) | def _initialize_distributed():
function _init_autoresume (line 282) | def _init_autoresume():
function _set_random_seed (line 291) | def _set_random_seed(seed_):
function write_args_to_tensorboard (line 305) | def write_args_to_tensorboard():
function initialize_wandb_experiment (line 314) | def initialize_wandb_experiment():
function _initialize_mem_buffs (line 331) | def _initialize_mem_buffs():
FILE: codegeex/megatron/learning_rates.py
class AnnealingLR (line 23) | class AnnealingLR(object):
method __init__ (line 26) | def __init__(
method get_lr (line 70) | def get_lr(self):
method step (line 116) | def step(self, increment, token_num=None):
method state_dict (line 127) | def state_dict(self):
method _check_and_set (line 140) | def _check_and_set(self, cls_value, sd_value, name):
method load_state_dict (line 155) | def load_state_dict(self, sd):
FILE: codegeex/megatron/memory.py
function allocate_mem_buff (line 24) | def allocate_mem_buff(name, numel, dtype, track_usage):
function get_mem_buff (line 31) | def get_mem_buff(name):
class MemoryBuffer (line 36) | class MemoryBuffer:
method __init__ (line 49) | def __init__(self, name, numel, dtype, track_usage):
method reset (line 78) | def reset(self):
method is_in_use (line 82) | def is_in_use(self):
method numel_in_use (line 86) | def numel_in_use(self):
method add (line 90) | def add(self, tensor):
method get_data (line 114) | def get_data(self):
method print_average_usage (line 121) | def print_average_usage(self):
class RingMemBuffer (line 134) | class RingMemBuffer:
method __init__ (line 137) | def __init__(self, name, num_buffers, numel, dtype, track_usage):
method get_next_buffer (line 145) | def get_next_buffer(self):
FILE: codegeex/megatron/merge_ckpt_parallel.py
function get_change_ckpt_args (line 13) | def get_change_ckpt_args(parser):
function main (line 43) | def main():
FILE: codegeex/megatron/microbatches.py
function build_num_microbatches_calculator (line 22) | def build_num_microbatches_calculator(args):
class NumMicroBatchesCalculator (line 70) | class NumMicroBatchesCalculator(ABC):
method __init__ (line 71) | def __init__(self):
method get (line 75) | def get(self):
method get_current_global_batch_size (line 78) | def get_current_global_batch_size(self):
method update (line 82) | def update(self, consumed_samples, consistency_check):
class ConstantNumMicroBatches (line 86) | class ConstantNumMicroBatches(NumMicroBatchesCalculator):
method __init__ (line 87) | def __init__(self, global_batch_size, micro_batch_size, data_parallel_...
method update (line 99) | def update(self, consumed_samples, consistency_check):
class RampupBatchsizeNumMicroBatches (line 103) | class RampupBatchsizeNumMicroBatches(NumMicroBatchesCalculator):
method __init__ (line 104) | def __init__(
method update (line 159) | def update(self, consumed_samples, consistency_check):
FILE: codegeex/megatron/mindspore_to_megatron.py
function get_change_ckpt_args (line 34) | def get_change_ckpt_args(parser):
function loadModelFromNp (line 53) | def loadModelFromNp(sd, args):
function loadEmbeddingFromNp (line 73) | def loadEmbeddingFromNp(npCkptPath, languageModel, vocabSize=52224):
function loadAttentionLayerFromNp (line 90) | def loadAttentionLayerFromNp(npCkptPath, transformer, layerID):
function loadQueryLayerFromNp (line 199) | def loadQueryLayerFromNp(npCkptPath, transformer):
function main (line 290) | def main():
FILE: codegeex/megatron/model/codegeex_model.py
class CodeGeeXModel (line 27) | class CodeGeeXModel(MegatronModule):
method __init__ (line 30) | def __init__(self, num_tokentypes=0, parallel_output=False):
method set_input_tensor (line 44) | def set_input_tensor(self, input_tensor):
method forward (line 48) | def forward(
method state_dict_for_save_checkpoint (line 99) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 108) | def load_state_dict(self, state_dict, strict=True):
function CrossEntropy (line 116) | def CrossEntropy(output, labels):
class CodeGeeXModelPipe (line 127) | class CodeGeeXModelPipe(PipelineModule, MegatronModule):
method __init__ (line 130) | def __init__(self, num_tokentypes=0, parallel_output=True):
FILE: codegeex/megatron/model/distributed.py
class MemoryBuffer (line 26) | class MemoryBuffer:
method __init__ (line 27) | def __init__(self, numel, dtype):
method zero (line 37) | def zero(self):
method get (line 41) | def get(self, shape, start_index):
class DistributedDataParallelBase (line 51) | class DistributedDataParallelBase(MegatronModule, ABC):
method __init__ (line 54) | def __init__(self, module):
method allreduce_gradients (line 60) | def allreduce_gradients(self):
method forward (line 63) | def forward(self, *inputs, **kwargs):
method state_dict (line 66) | def state_dict(self, destination=None, prefix="", keep_vars=False):
method state_dict_for_save_checkpoint (line 69) | def state_dict_for_save_checkpoint(
method load_state_dict (line 76) | def load_state_dict(self, state_dict, strict=True):
class DistributedDataParallel (line 80) | class DistributedDataParallel(DistributedDataParallelBase):
method __init__ (line 96) | def __init__(
method _make_param_hook (line 162) | def _make_param_hook(self, param):
method zero_grad_buffer (line 175) | def zero_grad_buffer(self):
method allreduce_gradients (line 182) | def allreduce_gradients(self):
FILE: codegeex/megatron/model/language_model.py
function get_shrink_embedding_gradient_alpha (line 29) | def get_shrink_embedding_gradient_alpha(iteration):
function parallel_lm_logits (line 46) | def parallel_lm_logits(input_, word_embeddings_weight, parallel_output, ...
function get_language_model (line 73) | def get_language_model(
class Embedding (line 100) | class Embedding(MegatronModule):
method __init__ (line 114) | def __init__(
method add_tokentype_embeddings (line 164) | def add_tokentype_embeddings(self, num_tokentypes):
method forward (line 180) | def forward(self, input_ids, position_ids, tokentype_ids=None):
method state_dict_for_save_checkpoint (line 196) | def state_dict_for_save_checkpoint(
method load_state_dict (line 214) | def load_state_dict(self, state_dict, strict=True):
class EmbeddingPipe (line 273) | class EmbeddingPipe(Embedding):
method forward (line 274) | def forward(self, inputs, **kwargs):
method word_embeddings_weight (line 302) | def word_embeddings_weight(self):
class QueryEmbedding (line 307) | class QueryEmbedding(MegatronModule):
method __init__ (line 321) | def __init__(self,
method add_tokentype_embeddings (line 360) | def add_tokentype_embeddings(self, num_tokentypes):
method forward (line 376) | def forward(self, position_ids, tokentype_ids=None):
method state_dict_for_save_checkpoint (line 391) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 406) | def load_state_dict(self, state_dict, strict=True):
class QueryEmbeddingPipe (line 448) | class QueryEmbeddingPipe(QueryEmbedding):
method forward (line 449) | def forward(self, inputs, **kwargs):
method word_embeddings_weight (line 476) | def word_embeddings_weight(self):
class TransformerLanguageModel (line 481) | class TransformerLanguageModel(MegatronModule):
method __init__ (line 500) | def __init__(self,
method set_input_tensor (line 537) | def set_input_tensor(self, input_tensor):
method forward (line 541) | def forward(
method state_dict_for_save_checkpoint (line 572) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 593) | def load_state_dict(self, state_dict, strict=True):
FILE: codegeex/megatron/model/module.py
function param_is_not_shared (line 31) | def param_is_not_shared(param):
class MegatronModule (line 35) | class MegatronModule(torch.nn.Module):
method __init__ (line 39) | def __init__(self, share_word_embeddings=True):
method state_dict_for_save_checkpoint (line 43) | def state_dict_for_save_checkpoint(
method word_embeddings_weight (line 50) | def word_embeddings_weight(self):
method initialize_word_embeddings (line 64) | def initialize_word_embeddings(self, init_method_normal):
function conversion_helper (line 120) | def conversion_helper(val, conversion):
function fp32_to_float16 (line 131) | def fp32_to_float16(val, float16_convertor):
function float16_to_fp32 (line 145) | def float16_to_fp32(val):
class Float16Module (line 159) | class Float16Module(MegatronModule):
method __init__ (line 160) | def __init__(self, module, args):
method forward (line 180) | def forward(self, *inputs, **kwargs):
method state_dict (line 188) | def state_dict(self, destination=None, prefix="", keep_vars=False):
method state_dict_for_save_checkpoint (line 191) | def state_dict_for_save_checkpoint(
method load_state_dict (line 198) | def load_state_dict(self, state_dict, strict=True):
FILE: codegeex/megatron/model/transformer.py
class ParallelMLP (line 55) | class ParallelMLP(MegatronModule):
method __init__ (line 64) | def __init__(
method forward (line 93) | def forward(self, hidden_states):
class ParallelSelfAttention (line 103) | class ParallelSelfAttention(MegatronModule):
method __init__ (line 110) | def __init__(self, init_method,
method forward (line 164) | def forward(
class ParallelTopQuerySelfAttention (line 325) | class ParallelTopQuerySelfAttention(MegatronModule):
method __init__ (line 332) | def __init__(self, init_method,
method forward (line 387) | def forward(
function bias_dropout_add (line 547) | def bias_dropout_add(x, bias, residual, prob, training):
function get_bias_dropout_add (line 554) | def get_bias_dropout_add(training):
function bias_dropout_add_fused_train (line 562) | def bias_dropout_add_fused_train(x, bias, residual, prob):
function bias_dropout_add_fused_inference (line 568) | def bias_dropout_add_fused_inference(x, bias, residual, prob):
class ParallelTransformerLayer (line 573) | class ParallelTransformerLayer(MegatronModule):
method __init__ (line 580) | def __init__(self, init_method,
method forward (line 619) | def forward(
class ParallelTransformerLayerPipe (line 695) | class ParallelTransformerLayerPipe(ParallelTransformerLayer):
method forward (line 716) | def forward(self, inputs, **kwargs):
class ParallelTopQueryLayer (line 732) | class ParallelTopQueryLayer(MegatronModule):
method __init__ (line 739) | def __init__(self, init_method,
method forward (line 776) | def forward(
class ParallelTopQueryLayerPipe (line 857) | class ParallelTopQueryLayerPipe(ParallelTopQueryLayer):
method forward (line 878) | def forward(self, inputs, **kwargs):
class ParallelTransformer (line 895) | class ParallelTransformer(MegatronModule):
method __init__ (line 898) | def __init__(self, init_method, output_layer_init_method):
method _get_layer_index (line 943) | def _get_layer_index(self, layer_number):
method _get_layer (line 950) | def _get_layer(self, layer_number):
method _checkpointed_forward (line 953) | def _checkpointed_forward(self, hidden_states, attention_mask):
method set_input_tensor (line 977) | def set_input_tensor(self, input_tensor):
method forward (line 987) | def forward(
FILE: codegeex/megatron/model/utils.py
function init_method_normal (line 22) | def init_method_normal(sigma):
function scaled_init_method_normal (line 31) | def scaled_init_method_normal(sigma, num_layers):
function attention_mask_func (line 41) | def attention_mask_func(attention_scores, attention_mask):
function get_linear_layer (line 47) | def get_linear_layer(rows, columns, init_method):
function fast_gelu (line 56) | def fast_gelu(x):
function gelu_impl (line 62) | def gelu_impl(x):
function openai_gelu (line 69) | def openai_gelu(x):
function erf_gelu (line 75) | def erf_gelu(x):
FILE: codegeex/megatron/mpu/cross_entropy.py
class _VocabParallelCrossEntropy (line 25) | class _VocabParallelCrossEntropy(torch.autograd.Function):
method forward (line 27) | def forward(ctx, vocab_parallel_logits, target):
method backward (line 92) | def backward(ctx, grad_output):
function vocab_parallel_cross_entropy (line 113) | def vocab_parallel_cross_entropy(vocab_parallel_logits, target):
FILE: codegeex/megatron/mpu/data.py
function _check_data_types (line 26) | def _check_data_types(keys, data, target_dtype):
function _build_key_size_numel_dictionaries (line 36) | def _build_key_size_numel_dictionaries(keys, data):
function broadcast_data (line 82) | def broadcast_data(keys, data, datatype):
FILE: codegeex/megatron/mpu/initialize.py
function is_unitialized (line 49) | def is_unitialized():
function initialize_model_parallel (line 54) | def initialize_model_parallel(
function model_parallel_is_initialized (line 184) | def model_parallel_is_initialized():
function get_model_parallel_group (line 195) | def get_model_parallel_group():
function get_tensor_model_parallel_group (line 201) | def get_tensor_model_parallel_group():
function get_pipeline_model_parallel_group (line 209) | def get_pipeline_model_parallel_group():
function get_data_parallel_group (line 217) | def get_data_parallel_group():
function get_embedding_group (line 223) | def get_embedding_group():
function set_tensor_model_parallel_world_size (line 229) | def set_tensor_model_parallel_world_size(world_size):
function set_pipeline_model_parallel_world_size (line 235) | def set_pipeline_model_parallel_world_size(world_size):
function get_tensor_model_parallel_world_size (line 241) | def get_tensor_model_parallel_world_size():
function get_model_parallel_world_size (line 249) | def get_model_parallel_world_size():
function get_pipeline_model_parallel_world_size (line 256) | def get_pipeline_model_parallel_world_size():
function set_tensor_model_parallel_rank (line 264) | def set_tensor_model_parallel_rank(rank):
function set_pipeline_model_parallel_rank (line 270) | def set_pipeline_model_parallel_rank(rank):
function get_tensor_model_parallel_rank (line 276) | def get_tensor_model_parallel_rank():
function get_model_parallel_rank (line 284) | def get_model_parallel_rank():
function get_pipeline_model_parallel_rank (line 291) | def get_pipeline_model_parallel_rank():
function is_pipeline_first_stage (line 299) | def is_pipeline_first_stage(ignore_virtual=False):
function is_pipeline_last_stage (line 310) | def is_pipeline_last_stage(ignore_virtual=False):
function get_virtual_pipeline_model_parallel_rank (line 327) | def get_virtual_pipeline_model_parallel_rank():
function set_virtual_pipeline_model_parallel_rank (line 333) | def set_virtual_pipeline_model_parallel_rank(rank):
function get_virtual_pipeline_model_parallel_world_size (line 339) | def get_virtual_pipeline_model_parallel_world_size():
function get_tensor_model_parallel_src_rank (line 345) | def get_tensor_model_parallel_src_rank():
function get_pipeline_model_parallel_first_rank (line 353) | def get_pipeline_model_parallel_first_rank():
function get_pipeline_model_parallel_last_rank (line 360) | def get_pipeline_model_parallel_last_rank():
function get_pipeline_model_parallel_next_rank (line 368) | def get_pipeline_model_parallel_next_rank():
function get_pipeline_model_parallel_prev_rank (line 377) | def get_pipeline_model_parallel_prev_rank():
function get_data_parallel_world_size (line 386) | def get_data_parallel_world_size():
function get_data_parallel_rank (line 391) | def get_data_parallel_rank():
function destroy_model_parallel (line 396) | def destroy_model_parallel():
FILE: codegeex/megatron/mpu/layers.py
function param_is_not_tensor_parallel_duplicate (line 49) | def param_is_not_tensor_parallel_duplicate(param):
function set_tensor_model_parallel_attributes (line 55) | def set_tensor_model_parallel_attributes(tensor, is_parallel, dim, stride):
function set_defaults_if_not_set_tensor_model_parallel_attributes (line 65) | def set_defaults_if_not_set_tensor_model_parallel_attributes(tensor):
function copy_tensor_model_parallel_attributes (line 74) | def copy_tensor_model_parallel_attributes(destination_tensor, source_ten...
function _initialize_affine_weight_gpu (line 83) | def _initialize_affine_weight_gpu(weight, init_method, partition_dim, st...
function _initialize_affine_weight_cpu (line 98) | def _initialize_affine_weight_cpu(
class VocabParallelEmbedding (line 141) | class VocabParallelEmbedding(torch.nn.Module):
method __init__ (line 152) | def __init__(self, num_embeddings, embedding_dim, init_method=init.xav...
method forward (line 211) | def forward(self, input_):
class ColumnParallelLinear (line 240) | class ColumnParallelLinear(torch.nn.Module):
method __init__ (line 264) | def __init__(
method forward (line 351) | def forward(self, input_):
class RowParallelLinear (line 367) | class RowParallelLinear(torch.nn.Module):
method __init__ (line 397) | def __init__(
method forward (line 483) | def forward(self, input_):
FILE: codegeex/megatron/mpu/mappings.py
function _reduce (line 26) | def _reduce(input_):
function _split (line 39) | def _split(input_):
function _gather (line 58) | def _gather(input_):
class _CopyToModelParallelRegion (line 82) | class _CopyToModelParallelRegion(torch.autograd.Function):
method symbolic (line 86) | def symbolic(graph, input_):
method forward (line 90) | def forward(ctx, input_):
method backward (line 94) | def backward(ctx, grad_output):
class _ReduceFromModelParallelRegion (line 98) | class _ReduceFromModelParallelRegion(torch.autograd.Function):
method symbolic (line 102) | def symbolic(graph, input_):
method forward (line 106) | def forward(ctx, input_):
method backward (line 110) | def backward(ctx, grad_output):
class _ScatterToModelParallelRegion (line 114) | class _ScatterToModelParallelRegion(torch.autograd.Function):
method symbolic (line 118) | def symbolic(graph, input_):
method forward (line 122) | def forward(ctx, input_):
method backward (line 126) | def backward(ctx, grad_output):
class _GatherFromModelParallelRegion (line 130) | class _GatherFromModelParallelRegion(torch.autograd.Function):
method symbolic (line 134) | def symbolic(graph, input_):
method forward (line 138) | def forward(ctx, input_):
method backward (line 142) | def backward(ctx, grad_output):
function copy_to_tensor_model_parallel_region (line 151) | def copy_to_tensor_model_parallel_region(input_):
function reduce_from_tensor_model_parallel_region (line 155) | def reduce_from_tensor_model_parallel_region(input_):
function scatter_to_tensor_model_parallel_region (line 159) | def scatter_to_tensor_model_parallel_region(input_):
function gather_from_tensor_model_parallel_region (line 163) | def gather_from_tensor_model_parallel_region(input_):
FILE: codegeex/megatron/mpu/random.py
function init_checkpointed_activations_memory_buffer (line 44) | def init_checkpointed_activations_memory_buffer():
function reset_checkpointed_activations_memory_buffer (line 72) | def reset_checkpointed_activations_memory_buffer():
function _set_cuda_rng_state (line 78) | def _set_cuda_rng_state(new_state, device=-1):
function split_tensor_into_1d_equal_chunks (line 112) | def split_tensor_into_1d_equal_chunks(tensor):
function gather_split_1d_tensor (line 121) | def gather_split_1d_tensor(tensor):
class CudaRNGStatesTracker (line 139) | class CudaRNGStatesTracker:
method __init__ (line 148) | def __init__(self):
method reset (line 154) | def reset(self):
method get_states (line 159) | def get_states(self):
method set_states (line 167) | def set_states(self, states):
method add (line 172) | def add(self, name, seed):
method fork (line 190) | def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME):
function get_cuda_rng_tracker (line 215) | def get_cuda_rng_tracker():
function model_parallel_cuda_manual_seed (line 220) | def model_parallel_cuda_manual_seed(seed):
class CheckpointFunction (line 265) | class CheckpointFunction(torch.autograd.Function):
method forward (line 274) | def forward(ctx, run_function, *args):
method backward (line 298) | def backward(ctx, *args):
function checkpoint (line 339) | def checkpoint(function, *args):
FILE: codegeex/megatron/mpu/utils.py
function ensure_divisibility (line 20) | def ensure_divisibility(numerator, denominator):
function divide (line 27) | def divide(numerator, denominator):
function split_tensor_along_last_dim (line 34) | def split_tensor_along_last_dim(tensor, num_partitions, contiguous_split...
class VocabUtility (line 54) | class VocabUtility:
method vocab_range_from_per_partition_vocab_size (line 60) | def vocab_range_from_per_partition_vocab_size(
method vocab_range_from_global_vocab_size (line 68) | def vocab_range_from_global_vocab_size(global_vocab_size, rank, world_...
FILE: codegeex/megatron/optimizer/__init__.py
function _get_params_for_weight_decay_optimization (line 26) | def _get_params_for_weight_decay_optimization(modules):
function get_megatron_optimizer (line 58) | def get_megatron_optimizer(model):
FILE: codegeex/megatron/optimizer/clip_grads.py
function clip_grad_norm_fp32 (line 29) | def clip_grad_norm_fp32(parameters, max_norm, norm_type=2):
function count_zeros_fp32 (line 126) | def count_zeros_fp32(parameters):
FILE: codegeex/megatron/optimizer/grad_scaler.py
class MegatronGradScaler (line 24) | class MegatronGradScaler(ABC):
method __init__ (line 25) | def __init__(self, initial_scale):
method scale (line 31) | def scale(self):
method inv_scale (line 35) | def inv_scale(self):
method update (line 39) | def update(self, found_inf):
method state_dict (line 43) | def state_dict(self):
method load_state_dict (line 47) | def load_state_dict(self, state_dict):
class ConstantGradScaler (line 51) | class ConstantGradScaler(MegatronGradScaler):
method update (line 52) | def update(self, found_inf):
method state_dict (line 55) | def state_dict(self):
method load_state_dict (line 58) | def load_state_dict(self, state_dict):
class DynamicGradScaler (line 62) | class DynamicGradScaler(MegatronGradScaler):
method __init__ (line 63) | def __init__(
method update (line 99) | def update(self, found_inf):
method state_dict (line 122) | def state_dict(self):
method load_state_dict (line 129) | def load_state_dict(self, state_dict):
FILE: codegeex/megatron/optimizer/optimizer.py
function _zero_grad_group_helper (line 33) | def _zero_grad_group_helper(group, set_to_none):
function _multi_tensor_copy_this_to_that (line 48) | def _multi_tensor_copy_this_to_that(this, that, overflow_buf=None):
class MegatronOptimizer (line 62) | class MegatronOptimizer(ABC):
method __init__ (line 63) | def __init__(
method get_parameters (line 74) | def get_parameters(self):
method clip_grad_norm (line 81) | def clip_grad_norm(self, clip_grad):
method count_zeros (line 85) | def count_zeros(self):
method zero_grad (line 90) | def zero_grad(self, set_to_none=True):
method get_loss_scale (line 94) | def get_loss_scale(self):
method scale_loss (line 98) | def scale_loss(self, loss):
method step (line 103) | def step(self):
method reload_model_params (line 107) | def reload_model_params(self):
method state_dict (line 116) | def state_dict(self):
method load_state_dict (line 120) | def load_state_dict(self, state_dict):
method _get_state (line 125) | def _get_state(self):
method _set_state (line 128) | def _set_state(self, value):
method _get_param_groups (line 136) | def _get_param_groups(self):
method _set_param_groups (line 139) | def _set_param_groups(self, value):
class Float16OptimizerWithFloat16Params (line 145) | class Float16OptimizerWithFloat16Params(MegatronOptimizer):
method __init__ (line 170) | def __init__(
method zero_grad (line 273) | def zero_grad(self, set_to_none=True):
method get_loss_scale (line 281) | def get_loss_scale(self):
method _copy_model_grads_to_main_grads (line 286) | def _copy_model_grads_to_main_grads(self):
method _unscale_main_grads_and_check_for_nan (line 303) | def _unscale_main_grads_and_check_for_nan(self):
method _get_model_and_main_params_data_float16 (line 332) | def _get_model_and_main_params_data_float16(self):
method _copy_main_params_to_model_params (line 343) | def _copy_main_params_to_model_params(self):
method _copy_model_params_to_main_params (line 350) | def _copy_model_params_to_main_params(self):
method reload_model_params (line 357) | def reload_model_params(self):
method step (line 361) | def step(self):
method state_dict (line 408) | def state_dict(self):
method load_state_dict (line 416) | def load_state_dict(self, state_dict):
class FP32Optimizer (line 453) | class FP32Optimizer(MegatronOptimizer):
method __init__ (line 454) | def __init__(
method zero_grad (line 464) | def zero_grad(self, set_to_none=True):
method get_loss_scale (line 469) | def get_loss_scale(self):
method step (line 474) | def step(self):
method reload_model_params (line 498) | def reload_model_params(self):
method state_dict (line 501) | def state_dict(self):
method load_state_dict (line 504) | def load_state_dict(self, state_dict):
FILE: codegeex/megatron/p2p_communication.py
function _communicate (line 24) | def _communicate(
function recv_forward (line 150) | def recv_forward(timers=None):
function recv_backward (line 168) | def recv_backward(timers=None):
function send_forward (line 186) | def send_forward(output_tensor, timers=None):
function send_backward (line 201) | def send_backward(input_tensor_grad, timers=None):
function send_forward_recv_backward (line 216) | def send_forward_recv_backward(output_tensor, timers=None):
function send_backward_recv_forward (line 234) | def send_backward_recv_forward(input_tensor_grad, timers=None):
function send_forward_recv_forward (line 252) | def send_forward_recv_forward(output_tensor, recv_prev, timers=None):
function send_backward_recv_backward (line 267) | def send_backward_recv_backward(input_tensor_grad, recv_next, timers=None):
function send_forward_backward_recv_forward_backward (line 282) | def send_forward_backward_recv_forward_backward(
FILE: codegeex/megatron/schedules.py
function get_forward_backward_func (line 31) | def get_forward_backward_func():
function forward_step (line 43) | def forward_step(forward_step_func, data_iterator, model, input_tensor, ...
function backward_step (line 72) | def backward_step(
function dummy_handler (line 113) | def dummy_handler():
function forward_backward_no_pipelining (line 120) | def forward_backward_no_pipelining(
function forward_backward_pipelining_with_interleaving (line 173) | def forward_backward_pipelining_with_interleaving(
function forward_backward_pipelining_without_interleaving (line 416) | def forward_backward_pipelining_without_interleaving(
FILE: codegeex/megatron/tokenizer/gpt2_tokenization.py
function lru_cache (line 33) | def lru_cache():
function bytes_to_unicode (line 54) | def bytes_to_unicode():
function get_pairs (line 81) | def get_pairs(word):
class GPT2Tokenizer (line 94) | class GPT2Tokenizer(object):
method from_pretrained (line 101) | def from_pretrained(
method __init__ (line 183) | def __init__(
method __len__ (line 212) | def __len__(self):
method set_special_tokens (line 215) | def set_special_tokens(self, special_tokens):
method bpe (line 230) | def bpe(self, token):
method tokenize (line 271) | def tokenize(self, text):
method convert_tokens_to_ids (line 282) | def convert_tokens_to_ids(self, tokens):
method convert_ids_to_tokens (line 307) | def convert_ids_to_tokens(self, ids, skip_special_tokens=False):
method encode (line 318) | def encode(self, text):
method decode (line 321) | def decode(self, tokens):
method save_vocabulary (line 328) | def save_vocabulary(self, vocab_path):
FILE: codegeex/megatron/tokenizer/tokenizer.py
function encode_whitespaces (line 26) | def encode_whitespaces(text: str, start_extra_id: int=10, max_len: int=10):
function decode_whitespaces (line 37) | def decode_whitespaces(text: str, start_extra_id: int=10, max_len: int=10):
function build_hgf_tokenizer (line 52) | def build_hgf_tokenizer(args):
function build_tokenizer (line 71) | def build_tokenizer(args):
function _vocab_size_with_padding (line 98) | def _vocab_size_with_padding(orig_vocab_size, args):
class AbstractTokenizer (line 118) | class AbstractTokenizer(ABC):
method __init__ (line 121) | def __init__(self, name):
method vocab_size (line 127) | def vocab_size(self):
method vocab (line 132) | def vocab(self):
method inv_vocab (line 138) | def inv_vocab(self):
method tokenize (line 143) | def tokenize(self, text):
method detokenize (line 146) | def detokenize(self, token_ids):
method cls (line 152) | def cls(self):
method sep (line 158) | def sep(self):
method pad (line 164) | def pad(self):
method eod (line 170) | def eod(self):
method mask (line 176) | def mask(self):
class _GPT2BPETokenizer (line 182) | class _GPT2BPETokenizer(AbstractTokenizer):
method __init__ (line 185) | def __init__(self, vocab_file, merge_file):
method vocab_size (line 195) | def vocab_size(self):
method vocab (line 199) | def vocab(self):
method inv_vocab (line 203) | def inv_vocab(self):
method tokenize (line 206) | def tokenize(self, text):
method detokenize (line 209) | def detokenize(self, token_ids):
method eod (line 213) | def eod(self):
class HgfTokenizerWrapper (line 217) | class HgfTokenizerWrapper(AbstractTokenizer):
method __init__ (line 220) | def __init__(
method tokenize (line 231) | def tokenize(self, text):
method detokenize (line 238) | def detokenize(self, token_ids):
method eod (line 245) | def eod(self):
method inv_vocab (line 249) | def inv_vocab(self):
method vocab (line 253) | def vocab(self):
method vocab_size (line 257) | def vocab_size(self):
FILE: codegeex/megatron/tools/collect_env.py
function main (line 7) | def main():
FILE: codegeex/megatron/tools/finetune_codegeex.py
function model_provider (line 19) | def model_provider(pre_process=True, post_process=True):
function get_batch (line 88) | def get_batch(data_iterator):
function get_batch_pipe (line 128) | def get_batch_pipe(data):
function loss_func (line 157) | def loss_func(loss_mask, output_tensor):
function valid_loss_func (line 182) | def valid_loss_func(loss_mask, output_tensor):
function forward_step (line 200) | def forward_step(data_iterator, model):
function valid_forward_step (line 215) | def valid_forward_step(data_iterator, model):
function train_valid_test_datasets_provider (line 230) | def train_valid_test_datasets_provider(train_val_test_num_samples):
FILE: codegeex/megatron/tools/pretrain_codegeex.py
function model_provider (line 20) | def model_provider(pre_process=True, post_process=True):
function get_batch (line 89) | def get_batch(data_iterator):
function get_batch_pipe (line 123) | def get_batch_pipe(data):
function loss_func (line 152) | def loss_func(loss_mask, output_tensor):
function forward_step (line 163) | def forward_step(data_iterator, model):
function train_valid_test_datasets_provider (line 178) | def train_valid_test_datasets_provider(train_val_test_num_samples):
function command_exists (line 197) | def command_exists(cmd):
FILE: codegeex/megatron/training.py
function print_datetime (line 70) | def print_datetime(string):
function pretrain (line 77) | def pretrain(
function update_train_iters (line 221) | def update_train_iters(args):
function get_model (line 250) | def get_model(model_provider_func):
function get_learning_rate_scheduler (line 353) | def get_learning_rate_scheduler(optimizer):
function setup_model_and_optimizer (line 396) | def setup_model_and_optimizer(model_provider_func):
function train_step (line 477) | def train_step(forward_step_func, data_iterator, model, optimizer, lr_sc...
function training_log (line 600) | def training_log(
function save_checkpoint_and_time (line 870) | def save_checkpoint_and_time(iteration, model, optimizer, lr_scheduler):
function train (line 882) | def train(
function evaluate (line 1022) | def evaluate(forward_step_func, data_iterator, model, verbose=False):
function evaluate_and_print_results (line 1090) | def evaluate_and_print_results(
function evaluate_and_print_results_gold (line 1157) | def evaluate_and_print_results_gold(
function cyclic_iter (line 1214) | def cyclic_iter(iter):
function build_train_valid_test_data_iterators (line 1220) | def build_train_valid_test_data_iterators(build_train_valid_test_dataset...
FILE: codegeex/megatron/utils.py
function unwrap_model (line 34) | def unwrap_model(model, module_instances=(torchDDP)):
function calc_params_l2_norm (line 49) | def calc_params_l2_norm(model):
function average_losses_across_data_parallel_group (line 81) | def average_losses_across_data_parallel_group(losses):
function report_memory (line 92) | def report_memory(name):
function print_params_min_max_norm (line 108) | def print_params_min_max_norm(optimizer, iteration):
function check_adlr_autoresume_termination (line 127) | def check_adlr_autoresume_termination(iteration, model, optimizer, lr_sc...
function get_ltor_masks_and_position_ids (line 145) | def get_ltor_masks_and_position_ids(
function get_parameters_in_billions (line 202) | def get_parameters_in_billions(model):
function flops_calculator (line 222) | def flops_calculator(model, args, iteration_time):
FILE: codegeex/mindspore/convertion_1p.py
function load_model (line 40) | def load_model(args_opt):
function export_mindir (line 188) | def export_mindir(model_predict, config):
function run_predict (line 206) | def run_predict(model_predict, config, args_opt, rank):
function main (line 246) | def main():
FILE: codegeex/mindspore/finetune.py
function set_weight_decay (line 57) | def set_weight_decay(params):
function add_checkpoint_callback_policy (line 72) | def add_checkpoint_callback_policy(args_param, callback, rank_id):
function set_parallel_context (line 102) | def set_parallel_context(args_opt):
function run_train (line 121) | def run_train(args_opt):
FILE: codegeex/mindspore/generation.py
function load_model (line 40) | def load_model(args_opt):
function export_mindir (line 177) | def export_mindir(model_predict, config):
function run_predict (line 195) | def run_predict(model_predict, config, args_opt, rank):
function main (line 249) | def main():
FILE: codegeex/mindspore/generation_1p.py
function load_model (line 40) | def load_model(args_opt):
function export_mindir (line 183) | def export_mindir(model_predict, config):
function run_predict (line 201) | def run_predict(model_predict, config, args_opt, rank):
function main (line 242) | def main():
FILE: codegeex/mindspore/generation_batch.py
function load_model (line 40) | def load_model(args_opt):
function export_mindir (line 175) | def export_mindir(model_predict, config):
function run_predict (line 193) | def run_predict(model_predict, config, args_opt, rank):
function main (line 298) | def main():
FILE: codegeex/mindspore/generation_finetune.py
function load_model (line 41) | def load_model(args_opt):
function export_mindir (line 178) | def export_mindir(model_predict, config):
function run_predict (line 196) | def run_predict(model_predict, config, args_opt, rank):
function main (line 285) | def main():
FILE: codegeex/mindspore/generation_humaneval.py
function load_model (line 42) | def load_model(args_opt):
function export_mindir (line 177) | def export_mindir(model_predict, config):
function run_predict (line 195) | def run_predict(model_predict, config, args_opt, rank):
function main (line 245) | def main():
FILE: codegeex/mindspore/generation_values.py
function load_model (line 39) | def load_model(args_opt):
function run_predict (line 174) | def run_predict(model_predict, config, args_opt, rank):
function main (line 208) | def main():
FILE: codegeex/mindspore/generation_values_1p.py
function load_model (line 40) | def load_model(args_opt):
function export_mindir (line 183) | def export_mindir(model_predict, config):
function run_predict (line 201) | def run_predict(model_predict, config, args_opt, rank):
function main (line 243) | def main():
FILE: codegeex/mindspore/save_1p_ckpt_from_8p_ckpt.py
function set_weight_decay (line 66) | def set_weight_decay(params):
function add_checkpoint_callback_policy (line 81) | def add_checkpoint_callback_policy(args_param, callback, rank_id):
function set_parallel_context (line 111) | def set_parallel_context(args_opt):
function download_ckpt (line 128) | def download_ckpt(args_opt, file_num, rank_num, rank_id):
function get_needed_model_parallel_list (line 147) | def get_needed_model_parallel_list(train_strategy_file, self_rank):
function transform_model_parallel (line 160) | def transform_model_parallel(restore_local_ckpt_file_list, train_strateg...
function run_transform_model_parallel_ckpt (line 213) | def run_transform_model_parallel_ckpt(args_opt):
FILE: codegeex/mindspore/save_8p_ckpt.py
function set_parallel_context (line 61) | def set_parallel_context(args_opt):
function download_ckpt (line 78) | def download_ckpt(args_opt, file_num, rank_num, rank_id):
function get_needed_opt_shard_list (line 95) | def get_needed_opt_shard_list(train_strategy_file, self_rank):
function transform_opt_shard (line 111) | def transform_opt_shard(restore_local_ckpt_file_list, train_strategy_fil...
function run_transform_opt_shard_ckpt (line 154) | def run_transform_opt_shard_ckpt(args_opt):
FILE: codegeex/mindspore/scripts/layer_norm.py
function get_op_support_info (line 43) | def get_op_support_info(input_x, input_gamma, input_beta,
function _division_sixteen (line 97) | def _division_sixteen(shape, begin_norm_axis):
function op_select_format (line 121) | def op_select_format(input_x, input_gamma, input_beta,
function to_frac_z_axis (line 273) | def to_frac_z_axis(ori_shape, ori_axis):
function _broadcast_nz (line 312) | def _broadcast_nz(tensor, shape):
function _check_vector_to_cube (line 330) | def _check_vector_to_cube(dtype, ori_shape_x, shape_x, begin_norm_axis, ...
function nz_non_aligned (line 354) | def nz_non_aligned(input_x, input_gamma, input_beta,
function layer_norm_compute_nz (line 433) | def layer_norm_compute_nz(input_x, input_gamma, input_beta,
function layer_norm_compute (line 571) | def layer_norm_compute(input_x, input_gamma, input_beta,
function is_support_nz_non_aligned (line 709) | def is_support_nz_non_aligned(ori_shape_x, begin_params_axis, impl_mode):
function layer_norm (line 726) | def layer_norm(input_x, input_gamma, input_beta,
function __dynamic_template_api (line 917) | def __dynamic_template_api(input_x, input_gamma, input_beta, output_y, o...
FILE: codegeex/mindspore/scripts/layer_norm_x_backprop_v2.py
function get_op_support_info (line 36) | def get_op_support_info(input_dy,
function _check_dynamic_format (line 87) | def _check_dynamic_format(shape_dy, shape_gamma, c_0):
function op_select_format (line 100) | def op_select_format(input_dy,
function _check_params (line 224) | def _check_params(params_map):
function _check_shape (line 247) | def _check_shape(params_map):
function _check_shape_mean (line 285) | def _check_shape_mean(shape_x, shape_mean):
function _check_shape_gamma (line 325) | def _check_shape_gamma(shape_x, shape_gamma):
function _broadcast_nz (line 351) | def _broadcast_nz(tensor, shape):
function _update_gamma_shape (line 366) | def _update_gamma_shape(shape_x, shape_gamma):
function _get_data_gm (line 398) | def _get_data_gm(shapes, dtype):
function _get_params (line 426) | def _get_params(shape_x, shape_mean, shape_gamma):
function _get_pd_xl (line 469) | def _get_pd_xl(data, shape_x):
function _get_pd_var_front (line 491) | def _get_pd_var_front(data, cast_dtype):
function _get_pd_var (line 519) | def _get_pd_var(data, params, shape_x, pd_xl, cast_dtype):
function _get_pd_mean (line 561) | def _get_pd_mean(params, pd_xl, pd_var, var_elta_2, sub_x_mean, cast_dty...
function _get_pd_x (line 600) | def _get_pd_x(data, params, shape_x, dtype, cast_dtype):
function _get_res (line 656) | def _get_res(data, params, shape_x, dtype, cast_dtype):
function _get_pds (line 688) | def _get_pds(data_dy, data_x, data_variance, data_mean, data_gamma, shap...
function layer_norm_x_backprop_v2_compute (line 748) | def layer_norm_x_backprop_v2_compute(input_dy,
function update_shape_nz (line 787) | def update_shape_nz(shape_x, shape_var, shape_gamma):
function _get_data_nz (line 847) | def _get_data_nz(param_nz, dtype):
function _get_pd_xl_nz (line 863) | def _get_pd_xl_nz(data, param_nz):
function _get_pd_var_front_nz (line 874) | def _get_pd_var_front_nz(data, cast_dtype):
function _get_pd_var_nz (line 889) | def _get_pd_var_nz(data, param_nz, pd_xl, cast_dtype):
function _get_pd_mean_nz (line 907) | def _get_pd_mean_nz(param_nz, pd_xl, pd_var, var_elta_2, sub_x_mean, cas...
function _get_pd_x_nz (line 920) | def _get_pd_x_nz(data, param_nz, dtype, cast_dtype):
function _get_res_nz (line 951) | def _get_res_nz(data, param_nz, dtype, cast_dtype):
function _get_pds_nz (line 961) | def _get_pds_nz(data_dy, data_x, data_variance, data_mean, data_gamma, p...
function layer_norm_x_back_nz_compute (line 994) | def layer_norm_x_back_nz_compute(data_dy, data_x, data_variance, data_me...
function layer_norm_x_backprop_v2 (line 1027) | def layer_norm_x_backprop_v2(input_dy,
FILE: codegeex/mindspore/src/adam.py
function _update_run_kernel (line 36) | def _update_run_kernel(opt, clip_value, beta1, beta2, eps, lr, weight_de...
function _check_param_value (line 53) | def _check_param_value(beta1, beta2, eps, prim_name):
class AdamWeightDecayOp (line 63) | class AdamWeightDecayOp(Optimizer):
method __init__ (line 136) | def __init__(self, params, learning_rate=1e-3, beta1=0.9, beta2=0.999,...
method construct (line 156) | def construct(self, gradients, clip_value):
method clone_param32 (line 182) | def clone_param32(self, prefix, init=None):
FILE: codegeex/mindspore/src/callbacks.py
class LossCallBack (line 31) | class LossCallBack(Callback):
method __init__ (line 37) | def __init__(
method step_end (line 61) | def step_end(self, run_context):
class EvalCallBack (line 92) | class EvalCallBack(Callback):
method __init__ (line 102) | def __init__(self, model, eval_dataset, ppl_metric, validation_loss, p...
method step_end (line 123) | def step_end(self, run_context):
class SaveCheckpointCallback (line 161) | class SaveCheckpointCallback(Callback):
method __init__ (line 162) | def __init__(self, cache_dir, bucket, local_rank=0, has_trained_epoch=...
method step_end (line 175) | def step_end(self, run_context):
method syn_files (line 183) | def syn_files(self):
FILE: codegeex/mindspore/src/code_tokenizer.py
function encode_whitespaces (line 8) | def encode_whitespaces(text, start_extra_id: int, max_len: int):
function decode_whitespaces (line 43) | def decode_whitespaces(text: str, start_extra_id: int, max_len: int):
class Code13BDictionary (line 58) | class Code13BDictionary(object):
method __init__ (line 59) | def __init__(
method _pad_to_vocab_size (line 87) | def _pad_to_vocab_size(self, vocab_size: int):
method _load_dict (line 94) | def _load_dict(self, dict_file: str):
method _add_symbol (line 103) | def _add_symbol(self, sym: str, count: int):
method __len__ (line 109) | def __len__(self):
method index (line 112) | def index(self, sym: str):
method string (line 115) | def string(self, idx: int):
method map_token (line 118) | def map_token(self, token: Union[int, str]):
method map_tokens (line 123) | def map_tokens(self, tokens):
method decode_tokens (line 126) | def decode_tokens(self, tokens):
class CodeTokenizer (line 131) | class CodeTokenizer(object):
method __init__ (line 132) | def __init__(
method encode_code (line 149) | def encode_code(self, code: str):
method decode_code (line 161) | def decode_code(self, input_ids):
FILE: codegeex/mindspore/src/dataset.py
function get_input_data_batch_slice_map (line 35) | def get_input_data_batch_slice_map(input_ids, eod_id, rank, dis, eod_res...
function create_dataset (line 80) | def create_dataset(batch_size, data_path, args_opt, device_num=1, rank=0...
function get_code_data_train (line 173) | def get_code_data_train(code_data_path, args_opt, process_fn=None, scale...
function get_code_data_eval (line 201) | def get_code_data_eval(code_data_path, args_opt, process_fn=None, scale=1):
FILE: codegeex/mindspore/src/dataset_finetune.py
function get_input_data_batch_slice_map (line 33) | def get_input_data_batch_slice_map(input_ids, loss_mask, eod_id, rank, d...
function create_dataset (line 82) | def create_dataset(batch_size, data_path, args_opt, device_num=1, rank=0...
function get_code_data (line 165) | def get_code_data(code_data_path, split, args_opt, process_fn=None, skip...
FILE: codegeex/mindspore/src/generate.py
function topk_fun (line 26) | def topk_fun(logits, topk=5):
function sampler (line 38) | def sampler(log_probs_revised, top_p, top_k_num, use_pynative=False):
function generate (line 84) | def generate(model, origin_inputs, config, verbose=False):
function generate_increment (line 170) | def generate_increment(model, origin_inputs, config, verbose=False):
FILE: codegeex/mindspore/src/generate_finetune.py
function topk_fun (line 28) | def topk_fun(logits, topk=5):
function sampler (line 35) | def sampler(log_probs_revised, top_p, top_k_num, use_pynative=False, bad...
function generate_increment (line 78) | def generate_increment(model, origin_inputs, origin_length, config, toke...
FILE: codegeex/mindspore/src/generate_greedy.py
function topk_fun (line 27) | def topk_fun(logits, topk=5):
function sampler (line 34) | def sampler(log_probs_revised, top_p, top_k_num, use_pynative=False, bad...
function generate_increment (line 45) | def generate_increment(model, origin_inputs, origin_length, config, toke...
FILE: codegeex/mindspore/src/generate_humaneval.py
function is_code_generation_finished (line 25) | def is_code_generation_finished(text: str):
function cleanup_text (line 37) | def cleanup_text(text: str):
function truncate_text (line 49) | def truncate_text(text: str):
function topk_fun (line 63) | def topk_fun(logits, topk=5):
function sampler (line 77) | def sampler(log_probs_revised, top_p, top_k_num, use_pynative=False, bad...
function generate_increment (line 130) | def generate_increment(model, origin_inputs, config, tokenizer, verbose=...
FILE: codegeex/mindspore/src/metrics.py
class PPLMetric (line 26) | class PPLMetric(Metric):
method __init__ (line 31) | def __init__(self, data_length):
method clear (line 40) | def clear(self):
method update (line 45) | def update(self, *inputs): # inputs
method eval (line 53) | def eval(self):
class ValidationLoss (line 67) | class ValidationLoss(Metric):
method __init__ (line 68) | def __init__(self, data_length):
method clear (line 77) | def clear(self):
method update (line 82) | def update(self, *inputs): # inputs
method eval (line 94) | def eval(self):
FILE: codegeex/mindspore/src/pangu_alpha.py
class EmbeddingLayer (line 38) | class EmbeddingLayer(nn.Cell):
method __init__ (line 41) | def __init__(self, config):
method construct (line 92) | def construct(
method get_word_embedding_weight (line 105) | def get_word_embedding_weight(self):
class QueryLayer (line 109) | class QueryLayer(TransformerEncoderLayer):
method __init__ (line 112) | def __init__(
method construct (line 144) | def construct(
class PanGuHead (line 220) | class PanGuHead(Cell):
method __init__ (line 232) | def __init__(
method construct (line 252) | def construct(self, state, embed):
function set_parallel_configure_for_layer (line 261) | def set_parallel_configure_for_layer(
class PanguAlpha_Model (line 296) | class PanguAlpha_Model(Cell):
method __init__ (line 299) | def __init__(self, config):
method construct (line 388) | def construct(self, input_ids,
method reshape_to_2d (line 419) | def reshape_to_2d(self, x):
method load_embedding_from_ckpt (line 427) | def load_embedding_from_ckpt(self, load_ckpt_path):
class PanguAlphaModel (line 461) | class PanguAlphaModel(nn.Cell):
method __init__ (line 474) | def __init__(self, config):
method construct (line 488) | def construct(self, input_ids, input_position, attention_mask,
class PanGUAlphaWithLoss (line 496) | class PanGUAlphaWithLoss(Cell):
method __init__ (line 508) | def __init__(self, config, network, loss):
method construct (line 527) | def construct(self, input_ids, input_position=None, attention_mask=None):
class EvalNet (line 547) | class EvalNet(nn.Cell):
method __init__ (line 561) | def __init__(self, backbone, generate=False, pad_token=6, seq_length=2...
method construct (line 575) | def construct(self, input_ids, current_index, init_reset=True, batch_v...
class LogitsNet (line 597) | class LogitsNet(nn.Cell):
method __init__ (line 610) | def __init__(self, backbone, generate=False, pad_token=6, seq_length=2...
method construct (line 624) | def construct(self, input_ids, init_reset=True, batch_valid_length=Non...
class PanGUAlphaWithFinetuneLoss (line 642) | class PanGUAlphaWithFinetuneLoss(Cell):
method __init__ (line 654) | def __init__(self, config, network, loss):
method construct (line 673) | def construct(self, input_ids, loss_mask, input_position, attention_ma...
FILE: codegeex/mindspore/src/pangu_alpha_config.py
class PanguAlphaConfig (line 21) | class PanguAlphaConfig:
method __init__ (line 26) | def __init__(self,
method __str__ (line 71) | def __str__(self):
function set_parse (line 80) | def set_parse(args_opt):
FILE: codegeex/mindspore/src/pangu_alpha_fp16_predict.py
class EmbeddingLayer (line 37) | class EmbeddingLayer(nn.Cell):
method __init__ (line 40) | def __init__(self, config):
method construct (line 89) | def construct(
method get_word_embedding_weight (line 101) | def get_word_embedding_weight(self):
class QueryLayer (line 105) | class QueryLayer(TransformerEncoderLayer):
method __init__ (line 108) | def __init__(
method construct (line 140) | def construct(
class PanGuHead (line 216) | class PanGuHead(Cell):
method __init__ (line 228) | def __init__(
method construct (line 248) | def construct(self, state, embed):
function set_parallel_configure_for_layer (line 257) | def set_parallel_configure_for_layer(
class PanguAlpha_Model (line 292) | class PanguAlpha_Model(Cell):
method __init__ (line 295) | def __init__(self, config):
method construct (line 383) | def construct(self, input_ids,
method reshape_to_2d (line 422) | def reshape_to_2d(self, x):
method load_embedding_from_ckpt (line 430) | def load_embedding_from_ckpt(self, load_ckpt_path):
class PanguAlphaModel (line 471) | class PanguAlphaModel(nn.Cell):
method __init__ (line 484) | def __init__(self, config):
method construct (line 499) | def construct(self, input_ids, input_position, attention_mask,
class PanGUAlphaWithLoss (line 511) | class PanGUAlphaWithLoss(Cell):
method __init__ (line 523) | def __init__(self, config, network, loss):
method construct (line 542) | def construct(self, input_ids, input_position=None, attention_mask=None):
class EvalNet (line 562) | class EvalNet(nn.Cell):
method __init__ (line 576) | def __init__(self, backbone, generate=False, pad_token=6, seq_length=2...
method construct (line 590) | def construct(self, input_ids, current_index, init_reset=True, batch_v...
FILE: codegeex/mindspore/src/pangu_alpha_wrapcell.py
function _clip_grad (line 37) | def _clip_grad(clip_type, clip_value, grad):
function tensor_grad_scale (line 72) | def tensor_grad_scale(scale, grad):
function tensor_grad_scale_pipeline (line 77) | def tensor_grad_scale_pipeline(scale, grad, accu_grad):
function tensor_shard_grad_scale_pipeline (line 87) | def tensor_shard_grad_scale_pipeline(scale, grad, accu_grad):
class PanguAlphaTrainOneStepWithLossScaleCell (line 94) | class PanguAlphaTrainOneStepWithLossScaleCell(TrainOneStepWithLossScaleC...
method __init__ (line 107) | def __init__(
method construct (line 130) | def construct(self, input_ids, input_position, attention_mask, layer_p...
class PanguAlphaTrainPipelineWithLossScaleCell (line 168) | class PanguAlphaTrainPipelineWithLossScaleCell(nn.Cell):
method __init__ (line 181) | def __init__(self, network, optimizer, config, scale_update_cell=None,...
method construct (line 221) | def construct(
FILE: codegeex/mindspore/src/pangu_alpha_wrapcell_finetune.py
function _clip_grad (line 37) | def _clip_grad(clip_type, clip_value, grad):
function tensor_grad_scale (line 72) | def tensor_grad_scale(scale, grad):
function tensor_grad_scale_pipeline (line 77) | def tensor_grad_scale_pipeline(scale, grad, accu_grad):
function tensor_shard_grad_scale_pipeline (line 87) | def tensor_shard_grad_scale_pipeline(scale, grad, accu_grad):
class PanguAlphaTrainOneStepWithLossScaleCell (line 94) | class PanguAlphaTrainOneStepWithLossScaleCell(TrainOneStepWithLossScaleC...
method __init__ (line 107) | def __init__(
method construct (line 130) | def construct(self, input_ids, loss_mask, input_position, attention_ma...
class PanguAlphaTrainPipelineWithLossScaleCell (line 168) | class PanguAlphaTrainPipelineWithLossScaleCell(nn.Cell):
method __init__ (line 181) | def __init__(self, network, optimizer, config, scale_update_cell=None,...
method construct (line 221) | def construct(
FILE: codegeex/mindspore/src/preprocess.py
function chunks (line 32) | def chunks(lst, n):
function package_file (line 38) | def package_file(it, n):
function clean_wikitext (line 53) | def clean_wikitext(string):
function tokenize_openwebtext (line 87) | def tokenize_openwebtext(tokenizer, iterator, seq_length, eot):
function tokenize_wiki (line 107) | def tokenize_wiki(tokenizer, file_path, seq_length, eot):
function tokenize_lambada (line 124) | def tokenize_lambada(tokenizer, file_path, seq_length, eot):
function task_unit (line 145) | def task_unit(iterator, tokenizer, seq_length, eot, parallel_writer=True):
FILE: codegeex/mindspore/src/sat_dataset.py
class Dataset (line 8) | class Dataset(ABC):
method __len__ (line 10) | def __len__(self):
method __getitem__ (line 14) | def __getitem__(self, idx):
class LMDBDataset (line 18) | class LMDBDataset(Dataset):
method __init__ (line 19) | def __init__(self, path, process_fn=None):
method __len__ (line 38) | def __len__(self):
method __getitem__ (line 41) | def __getitem__(self, idx):
class PadDataset (line 56) | class PadDataset(Dataset):
method __init__ (line 57) | def __init__(self, dataset, seq_len, eod_id):
method __len__ (line 62) | def __len__(self):
method __getitem__ (line 65) | def __getitem__(self, idx):
class BinaryDataset (line 72) | class BinaryDataset(Dataset):
method __init__ (line 73) | def __init__(
method __len__ (line 100) | def __len__(self):
method __getitem__ (line 103) | def __getitem__(self, index):
class TSVDataset (line 108) | class TSVDataset(Dataset):
method __init__ (line 109) | def __init__(self, path, process_fn, with_heads=True, **kwargs):
method __len__ (line 118) | def __len__(self):
method __getitem__ (line 121) | def __getitem__(self, index):
class ConcatDataset (line 125) | class ConcatDataset(Dataset):
method cumsum (line 136) | def cumsum(sequence, weights):
method __init__ (line 144) | def __init__(self, datasets, weights=None, skip_num=0, **kwargs):
method __len__ (line 158) | def __len__(self):
method __getitem__ (line 161) | def __getitem__(self, idx):
class RandomMappingDataset (line 175) | class RandomMappingDataset(Dataset):
method __init__ (line 181) | def __init__(self, ds):
method __len__ (line 185) | def __len__(self):
method __getitem__ (line 188) | def __getitem__(self, index):
class BlockedSplitDataset (line 197) | class BlockedSplitDataset(Dataset):
method __init__ (line 204) | def __init__(self, ds, indices, block_size):
method __len__ (line 214) | def __len__(self):
method __getitem__ (line 217) | def __getitem__(self, index):
class SubsetDataset (line 223) | class SubsetDataset(Dataset):
method __init__ (line 224) | def __init__(self, ds, start, length):
method __len__ (line 230) | def __len__(self):
method __getitem__ (line 233) | def __getitem__(self, idx):
function split_train_val_test (line 239) | def split_train_val_test(ds, split=[0.99, 0.01, 0.0], seed=None):
FILE: codegeex/mindspore/src/tokenization_jieba.py
class JIEBATokenizer (line 25) | class JIEBATokenizer():
method __init__ (line 30) | def __init__(self, vocab_file, model_file, max_len=None):
method vocab_size (line 49) | def vocab_size(self):
method __len__ (line 52) | def __len__(self):
method eod (line 56) | def eod(self):
method tokenize (line 59) | def tokenize(self, text):
method convert_tokens_to_ids (line 65) | def convert_tokens_to_ids(self, tokens):
method convert_ids_to_tokens (line 68) | def convert_ids_to_tokens(self, ids):
method encode (line 71) | def encode(self, text):
method decode (line 75) | def decode(self, tokens):
FILE: codegeex/mindspore/src/utils.py
class FP32StateAdamWeightDecay (line 40) | class FP32StateAdamWeightDecay(AdamWeightDecay):
method __init__ (line 49) | def __init__(self, params, learning_rate=1e-3, beta1=0.9, beta2=0.999,...
method clone_state (line 59) | def clone_state(self, parameter_tuple, prefix, init):
function _get_square_sum (line 80) | def _get_square_sum(grad, value):
function _apply_global_norm (line 90) | def _apply_global_norm(enable_grad_fp16, clip_norm, global_norm, grad):
function _get_model_parallel_group (line 98) | def _get_model_parallel_group(mp):
function _get_pipeline_group (line 118) | def _get_pipeline_group():
class GlobalNorm (line 136) | class GlobalNorm(nn.Cell):
method __init__ (line 141) | def __init__(self, params, config):
method construct (line 193) | def construct(self, grads):
class ClipByGlobalNorm (line 206) | class ClipByGlobalNorm(nn.Cell):
method __init__ (line 213) | def __init__(self, params, config, clip_norm=1.0):
method construct (line 223) | def construct(self, grads):
class LearningRate (line 232) | class LearningRate(LearningRateSchedule):
method __init__ (line 237) | def __init__(self,
method construct (line 261) | def construct(self, global_step):
function add_inference_params (line 278) | def add_inference_params(opt):
function add_training_params (line 315) | def add_training_params(opt):
function add_retrain_params (line 419) | def add_retrain_params(opt):
function get_args (line 461) | def get_args(inference=False):
function download_data (line 587) | def download_data(src_data_url, tgt_data_path, rank):
FILE: codegeex/mindspore/train.py
function set_weight_decay (line 59) | def set_weight_decay(params):
function add_checkpoint_callback_policy (line 74) | def add_checkpoint_callback_policy(args_param, callback, rank_id):
function set_parallel_context (line 104) | def set_parallel_context(args_opt):
function run_train (line 121) | def run_train(args_opt):
function restore_checkpoint (line 333) | def restore_checkpoint(args_param, sink_size, dataset, model, network, e...
function get_exception_checkpoints (line 376) | def get_exception_checkpoints(args_param):
function check_exception_checkpoints (line 408) | def check_exception_checkpoints(ckpt_file_list):
function restore_exception_checkpoint (line 425) | def restore_exception_checkpoint(args_param, sink_size, dataset, model, ...
function set_pipeline_parallel_context (line 487) | def set_pipeline_parallel_context(args_opt):
function run_train_pipeline (line 506) | def run_train_pipeline(args_opt):
FILE: codegeex/oneflow/codegeex_model.py
function fast_gelu (line 7) | def fast_gelu(x):
class MLP (line 14) | class MLP(torch.nn.Module):
method __init__ (line 22) | def __init__(
method forward (line 42) | def forward(self, hidden_states):
class SelfAttention (line 52) | class SelfAttention(torch.nn.Module):
method __init__ (line 58) | def __init__(
method forward (line 85) | def forward(
class TopQuerySelfAttention (line 279) | class TopQuerySelfAttention(torch.nn.Module):
method __init__ (line 285) | def __init__(
method forward (line 312) | def forward(
class TransformerLayer (line 494) | class TransformerLayer(torch.nn.Module):
method __init__ (line 500) | def __init__(
method forward (line 530) | def forward(
class TopQueryLayer (line 573) | class TopQueryLayer(torch.nn.Module):
method __init__ (line 579) | def __init__(
method forward (line 607) | def forward(
class Transformer (line 657) | class Transformer(torch.nn.Module):
method __init__ (line 660) | def __init__(
method _get_layer_index (line 698) | def _get_layer_index(self, layer_number):
method _get_layer (line 701) | def _get_layer(self, layer_number):
method forward (line 704) | def forward(
method state_dict_for_save_checkpoint (line 767) | def state_dict_for_save_checkpoint(
class Embedding (line 773) | class Embedding(torch.nn.Module):
method __init__ (line 782) | def __init__(
method forward (line 802) | def forward(self, input_ids, position_ids):
method state_dict_for_save_checkpoint (line 810) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 823) | def load_state_dict(self, state_dict, strict=True):
class QueryEmbedding (line 852) | class QueryEmbedding(torch.nn.Module):
method __init__ (line 861) | def __init__(
method forward (line 878) | def forward(self, position_ids):
method state_dict_for_save_checkpoint (line 884) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 895) | def load_state_dict(self, state_dict, strict=True):
class TransformerLanguageModel (line 911) | class TransformerLanguageModel(torch.nn.Module):
method __init__ (line 926) | def __init__(
method forward (line 959) | def forward(
method state_dict_for_save_checkpoint (line 986) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 1003) | def load_state_dict(self, state_dict, strict=True):
class CodeGeeXModel (line 1039) | class CodeGeeXModel(torch.nn.Module):
method __init__ (line 1042) | def __init__(
method forward (line 1059) | def forward(
method state_dict_for_save_checkpoint (line 1088) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 1097) | def load_state_dict(self, state_dict, strict=True):
FILE: codegeex/oneflow/inference.py
function get_ltor_masks_and_position_ids (line 12) | def get_ltor_masks_and_position_ids(
function get_batch (line 67) | def get_batch(
function top_k_logits (line 87) | def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")):
function pad_batch (line 117) | def pad_batch(batch, pad_id, seq_length):
function forward_step (line 127) | def forward_step(
function get_token_stream (line 158) | def get_token_stream(
function switch (line 212) | def switch(val1, val2, boolean):
function sample_sequence_batch (line 217) | def sample_sequence_batch(
FILE: codegeex/paddle/codegeex_model.py
function fast_gelu (line 6) | def fast_gelu(x):
class MLP (line 11) | class MLP(paddle.nn.Layer):
method __init__ (line 20) | def __init__(
method forward (line 40) | def forward(self, hidden_states):
class SelfAttention (line 50) | class SelfAttention(paddle.nn.Layer):
method __init__ (line 57) | def __init__(
method forward (line 84) | def forward(
class TopQuerySelfAttention (line 226) | class TopQuerySelfAttention(paddle.nn.Layer):
method __init__ (line 233) | def __init__(
method forward (line 260) | def forward(
class TransformerLayer (line 400) | class TransformerLayer(paddle.nn.Layer):
method __init__ (line 407) | def __init__(
method forward (line 437) | def forward(
class TopQueryLayer (line 478) | class TopQueryLayer(paddle.nn.Layer):
method __init__ (line 485) | def __init__(
method forward (line 513) | def forward(
class Transformer (line 563) | class Transformer(paddle.nn.Layer):
method __init__ (line 566) | def __init__(
method _get_layer_index (line 604) | def _get_layer_index(self, layer_number):
method _get_layer (line 607) | def _get_layer(self, layer_number):
method forward (line 610) | def forward(
method state_dict_for_save_checkpoint (line 672) | def state_dict_for_save_checkpoint(
class Embedding (line 678) | class Embedding(paddle.nn.Layer):
method __init__ (line 688) | def __init__(
method forward (line 708) | def forward(self, input_ids, position_ids):
method state_dict_for_save_checkpoint (line 716) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method set_state_dict (line 729) | def set_state_dict(self, state_dict, use_structured_name=True):
class QueryEmbedding (line 758) | class QueryEmbedding(paddle.nn.Layer):
method __init__ (line 768) | def __init__(
method forward (line 785) | def forward(self, position_ids):
method state_dict_for_save_checkpoint (line 791) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method set_state_dict (line 802) | def set_state_dict(self, state_dict, use_structured_name=True):
class TransformerLanguageModel (line 818) | class TransformerLanguageModel(paddle.nn.Layer):
method __init__ (line 834) | def __init__(
method forward (line 867) | def forward(
method state_dict_for_save_checkpoint (line 894) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method set_state_dict (line 911) | def set_state_dict(self, state_dict, use_structured_name=True):
class CodeGeeXModel (line 947) | class CodeGeeXModel(paddle.nn.Layer):
method __init__ (line 950) | def __init__(
method forward (line 967) | def forward(
method state_dict_for_save_checkpoint (line 996) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method set_state_dict (line 1005) | def set_state_dict(self, state_dict, use_structured_name=True):
FILE: codegeex/paddle/inference.py
function get_ltor_masks_and_position_ids (line 12) | def get_ltor_masks_and_position_ids(
function get_batch (line 67) | def get_batch(
function top_k_logits (line 87) | def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")):
function pad_batch (line 117) | def pad_batch(batch, pad_id, seq_length):
function forward_step (line 127) | def forward_step(
function get_token_stream (line 158) | def get_token_stream(
function switch (line 212) | def switch(val1, val2, boolean):
function sample_sequence_batch (line 217) | def sample_sequence_batch(
FILE: codegeex/paddle/pt_to_pdparams.py
function WalkDict (line 15) | def WalkDict(x):
function parse_opt (line 29) | def parse_opt():
function main (line 47) | def main(opt):
FILE: codegeex/quantization/quantize.py
class W8A16Linear (line 9) | class W8A16Linear(torch.autograd.Function):
method forward (line 11) | def forward(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: to...
method backward (line 23) | def backward(ctx, grad_output: torch.Tensor):
class QuantizedLinear (line 32) | class QuantizedLinear(torch.nn.Module):
method __init__ (line 33) | def __init__(
method forward (line 68) | def forward(self, input_):
class QuantizedColumnParallelLinear (line 77) | class QuantizedColumnParallelLinear(ColumnParallelLinear):
method __init__ (line 78) | def __init__(
method forward (line 118) | def forward(self, input_):
class QuantizedRowParallelLinear (line 136) | class QuantizedRowParallelLinear(RowParallelLinear):
method __init__ (line 137) | def __init__(
method forward (line 177) | def forward(self, input_):
function quantize (line 196) | def quantize(model, weight_bit_width, backend="torch"):
FILE: codegeex/quantization/quantize_oneflow.py
function _pack_int8_to_int4 (line 5) | def _pack_int8_to_int4(x):
function _quantize (line 16) | def _quantize(num_bits, symmetric, x, group_dim, group_size, quant_type):
class QuantizedLinear (line 53) | class QuantizedLinear(torch.nn.Module):
method __init__ (line 54) | def __init__(
method forward (line 89) | def forward(self, input_):
function quantize_oneflow (line 103) | def quantize_oneflow(model, weight_bit_width):
FILE: codegeex/tokenizer/tokenizer.py
function encode_whitespaces (line 7) | def encode_whitespaces(text: str, start_extra_id: int, max_len: int):
function decode_whitespaces (line 18) | def decode_whitespaces(text: str, start_extra_id: int, max_len: int):
class CodeGeeXTokenizer (line 33) | class CodeGeeXTokenizer(object):
method __init__ (line 34) | def __init__(
method encode_code (line 51) | def encode_code(self, code: str):
method decode_code (line 58) | def decode_code(self, input_ids):
FILE: codegeex/torch/codegeex_model.py
function fast_gelu (line 7) | def fast_gelu(x):
class MLP (line 12) | class MLP(torch.nn.Module):
method __init__ (line 21) | def __init__(
method forward (line 41) | def forward(self, hidden_states):
class SelfAttention (line 51) | class SelfAttention(torch.nn.Module):
method __init__ (line 58) | def __init__(
method forward (line 85) | def forward(
class TopQuerySelfAttention (line 227) | class TopQuerySelfAttention(torch.nn.Module):
method __init__ (line 234) | def __init__(
method forward (line 261) | def forward(
class TransformerLayer (line 401) | class TransformerLayer(torch.nn.Module):
method __init__ (line 408) | def __init__(
method forward (line 438) | def forward(
class TopQueryLayer (line 479) | class TopQueryLayer(torch.nn.Module):
method __init__ (line 486) | def __init__(
method forward (line 514) | def forward(
class Transformer (line 564) | class Transformer(torch.nn.Module):
method __init__ (line 567) | def __init__(
method _get_layer_index (line 605) | def _get_layer_index(self, layer_number):
method _get_layer (line 608) | def _get_layer(self, layer_number):
method forward (line 611) | def forward(
method state_dict_for_save_checkpoint (line 673) | def state_dict_for_save_checkpoint(
class Embedding (line 679) | class Embedding(torch.nn.Module):
method __init__ (line 689) | def __init__(
method forward (line 709) | def forward(self, input_ids, position_ids):
method state_dict_for_save_checkpoint (line 717) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 730) | def load_state_dict(self, state_dict, strict=True):
class QueryEmbedding (line 759) | class QueryEmbedding(torch.nn.Module):
method __init__ (line 769) | def __init__(
method forward (line 786) | def forward(self, position_ids):
method state_dict_for_save_checkpoint (line 792) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 803) | def load_state_dict(self, state_dict, strict=True):
class TransformerLanguageModel (line 819) | class TransformerLanguageModel(torch.nn.Module):
method __init__ (line 835) | def __init__(
method forward (line 868) | def forward(
method state_dict_for_save_checkpoint (line 895) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 912) | def load_state_dict(self, state_dict, strict=True):
class CodeGeeXModel (line 948) | class CodeGeeXModel(torch.nn.Module):
method __init__ (line 951) | def __init__(
method forward (line 968) | def forward(
method state_dict_for_save_checkpoint (line 997) | def state_dict_for_save_checkpoint(self, destination=None, prefix='',
method load_state_dict (line 1006) | def load_state_dict(self, state_dict, strict=True):
FILE: codegeex/torch/get_ckpt_qkv.py
function main (line 9) | def main():
FILE: codegeex/torch/inference.py
function get_ltor_masks_and_position_ids (line 12) | def get_ltor_masks_and_position_ids(
function get_batch (line 67) | def get_batch(
function top_k_logits (line 87) | def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")):
function pad_batch (line 117) | def pad_batch(batch, pad_id, seq_length):
function forward_step (line 127) | def forward_step(
function get_token_stream (line 158) | def get_token_stream(
function switch (line 212) | def switch(val1, val2, boolean):
function sample_sequence_batch (line 217) | def sample_sequence_batch(
FILE: deployment/server_gradio.py
function model_provider (line 14) | def model_provider(args):
function add_code_generation_args (line 27) | def add_code_generation_args(parser):
function main (line 85) | def main():
FILE: tests/test_inference.py
function model_provider (line 12) | def model_provider(args):
function add_code_generation_args (line 26) | def add_code_generation_args(parser):
function main (line 118) | def main():
FILE: tests/test_inference_megatron.py
function set_random_seed (line 17) | def set_random_seed(seed):
function model_provider (line 26) | def model_provider(pre_process=True, post_process=True):
function add_code_generation_args (line 36) | def add_code_generation_args(parser):
function main (line 136) | def main():
FILE: tests/test_inference_oneflow.py
function model_provider (line 17) | def model_provider(args):
function add_code_generation_args (line 31) | def add_code_generation_args(parser):
function main (line 119) | def main():
FILE: tests/test_inference_paddle.py
function model_provider (line 15) | def model_provider(args):
function add_code_generation_args (line 41) | def add_code_generation_args(parser):
function main (line 129) | def main():
Condensed preview — 189 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (3,317K chars).
[
{
"path": ".gitmodules",
"chars": 169,
"preview": "[submodule \"vscode-extension/codegeex-vscode-extension\"]\n\tpath = vscode-extension/codegeex-vscode-extension\n\turl = git@g"
},
{
"path": "LICENSE",
"chars": 11357,
"preview": " Apache License\n Version 2.0, January 2004\n "
},
{
"path": "MODEL_LICENSE",
"chars": 2327,
"preview": "The CodeGeeX License\n\n1. Definitions\n\n“Licensor” means the CodeGeeX Model Team that distributes its Software.\n\n“Software"
},
{
"path": "README.md",
"chars": 19570,
"preview": "<img src=\"resources/logo/codegeex_logo.png\">\n\n<p align=\"center\">\n 🏠 <a href=\"https://codegeex.cn\" target=\"_blank\">Hom"
},
{
"path": "README_zh.md",
"chars": 11404,
"preview": "<img src=\"resources/logo/codegeex_logo.png\">\n\n<p align=\"center\">\n 🏠 <a href=\"https://models.aminer.cn/codegeex/zh-CN\""
},
{
"path": "api/README_zh.md",
"chars": 569,
"preview": "\n\n# 创建CodeGeeX API\n\n使用[天启 · API开放平台](https://tianqi.aminer.cn/open/"
},
{
"path": "api/codegeex-api-example-java/pom.xml",
"chars": 4655,
"preview": "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<project xmlns=\"http://maven.apache.org/POM/4.0.0\"\n xmlns:xsi=\"http://www"
},
{
"path": "api/codegeex-api-example-java/src/main/java/cn/aminer/codegeex/example/CodeGenerationExample.java",
"chars": 2710,
"preview": "package cn.aminer.codegeex.example;\n\nimport cn.aminer.codegeex.example.pojo.Payload;\nimport com.fasterxml.jackson.databi"
},
{
"path": "api/codegeex-api-example-java/src/main/java/cn/aminer/codegeex/example/pojo/Payload.java",
"chars": 710,
"preview": "package cn.aminer.codegeex.example.pojo;\n\nimport com.fasterxml.jackson.annotation.JsonIgnoreProperties;\nimport com.faste"
},
{
"path": "api/codegeex-api-example-python/generation_example.py",
"chars": 1125,
"preview": "# encoding:utf-8\n\nimport json\n\nimport requests\n\n'''\nCode Generation\n'''\nAPI_KEY = \"\" # Get from Tianqi console. 从控制台获取\n"
},
{
"path": "codegeex/__init__.py",
"chars": 2146,
"preview": "import copy\n\nfrom typing import *\nfrom codegeex.tokenizer import CodeGeeXTokenizer\nfrom codegeex.torch.inference import "
},
{
"path": "codegeex/benchmark/README.md",
"chars": 4525,
"preview": "# HumanEval-X: A new benchmark for Multilingual Program Synthesis\n\n🌐 <a href=\"README_zh.md\" target=\"_blank\">中文</a>\n\nHuma"
},
{
"path": "codegeex/benchmark/README_zh.md",
"chars": 2433,
"preview": "# HumanEval-X: 多语言代码生成基准\n\n🌐 <a href=\"README.md\" target=\"_blank\">English</a>\n\n为了更好地评测代码生成模型的多语言生成能力，我们构建了一个新基准HumanEval-X"
},
{
"path": "codegeex/benchmark/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "codegeex/benchmark/evaluate_humaneval_x.py",
"chars": 8436,
"preview": "import os\nimport sys\nimport fire\nimport json\nimport gzip\nimport regex\nimport numpy as np\n\nfrom typing import *\nfrom tqdm"
},
{
"path": "codegeex/benchmark/execution.py",
"chars": 21844,
"preview": "import contextlib\nimport faulthandler\nimport io\nimport multiprocessing\nimport os\nimport platform\nimport signal\nimport ra"
},
{
"path": "codegeex/benchmark/gather_output.py",
"chars": 1891,
"preview": "import os\nimport sys\nimport fire\nimport glob\n\n\ndef gather_output(\n output_dir: str = \"./output\",\n output_prefix: s"
},
{
"path": "codegeex/benchmark/humaneval-x/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "codegeex/benchmark/humaneval-x/cpp/evaluation/test.cpp",
"chars": 1681,
"preview": "/*\nInput to this function is a string containing multiple groups of nested parentheses. Your goal is to\nseparate those g"
},
{
"path": "codegeex/benchmark/humaneval-x/evaluate_humaneval_x.py",
"chars": 8436,
"preview": "import os\nimport sys\nimport fire\nimport json\nimport gzip\nimport regex\nimport numpy as np\n\nfrom typing import *\nfrom tqdm"
},
{
"path": "codegeex/benchmark/humaneval-x/generate_humaneval_x.py",
"chars": 12449,
"preview": "import argparse\nimport logging\nimport os\nimport random\nimport socket\nimport time\nfrom typing import *\n\nimport torch\nimpo"
},
{
"path": "codegeex/benchmark/humaneval-x/go/evaluation/go.mod",
"chars": 261,
"preview": "module humanEval\n\ngo 1.18\n\nrequire (\n\tgithub.com/go-openapi/inflect v0.19.0\n\tgithub.com/stretchr/testify v1.8.0\n)\n\nrequi"
},
{
"path": "codegeex/benchmark/humaneval-x/go/evaluation/go.sum",
"chars": 1520,
"preview": "github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=\ngithub.com/davecgh/go-spew v1.1"
},
{
"path": "codegeex/benchmark/humaneval-x/translate_humaneval_x.py",
"chars": 13329,
"preview": "import os\nimport zmq\nimport time\nimport torch\nimport random\nimport socket\nimport logging\nimport argparse\n\nfrom typing im"
},
{
"path": "codegeex/benchmark/inspect_result.py",
"chars": 12118,
"preview": "import os\nimport sys\nimport fire\nimport json\nimport glob\nimport numpy as np\nimport pandas as pd\n\nfrom collections import"
},
{
"path": "codegeex/benchmark/metric.py",
"chars": 2039,
"preview": "# Copyright (c) OpenAI (https://openai.com)\n\n# Permission is hereby granted, free of charge, to any person obtaining a c"
},
{
"path": "codegeex/benchmark/rust/Cargo.toml",
"chars": 213,
"preview": "[package]\nname = \"rust\"\nversion = \"0.1.0\"\nedition = \"2021\"\n\n# See more keys and their definitions at https://doc.rust-la"
},
{
"path": "codegeex/benchmark/utils.py",
"chars": 5826,
"preview": "import os\nfrom typing import *\nfrom codegeex.data.data_utils import stream_jsonl, LANGUAGE_TAG\n\n\nIMPORT_HELPER = {\n \""
},
{
"path": "codegeex/data/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "codegeex/data/data_utils.py",
"chars": 4844,
"preview": "import os\nimport gzip\nimport json\n\nfrom typing import *\n\nLANGUAGE_TAG = {\n \"c\" : \"// language: C\",\n \"c+"
},
{
"path": "codegeex/data/process_pretrain_dataset.py",
"chars": 4456,
"preview": "import os\nimport glob\nimport fire\nimport torch\nimport multiprocessing\n\nfrom typing import *\nfrom tqdm.auto import tqdm\nf"
},
{
"path": "codegeex/data/processor.py",
"chars": 6112,
"preview": "from typing import *\nfrom time import perf_counter\n\nfrom codegeex.data.data_utils import sliding_window\nfrom codegeex.da"
},
{
"path": "codegeex/data/types.py",
"chars": 304,
"preview": "from typing import *\nfrom dataclasses import dataclass\n\n\n@dataclass\nclass PromptSample:\n prompt: str\n code: str\n "
},
{
"path": "codegeex/docker/Dockerfile",
"chars": 2872,
"preview": "FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime\n\nRUN cp /etc/apt/sources.list /etc/apt/sources.list.bak \\\n && sed"
},
{
"path": "codegeex/kernels/__init__.py",
"chars": 2993,
"preview": "import pkg_resources\nimport torch\nimport ctypes\n\nfrom typing import List\nfrom cpm_kernels.kernels.base import LazyKernel"
},
{
"path": "codegeex/megatron/__init__.py",
"chars": 1658,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/arguments.py",
"chars": 48485,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/checkpointing.py",
"chars": 20400,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/code_generation_utils.py",
"chars": 41383,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/convert_ckpt_parallel.py",
"chars": 4958,
"preview": "\"\"\"Get model parallel partitions.\"\"\"\n\nimport os\nimport torch\nimport argparse\n\n\ndef get_change_ckpt_args(parser):\n \"\"\""
},
{
"path": "codegeex/megatron/data/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "codegeex/megatron/data/blendable_dataset.py",
"chars": 2144,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/data/data_samplers.py",
"chars": 6567,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/data/dataset_utils.py",
"chars": 19607,
"preview": "# coding=utf-8\n# Copyright 2018 The Google AI Language Team Authors, and NVIDIA.\n#\n# Licensed under the Apache License, "
},
{
"path": "codegeex/megatron/data/helpers.cpp",
"chars": 27148,
"preview": "/*\n coding=utf-8\n Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n\n Licensed under the Apache License, Ver"
},
{
"path": "codegeex/megatron/data/indexed_dataset.py",
"chars": 18665,
"preview": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n#"
},
{
"path": "codegeex/megatron/data/prompt_dataset.py",
"chars": 11826,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/enums.py",
"chars": 830,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/global_vars.py",
"chars": 8254,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/inference.py",
"chars": 11215,
"preview": "import copy\nimport json\nimport random\nimport traceback\nfrom typing import *\n\nimport numpy\nimport torch\nimport zmq\n\nfrom "
},
{
"path": "codegeex/megatron/initialize.py",
"chars": 11502,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/learning_rates.py",
"chars": 6764,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/memory.py",
"chars": 5131,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/merge_ckpt_parallel.py",
"chars": 7734,
"preview": "\"\"\"Merge model parallel partitions into a single checkpoint.\"\"\"\n\nimport os\nimport torch\nimport random\n\nfrom codegeex.meg"
},
{
"path": "codegeex/megatron/microbatches.py",
"chars": 6959,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/mindspore_to_megatron.py",
"chars": 13299,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/model/__init__.py",
"chars": 828,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/model/codegeex_model.py",
"chars": 7863,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/model/distributed.py",
"chars": 8428,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/model/language_model.py",
"chars": 25257,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/model/module.py",
"chars": 7501,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/model/transformer.py",
"chars": 40825,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/model/utils.py",
"chars": 2276,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/mpu/__init__.py",
"chars": 3268,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/mpu/cross_entropy.py",
"chars": 4729,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/mpu/data.py",
"chars": 3999,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/mpu/initialize.py",
"chars": 15334,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/mpu/layers.py",
"chars": 18743,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/mpu/mappings.py",
"chars": 4432,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/mpu/random.py",
"chars": 12684,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/mpu/utils.py",
"chars": 2649,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/optimizer/__init__.py",
"chars": 4435,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/optimizer/clip_grads.py",
"chars": 5513,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/optimizer/grad_scaler.py",
"chars": 4367,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/optimizer/optimizer.py",
"chars": 19191,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/p2p_communication.py",
"chars": 10304,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/schedules.py",
"chars": 19276,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/tokenizer/__init__.py",
"chars": 666,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/tokenizer/gpt2_tokenization.py",
"chars": 14146,
"preview": "# coding=utf-8\n# Copyright 2018 The Open AI Team Authors and The HuggingFace Inc. team.\n#\n# Licensed under the Apache Li"
},
{
"path": "codegeex/megatron/tokenizer/tokenizer.py",
"chars": 7502,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/tools/collect_env.py",
"chars": 513,
"preview": "import os\n\n\nENV_NAMES = [\"CUDA_HOME\", \"LD_LIBRARY_PATH\", \"PATH\", \"TORCH_EXTENSIONS_DIR\", \"CUDA_LAUNCH_BLOCKING\"]\n\n\ndef m"
},
{
"path": "codegeex/megatron/tools/finetune_codegeex.py",
"chars": 11489,
"preview": "import os\nimport torch\nimport logging\n\nlogging.getLogger(\"torch\").setLevel(logging.WARNING)\n\nimport deepspeed\nfrom deeps"
},
{
"path": "codegeex/megatron/tools/pretrain_codegeex.py",
"chars": 6970,
"preview": "import os\nimport subprocess\nimport torch\nimport logging\n\nlogging.getLogger(\"torch\").setLevel(logging.WARNING)\n\nimport de"
},
{
"path": "codegeex/megatron/training.py",
"chars": 51481,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/megatron/utils.py",
"chars": 8578,
"preview": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Ve"
},
{
"path": "codegeex/mindspore/configs/13B.sh",
"chars": 1006,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/configs/13B_128p_save_1p.sh",
"chars": 1111,
"preview": "#!/bin/bash\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the pa"
},
{
"path": "codegeex/mindspore/configs/13B_128p_save_8p_ckpt.sh",
"chars": 1104,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/configs/13B_1p_to_torch.sh",
"chars": 1072,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/configs/13B_finetune.sh",
"chars": 1013,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/configs/13B_generate.sh",
"chars": 1254,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/configs/13B_generate_1p.sh",
"chars": 1296,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/configs/13B_generate_1p_values.sh",
"chars": 1274,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/configs/13B_generate_finetune.sh",
"chars": 1324,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/configs/13B_generate_humaneval.sh",
"chars": 1317,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/configs/13B_generate_values.sh",
"chars": 1293,
"preview": "#!/bin/bash\n\nscript_path=$(realpath $BASH_SOURCE)\nscript_dir=$(dirname $script_path)\n\nCODE_DATA_DIR=\"\" # TODO: set the p"
},
{
"path": "codegeex/mindspore/convertion_1p.py",
"chars": 12417,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/finetune.py",
"chars": 16464,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/generation.py",
"chars": 13490,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/generation_1p.py",
"chars": 12205,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/generation_batch.py",
"chars": 15630,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/generation_finetune.py",
"chars": 13370,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/generation_humaneval.py",
"chars": 11793,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/generation_values.py",
"chars": 10261,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/generation_values_1p.py",
"chars": 12181,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/save_1p_ckpt_from_8p_ckpt.py",
"chars": 11504,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/save_8p_ckpt.py",
"chars": 8623,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/scripts/custom_tune_bank_new/Ascend910ProA/cube/repository_ascend910ProA_matmul.json",
"chars": 33383,
"preview": "{\n \"shape\": {\n \"op_type\": \"matmul\",\n \"A_shape\": [\n 1,\n 320,\n 1024,\n 16,\n 16\n ],\n \""
},
{
"path": "codegeex/mindspore/scripts/custom_tune_bank_new/Ascend910ProA/vector/Ascend910ProA_AiCore_32_v001_20220509_200939_588817.json",
"chars": 28204,
"preview": "{\n \"[[10, [16384, 6528], [], 11, []], [0, [16384, 6528], [], 11, [0]], [29, [16384, 6528], [], 11, [0]], [0, [1, 1], []"
},
{
"path": "codegeex/mindspore/scripts/layer_norm.py",
"chars": 46190,
"preview": "# Copyright 2019 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/scripts/layer_norm_x_backprop_v2.py",
"chars": 42531,
"preview": "# Copyright 2019 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/scripts/ma-pre-start.sh",
"chars": 5459,
"preview": "#!/bin/bash\n\nsource ~/.bashrc\necho \"Start to intall the run package\"\nWORK_DIR=start_1.7\nRUN_DIR=run\nmindspore_file=minds"
},
{
"path": "codegeex/mindspore/scripts/run_modelarts.py",
"chars": 1426,
"preview": "import argparse\nimport os\nimport time\nfrom pathlib import Path\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\""
},
{
"path": "codegeex/mindspore/scripts/run_modelarts_gen_finetune.py",
"chars": 1599,
"preview": "import argparse\nimport os\nimport time\nfrom pathlib import Path\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\""
},
{
"path": "codegeex/mindspore/scripts/run_modelarts_gen_humaneval_x.py",
"chars": 1577,
"preview": "import argparse\nimport os\nimport time\nfrom pathlib import Path\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\""
},
{
"path": "codegeex/mindspore/src/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "codegeex/mindspore/src/adam.py",
"chars": 11677,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/callbacks.py",
"chars": 7706,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/code_tokenizer.py",
"chars": 5577,
"preview": "from typing import *\n\nimport numpy as np\nfrom transformers import AutoTokenizer\nfrom transformers.models.gpt2 import GPT"
},
{
"path": "codegeex/mindspore/src/dataset.py",
"chars": 10092,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/dataset_finetune.py",
"chars": 8035,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/generate.py",
"chars": 10898,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/generate_finetune.py",
"chars": 8023,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/generate_greedy.py",
"chars": 6201,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/generate_humaneval.py",
"chars": 10346,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/metrics.py",
"chars": 3368,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/pangu_alpha.py",
"chars": 31502,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/pangu_alpha_config.py",
"chars": 7061,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/pangu_alpha_fp16_predict.py",
"chars": 28164,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/pangu_alpha_wrapcell.py",
"chars": 10968,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/pangu_alpha_wrapcell_finetune.py",
"chars": 11030,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/preprocess.py",
"chars": 9221,
"preview": "# Copyright 2020 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/src/sat_dataset.py",
"chars": 9522,
"preview": "import pickle\nfrom abc import ABC, abstractmethod\nfrom bisect import bisect_right\n\nimport numpy as np\n\n\nclass Dataset(AB"
},
{
"path": "codegeex/mindspore/src/tokenization_jieba.py",
"chars": 2431,
"preview": "# coding=utf-8\n# Copyright 2018 The Open AI Team Authors and The HuggingFace Inc. team.\n#\n# Licensed under the Apache Li"
},
{
"path": "codegeex/mindspore/src/utils.py",
"chars": 27356,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/mindspore/train.py",
"chars": 30339,
"preview": "# Copyright 2021 Huawei Technologies Co., Ltd\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you "
},
{
"path": "codegeex/oneflow/__init__.py",
"chars": 41,
"preview": "from .codegeex_model import CodeGeeXModel"
},
{
"path": "codegeex/oneflow/codegeex_model.py",
"chars": 43244,
"preview": "import math\nimport oneflow as torch\nimport oneflow.nn.functional as F\nfrom oneflow.nn.parameter import Parameter\nfrom .."
},
{
"path": "codegeex/oneflow/inference.py",
"chars": 10828,
"preview": "import copy\nimport json\nimport os\nimport time\nfrom typing import *\n\nimport oneflow as torch\nimport oneflow.nn.functional"
},
{
"path": "codegeex/paddle/__init__.py",
"chars": 41,
"preview": "from .codegeex_model import CodeGeeXModel"
},
{
"path": "codegeex/paddle/codegeex_model.py",
"chars": 37450,
"preview": "import math\nimport paddle\nimport paddle.nn.functional as F\n\n\ndef fast_gelu(x):\n \"\"\"Mindspore's fast gelu implementati"
},
{
"path": "codegeex/paddle/inference.py",
"chars": 10915,
"preview": "import copy\nimport json\nimport os\nimport time\nfrom typing import *\n\nimport paddle\nimport paddle.nn.functional as F\nfrom "
},
{
"path": "codegeex/paddle/pt_to_pdparams.py",
"chars": 1225,
"preview": "import argparse\nimport paddle\nimport torch\n\nlinear_layer = [\n \"mlp.dense_h_to_4h\",\n \"mlp.dense_4h_to_h\",\n \"atte"
},
{
"path": "codegeex/quantization/__init__.py",
"chars": 174,
"preview": "from .quantize import quantize\ntry:\n from .quantize_oneflow import quantize_oneflow\n from .quantize_oneflow import"
},
{
"path": "codegeex/quantization/quantize.py",
"chars": 15139,
"preview": "import torch\n\nfrom torch.nn.parameter import Parameter\nfrom codegeex.kernels import extract_weight_to_half\nfrom codegeex"
},
{
"path": "codegeex/quantization/quantize_oneflow.py",
"chars": 7442,
"preview": "import numpy as np\nimport oneflow as torch\nfrom oneflow.nn.parameter import Parameter\n\ndef _pack_int8_to_int4(x):\n n"
},
{
"path": "codegeex/tokenizer/__init__.py",
"chars": 40,
"preview": "from .tokenizer import CodeGeeXTokenizer"
},
{
"path": "codegeex/tokenizer/added_tokens.json",
"chars": 4040,
"preview": "{\"<|extratoken_14|>\": 50270, \"<|extratoken_121|>\": 50377, \"<|extratoken_3|>\": 50259, \"<|extratoken_25|>\": 50281, \"<|extr"
},
{
"path": "codegeex/tokenizer/merges.txt",
"chars": 420572,
"preview": "#version: 0.2\nĠ t\nĠ a\nh e\ni n\nr e\no n\nĠt he\ne r\nĠ s\na t\nĠ w\nĠ o\ne n\nĠ c\ni t\ni s\na n\no r\ne s\nĠ b\ne d\nĠ f\nin g\nĠ p\no u\nĠa "
},
{
"path": "codegeex/tokenizer/special_tokens_map.json",
"chars": 358,
"preview": "{\"bos_token\": {\"content\": \"<|endoftext|>\", \"single_word\": false, \"lstrip\": false, \"rstrip\": false, \"normalized\": true}, "
},
{
"path": "codegeex/tokenizer/tokenizer.py",
"chars": 2268,
"preview": "import numpy as np\nfrom typing import *\nfrom transformers import AutoTokenizer\nfrom transformers.models.gpt2 import GPT2"
},
{
"path": "codegeex/tokenizer/tokenizer_config.json",
"chars": 620,
"preview": "{\"unk_token\": {\"content\": \"<|endoftext|>\", \"single_word\": false, \"lstrip\": false, \"rstrip\": false, \"normalized\": true, \""
},
{
"path": "codegeex/tokenizer/vocab.json",
"chars": 1042301,
"preview": "{\"!\": 0, \"\\\"\": 1, \"#\": 2, \"$\": 3, \"%\": 4, \"&\": 5, \"'\": 6, \"(\": 7, \")\": 8, \"*\": 9, \"+\": 10, \",\": 11, \"-\": 12, \".\": 13, \"/"
},
{
"path": "codegeex/torch/__init__.py",
"chars": 41,
"preview": "from .codegeex_model import CodeGeeXModel"
},
{
"path": "codegeex/torch/codegeex_model.py",
"chars": 36972,
"preview": "import math\nimport torch\nimport torch.nn.functional as F\nfrom torch.nn.parameter import Parameter\n\n\ndef fast_gelu(x):\n "
},
{
"path": "codegeex/torch/get_ckpt_qkv.py",
"chars": 2782,
"preview": "import os\nimport sys\nimport torch\nimport random\nimport argparse\nimport numpy as np\n\n\ndef main():\n parser = argparse.A"
},
{
"path": "codegeex/torch/inference.py",
"chars": 10815,
"preview": "import copy\nimport json\nimport os\nimport time\nfrom typing import *\n\nimport torch\nimport torch.nn.functional as F\nfrom da"
},
{
"path": "configs/codegeex_13b.sh",
"chars": 561,
"preview": "# CodeGeeX-13B configuration\n\nCHECKPOINT_PATH=\"<path where you put the checkpoint (e.g., XXX/codegeex_13b.pt)>\"\n\nMODEL_A"
},
{
"path": "configs/codegeex_13b_paddle.sh",
"chars": 574,
"preview": "# CodeGeeX-13B paddle configuration\n\nCHECKPOINT_PATH=\"<path where you put the checkpoint (e.g., XXX/codegeex_13b.pdparam"
},
{
"path": "configs/codegeex_13b_parallel.sh",
"chars": 686,
"preview": "# CodeGeeX-13B parallel configuration\n# Parallel checkpoints are named under the format \"mp_rank_0{i}_model_states.pt\", "
},
{
"path": "deployment/example_inputs.jsonl",
"chars": 3170,
"preview": "{\"code\": \"# Write a function that returns the sum of the numbers from 1 to n.\\n# For example, if n is 5, then the functi"
},
{
"path": "deployment/server_gradio.py",
"chars": 6970,
"preview": "import json\nimport torch\nimport argparse\nimport gradio as gr\n\nimport codegeex\nfrom codegeex.torch import CodeGeeXModel\nf"
},
{
"path": "requirements.txt",
"chars": 180,
"preview": "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"
},
{
"path": "resources/zh/wechat.md",
"chars": 162,
"preview": "<div align=\"center\">\n<img src=join_wechat.png width=\"60%\"/>\n\n<p> 扫码关注公众号加入「CodeGeeX交流群」 </p>\n<p> Scan the QR code to joi"
},
{
"path": "scripts/convert_ckpt_parallel.sh",
"chars": 1046,
"preview": "# This script is used to convert checkpoint model parallel partitions.\n\nLOAD_CKPT_PATH=$1 # Path to weights in .pt form"
},
{
"path": "scripts/convert_mindspore_to_megatron.sh",
"chars": 730,
"preview": "# This script is used to convert mindspore checkpoint to the megatron format.\n\nNPY_CKPT_PATH=$1 # Path to Mindspore exp"
},
{
"path": "scripts/evaluate_humaneval_x.py",
"chars": 2095,
"preview": "import argparse\nimport os\nfrom pathlib import Path\nfrom codegeex.benchmark.evaluate_humaneval_x import evaluate_function"
},
{
"path": "scripts/evaluate_humaneval_x.sh",
"chars": 1139,
"preview": "# This script is for evaluating the functional correctness of the generated codes of HumanEval-X.\n\nINPUT_FILE=$1 # Path"
},
{
"path": "scripts/finetune_codegeex.sh",
"chars": 3218,
"preview": "SCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\n\n# ====== Environmen"
},
{
"path": "scripts/gather_output.sh",
"chars": 561,
"preview": "# This script is used to gather the distributed outputs of different ranks.\n\nOUTPUT_DIR=$1\nOUTPUT_PREFIX=$2\nIF_REMOVE_RA"
},
{
"path": "scripts/generate_humaneval_x.sh",
"chars": 2547,
"preview": "# This script is used to generate solutions of HumanEval-X.\n\nLANGUAGE=$1 # Target programming language, currently suppo"
},
{
"path": "scripts/pretrain_codegeex.sh",
"chars": 3220,
"preview": "SCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR=$(dirname \"$SCRIPT_PATH\")\nMAIN_DIR=$(dirname \"$SCRIPT_DIR\")\n\n# ====== Environmen"
},
{
"path": "scripts/process_pretrain_dataset.sh",
"chars": 578,
"preview": "# Process dataset for CodeGeeX pretraining\n\nDATASET_PATH=$1\nOUTPUT_PATH=$2\nLANGUAGE=$3\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCR"
},
{
"path": "scripts/test_inference.sh",
"chars": 877,
"preview": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR"
},
{
"path": "scripts/test_inference_oneflow.sh",
"chars": 885,
"preview": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR"
},
{
"path": "scripts/test_inference_oneflow_quantized.sh",
"chars": 887,
"preview": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR"
},
{
"path": "scripts/test_inference_paddle.sh",
"chars": 891,
"preview": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR"
},
{
"path": "scripts/test_inference_parallel.sh",
"chars": 1224,
"preview": "# This script is used to test the inference of CodeGeeX.\n\nMP_SIZE=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT"
},
{
"path": "scripts/test_inference_quantized.sh",
"chars": 879,
"preview": "# This script is used to test the inference of CodeGeeX.\n\nGPU=$1\nPROMPT_FILE=$2\n\nSCRIPT_PATH=$(realpath \"$0\")\nSCRIPT_DIR"
},
{
"path": "scripts/translate_humaneval_x.sh",
"chars": 2991,
"preview": "# This script is used to translate solutions of HumanEval-X.\n\nLANG_SRC_TYPE=$1 # Source programming language, currently"
},
{
"path": "setup.py",
"chars": 645,
"preview": "from setuptools import setup, find_packages\n\nsetup(\n name=\"codegeex\",\n py_modules=[\"codegeex\"],\n version=\"1.0\","
},
{
"path": "tests/test_inference.py",
"chars": 4799,
"preview": "import time\nimport torch\nimport argparse\nimport numpy as np\n\nimport codegeex\nfrom codegeex.torch import CodeGeeXModel\nfr"
},
{
"path": "tests/test_inference_megatron.py",
"chars": 6561,
"preview": "import copy\nimport time\nimport torch\nimport numpy as np\n\nfrom codegeex.megatron import get_tokenizer, get_args, print_ra"
},
{
"path": "tests/test_inference_oneflow.py",
"chars": 5951,
"preview": "\nimport os\nimport copy\nimport time\nimport oneflow as torch\nimport random\nimport argparse\nimport numpy as np\n\nfrom codege"
},
{
"path": "tests/test_inference_paddle.py",
"chars": 6433,
"preview": "\nimport os\nimport copy\nimport time\nimport paddle\nimport random\nimport argparse\nimport numpy as np\n\nfrom codegeex.paddle."
},
{
"path": "tests/test_prompt.txt",
"chars": 403,
"preview": "code translation\nJava:\npublic class Solution {\n public static boolean hasCloseElements(int[] nums, int threshold) {\n "
},
{
"path": "vscode-extension/README.md",
"chars": 8389,
"preview": "\n\n🌐 <a href=\"https://github.com/THUDM/CodeGeeX/blob/main/vscode-ext"
},
{
"path": "vscode-extension/README_zh.md",
"chars": 4902,
"preview": "\n\n🌐 <a href=\"https://github.com/THUDM/CodeGeeX/blob/main/vscode-ext"
}
]
// ... and 1 more files (download for full content)
About this extraction
This page contains the full source code of the zai-org/CodeGeeX GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 189 files (2.9 MB), approximately 773.5k tokens, and a symbol index with 1224 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.