如何寫最高端的代碼？Facebook教你怎樣用機器學習做最美的代碼搜索工具

選自Facebook AI Blog

作者：Sonia Kim、Hongyu Li、Satish Chandra

機器之心編譯

參與：路、一鳴、思源

如何基於文本查詢快速獲取代碼示例，對於工程師而言是一個很影響效率的事兒。Facebook 最近提出了新型代碼搜索工具——神經代碼搜索（NCS）和 UNIF，分別基於無監督和監督的方式提供快速高效的代碼檢索。

當工程師能夠輕鬆獲取代碼示例，指導其完成特定編程任務時，他們的工作效率會顯著提高。例如，對於「如何以編程方式關閉或隱藏安卓軟鍵盤？」這類問題，工程師可以從 Stack Overflow 等常用網站上獲取可用信息。但是當問題涉及專有代碼或 API（或者用不常用編程語言寫的代碼）時，工程師需要不同的解決方案，因為在常用論壇上可能找不到這方面的答案。

為了解決這個需求，Facebook 開發了一個代碼搜索工具——神經代碼搜索（Neural Code Search，NCS），該工具使用自然語言處理（NLP）和信息檢索（IR）技術直接處理源代碼文本。該工具接收自然語言作為查詢（query），並返回從代碼庫中直接檢索到的相關代碼段。這裡的前提是能夠獲取大型代碼庫，從而更有可能搜索到與開發者提出的查詢相關的代碼段。

Facebook 介紹了兩種可用模型：

NCS：結合自然語言處理和信息檢索技術的無監督模型；

UNIF：NCS 的擴展，當訓練過程中有好的監督數據時，UNIF 使用監督神經網路模型來提升性能。

利用 Facebook AI 開源工具（包括 fastText、FAISS、PyTorch），NCS 和 UNIF 將自然語言查詢和代碼段表示為向量，然後訓練一個網路，使語義相近的代碼段和查詢的向量表示在向量空間中彼此靠近。使用這些模型，我們能夠從代碼庫中直接尋找代碼段，從而高效解決工程師的問題。為了評估 NCS 和 UNIF，Facebook 使用了新創建的數據集（包含 Stack Overflow 上的公開查詢和對應的代碼段答案）。結果表明，這兩個模型可以正確回答該數據集中的問題，如：

如何關閉／隱藏安卓軟鍵盤？

如何在安卓中將點陣圖轉換為可畫的？

如何刪除一整個文件夾及其內容？

如何處理 back button？

NCS 的性能表明，相對簡單的方法可以很好地處理源代碼；UNIF 的性能表明，簡單的監督學習方法可以在有可用標註數據時，提供顯著的額外收益。當這些模型與其他 Facebook 構建系統（如 Aroma 和 Getafix）結合時，這個項目可以為工程師提供可擴展且不斷增長的 ML 工具包，幫助他們更高效地寫代碼、管理代碼。

NCS 如何使用嵌入向量

NCS 模型使用嵌入（連續向量表示）來捕捉程序語義（即代碼段的意圖）。當進行恰當計算時，這些嵌入能夠將語義相近的實體在向量空間中拉近距離。

如下圖示例所示，關於如何關閉／隱藏安卓軟鍵盤有兩個不同的方法。由於它們共享類似的語義，因此即使它們代碼行不完全相同，它們在向量空間中的位置彼此接近。

上圖表明語義相近的代碼段在向量空間中距離較為接近。

Facebook 使用這個概念構建 NCS 模型。總體而言，在模型生成過程中，每個代碼段以方法級粒度嵌入到向量空間中。一旦模型構建完成，給定的查詢語句就會映射到同樣的向量空間中，那麼向量距離就可用於估計代碼段與查詢的相關程度。

下圖展示了模型生成和搜索檢索流程：

上圖展示了 NCS 的整個模型生成和搜索檢索過程。

模型生成

為了生成模型，NCS 必須抽取單詞，構建詞嵌入，然後構建文檔嵌入。（這裡「文檔」指方法體（method body）。）

抽取單詞

