是不是有一天想象著讓程式碼自動補全，今天他來了！！！

程式零世界發表於2020-07-30

在這裡插入圖片描述

作者：熊唯，黃飛，騰訊 PCG/QQ研發中心/CV應用研究組

AI 如果真的可以寫程式碼了，程式設計師將何去何從？近幾年，NLP 領域的生成式任務有明顯的提升，那通過 AI
我們可以讓程式碼自動完成後續補全嗎？本文主要介紹瞭如何使用 GPT2 框架實現程式碼自動補全的功能。

如果 AI 真的可以自己寫程式碼了，程式設計師將何去何從？

我去年做過一個程式碼補全的小功能，打包為 androidStudio 外掛，使用效果如下：

在這裡插入圖片描述

程式碼補全模型預測出的結果有時的確會驚嚇到我，這也能學到~? 那如果給它見識了全世界的優秀程式碼，再給足夠量級引數和優秀的模型框架，真的可以實現需求作為輸入，直接輸出程式碼嗎?

"我的需求講完了，你的程式碼呢?" 希望可以看到這一天。

程式碼補齊功能有其他優秀外掛也已實現，比如 tabnine，Kite 和國產的 aixcoder。本文主要介紹下程式碼補全功能需要實現的整套流程。主要包括資料，演算法和工程。

資料

眾所周知，演算法工程師大部分時間都在處理資料。

深度學習是使用大資料訓練模型的一個過程，資料是很重要的一個模組。人是會累的，休息不好還導致記憶不好。AI 是你給多少資料它就能儲存接收多少資料，學不到資訊那是人的錯，給的資料不好或者演算法設計不好。所以我們先儘可能多的準備好訓練資料。

1、資料採集
本文的目的是程式碼補全，訓練資料就是程式碼段。考慮到每種語言風格和語法都不一致，所以單個模型只針對一種程式碼語言。

我使用的訓練資料主要來源於 GitHub，編寫了一個簡單的爬蟲程式碼，指定語言後根據 stars 的排序下載工程。

Github 的 search API 官方地址：

https://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 部分，自迴歸語言模型，適合生成式任務。

Transformer框架圖

GPT2和BTER框架示意圖

程式碼補全功能就是基於 GPT2 框架，OPenAI 官方提供了多套 GPT2 預訓練模型：

官方提供GPT2引數

作為一個經常要把模型部署到移動端的 CVer，看到這個引數級別，我選擇最小的模型進行 finetune。

對於 GPT 演算法，下面這篇文章講的很好，感興趣同學可以看看：

https://zhuanlan.zhihu.com/p/137350403

GPT2的預測過程

本文在訓練中使用 512 個上文，預測到回車符為終止。模型網路使用超參:12 個層，768 個隱藏節點，12 個 heads，採用了 uber 的 Horovod 分散式框架進行訓練。

infer 階段採用 beam-search 會導致整個預測過程特別耗時，所以參考了https://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 = http://ip:8080/plugin?data="輸入"

其中 model_infer 函式需要實現模型的 infer 前向計算邏輯，從請求中獲取 data 欄位作為輸入，infer 預測的結果列表作為輸出返回給呼叫方。

經過上面的工作，我們已經提供了一個服務介面，返回我們程式碼補全的預測結果。

外掛編寫

最後一步就是如何在 IDE 上使用功能了。我們要開發 AS 的外掛，需要使用 IntelliJ，首先需要在本機安裝配置 IntelliJ IDEA

下載地址:

https://www.jetbrains.com/idea/download/

社群版原始碼:

https://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 = "http://ip:8080/plugin?data=";
// 獲取當前編輯位置文字
PsiFile str = position.getContainingFile();
// 根據模型上文限制獲取程式碼端
String data = getContentCode();
String url = baseUrl+data;
// 傳送請求
String result = HttpUtils.doGet(url);
// 後處理邏輯，在提示框顯示預測結果
show()

最終呈現形式：

在這裡插入圖片描述