其他分享
首页 > 其他分享> > Code Search

Code Search

作者:互联网

简介

本次作业的主题是 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 模型实现来训练一个可以在句子中预测下一个词的模型,这里的训练语料可以来自和我们任务相关的数据,比如 StackOverflow 上面的问答,这里为了简单使用训练集中的 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 作为一种约束,返回满足这种输入和输出或近似满足的代码片段,和基于自然语言搜索的方法结合应该会有更好的效果。

这次和结对编程的队友的合作主要是比较独立的分工,我们各自完成了不同的调研和复现,互相交流之后学习到了很多经验。

 

标签:Search,Code,训练,代码,docstring,encoder,模型
来源: https://www.cnblogs.com/jiyan-he/p/11666791.html