NCS 從源代碼中抽取單詞，並執行分詞，生成詞的線性序列。

為了生成能表示方法體的向量，Facebook 將源代碼看作文本，從以下句法類中抽取單詞：方法名稱、方法調用、枚舉值、字元串文本和注釋。然後基於標準英語規範（如空格、標點）和代碼相關標點（如下劃線命名法和駝峰命名法）執行分詞。例如上圖中，對於方法體（或方法名）「pxToDp」，源代碼可作為以下單詞的集合：「converts pixel in dp px to dp get resources get display metrics」。

對於代碼庫中的每個方法體，我們都可以用這種方法對源代碼執行分詞，並為每個詞學習一個嵌入。之後，從每個方法體中抽取的單詞列表類似一個自然語言文檔。

構建詞嵌入

Facebook 使用 fastText 為辭彙語料庫中的所有單詞構建詞嵌入。fastText 使用一個兩層神經網路計算向量表示，該網路可以在大型語料庫上以無監督方式訓練。具體來講，Facebook 使用 skip-gram 模型，利用目標 token 的嵌入來預測固定窗口大小中的上下文 token 的嵌入。如上例所示，給出 token「dp」的嵌入、窗口大小為 2，skip-gram 模型學習預測 token「pixel,」、「in,」、「px,」、「to.」的嵌入。其目標是學習嵌入矩陣，其中 |V_c| 表示語料庫大小，d 表示詞嵌入的維度，T 的第k 行表示 V_c 中第 k 個單詞的嵌入。

在該矩陣中，如果兩個向量表示對應的單詞經常出現在相似語境，則這兩個向量表示距離較近。Facebook 使用該命題的逆命題幫助定義語義關係：向量距離接近的單詞應該語義相關性較高。這在自然語言處理中叫作「分布假設」（distributional hypothesis），源文本同樣適用於這一概念。

構建文檔嵌入向量

下一步是利用方法體中的單詞將其總體意圖表達出來。為此，研究人員計算了方法體中所有詞語的詞嵌入向量的加權平均值。這被稱為是文檔嵌入。

公式中，d 表示方法體的詞語集合，v_w 是詞 w 的詞嵌入，使用 fastText 處理。C 是包含所有文檔的語料，u 是歸一化函數。

研究人員使用詞頻-逆文檔頻率（TF-IDF），為給定文檔中的給定詞語分配權重。這一步驟旨在強調文檔中最為重要的詞語，即如果一個詞在文檔中出現頻率很高，則它的權重很高，但如果它在語料庫中的多個文檔中都有出現，則該辭彙會受到懲罰。

在這一步之後，研究人員給語料庫中每個方法體的文檔向量表示一個索引數字。模型生成就完成了。

搜索檢索

搜索查詢可以用自然語言表達，例如「關閉/隱藏軟鍵盤」或「如何建立一個沒有標題的對話框」。研究人員使用同樣的方式對查詢和源代碼執行分詞，且使用同樣的 fastText 詞嵌入矩陣 T。研究人員簡單地計算了詞向量表示的平均值，建立一個查詢語句的文檔嵌入，詞表外的詞被捨棄。研究使用標準的相似度搜索演算法 FAISS，用於尋找和查詢的餘弦相似度最接近的文檔向量，並返回 top n 個結果。（此外還有一些訓練後排序，詳情參見：https://dl.acm.org/citation.cfm?id=3211353）

兩個方法體和查詢被映射在相同的向量空間中，且位置較為接近。這說明查詢和這兩個方法體在語義上相似，且相關。

實驗結果

研究人員使用 Stack Overflow 問題測試 NCS 的性能。使用標題作為查詢，以及答案中的一個代碼段作為理想的代碼回答。給定一個查詢，研究人員評價模型能否在 top 1、5 或 10 個回答上從 GitHub 倉庫集合中抽取正確的代碼段。研究人員還提供了平均排序倒數（MRR）分數，用于衡量 NCS 能否在第 n 個結果返回正確的答案。

