🏠 主页｜🛠 插件 VS Code, Jetbrains｜🤗 模型下载｜📄 论文｜👋 加入微信开发者交流群

在 English 中阅读
用 日本語 阅读
以 Français 阅读

⭐️ 最新一代 CodeGeeX4 模型已经正式开源。 The newest CodeGeeX4 has been released.

CodeGeeX2: 更强大的多语言代码生成模型

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编程助手，支持中英文对话解决各种编程问题，包括且不限于代码解释、代码翻译、代码纠错、文档生成等，帮助程序员更高效开发。
• 更开放的协议：CodeGeeX2-6B 权重对学术研究完全开放，填写登记表申请商业使用。

使用教程

• 快速开始
• 推理教程（多卡推理，加速推理，多平台推理等）

AI编程助手

我们开发了支持 VS Code、 IntelliJ IDEA、PyCharm、GoLand、WebStorm、Android Studio 等IDE的 CodeGeeX 插件。在插件中，可以更直接地体验到 CodeGeeX2 模型在代码生成与补全、添加注释、代码翻译及技术问答方面的能力为开发效率带来的提升。欢迎在IDE中下载 CodeGeeX 插件获得更加全面的AI编程体验，详情见CodeGeeX主页。

快速开始

使用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)

Model	Pass@1	Pass@10	Pass@100
CodeGen-16B-multi	19.2	34.6	55.2
CodeGeeX-13B	22.9	39.6	60.9
Codex-12B	28.8	46.8	72.3
CodeT5Plus-16B-mono	30.9	51.6	76.7
Code-Cushman-001	33.5	54.3	77.4
LLaMA-65B	23.7	-	79.3
LLaMA2-70B	29.9	-	-
CodeGen2.5-7B-mono	33.4	58.4	82.7
StarCoder-15B	33.2	61.0	84.7
CodeGeeX2-6B	35.9	62.6	88.3

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)

模型	Python	C++	Java	JavaScript	Go	Rust	Overall
CodeGen-16B-multi	19.2	18.1	15.0	18.4	13.0	1.8	14.2
CodeGeeX-13B	22.9	17.1	20.0	17.6	14.4	4.3	16.0
Replit-code-v1-3B	22.0	20.1	20.1	20.1	12.2	8.6	17.2
CodeGen2.5-7B-multi	30.6	24.3	29.0	27.5	18.9	20.1	25.1
StarCoder-15B	35.5	28.2	31.5	33.2	21.3	17.8	27.9
CodeGeeX2-6B	35.9	29.3	30.8	32.2	22.5	18.1	28.1

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

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

DS1000 (Pass@1)

Model	Matplotlib	Numpy	Pandas	Pytorch	SciPy	Scikit-learn	TensorFlow	Overall
# Samples	155	220	291	68	106	115	45	1000
CodeGen-16B-Mono	31.7	10.9	3.4	7.0	9.0	10.8	15.2	11.7
code-cushman-001	40.7	21.8	7.9	12.4	11.3	18.0	12.2	18.1
Codex-001	41.8	26.6	9.4	9.7	15.0	18.5	17.2	20.2
CodeGeeX2-6B	40.5	25.5	14.5	17.3	19.3	24.0	23.0	23.1
StarCoder-15B	51.7	29.7	11.4	21.4	20.2	29.5	24.5	26.0
Codex-002	57.0	43.1	26.5	41.8	31.8	44.8	39.3	39.2

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

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

量化推理性能

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

量化

Model	FP16/BF16	INT8	INT4
CodeGeeX-13B	26.9 GB	14.7 GB	-
CodeGeeX2-6B	13.1 GB	8.2 GB	5.5 GB

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

推理

Model	推理速度 (字符/秒)
CodeGeeX-13B	32
CodeGeeX2-6B	94

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

协议

本仓库的代码依照 Apache-2.0 协议开源，模型的权重的使用则需要遵循 Model License。CodeGeeX2-6B 权重对学术研究完全开放，填写登记表申请商业使用。

引用

如果觉得我们的工作有帮助，欢迎引用以下论文：

@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}
}

CodeGeeX2 快速上手指南

CodeGeeX2 是一款强大的多语言代码生成模型，基于 ChatGLM2 架构，支持超过 100 种编程语言。本指南将帮助您快速在本地部署并运行 CodeGeeX2-6B 模型。

环境准备

系统要求

• 操作系统: Linux (推荐), Windows, macOS
• GPU: 支持 CUDA 的 NVIDIA 显卡
  • FP16/BF16 推理: 建议显存 ≥ 14GB (如 RTX 3090/4090)
  • INT4 量化推理: 仅需显存 ≥ 6GB (如 RTX 2060/3060 等)
• Python: 3.8 及以上版本

前置依赖

请确保已安装以下基础库：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install transformers>=4.30.0
pip install accelerate
pip install sentencepiece
pip install protobuf

提示：国内用户可使用清华或阿里镜像源加速安装，例如：pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <package_name>

安装步骤

CodeGeeX2 可直接通过 Hugging Face transformers 库调用，无需额外克隆整个仓库即可进行基础推理。若需运行官方 Demo 或进行多卡部署，可克隆仓库：

git clone https://github.com/THUDM/CodeGeeX2.git
cd CodeGeeX2
pip install -r requirements.txt

模型权重将在首次运行时自动从 Hugging Face 下载。

国内加速：如果访问 Hugging Face 较慢，建议配置镜像或使用国内镜像站下载模型后本地加载。 环境变量设置示例：

export HF_ENDPOINT=https://hf-mirror.com

基本使用

以下是最简单的 Python 代码生成示例。我们将加载模型并生成一个冒泡排序函数。

1. 基础推理脚本

创建文件 run_codegeex2.py 并填入以下代码：

from transformers import AutoTokenizer, AutoModel

# 加载分词器和模型
# trust_remote_code=True 是必须的，因为模型使用了自定义代码
tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True)

# 加载模型到 GPU
# 如果显存不足或不支持 bfloat16，请使用 .half().cuda()
model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True, device='cuda')
# model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True).half().cuda() 
model = model.eval()

# 构造 Prompt
# 注意：为了获得最佳效果，请在开头添加语言标签 "# language: Python"
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)

2. 运行脚本

在终端执行：

python run_codegeex2.py

3. 预期输出

模型将输出类似以下的 Python 代码：

# 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]))

💡 关键提示

• Prompt 格式：务必在注释中包含 # language: <语言名> 标签（如 # language: C++），这能显著提升生成准确率。
• 显存优化：如果您的显卡显存小于 13GB，建议使用 INT4 量化版本或参考官方文档使用 chatglm.cpp 进行加速推理。
• 多卡部署：如需在多张 GPU 上运行，请参考仓库中 demo/gpus.py 的实现逻辑加载模型。