Neural code search

Published in

翻滾吧！駭客女孩！

9 min readMay 20, 2020

只要是身為開發者，總免不了利用 google 搜尋技巧，或在 stackoverflow 上提問。然而 Google 搜尋和 Stack overflow 為了方便開發者快速搜尋，通常會仰賴提問者提供的敘述和 tag，而非造成問題的 code snippet。我們稱呼對造成問題的 code snippet 相關的 tag 和以自然語言寫成的敘述為 metadata。Google 搜尋和 Stackoverflow 主要使用 metadata 作為 index 加速搜尋的速度，也避免處理問題 code snippet。然而這些 metadata 對開發者的問題只是觸及冰山一角，因為缺乏完整捕捉 code snippet 的能力，以至於無法辨識問題 code snippet 所達成的功能，而開發者也無法藉由閱讀原始碼樣本，來增進問題的了解。

有鑑於此，Facebook AI Research（FAIR）就應用 NLP 的方式來捕捉 code snippet 內容的能力並併入 Information Retrieval （IR）的方法來進行搜尋。他們建立了一個原始碼搜尋系統，在這個系統中使用者會輸入一個問題敘述，而系統會回傳一個和問題敘述相關的 code snippet。

FAIR 提出了兩個方法來完成原始碼搜尋系統，分別是：NCS 和 UNIF。這兩個方法分別發表在不同的論文中。NCS 是一個非監督式學習的方法，而 UNIF 則以監督式學習的方式擴充 NCS 方法。兩者都需要以稠密的浮點數向量，來表示該原始碼的語義。也仰賴一個相似度搜尋的演算法，進行資料庫搜尋比對而找出相似的 code snippet。

NCS

在 NCS 中，他們使用 FastText 架構去訓練 code token-based 的 embedding。雖說是 code token 但實際的流程和一般處理自然語言差異不大，比如在產生所謂的 code token，使用的是處理英文文章單字的 tokenizer ，使用 FastText架構下的 skip-gram 模型建立 code token-based embedding matrix，取出該 token 的向量代表。在做 IR 的方式也與自然語言文件搜尋相似，使用 TF-IDF （Term Frequency — Inverse Document Frequency）來給予文件權重，並以平均文件中 tokens 的向量表示法，來計算文件的向量表示法。最後採用 cosine similarity 來衡量兩個文件的相似度，只不過在 NCS 中，文件的定義是一個 method。

事實上，NCS 在處理原始碼所用的 Tokenize 的方法比起英語所用的 space tokenizer 要更為細緻。NCS 在取得一個完整的原始碼後，會先就程式文法架構，將每一個原始碼中的 method 拆解成一個包含數個 tokens 的 list。Tokenize 的程序包括取出 method 名稱，呼叫 method 所傳入的引數，enums，字串常數和註釋等。取出這些文法結構元件後，再以自然語言的方式來做 Tokenize。最後，透過標準的文件比較演算法， FAISS ，來完成 method-level 的文件的比較，並依相似程度來做排序。

雖然 NCS 表現比傳統使用 TF-IDF 的搜尋架構 BM25 為佳，但是 NCS 的問題在於這個 framework，會將 code token 置於和英文單詞 token 同樣的向量空間中。但是這樣的結果是不恰當的。其中一個問題就是：若有一個 code snippet 的功能是請求尚未被配置的記憶體，以進行記憶體的配置（如下圖）。

那麼以「獲得尚未配置的內部記憶體」（“Get free space on internal memory”）為搜尋關鍵字，透過 NCS 方法學習的搜尋演算法，找尋到 code snippets 無法包含上面的 code snippet（未包含重要 tokens）。因為兩者的 vector space 各自對應不同功能領域，然而非監督式學習的 NCS 只學習到其中一個 domain（原始碼或英文）。這個時候，我們就需要使用 UNIF 架構。

UNIF

UNIF 使用監督式學習來產生兩個不同的 domains 的 embedding：一個是開發者所使用的問題敘述，另一個是所對應的 code snippet。兩者各自有自己的處理 pipeline，和 NCS 的混合訓練不同的是，最後會各自產生不同的 embeddings：code embedding 和 NLP query embedding。為了產生兩個 embeddings，UNIF 會以不同的機制訓練：

