Skip to content

[Machine Learning] CodeBERT 使用方法與介紹(附範例程式碼)

Last Updated on 2023-12-30 by Clay

簡介

CodeBERT 這一預訓練模型(pre-trained model)提出自《CodeBERT: A Pre-Trained Model for Programming and Natural Languages》。

與大部分 BERT 預訓練模型相同,都可以將『輸入的資訊』來來產生一組『特徵』,以此來進行下游任務的應用。

而在 CodeBERT 所提出適用的下游任務中,就包含了程式碼檢索程式碼生成...... 等等與程式語言處理相關的任務。我想這就是把此預訓練模型取名叫 CodeBERT 所要強調的。

CodeBERT 論文出處:https://arxiv.org/abs/2002.08155


實驗設計

論文中提到的設定了四種不同的實驗來測試模型的有效性:

  • 自然語言程式碼檢索
  • NL-PL 探測
  • 程式碼文件生成
  • 在未參與預訓練的程式碼語言上的效果


自然語言程式碼檢索

給定一段自然語言文字敘述,從眾多程式碼中找到語義上最相關的程式碼。

為了能夠與不同的檢索方法進行比較,選擇了 CodeSearchNet (https://github.com/github/CodeSearchNet) 這個資料集進行訓練,並使用 MRR 作為評價指標。

平均倒數排名Mean Reciprocal Rank, MRR)是一種評估資料檢索回傳結果品質的評估指標。假設查詢回傳了 數個結果,標準答案出現得越早分數越高,都沒有出現正確答案則分數為 0,最後將實驗次數平均,得到的即為 MRR 指標的評估分數。


NL-PL Probing

