简介
本次作业的主题是 Code Search，即使用更自然的方式搜索代码，是对传统代码搜索的一种改进。一种最自然的方法去搜索代码则是使用自然语言去搜索代码，描述需要的代码的功能，搜索系统返回与此类似的代码片段或函数。
本次结对编程我和队友分别调研复现了两个不同的 Code Search 方法，在本博客中我主要记录和描述 Hamel’s model (https://github.com/hamelsmu/code_search)的一些原理和我的观点。
这个 model 的主要框架如图：
我按照框架的顺序分为五个小节介绍，分别是：数据预处理、训练 Seq2Seq Function Summarizer 模型、训练语言模型、训练代码到句子模型、创建索引和搜索引擎，下面是具体内容：
数据预处理
首先下载原始数据：
在我的服务器上花了约 5 分钟下载了全部的数据，为了避免之后重新下载，我们把处理好的 dataframe 进行保存。
可以看到 csv 格式的 1241664 条原始数据约 5G，加载到 pandas 中要占用 10G 左右的内存。
然后对代码进行 docstring 提取处理，将原始代码解析为 (code, docstring) 的 pair，这里的提取方法主要是利用 Python 3 中的 ast 包，解析代码中的 docstring：
然后根据我服务器的情况，调整 CPU 核数，开始并行处理代码：
可以从监控中看到并行处理时，所有的核心都被完全利用了内存占用约变为两倍。
在处理应该快要结束时，我遇到了以下错误：
我不想深究这里面到底出了什么问题，少利用这些造成异常数据也不会对训练造成太大的影响，只需要把捕获异常的范围扩大，我做了以下改动：
修改之后花费大概 11 分钟，终于得到的处理好的 docstring pair：
下一步是将 pairs 中的数据展开，没什么好说的，大概花了 10 分钟左右，最后得到：
下面的一些步骤是删除重复条目，将有 docstring 和没有 docstring 的数据分开存放，训练时只能利用有 docstring 的，以及在「同一个 repo 中代码有潜在的相似风格」这个假设下，我们在划分训练集和验证集、测试集的时候最好不要出现同一个 repo 下的数据出现在不同集中的情况。这些步骤都较为简单，在我的机器上一共花费 3 分钟不到，过程这里略过，只展示最后的训练集的部分结果：
最后将训练集，测试集，验证集分别保存，得到预处理完成的数据共 2.6 G：
训练 Seq2Seq Function Summarizer 模型
这一步中我们使用 Seq2Seq 训练一个模型，利用函数代码预测 docstring（或理解为函数的描述），这个预测模型其实相当于代码的 encoder 可供后续使用。
根据 https://towardsdatascience.com/how-to-create-data-products-that-are-magical-using-sequence-to-sequence-models-703f86a231f8 ，对数据 tokenize 之后构建 seq2seq 模型：
（由于我没有找到足够的 GPU 来训练，下面的部分训练过程参数用作者公布的预训练的参数继续完成实验）
构建 seq2seq 网络：
训练完成之后，通过一些例子直观感受模型的效果，如：
在这个例子中我们模型预测的结果要比真实的结果更具体化，但是没有完全偏离意思。
在这个函数比较复杂的例子中，我们模型的输出就有一点不尽人意，从这个例子来看，一方面这个函数的 docstring 并不是很自然的方式，而是一种标准模板，我们的模型也没有抓住 socket 这个重要的名词的信息，这是很致命的错误。
当然对于一些比较小的函数，我们的模型已经表现得比较好，甚至比原来得描述更合理一些。
在训练之后，将 seq2seq 模型保存备用。
训练语言模型
在这一步中我们使用 fast.ai 的 AWD LSTM 模型实现来训练一个可以在句子中预测下一个词的模型，这里的训练语料可以来自和我们任务相关的数据，比如 * 上面的问答，这里为了简单使用训练集中的 docstring 来训练这个模型。
训练完成后我们将得到一个 sentence encoder，并且基于这个 encoder 我们利用 nmslib 包，可以将数据预索引，以达到可以在 embedding 空间内快速搜索相邻向量，然后可以手动查看训练的效果
可以看到在 embedding 空间中相似的向量在自然语言中确实表现出了很强的相似性。
这里有一个坑需要注意，原作者代码中构建搜索引擎用的 test_raw 作为搜索引擎的 ref_data 应该是错的，因为我们训练的 nmslib 的 index 是基于 train 数据，如果不改显然是会 IndexError 的（因为 test 数据比 train 数据少，要不然这个问题还不好被发现），所以应该做下面的修改！
在这一步完成后，保存语言模型备用。
训练代码到句子模型
之前我们训练了一个 Seq2Seq 模型可以根据代码内容预测 docstring（或称之为函数描述），这一步我们要利用前面训练的 encoder 和 decoder 和前面训练的 docstring embedding 空间中 fine tune 模型，要完成这一任务大概需要以下几步：
加载之前的 Seq2Seq 模型的 encoder，固定参数（将 trainable 设为 False）；
在 encoder 之上添加新的全连接层；
根据 (code, docstring-embedding) 训练新的模型；
接触参数固定，继续训练（将 trainable 设为 True）；
然后将所有的代码预编码备用；
新的模型比较简单：
就是在 encoder 的输出加了一些 layer，一开始时，要将 encoder 参数固定。
训练一段时间之后，再取消参数固定，继续训练几个 epoch。
最后利用这个模型将全部代码都 encode 并保存备用。
创建索引和搜索引擎
在这一步中，我们使用 nmslib 创建代码搜索需要的索引，为了快速查询，我们不能在 query 来临的时候再查询相邻的向量，所以需要将之前步骤代码的向量进行预处理，使得之后能进行快速的查询。
创建索引的过程需要大量的内存（~50 G），最后预构建的索引也有约 8 G 大小。
此时很容易遇到 Cuda 内存不够的情况，可能需要服务器清理一下内存，比如关闭之前正在运行的 python。
最后根据处理好的 8 G 索引文件，我们就可以构建搜索引擎了，为了更加易用，我们这里还从原始数据中加载了 GitHub repo 中文件的地址和代码行数：
至此就可以测试整个 Code Search 了，如：
总结
通过复现以上 Code Search 模型我学到了很多新的概念，也接触到了一些很实用的模型实现和好用的包，在之后我自己的科研过程中应该也能派上用场。
我复现的这个模型在数据预处理上做的仍然不够精细，一个很容易想到的改进是，将代码的 AST 信息利用起来，做代码结构的 embedding，或者是在 AST 中分析代码，把字面常量占比多的代码设置为更低的权重，把 docstring 中一些固定写法的东西去除或转换成其他特征）（如“@author”）。
另外我觉得 Code Search 从 NLP 的各种方法出发也不是唯一的解决方案，我们（开发者）感到自然的代码搜索方法还有很多，不一定非要是自然语言，比如：我们可以将 testcase 作为一种约束，返回满足这种输入和输出或近似满足的代码片段，和基于自然语言搜索的方法结合应该会有更好的效果。
这次和结对编程的队友的合作主要是比较独立的分工，我们各自完成了不同的调研和复现，互相交流之后学习到了很多经验。