作者丨熊唯、黃飛?
來源丨騰訊技術(shù)工程
AI 如果真的可以寫代碼了，程序員將何去何從？近幾年，NLP 領(lǐng)域的生成式任務(wù)有明顯的提升，那通過 AI 我們可以讓代碼自動完成后續(xù)補(bǔ)全嗎？本文主要介紹了如何使用 GPT2 框架實現(xiàn)代碼自動補(bǔ)全的功能。
我去年做過一個代碼補(bǔ)全的小功能，打包為 androidStudio 插件，使用效果如下：代碼補(bǔ)全模型預(yù)測出的結(jié)果有時的確會驚嚇到我，這也能學(xué)到~? 那如果給它見識了全世界的優(yōu)秀代碼，再給足夠量級參數(shù)和優(yōu)秀的模型框架，真的可以實現(xiàn)需求作為輸入，直接輸出代碼嗎?"我的需求講完了，你的代碼呢?" 希望可以看到這一天。代碼補(bǔ)齊功能有其他優(yōu)秀插件也已實現(xiàn)，比如 tabnine，Kite 和國產(chǎn)的 aixcoder。本文主要介紹下代碼補(bǔ)全功能需要實現(xiàn)的整套流程。主要包括數(shù)據(jù)，算法和工程。數(shù)據(jù)
眾所周知，算法工程師大部分時間都在處理數(shù)據(jù)。深度學(xué)習(xí)是使用大數(shù)據(jù)訓(xùn)練模型的一個過程，數(shù)據(jù)是很重要的一個模塊。人是會累的，休息不好還導(dǎo)致記憶不好。AI 是你給多少數(shù)據(jù)它就能存儲接收多少數(shù)據(jù)，學(xué)不到信息那是人的錯，給的數(shù)據(jù)不好或者算法設(shè)計不好。所以我們先盡可能多的準(zhǔn)備好訓(xùn)練數(shù)據(jù)。1、數(shù)據(jù)采集
本文的目的是代碼補(bǔ)全，訓(xùn)練數(shù)據(jù)就是代碼段?？紤]到每種語言風(fēng)格和語法都不一致，所以單個模型只針對一種代碼語言。我使用的訓(xùn)練數(shù)據(jù)主要來源于 GitHub，編寫了一個簡單的爬蟲代碼，指定語言后根據(jù) stars 的排序下載工程。Github 的 search API 官方地址：https://developer.github.com/v3/search/2、數(shù)據(jù)清理
直接下載的數(shù)據(jù)肯定是不能直接用的，我們還需要對數(shù)據(jù)進(jìn)行清理。首先，我們的訓(xùn)練數(shù)據(jù)只需要工程中的代碼文件，以 java 工程為例，我們只保留.java 結(jié)尾的文件，其他文件可剔除。其次，我的代碼補(bǔ)全目標(biāo)是代碼段，不針對注釋功能。而且對于代碼補(bǔ)全訓(xùn)練時，我們是會給定一定范圍的上文，如果存在注釋段會占用有效代碼信息。另外注釋除英文外其他字符不在我的訓(xùn)練 vocab 范圍內(nèi)，所以需要對代碼中注釋和日志進(jìn)行清理。/*?注釋文本*/
/**
注釋段落
*/
//?注釋文本
?code?//注釋
經(jīng)過以上數(shù)據(jù)清理后，得到純代碼數(shù)據(jù)。3、數(shù)據(jù)編碼
得到了訓(xùn)練數(shù)據(jù)后還需要把代碼文本進(jìn)行編碼。本文使用的是 bpe（byte pair encoder）字節(jié)對編碼，主要為了數(shù)據(jù)壓縮。bpe 簡單理解為將一個單詞再拆分為多個字母組合，比如 tencent 拆分為 ten-cent，這些組合方式則是根據(jù)大量數(shù)據(jù)，統(tǒng)計頻率得到。由于我們期待的代碼補(bǔ)全功能是在行首輸入幾個字母，根據(jù)上文預(yù)期出本行內(nèi)容。假設(shè) tensorflow 這個 token 被編碼對應(yīng)到一個 id，那我希望輸入 ten 就輸出 tensorflow 是無法實現(xiàn)的。所以在訓(xùn)練過程中，我會隨機(jī)把 token 打斷，比如將 tensorflow 打斷為 t-en-sor-flow 進(jìn)行編碼，打斷原則是被切分的部分一定要在詞匯表中。數(shù)據(jù)編碼后，代碼的每個 token 被編碼為 1~N 個 id。模型預(yù)測到的 id 反編碼為 token 即可?；剀嚪J(rèn)為是預(yù)測的終止符。經(jīng)過以上處理，我們就準(zhǔn)備好了訓(xùn)練數(shù)據(jù)，下面就可以進(jìn)行算法部分了。模型算法
在騰訊文檔的錯別字糾錯需求中，我們采用了基于 LSTM 的 seq2seq 以及 facebook 提出的基于 CNN 的 seq2seq，可以得到不錯的糾錯效果。直到 NLP 出現(xiàn)了一個"網(wǎng)紅"--BERT，采用后精度直接提升 8 個點左右，不虧是 google。下面先簡單介紹下 bert 和 gpt2。2017 年中 google 提出了 Transformer 結(jié)構(gòu)。不用 rnn，不用 cnn，提出 attention is all you need。2018 年 openAI 采用了 transformers 結(jié)構(gòu)在 18 年發(fā)布了 GPT。同年 google AI Language 發(fā)布了 bert 論文，提出的 BERT 模型在 11 個 NLP 任務(wù)上刷新了記錄。2019 年 openAI 又推出了 GPT-2 模型。Bidirectional Encoder Representation from Transformers：基于 transformers 框架的 encoder 部分，自編碼語言模型，適合 N-1（比如句子分類），N-N（比如詞性標(biāo)注）的任務(wù)，但是它并不適合做生成任務(wù)。基于 transformers 的 decoder 部分，自回歸語言模型，適合生成式任務(wù)。
Transformer框架圖
GPT2和BTER框架示意圖代碼補(bǔ)全功能就是基于 GPT2 框架，OPenAI 官方提供了多套 GPT2 預(yù)訓(xùn)練模型：
官方提供GPT2參數(shù)作為一個經(jīng)常要把模型部署到移動端的 CVer，看到這個參數(shù)級別，我選擇最小的模型進(jìn)行 finetune。對于 GPT 算法，下面這篇文章講的很好，感興趣同學(xué)可以看看：https://zhuanlan.zhihu.com/p/137350403??
GPT2的預(yù)測過程本文在訓(xùn)練中使用 512 個上文，預(yù)測到回車符為終止。模型網(wǎng)絡(luò)使用超參:12 個層，768 個隱藏節(jié)點，12 個 heads，采用了 uber 的 Horovod 分布式框架進(jìn)行訓(xùn)練。infer 階段采用 beam-search 會導(dǎo)致整個預(yù)測過程特別耗時，所以參考了https://arxiv.org/abs/1904.09751論文，采用 top-k sampling，每次預(yù)測 top3 的結(jié)果再通過概率閾值過濾后作為最終候選輸出。輸入一段代碼，預(yù)測出后續(xù)代碼，以回車符截止。訓(xùn)練出模型后，還要把模型應(yīng)用起來，所以還需要一些工程工作需要實現(xiàn)。代碼補(bǔ)全功能，最合適的應(yīng)用場景就是上 IDE。nlp 模型不太適合在本機(jī)部署，最終選擇了在 GPU 機(jī)器上部署模型，然后終端通過 http 請求獲取預(yù)測文本顯示的方案。Flask 是一個 Web 應(yīng)用程序框架，靈活，輕便，容易上手。本文簡單介紹如何利用 flask 啟動一個 web 服務(wù)，以及如何訪問和調(diào)用我們的功能接口。首先我們創(chuàng)建一個 conda 環(huán)境：conda?create?-n?flask?python=3.6
source?activate?flask
pip?install?flask
from?flask?import?Flask
from?flask?import?request
app?=?Flask()
#?route把一個函數(shù)綁定到對應(yīng)的?url?上
@app.route("/plugin",methods=['GET',])
def?send():
????data?=?request.args.get('data')
???#?模型預(yù)測邏輯
????out?=??model_infer(data)
????return?out
if?__name__?==?'__main__':
????app.run(host='0.0.0.0',port=8080,?debug=False)
執(zhí)行 run.py 代碼，后臺服務(wù)開啟運(yùn)行：url?=?http://ip:8080/plugin?data="輸入"
其中 model_infer 函數(shù)需要實現(xiàn)模型的 infer 前向計算邏輯，從請求中獲取 data 字段作為輸入，infer 預(yù)測的結(jié)果列表作為輸出返回給調(diào)用方。經(jīng)過上面的工作，我們已經(jīng)提供了一個服務(wù)接口，返回我們代碼補(bǔ)全的預(yù)測結(jié)果。最后一步就是如何在 IDE 上使用功能了。我們要開發(fā) AS 的插件，需要使用 IntelliJ，首先需要在本機(jī)安裝配置 IntelliJ IDEAhttps://www.jetbrains.com/idea/download/https://github.com/JetBrains/intellij-community好用的插件可以節(jié)省程序員很多時間，在插件實現(xiàn)時，我還添加了一個小的 git-blame 功能，實時查看指定行的 git 提交人，對于手 Q 這種多人合作的工作，比較實用。大家也可以通過 IntelliJ 自己開發(fā)一些常用功能。
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
????????//獲得當(dāng)前本地代碼根目錄
????????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;
????????}
????????//獲取當(dāng)前編輯文檔的目錄
????????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);
????????}
????}
}
本文的代碼補(bǔ)全插件主要代碼邏輯為調(diào)用上一步后臺部署的請求。//?請求url格式（和flask接口一致）
String?baseUrl?=?"http://ip:8080/plugin?data=";
//?獲取當(dāng)前編輯位置文本
PsiFile?str?=?position.getContainingFile();
//?根據(jù)模型上文限制獲取代碼端
String?data?=?getContentCode();
String?url?=?baseUrl+data;
//?發(fā)送請求
String?result?=?HttpUtils.doGet(url);
//?后處理邏輯，在提示框顯示預(yù)測結(jié)果
show()
可以看出，模型的預(yù)計結(jié)果還是不錯的~以上為代碼補(bǔ)全功能的實現(xiàn)和應(yīng)用，算是 AI 自動寫代碼的一小步。AI 能否自己寫代碼，達(dá)到疑犯追蹤里 TM 那種水平，我不敢說一定不可能，但以我目前的認(rèn)知是實現(xiàn)不了，畢竟寫代碼的是程序員，給算法喂數(shù)據(jù)的是程序員，算法設(shè)計還是程序員，AI 連幫人類解 bug 的功能都還不出現(xiàn)！[1] https://arxiv.org/abs/1706.03762[2] https://arxiv.org/abs/1810.04805[3] https://github.com/openai/gpt-2[4] https://arxiv.org/abs/1904.09751
添加極市小助手微信（ID : cv-mart），備注：姓名-學(xué)校/公司-研究方向-城市（如：小極-北大-目標(biāo)檢測-深圳），即可申請加入極市目標(biāo)檢測/圖像分割/工業(yè)檢測/人臉/醫(yī)學(xué)影像/3D/SLAM/自動駕駛/超分辨率/姿態(tài)估計/ReID/GAN/圖像增強(qiáng)/OCR等技術(shù)交流群，更有每月大咖直播分享、真實項目需求對接、求職內(nèi)推、算法競賽、干貨資訊匯總、行業(yè)技術(shù)交流，一起來讓思想之光照的更遠(yuǎn)吧~
覺得有用麻煩給個在看啦~??