在 287 個問題中，NCS 能在 top 10 個結果內正確回答 175 個問題，大約是整體數據集的 60%。研究人員同時對比了 NCS 和其他傳統信息檢索演算法的表現，如 BM25。如下圖所示，NCS 的性能優於 BM25。

下面是 NCS 回答很好的一個問題示例：「從 app 中打開安卓市場」，NCS 返回的第一個答案如下所示：

UNIF：監督式方法

NCS 的關鍵在於詞嵌入。由於 NCS 是無監督模型，它有很多優勢。首先，它可以直接從搜索語料中學習，訓練很快、很簡單。NCS 假設查詢中的詞和源代碼中抽取的詞有著相同的域，因為查詢和代碼段被映射在相同的向量空間中。然而，事實不一定總是這樣。

例如，「Get free space on internal memory」里，沒有一個詞出現在以下代碼段中。研究人員想實現的是將查詢詞「free space」映射到下列代碼中的單詞「available」。

基於包含 14005 個 Stack Overflow 帖子的數據集，Facebook 研究人員分析了查詢詞和源代碼中詞語的重疊情況。他們發現，查詢中有 13,972 個唯一字，其中只有不到一半（6072 個詞）在源代碼中出現了。這說明，如果一個查詢包含源代碼沒有的詞，則 NCS 模型無法有效地檢索正確的方法。這一結果促使研究人員進一步探索監督學習模型，以將查詢詞映射到源代碼中。

Facebook 研究人員決定嘗試 UNIF。這是一個監督模型，基於 NCS 進行了微小的擴展，用於連接自然語言和源代碼。在這個模型中，研究人員使用監督學習訓練詞嵌入矩陣 T，生成兩個嵌入矩陣 T_c 和 T_q，分別對應代碼 token 和查詢 token。此外，研究人員還用基於注意力的權重機制替換了代碼 token 嵌入向量的 TF-IDF 權重機制。

UNIF 模型工作原理

研究人員在與訓練 NCS 同樣的 (c, q) 數據點集合中訓練 UNIF，c、q 分別表示代碼 token 和查詢 token。模型架構表示如下：

T_q ∈ R^|V_q|×d 和 T_c ∈ R^|Vc |×d 是兩個嵌入矩陣，分別映射每個自然語言描述單詞和代碼 token 到向量，向量長度為 d（V_q 是自然語言查詢語料，V_c 是代碼語料）。這兩個矩陣使用同樣的初始權重 T 初始化，並分別在訓練中進行修改（對應 fastText）。

為了將離散的代碼 token 向量結合為一個文檔向量，研究人員使用注意力機制計算加權平均值。注意力權重 a_c ∈ R^d，是一個 d 維向量，從訓練中學得，並代替了原有的 TF-IDF 權重。給定一個代碼詞袋嵌入向量 ，每個 e_i 的注意力權重 a_i 如下：

文檔向量通過和注意力權重相乘後的詞嵌入向量相加得到：

為了創建查詢文檔向量 e_q，研究人員求得每個查詢詞嵌入向量的平均，這和 NCS 的方法相同。在訓練過程中，模型從標準的反向傳播中學習參數 T_q、T_c、a_c。

UNIF 網路結構圖示。

檢索工作和 NCS 一致。給定一個查詢，利用上述方法將其表示為文檔向量，然後使用 FAISS 尋找和查詢的餘弦相似度最相近的文檔向量。（原則上，UNIF 可以從後處理排序中獲益，正如 NCS 那樣。）

UNIF 和 NCS 的效果對比

研究人員對比了 NCS 和 UNIF 在 Stack Overflow 評測數據集上的表現。評價方法是模型能否用 top1、top5、top10 結果正確地解答此查詢，評價使用 MRR 分數。下表說明，相比 NCS，UNIF 顯著提高了回答問題的數量。

如此驚人的搜索性能說明，如果提供合適的訓練語料，監督模型可以完成這項任務。例如，搜索查詢是「如何退出應用並展示主菜單？」，NCS 會返回：

而 UNIF 提供了更相關的代碼：

另一個例子是提問：「如何獲得 ActionBar 的高度？」，NCS 模型會返回：

UNIF 返回相關的代碼：

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！