固定模型參數,研究 CodeBERT 究竟學習了什麼類型的知識。(此為本人比較沒看懂的實驗


程式碼文件生成(有參與、未參與預訓練的兩組)

這是近年來很倡導的良好習慣:開發與文件介紹同步撰寫

試想如果你製作的系統、API 有個良好的文件說明,這樣就可以免去了後人每次遇到不清楚的部分就得反覆去追蹤原始碼。也就是說,這份文件清楚到可以不用閱讀原始碼就能理解該如何操作系統。

而 CodeBERT 這個任務就是打算透過程式碼來自動產生該程式的說明文件

預訓練模型總共使用了六種程式語言:

  • Go
  • Java
  • JavaScript
  • PHP
  • Python
  • Ruby

而未參與預訓練的程式語言為 C#

選擇了 Codenn 資料集(https://github.com/sriniiyer/codenn)。


論文總結

CodeBERT 是第一個同時面向程式語言自然語言的預訓練模型,並在實驗中的下游任務中取得了 SOTA (state-of-the-art)效果。

論文中提出說加入程式語言的 AST 結構資訊是一個很有潛力的研究方向,據我所知他們後來也在 GraphCodeBERT 的論文中實現了這個研究方向的承諾。


CodeBERT 使用方法(Code Documentation Generation)

詳細的使用方法請參閱 CodeBERT 的論文以及 GitHub 專案(https://github.com/guoday/CodeBERT),在這裡我簡單介紹如何使用 CodeBERT、並以程式碼文件生成Code Documentation Generation)這個任務作為範例。


安裝相關套件

pip3 install torch==1.4.0
pip3 install transformers==2.5.0
pip3 install filelock



資料前處理

專案中對於資料集的處理如下:

  • 刪除了程式碼中的註釋
  • 刪除了程式碼無法解析(parse)成抽象語法樹Abstract Syntax Tree, ATS)的資料
  • 刪除了文件中的 # 令牌 < 3 或 > 256 的資料
  • 刪除了非英文文件的資料

你可以從這個網站中下載資料,或是使用以下指令:

pip3 install gdown
mkdir -p data/code2nl
cd data/code2nl
gdown https://drive.google.com/uc?id=1rd2Tc6oUWBo7JouwexW3ksQ0PaOhUr6h
unzip Cleaned_CodeSearchNet.zip
rm Cleaned_CodeSearchNet.zip
cd ../..


如果使用 tree 指令查看資料架構的話,應如下所示:

tree data/code2nl/CodeSearchNet/


Output:

data/code2nl/CodeSearchNet/
├── go
│   ├── test.jsonl
│   ├── train.jsonl
│   └── valid.jsonl
├── java
│   ├── test.jsonl
│   ├── train.jsonl
│   └── valid.jsonl
├── javascript
│   ├── test.jsonl
│   ├── train.jsonl
│   └── valid.jsonl
├── php
│   ├── test.jsonl
│   ├── train.jsonl
│   └── valid.jsonl
├── python
│   ├── test.jsonl
│   ├── train.jsonl
│   └── valid.jsonl
└── ruby
    ├── test.jsonl
    ├── train.jsonl
    └── valid.jsonl

6 directories, 18 files


可以發現,每個語言都分別有著訓練、驗證、測試等三種不同的資料,一共十八個檔案。


程式執行

前去:https://github.com/microsoft/CodeBERT/tree/master/CodeBERT/code2nlrun.pybleu.pymodel.py 三份程式放在 data/code2nl 資料夾中。

接著執行以下程式碼,原本官方說明中的 batch_size=128,但是對大部分單人使用的 GPU 來說記憶體負荷不了,所以擅自改成了 batch_size=4

lang=php #programming language
beam_size=10
batch_size=4
source_length=256
target_length=128
output_dir=model/$lang
data_dir=../data/code2nl/CodeSearchNet
dev_file=$data_dir/$lang/valid.jsonl
test_file=$data_dir/$lang/test.jsonl
test_model=$output_dir/checkpoint-best-bleu/pytorch_model.bin #checkpoint for test

python run.py --do_test --model_type roberta --model_name_or_path microsoft/codebert-base --load_model_path $test_model --dev_filename $dev_file --test_filename $test_file --output_dir $output_dir --max_source_length $source_length --max_target_length $target_length --beam_size $beam_size --eval_batch_size $batch_size



如何使用 CodeBERT

如果想要看如何使用 CodeBERT 轉換出程式碼特徵向量,可以參考下方範例程式碼:

from transformers import AutoTokenizer, AutoModel
import torch


# Init
tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
model = AutoModel.from_pretrained("microsoft/codebert-base")


# Tokenization 
nl_tokens = tokenizer.tokenize("return maximum value")
code_tokens = tokenizer.tokenize("def max(a,b): if a>b: return a else return b")
tokens = [tokenizer.cls_token]+nl_tokens+[tokenizer.sep_token]+code_tokens+[tokenizer.sep_token]


# Convert tokens to ids
tokens_ids = tokenizer.convert_tokens_to_ids(tokens)
context_embeddings = model(torch.tensor(tokens_ids)[None,:])[0]


# Print
print(context_embeddings)


Output:

tensor([[-0.1423,  0.3766,  0.0443,  ..., -0.2513, -0.3099,  0.3183],
        [-0.5739,  0.1333,  0.2314,  ..., -0.1240, -0.1219,  0.2033],
        [-0.1579,  0.1335,  0.0291,  ...,  0.2340, -0.8801,  0.6216],
        ...,
        [-0.4042,  0.2284,  0.5241,  ..., -0.2046, -0.2419,  0.7031],
        [-0.3894,  0.4603,  0.4797,  ..., -0.3335, -0.6049,  0.4730],
        [-0.1433,  0.3785,  0.0450,  ..., -0.2527, -0.3121,  0.3207]],
       grad_fn=<SelectBackward>)


如果想要參考更多 HuggingFace 團隊的 transformers 套件使用方法,可以參考我之前寫過的:[PyTorch] 如何使用 Hugging Face 所提供的 Transformers —— 以 BERT 為例


References


Read More

Leave a Reply