CodeGeeX2 是多语言代码生成模型 CodeGeeX (KDD’23) 的第二代模型。不同于一代 CodeGeeX（完全在国产华为昇腾芯片平台训练） ，CodeGeeX2 是基于 ChatGLM2 架构加入代码预训练实现，得益于 ChatGLM2 的更优性能，CodeGeeX2 在多项指标上取得性能提升（+107% > CodeGeeX；仅60亿参数即超过150亿参数的 StarCoder-15B 近10%），更多特性包括：

• 更强大的代码能力：基于 ChatGLM2-6B 基座语言模型，CodeGeeX2-6B 进一步经过了 600B 代码数据预训练，相比一代模型，在代码能力上全面提升，HumanEval-X 评测集的六种编程语言均大幅提升 (Python +57%, C++ +71%, Java +54%, JavaScript +83%, Go +56%, Rust +321%)，在Python上达到 35.9% 的 Pass@1 一次通过率，超越规模更大的 StarCoder-15B。
• 更优秀的模型特性：继承 ChatGLM2-6B 模型特性，CodeGeeX2-6B 更好支持中英文输入，支持最大 8192 序列长度，推理速度较一代 CodeGeeX-13B 大幅提升，量化后仅需6GB显存即可运行，支持轻量级本地化部署。
• 更全面的AI编程助手：CodeGeeX插件（VS Code, Jetbrains）后端升级，支持超过100种编程语言，新增上下文补全、跨文件补全等实用功能。结合 Ask CodeGeeX 交互式AI编程助手，支持中英文对话解决各种编程问题，包括且不限于代码解释、代码翻译、代码纠错、文档生成等，帮助程序员更高效开发。

快速开始

使用transformers快速调用CodeGeeX2-6B：

from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True) model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True, device='cuda') model = model.eval（) # remember adding a language tag for better performance prompt = "# language: Python\n# write a bubble sort function\n" inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device) outputs = model.generate(inputs, max_length=256, top_k=1) response = tokenizer.decode(outputs[0]) >>> print(response) # language: Python # write a bubble sort function def bubble_sort(list): for i in range(len(list) - 1): for j in range(len(list) - 1): if list[j] > list[j + 1]: list[j], list[j + 1] = list[j + 1], list[j] return list print(bubble_sort([5, 2, 1, 8, 4]))

启动 Gradio DEMO：

python ./demo/run_demo.py

usage: run_demo.py [-h] [--model-path MODEL_PATH] [--example-path EXAMPLE_PATH] [--quantize QUANTIZE]
                   [--chatglm-cpp] [--fastllm] [--n-gpus N_GPUS] [--gpu GPU] [--cpu] [--auth] [--username yourname]
                   [--password yourpassword]
                   [--port PORT] [--listen ADDRESS]

# 若要启用身份验证，请先启用--auth，然后定义--username与--password，如：
python run_demo.py --auth --username user --password password  # 若要监听所有地址请指定 --listen 0.0.0.0

支持使用 ChatGLM.cpp 量化推理加速：

python ./demo/run_demo.py --quantize 4 --chatglm-cpp

启动FAST API:

python ./demo/fastapicpu.py
usage: fastapicpu.py [-h] [--model-path MODEL_PATH] [--listen ADDRESS] [--port PORT] [--workders NUM] [--cpu] [--half] [--quantize QUANTIZE] [--chatglm-cpp]
# --cpu启用cpu --half启用.half()

支持使用 ChatGLM.cpp 量化推理加速，同样添加 --quantize 4 --chatglm-cpp 参数即可。

API使用示例

curl -X POST "http://127.0.0.1:7860" \
    -H 'Content-Type: application/json' \
    -d '{"lang": "Python", "prompt": "# Write a quick sort function"}'

❗️请注意：

• CodeGeeX2-6B 是一个基座代码生成模型，不具备聊天能力。请前往插件中体验更全面的 Ask CodeGeeX 聊天功能。
• 在使用 CodeGeeX2-6B 的补全功能时，输入prompt需要遵循特定的格式以获得最好的效果。比如需要在开头加入编程语言标签（# language: Python，请查看完整语言列表），以注释的形式写prompt等。参考run_demo.py中的处理。
• 如果显卡不支持bfloat16格式，将会输出错误的内容，需要将模型转换成float16格式： model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True).half().cuda()
• 如果需要使用多显卡加载模型,可以将以下代码： tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True) model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True, device='cuda') model = model.eval（) 替换为 def get_model(): tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True) from gpus import load_model_on_gpus # gpus文件在demo文件夹中 model = load_model_on_gpus("THUDM/codegeex2-6b", num_gpus=2) model = model.eval（) return tokenizer, model tokenizer, model = get_model()

代码能力评测

CodeGeeX2 作为一个多语言代码生成基座模型，代码能力较上一代大幅提升，以下是在 HumanEval，HumanEval-X, DS1000 基准上的评测结果（评价指标 Pass@k 定义与论文中一致）：

HumanEval (Pass@1,10,100)

Pass@1 使用 n=20, t=0.2, top_p=0.95；Pass@10,Pass@100 使用 n=200, t=0.8, top_p=0.95。

HumanEval-X (Pass@1)

Pass@1 使用 n=20, t=0.2, top_p=0.95。

以上结果可使用脚本scripts/run_humanevalx.sh复现。环境配置和说明参见评测环境。

DS1000 (Pass@1)

Pass@1 使用 n=40, t=0.2, top_p=0.5。

以上结果可使用DS1000评测代码复现。

量化推理性能

CodeGeeX2 与上一代相比，对部署更加友好。得益于使用 Multi-Query Attention 和 Flash Attention，推理速度更快，且量化后仅需6GB显存即可运行：

量化

基于 PyTorch 2.0 测试，利用torch.nn.functional.scaled_dot_product_attention实现高效的 Attention 计算。

推理

batch_size=1, max_length=2048，均使用加速框架，测试硬件为GeForce RTX-3090。