重现baseline
我们选择重现CODEnn模型(论文:Deep Code Search),因为它结构简单、端到端可训练,且相比其它方法拥有较高的性能。
Baseline原理
为了根据给定的query(文本)查询相关的代码,需要计算文本和所有代码之间的相似度,从而选择相似度最高的k个代码作为查询结果。由于文本与代码为异构数据,需要将它们编码到统一的embedding space中。CODEnn是一种端到端模型,用两个encoder分别将文本和代码编码到同一高维空间中,并用cosine similarity计算文本与代码间的相似度。其整体结构可以用下图简要概括:
完整结构图如下:
RNN
RNN经常用于编码序列数据,能够有效捕捉序列中不同时刻之间的关联性。其结构如下:
RNN按顺序处理顺序输入数据,用hidden state(图中\(h\))保存上下文信息,从而能够处理不同时刻之间的依赖关系并对序列进行编码。
Code Embedding Network
Code embedding network对代码的三个方面分别编码:method名字、调用的API序列、代码tokens。网络结构如下:
其中:
- method名字可以表示为单词序列 \(M = w_1, w_2, ..., w_{N_M}\),经过RNN模块为每个时间产生一个embedding,再经过max pooling模块,得到序列的embedding向量\(\vec{m}\)。
- API序列 \(A = a_1, a_2, ..., a_{N_A}\) 可以通过AST结构获得。类似method名字,经过RNN模块与max pooling模块编码为向量\(\vec{a}\)。
- 代码中的tokens表示为集合 \(\Gamma = \{ \tau_1, \tau_2, ..., \tau_{N_\Gamma} \}\),因为代码中 tokens 较多、可能重复,不应该表示为序列。因此,\(\Gamma\)不使用RNN编码,直接经过MLP模块与max pooling模块编码为向量\(\vec{t}\)。
最终的code embedding为\(\vec{c} = \text{tanh}(W^C [\vec{m}, \vec{a}, \vec{t}])\)。
Text Embedding Network
文本(query内容,或对代码的描述)可以表示为单词序列\(D = w_1, w_2, ..., w_{N_D}\),因此也使用RNN与max pooling模块编码为向量\(\vec{d}\)。
相似度计算
本文采用余弦相似度作为向量相似度指标:
\[sim(\vec{c}, \vec{d}) = cos(\vec{c}, \vec{d}) = \frac{\vec{d}^T \vec{d}}{\|\vec{c}\| \|\vec{d}\|}\]
比起L2距离,余弦相似度能够不受向量尺度影响,而更好地表征两个向量之间的相似度。详见Cosine Similarity – Understanding the math and how it works (with python codes)。
训练
训练过程中,每个样本包含一段代码\(\vec{c}\)、代码的正确描述\(\vec{d}_+\)和随机采样的一段错误描述\(\vec{d}_-\)。采用triplet loss的思路设计损失函数:
\[L(\vec{c}, \vec{d}) = \max (0, \epsilon - cos (\vec{c}, \vec{d}_+) + cos (\vec{c}, \vec{d}_-) ) \]
优缺点
优点:
- 将代码与文本两种不同domain的数据编码到同一个高维空间里,用端到端训练优化encoder、