是不是有一天想像着讓代碼自動補全,今天他來了!!!
作者:熊唯,黃飛 ,騰訊 PCG/QQ研發中心/CV應用研究組
AI 如果真的可以寫代碼了,程序員將何去何從?近幾年,NLP 領域的生成式任務有明顯的提升,那通過 AI
我們可以讓代碼自動完成後續補全嗎?本文主要介紹了如何使用 GPT2 框架實現代碼自動補全的功能。
如果 AI 真的可以自己寫代碼了,程序員將何去何從?
我去年做過一個代碼補全的小功能,打包為 androidStudio 插件,使用效果如下:
代碼補全模型預測出的結果有時的確會驚嚇到我,這也能學到~? 那如果給它見識了全世界的優秀代碼,再給足夠量級參數和優秀的模型框架,真的可以實現需求作為輸入,直接輸出代碼嗎?
“我的需求講完了,你的代碼呢?” 希望可以看到這一天。
代碼補齊功能有其他優秀插件也已實現,比如 tabnine,Kite 和國產的 aixcoder。本文主要介紹下代碼補全功能需要實現的整套流程。主要包括數據,算法和工程。
數據
眾所周知,算法工程師大部分時間都在處理數據。
深度學習是使用大數據訓練模型的一個過程,數據是很重要的一個模塊。人是會累的,休息不好還導致記憶不好。AI 是你給多少數據它就能存儲接收多少數據,學不到信息那是人的錯,給的數據不好或者算法設計不好。所以我們先儘可能多的準備好訓練數據。
1、數據採集
本文的目的是代碼補全,訓練數據就是代碼段。考慮到每種語言風格和語法都不一致,所以單個模型只針對一種代碼語言。
我使用的訓練數據主要來源於 GitHub,編寫了一個簡單的爬蟲代碼,指定語言後根據 stars 的排序下載工程。
Github 的 search API 官方地址:
//developer.github.com/v3/search/
2、數據清理
直接下載的數據肯定是不能直接用的,我們還需要對數據進行清理。
首先,我們的訓練數據只需要工程中的代碼文件,以 java 工程為例,我們只保留.java 結尾的文件,其他文件可剔除。
其次,我的代碼補全目標是代碼段,不針對注釋功能。而且對於代碼補全訓練時,我們是會給定一定範圍的上文,如果存在注釋段會佔用有效代碼信息。另外注釋除英文外其他字符不在我的訓練 vocab 範圍內,所以需要對代碼中注釋和日誌進行清理。
1.刪除代碼行中存在除符號和英文外的字符
2.刪除日誌行
3.刪除注釋行,主要針對以下格式
/* 注釋文本*/
/**
注釋段落
*/
// 注釋文本
code //注釋
經過以上數據清理後,得到純代碼數據。
3、數據編碼
得到了訓練數據後還需要把代碼文本進行編碼。本文使用的是 bpe(byte pair encoder)位元組對編碼,主要為了數據壓縮。bpe 簡單理解為將一個單詞再拆分為多個字母組合,比如 tencent 拆分為 ten-cent,這些組合方式則是根據大量數據,統計頻率得到。由於我們期待的代碼補全功能是在行首輸入幾個字母,根據上文預期出本行內容。
假設 tensorflow 這個 token 被編碼對應到一個 id,那我希望輸入 ten 就輸出 tensorflow 是無法實現的。所以在訓練過程中,我會隨機把 token 打斷,比如將 tensorflow 打斷為 t-en-sor-flow 進行編碼,打斷原則是被切分的部分一定要在詞彙表中。數據編碼後,代碼的每個 token 被編碼為 1~N 個 id。模型預測到的 id 反編碼為 token 即可。回車符認為是預測的終止符。經過以上處理,我們就準備好了訓練數據,下面就可以進行算法部分了。
模型算法
眾所周知,算法工程師大部分時間都在研究算法。
在騰訊文檔的錯別字糾錯需求中,我們採用了基於 LSTM 的 seq2seq 以及 facebook 提出的基於 CNN 的 seq2seq,可以得到不錯的糾錯效果。直到 NLP 出現了一個”網紅”–BERT,採用後精度直接提升 8 個點左右,不虧是 google。下面先簡單介紹下 bert 和 gpt2。
BERT 和 GPT2
2017 年中 google 提出了 Transformer 結構。不用 rnn,不用 cnn,提出 attention is all you need。2018 年 openAI 採用了 transformers 結構在 18 年發佈了 GPT。同年 google AI Language 發佈了 bert 論文,提出的 BERT 模型在 11 個 NLP 任務上刷新了記錄。2019 年 openAI 又推出了 GPT-2 模型。。
BERT(Bidirectional Encoder Representation from Transformers)是基於 transformers 框架的 encoder 部分,自編碼語言模型,適合 N-1(比如句子分類),N-N(比如詞性標註)的任務,但是它並不適合做生成任務。
GPT(Generative Pre-Training)基於 transformers 的 decoder 部分,自回歸語言模型,適合生成式任務。
代碼補全功能就是基於 GPT2 框架,OPenAI 官方提供了多套 GPT2 預訓練模型:
作為一個經常要把模型部署到移動端的 CVer,看到這個參數級別,我選擇最小的模型進行 finetune。
對於 GPT 算法,下面這篇文章講的很好,感興趣同學可以看看:
//zhuanlan.zhihu.com/p/137350403
本文在訓練中使用 512 個上文,預測到回車符為終止。模型網絡使用超參:12 個層,768 個隱藏節點,12 個 heads,採用了 uber 的 Horovod 分佈式框架進行訓練。
infer 階段採用 beam-search 會導致整個預測過程特別耗時,所以參考了//arxiv.org/abs/1904.09751論文,採用 top-k sampling,每次預測 top3 的結果再通過概率閾值過濾後作為最終候選輸出。
最終 infer 效果:
輸入一段代碼,預測出後續代碼,以回車符截止。
工程
眾所周知,算法工程師大部分時間都在做工程。
訓練出模型後,還要把模型應用起來,所以還需要一些工程工作需要實現。代碼補全功能,最合適的應用場景就是上 IDE。nlp 模型不太適合在本機部署,最終選擇了在 GPU 機器上部署模型,然後終端通過 http 請求獲取預測文本顯示的方案。
後台部署
Flask 是一個 Web 應用程序框架,靈活,輕便,容易上手。本文簡單介紹如何利用 flask 啟動一個 web 服務,以及如何訪問和調用我們的功能接口。首先我們創建一個 conda 環境:
conda create -n flask python=3.6
source activate flask
pip install flask
代碼中增加一個接口函數:
from flask import Flask
from flask import request
app = Flask()
# route把一個函數綁定到對應的 url 上
@app.route("/plugin",methods=['GET',])
def send():
data = request.args.get('data')
# 模型預測邏輯
out = model_infer(data)
return out
if __name__ == '__main__':
app.run(host='0.0.0.0',port=8080, debug=False)
執行 run.py 代碼,後台服務開啟運行:
客戶端請求:
url = //ip:8080/plugin?data="輸入"
其中 model_infer 函數需要實現模型的 infer 前向計算邏輯,從請求中獲取 data 字段作為輸入,infer 預測的結果列表作為輸出返回給調用方。
經過上面的工作,我們已經提供了一個服務接口,返回我們代碼補全的預測結果。
插件編寫
最後一步就是如何在 IDE 上使用功能了。我們要開發 AS 的插件,需要使用 IntelliJ,首先需要在本機安裝配置 IntelliJ IDEA
下載地址:
//www.jetbrains.com/idea/download/
社區版源碼:
//github.com/JetBrains/intellij-community
好用的插件可以節省程序員很多時間,在插件實現時,我還添加了一個小的 git-blame 功能,實時查看指定行的 git 提交人,對於手 Q 這種多人合作的工作,比較實用。大家也可以通過 IntelliJ 自己開發一些常用功能。
gitBlame 的主要代碼:
public class GitBlame extends AnAction {
private void showPopupBalloon(final Editor editor, final String result) {
ApplicationManager.getApplication().invokeLater(new Runnable() {
public void run() {
JBPopupFactory factory = JBPopupFactory.getInstance();
factory.createHtmlTextBalloonBuilder(result, null, new JBColor(new Color(186, 238, 186), new Color(73, 117, 73)), null)
.setFadeoutTime(5000)
.createBalloon()
.show(factory.guessBestPopupLocation(editor), Balloon.Position.below);
}
});
}
@Override
public void actionPerformed(AnActionEvent e) {
// TODO: insert action logic here
//獲得當前本地代碼根目錄
String base_path = e.getProject().getBasePath();
String file_path = e.getProject().getProjectFilePath();
//獲取編輯mEditor
final Editor mEditor = e.getData(PlatformDataKeys.EDITOR);
if (null == mEditor) {
return;
}
SelectionModel model = mEditor.getSelectionModel();
final String selectedText = model.getSelectedText();
if (TextUtils.isEmpty(selectedText)) {
return;
}
//獲取當前編輯文檔的目錄
PsiFile mPsifile = e.getData(PlatformDataKeys.PSI_FILE);
VirtualFile file = mPsifile.getContainingFile().getOriginalFile().getVirtualFile();
if (file != null && file.isInLocalFileSystem()) {
file_path = file.getCanonicalPath();
}
//gitkit工具
JGitUtil gitKit = new JGitUtil();
String filename = file_path.replace(base_path+"/","");
//得到blame信息
int line_index = mEditor.getSelectionModel().getSelectionStartPosition().getLine();
String blame_log = gitKit.git_blame(base_path,filename,line_index);
//展示
if (!blame_log.isEmpty()){
showPopupBalloon(mEditor, blame_log);
}
}
}
本文的代碼補全插件主要代碼邏輯為調用上一步後台部署的請求。
// 請求url格式(和flask接口一致)
String baseUrl = "//ip:8080/plugin?data=";
// 獲取當前編輯位置文本
PsiFile str = position.getContainingFile();
// 根據模型上文限制獲取代碼端
String data = getContentCode();
String url = baseUrl+data;
// 發送請求
String result = HttpUtils.doGet(url);
// 後處理邏輯,在提示框顯示預測結果
show()
最終呈現形式:
可以看出,模型的預計結果還是不錯的~
以上為代碼補全功能的實現和應用,算是 AI 自動寫代碼的一小步。
AI 能否自己寫代碼,達到疑犯追蹤里 TM 那種水平,我不敢說一定不可能,但以我目前的認知是實現不了,畢竟寫代碼的是程序員,給算法喂數據的是程序員,算法設計還是程序員,AI 連幫人類解 bug 的功能都還不出現!\
參考資料:
[1] //arxiv.org/abs/1706.03762
[2] //arxiv.org/abs/1810.04805
[4] //arxiv.org/abs/1904.09751