Query 敘述（NLP）為平均所有 tokens vector，而 code embedding 則是替換 TF-IDF 改使用 attention 機制來訓練。這兩個 embedding 最後會將 query 和 code snippet 映射到共享的向量空間中做比較，而這個共享的向量空間會透過前所述的監督式學習來完成，也就是建立 code 和 query 的配對組合，當作訓練集，以確保產生的兩種 embeddings 。他們除了能捕捉各自的 domain 的表示法，亦能符合最佳的映射關聯。

在 UNIF 的論文中，除了非監督的 NCS，又分別比較了兩個監督式的類神經網路架構的原始碼搜雲模型，他們分別是 Code Description Embedding Neural Network (CODEnn) 和 Semantic Code Search（SCS）。相較於 UNIF 使用 bag-of-words （skip-gram），這兩個方法都採用較為複雜多個 RNN sequences 的架構。

相較於 UNIF 只產生兩個 embedding matrices，CODEnn 至少產生 4 個不同的 embeddings（見上圖），分別是：Method name, API call sequence， bag-of-word code snippet tokens 和 query（NL Desc / Query Tokens）。除了code snippet tokens 採用 Feedforward network 外，每一個 embedding 都有自己的 bi-directional LSTM trainer，最後使用一個 max-pooling 的方法對三種 code relevant embeddings 做 downsampling，並以 Feedforward （MLP）的方式來結合三種不同的 embeddings 產生最後的 code embedding。這個 code embedding 會與同樣使用 RNN sequence encoder 產生的 query embedding，成為量測相似 cosine similarity 分數的輸入。

而使用 seq2seq 架構的 SCS，除了使用 code snippet，還引入了 Docstring 來學習與 query 同一個 domain 的自然語言模型。除此之外，與其學習 code 和 query 共享的向量空間，SCS 主要學習從 code 映射到 docstring 的向量空間，並與自然語言模型產生的 embedding 做相似比較。在這個搜尋系統中，有三個主要的模組，這三個模組可以分開獨立存在，他們分別是：

Train code-to-docstring model: 在這個模組中，主要是學習 code 對 docstring 的映射關係。如上圖所示，使用 sequence-to-sequence 的架構來訓練配對 code tokens 和 docstring 輸出。 Encoder 會以 code tokens 作為輸入，而相對應的 docstring 則經由 decoder 輸出。
Train language model: 在這個模組，主要是藉由 RNN encoder 訓練一個語言模型，會將輸入的 docstring 或 query tokens 映射到語言模型的向量空間中。該語言模型亦可以比較 query 和數個自然語言輸入，輸出 cosine similarity。
Train code-to-NL embedding model: 在這個模組內，會使用模組一中的sequence-to-sequence 的架構，但與其使用 docstring sequence 輸出，該方法會採用 docstring embedding 的輸出，並作為搜尋的 index。

結果：

在 UNIF 的論文中使用下列資料及來做 benchmarking，他們分別是：

Java-50：這組 dataset 最早由 CODEnn 論文提出，包含了在 stackoverflow 前 50 個最受歡迎，關於 Java 的問題。這些問題，必須要在回答中提供原始碼。
Android 287：這組 dataset 最早由 NCS 論文提出，包含了在 stackoverflow 中包含 Android 和 Java 在 tag 中。除此之外，必須要有包括原始碼且被按讚的回答，該原始碼還必須與 GitHub Android corpus 相匹配。模型的訓練，則使用 CODEnn-Java-Train

而模型的訓練，則使用 CODEnn-Java-Tran 和 GitHub-Android-Train。關於更多訓練資料集的準備，請見參考資料。

下圖中的數字是有多少問題被正確的回答。在 Benchmark queries 欄位中，是評估資料集的名字，名字中帶有的數字為該資料集擁有問題的個數，細節可看前文。 Model 欄位則是評估哪一個模型。而 Answered@1, 5, 10，則是模型分別回傳 top 1，top 5 和 top 10 的答案們，其中包含了真實答案。