CodeRetriever: A Large Scale Contrastive Pre-Training Method for Code Search

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 2898 - 2910 December 7-11, 2022 ©2022 Association for Computational Linguistics CodeRetriever: Large-scale Contrastive Pre-training for Code Search Xiaonan Li1∗, Yeyun Gong2, Yelong Shen2, Xipeng Qiu1†, Hang Zhang2, Bolun Yao2, Weizhen Qi2, Daxin Jiang2, Weizhu Chen2, Nan Duan2 1 Shanghai Key Laboratory of Intelligent Information Processing, Fudan University 1 School of Computer Science, Fudan University 2Microsoft 1{lixn20, xpqiu}@fudan.edu.cn, 2{yegong, yeshe, v-zhhang, yaobolun, weizhen djiang, wzchen, nanduan}@microsoft.com Abstract In this paper, we propose the CodeRetriever model, which learns the function-level code semantic representations through large-scale code-text contrastive pre-training. We adopt two contrastive learning schemes in CodeRe- triever: unimodal contrastive learning and bi- modal contrastive learning. For unimodal contrastive learning, we design an unsuper- vised learning approach to build semantic- related code pairs based on the documenta- tion and function name. For bimodal con- trastive learning, we leverage the documenta- tion and in-line comments of code to build code-text pairs. Both contrastive objectives can fully leverage large-scale code corpus for pre-training. Extensive experimental results show that CodeRetriever achieves new state- of-the-art with significant improvement over existing code pre-trained models, on eleven domain/language-specific code search tasks with six programming languages in different code granularity (function-level, snippet-level and statement-level). These results demonstrate the effectiveness and robustness of CodeRe- triever. The codes and resources are avail- able at https://github.com/microsoft/ AR2/tree/main/CodeRetriever. 1 Introduction Code search aims to retrieve functionally relevant code given a natural language query to boost de- velopers’ productivity (Parvez et al., 2021; Hu- sain et al., 2019). Recently, it has been shown that code pre-training techniques, such as Code- BERT (Feng et al., 2020) and GraphCodeBERT (Guo et al., 2021), could significantly improve code search performance via self-supervised pre-training using large-scale code corpus (Husain et al., 2019). However, existing code pre-training approaches usually adopt (masked) language modeling as the Asia.∗Work is done during internship at Microsoft Research †Corresponding author. Doc: Return the Fibonacci number. Code: def Fibonacci(n): if n == 0: return 0 elif n in [1,2]: return 1 return \ Fibonacci(n-1)+Fibonacci(n-2) Doc: Get the Fibonacci number. Code: def Fibonacci_Number(index): cache = [0]*(index+1) cache[0] = 0 cache[1] = 1 cache[2] = 1 for i in range(3,index+1): cache[i] = \ cache[i-1] + cache[i-2] return cache[index] (a) Fibonacci Doc: Sort the input array into ascending order. Code: def bubbleSort(arr): n = len(arr) for i in range(n): for j in range(0, n-i-1): # if adjacent elements appear in # descending order, swap them. if arr[j] > arr[j+1]: arr[j], arr[j+1] = arr[j+1], arr[j] (b) BubbleSort Figure 1: Code examples. (a) Two different imple- mentations of Fibonacci number algorithm; (b) Docu- mentation, in-line comment, and code in BubbleSort implementation. training objective which targets on learning to pre- dict (masked) tokens in a given code context (Feng et al., 2020; Guo et al., 2021; Ahmad et al., 2021; Wang et al., 2021b). However, this token-based approach generally results in poor code semantic representations due to two reasons. The first one is the anisotropy representation issue. As …

Content truncated to 500 words