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 模型在代碼生成與補全、添加註釋、代碼翻譯及技術問答方面的能力為開發效率帶來的提升。詳情見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, 4, 6, 1, 3]))

代碼能力評測

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。