CN118035424A - 一种代码搜索方法、装置、电子设备及存储介质 - Google Patents

Abstract

本申请提供了一种代码搜索方法、装置、电子设备及存储介质，采用预训练编码器编码输入的代码与查询文本得到嵌入向量，由动量对比学习和跨模态匹配单元微调代码搜索任务；动量对比学习单元针对嵌入向量构造正样本与负样本，最大化正样本间相似度并最小化负样本间相似度；跨模态匹配单元以编码器输出的代码隐状态和查询隐状态为输入，判断是否在语义上匹配进行辅助学习；根据跨模态匹配与动量对比对应的损失函数确定目标损失函数训练代码搜索模型。在微调阶段改善负样本选择范围局限于单一批次的缺点，通过维护队列实现跨批次选择，学习更好的嵌入空间，跨模态匹配辅助任务通过交叉注意力融合代码和查询文本的信息，在嵌入层面弥补两者的语义差距。

Description

一种代码搜索方法、装置、电子设备及存储介质

技术领域

本申请涉及代码搜索技术领域，具体而言，涉及一种代码搜索方法、装置、电子设备及存储介质。

背景技术

代码搜索（code search）又称代码检索（code retrieval），是指输入自然语言查询，输出满足查询意图的代码片段的技术，开发人员通过学习和重用检索到的代码可以节省软件开发和维护的时间及精力，继而推动项目和系统的高效开发，正因如此，在软件工程和程序理解领域，代码搜索逐渐成为研究热点。

近年来，随着神经网络在软件工程等领域的广泛应用，基于深度学习特别是基于预训练模型的代码搜索也迅速成为研究热点，但是现有的代码预训练模型忽视了具体下游任务的微调优化，导致应用到代码搜索任务上时效果并非理想，此外，现有方法仅通过语义向量间的余弦相似度实现代码和自然语言的简单交互，难以弥补两者之间的语义鸿沟。

发明内容

本申请至少提供一种代码搜索方法、装置、电子设备及存储介质，可以在微调阶段通过动量对比学习改善普通对比学习的负样本选择范围局限于单一批次内的缺点, 通过维护队列以实现跨批次选择, 从而学习更好的嵌入空间，同时跨模态匹配辅助任务通过交叉注意力融合代码和查询的信息，在嵌入层面弥补两者的语义差距。

本申请提供了一种代码搜索方法，应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，所述方法包括：

获取查询文本以及源代码，通过所述预训练编码器编码所述查询文本对应的文本嵌入向量和动量文本嵌入向量，以及所述源代码对应的代码嵌入向量和动量代码嵌入向量；

采用所述动量对比学习单元，根据所述动量文本嵌入向量以及所述动量代码嵌入向量，分别构建所述代码嵌入向量与所述文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数；

将所述文本嵌入向量和所述代码嵌入向量输入至所述跨模态匹配单元，确定对应的融合嵌入表征，并根据所述融合嵌入表征确定所述跨模态匹配单元对应的跨模态匹配损失函数；

根据所述跨模态匹配损失函数以及所述动量对比损失函数确定目标损失函数，根据所述目标损失函数训练所述代码搜索模型；

将目标查询文本以及待搜索源代码输入至训练好的所述代码搜索模型，确定与所述目标查询文本匹配的目标源代码。

可选的，所述预训练编码器包括代码编码器、查询文本编码器、动量代码编码器以及动量查询文本编码器；

所述代码编码器用于编码所述代码嵌入向量，所述动量代码编码器用于编码所述动量代码嵌入向量；

所述查询文本编码器用于编码所述文本嵌入向量，所述动量查询文本编码器用于编码所述动量文本嵌入向量；

其中，所述代码编码器与所述查询文本编码器为利用反向传播进行梯度更新的编码器；

所述动量代码编码器与动量查询文本编码器为利用预设动量更新规则进行参数更新的编码器。

可选的，基于以下步骤构建正样本和负样本，具体包括：

根据所述动量文本嵌入向量，构建文本嵌入队列；

在代码-查询文本对比学习中，针对每个所述代码嵌入向量，正样本为该代码嵌入向量在数据集当前批次中有配对关系的所述动量文本嵌入向量，负样本为无配对关系的动量文本嵌入向量和所述文本嵌入队列中的全部嵌入向量；

根据所述动量代码嵌入向量，构建代码嵌入队列；

在查询文本-代码对比学习中，针对每个所述文本嵌入向量，正样本为该文本嵌入向量在数据集当前批次中有配对关系的所述动量代码嵌入向量，负样本为无配对关系的动量代码嵌入向量和所述代码嵌入队列中的全部嵌入向量。

可选的，所述动量对比学习单元还用于：

针对所述动量对比学习单元输出的每个动量嵌入向量，确定该动量嵌入向量与对应的正样本与全部负样本之间的相似度；

将所述相似度进行Softmax处理后确定为动量蒸馏的伪标签，作为参数更新编码器的辅助监督信号。

可选的，基于以下步骤确定所述动量对比损失函数：

确定所述代码-查询文本对比学习在动量对比学习中的第一对比损失函数，以及所述查询文本-代码对比学习在动量对比学习中的第二对比损失函数；

针对所述代码-查询文本对比学习，确定所述代码嵌入向量与对应的正样本与全部负样本之间的第一相似度预测标签；

根据所述动量代码嵌入向量与对应的正样本与全部负样本之间的所述伪标签，以及所述第一相似度预测标签确定所述代码-查询文本对比学习对应的第一动量蒸馏损失函数；

根据所述第一对比损失函数、所述第一动量蒸馏损失函数以及预设权重超参数，确定代码-查询文本对比损失函数；

针对所述查询文本-代码对比学习，确定所述文本嵌入向量与对应的正样本与全部负样本之间的第二相似度预测标签；

根据所述动量文本嵌入向量与对应的正样本与全部负样本之间的所述伪标签，以及所述第二相似度预测标签确定所述查询文本-代码对比学习对应的第二动量蒸馏损失函数；

根据所述第二对比损失函数、所述第二动量蒸馏损失函数以及预设权重超参数，确定查询文本-代码对比损失函数；

聚合所述代码-查询文本对比损失函数以及所述查询文本-代码对比损失函数，确定所述动量对比损失函数。

可选的，所述跨模态匹配单元采用Transformer编码器架构的代码预训练模型的最后6层；

其中，在多头注意力子层和前馈层之间增加交叉注意力子层。

本申请还提供一种代码搜索装置，应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，所述装置包括：

嵌入向量编码模块，用于获取查询文本以及源代码的嵌入向量，通过所述预训练编码器编码所述查询文本对应的文本嵌入向量和动量文本嵌入向量，以及所述源代码对应的代码嵌入向量和动量代码嵌入向量；

动量对比学习模块，用于采用所述动量对比学习模块，根据所述动量文本嵌入向量以及所述动量代码嵌入向量，分别构建所述代码嵌入向量与所述文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数；

跨模态匹配模块，用于将所述文本嵌入向量和所述代码嵌入向量输入至所述跨模态匹配模块，确定对应的融合嵌入表征，并根据所述融合嵌入表征确定所述跨模态匹配模块对应的跨模态匹配损失函数；

训练模块，用于根据所述跨模态匹配损失函数以及所述动量对比损失函数确定目标损失函数，根据所述目标损失函数训练所述代码搜索模型；

搜索模块，用于将目标查询文本以及待搜索源代码输入至训练好的所述代码搜索模型，确定与所述目标查询文本匹配的目标源代码。

本申请还提供一种电子设备，包括：处理器、存储器和总线，所述存储器存储有所述处理器可执行的机器可读指令，当电子设备运行时，所述处理器与所述存储器之间通过总线通信，所述机器可读指令被所述处理器执行时执行上述代码搜索方法，或上述代码搜索方法中任一种可能的实施方式中的步骤。

本申请还提供一种计算机可读存储介质，该计算机可读存储介质上存储有计算机程序，该计算机程序被处理器运行时执行上述代码搜索方法，或上述代码搜索方法中任一种可能的实施方式中的步骤。

本申请还提供一种计算机程序产品，包括计算机程序/指令，该计算机程序、指令被处理器执行时实现上述代码搜索方法，或上述代码搜索方法中任一种可能的实施方式中的步骤。

本申请提供的一种代码搜索方法、装置、电子设备及存储介质，应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，根据获取的查询文本以及源代码，通过预训练编码器编码查询文本对应的文本嵌入向量和动量文本嵌入向量，以及源代码对应的代码嵌入向量和动量代码嵌入向量；采用动量对比学习单元，根据动量文本嵌入向量以及动量代码嵌入向量，分别构建代码嵌入向量与文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数；将文本嵌入向量和代码嵌入向量输入至跨模态匹配单元，确定对应的融合嵌入表征，并根据融合嵌入表征确定跨模态匹配单元对应的跨模态匹配损失函数；根据跨模态匹配损失函数以及动量对比损失函数确定目标损失函数，根据目标损失函数训练代码搜索模型；将目标查询文本以及待搜索源代码输入至训练好的代码搜索模型，确定与目标查询文本匹配的目标源代码。可以在微调阶段通过动量对比学习改善普通对比学习的负样本选择范围局限于单一批次内的缺点, 通过维护队列以实现跨批次选择, 从而学习更好的嵌入空间，同时跨模态匹配辅助任务通过交叉注意力融合代码和查询的信息，在嵌入层面弥补两者的语义差距。

为使本申请的上述目的、特征和优点能更明显易懂，下文特举较佳，并配合所附附图，作详细说明如下。

附图说明

为了更清楚地说明本申请实施例的技术方案，下面将对实施例中所需要使用的附图作简单地介绍，此处的附图被并入说明书中并构成本说明书中的一部分，这些附图示出了符合本申请的实施例，并与说明书一起用于说明本申请的技术方案。应当理解，以下附图仅示出了本申请的某些实施例，因此不应被看作是对范围的限定，对于本领域普通技术人员来讲，在不付出创造性劳动的前提下，还可以根据这些附图获得其他相关的附图。

图1示出了本申请实施例所提供的一种代码搜索方法的流程图；

图2示出了本申请实施例所提供的一种构建正样本和负样本方法的流程图；

图3示出了本申请实施例所提供的一种代码搜索装置的示意图；

图4示出了本申请实施例所提供的一种电子设备的示意图。

具体实施方式

为使本申请实施例的目的、技术方案和优点更加清楚，下面将结合本申请实施例中附图，对本申请实施例中的技术方案进行清楚、完整地描述，显然，所描述的实施例仅仅是本申请一部分实施例，而不是全部的实施例。通常在此处附图中描述和示出的本申请实施例的组件可以以各种不同的配置来布置和设计。因此，以下对在附图中提供的本申请的实施例的详细描述并非旨在限制要求保护的本申请的范围，而是仅仅表示本申请的选定实施例。基于本申请的实施例，本领域技术人员在没有做出创造性劳动的前提下所获得的所有其他实施例，都属于本申请保护的范围。

应注意到：相似的标号和字母在下面的附图中表示类似项，因此，一旦某一项在一个附图中被定义，则在随后的附图中不需要对其进行进一步定义和解释。

本文中术语“和/或”，仅仅是描述一种关联关系，表示可以存在三种关系，例如，A和/或B，可以表示：单独存在A，同时存在A和B，单独存在B这三种情况。另外，本文中术语“至少一种”表示多种中的任意一种或多种中的至少两种的任意组合，例如，包括A、B、C中的至少一种，可以表示包括从A、B和C构成的集合中选择的任意一个或多个元素。

经研究发现，近年来，随着神经网络在软件工程等领域的广泛应用，基于深度学习特别是基于预训练模型的代码搜索也迅速成为研究热点，但是现有的代码预训练模型忽视了具体下游任务的微调优化，导致应用到代码搜索任务上时效果并非理想，此外，现有方法仅通过语义向量间的余弦相似度实现代码和自然语言的简单交互，难以弥补两者之间的语义鸿沟。

基于上述研究，本申请提供了一种代码搜索方法、装置、电子设备及存储介质，应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，根据获取的查询文本以及源代码，通过预训练编码器编码查询文本对应的文本嵌入向量和动量文本嵌入向量，以及源代码对应的代码嵌入向量和动量代码嵌入向量；采用动量对比学习单元，根据动量文本嵌入向量以及动量代码嵌入向量，分别构建代码嵌入向量与文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数；将文本嵌入向量和代码嵌入向量输入至跨模态匹配单元，确定对应的融合嵌入表征，并根据融合嵌入表征确定跨模态匹配单元对应的跨模态匹配损失函数；根据跨模态匹配损失函数以及动量对比学习单元对应的动量对比损失函数确定目标损失函数，根据目标损失函数训练代码搜索模型；将目标查询文本以及待搜索源代码输入至训练好的代码搜索模型，确定与目标查询文本匹配的目标源代码。可以在微调阶段通过动量对比学习改善普通对比学习的负样本选择范围局限于单一批次内的缺点, 通过维护队列以实现跨批次选择, 从而学习更好的嵌入空间，同时跨模态匹配辅助任务通过交叉注意力融合代码和查询的信息，在嵌入层面弥补两者的语义差距。

为便于对本实施例进行理解，首先对本申请实施例所公开的一种代码搜索方法进行详细介绍，本申请实施例所提供的代码搜索方法的执行主体一般为具有一定计算能力的计算机设备，该计算机设备例如包括：终端设备或服务器或其它处理设备，终端设备可以为用户设备（User Equipment，UE）、移动设备、用户终端、终端、蜂窝电话、无绳电话、个人数字助理（Personal Digital Assistant，PDA）、手持设备、计算设备、车载设备、可穿戴设备等。在一些可能的实现方式中，该代码搜索方法可以通过处理器调用存储器中存储的计算机可读指令的方式来实现。

参见图1所示，为本申请实施例提供的一种代码搜索方法的流程图，应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，所述方法包括步骤S101~S105，其中：

S101、获取查询文本以及源代码，通过所述预训练编码器编码所述查询文本对应的文本嵌入向量和动量文本嵌入向量，以及所述源代码对应的代码嵌入向量和动量代码嵌入向量。

在具体实施中，代码搜索任务可以用形式化的方法定义。代码搜索首先分别用代码的嵌入模型和自然语言查询的嵌入模型, 将源代码和查询文本转换为对应的嵌入向量。得到两种模态的向量后, 需要计算二者的相似度，相似度越高，说明该代码与查询文本的匹配程度越高，从而越有可能是符合要求的目标代码。在得到所有代码对应的相似度后，根据相似度大小降序排序，返回排序结果即为搜索结果。

在本申请实施例中，提供一种代码搜索微调阶段的代码搜索模型，分别为预训练编码器、动量对比学习模块和跨模态匹配模块，预训练编码器分为代码编码器和查询编码器，每种编码器存在对应的动量版本，总共形成四种编码器，它们对输入的代码或自然语言查询进行编码，得到嵌入向量。

具体的，预训练编码器包括代码编码器、查询文本编码器、动量代码编码器以及动量查询文本编码器；代码编码器用于编码代码嵌入向量，动量代码编码器用于编码动量代码嵌入向量；查询文本编码器用于编码文本嵌入向量，动量查询文本编码器用于编码动量文本嵌入向量；其中，代码编码器与查询文本编码器为利用反向传播进行梯度更新的编码器；动量代码编码器与动量查询文本编码器为利用预设动量更新规则进行参数更新的编码器。

这里，预训练编码器将输入的代码和自然语言查询编码为嵌入向量，便于后续处理，代码搜索模型中包含四种结构相同且初始化参数相同的编码器，分别为代码编码器、查询文本编码器、动量代码编码器以及动量查询文本编码器。前两者是利用反向传播进行梯度更新的编码器，而后两者是利用动量更新规则进行参数更新的编码器。

可选的，动量更新规则可以定义为：

其中，代表代码编码器或查询文本编码器的参数；/>代表对应的动量代码编码器或动量查询文本编码器的参数；/>为动量系数，值范围为/>。

这里，两种动量编码器在每次训练迭代中不通过梯度更新，新参数为自身参数和对应梯度更新编码器参数的加权和，权值大小由动量系数控制。通过设置一个较大的动量系数使动量编码器缓慢更新，使得短时间内动量编码器变化较小，编码出的向量具有较高的一致性，从而为后续动量对比学习带来益处。

作为一种可能的实施方式，代码搜索模型使用UniXcoder的预训练参数初始化全部四种编码器，在训练过程中，代码和查询编码器共享参数，动量代码编码器和动量查询编码器共享参数。

示例性的，给定一对经过分词处理后的源代码-查询文本文本样本 (，)，其中，/>，/>，/>和/>分别为代码和查询的词元（token）数。依照UniXcoder的设置，构造的输入形式为/>= {[CLS]，[ENC]，[SEP]，/>，[SEP]} 和/>= {[CLS]， [ENC]，[SEP]，/>，[SEP]}。其中，[CLS]是序列开始标记，[SEP]是分隔标记，[ENC]是设置UniXcoder为encoder-only模式的标记。/>被输入代码编码器/>和动量代码编码器/>，/>被输入查询编码器/>和动量查询编码器/>，分别得到编码后的隐状态，即为代码嵌入向量和动量代码嵌入向量、文本嵌入向量和动量文本嵌入向量。

S102、采用所述动量对比学习单元，根据所述动量文本嵌入向量以及所述动量代码嵌入向量，分别构建所述代码嵌入向量与所述文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数。

在具体实施中，对比学习是一类有助于学习更好嵌入向量的方法，对比学习的基本思想是将语义匹配的正样本之间的距离拉近，同时将语义不匹配的负样本之间的距离拉远，其核心在于如何构造正样本和负样本。

然而，若只在同一批次中选取负样本，会受限于批次的大小，使得模型只在一个较小范围内学习负样本的特征，导致训练的不充分性。在代码搜索的推理阶段，每一个查询都需要和代码库中所有的代码计算相似度，尝试判别它们是否为需要被检索的代码，这要求模型能够识别大量的负样本，为了弥补训练和推理时的差距，在本申请实施例中使用动量对比学习的方式，引入动量编码器和队列。

具体的，参见图2所示，为本申请实施例提供的一种构建正样本和负样本方法的流程图，所述方法包括步骤S1021~S1024，其中：

S1021、根据所述动量文本嵌入向量，构建文本嵌入队列。

S1022、在代码-查询文本对比学习中，针对每个所述代码嵌入向量，正样本为该代码嵌入向量在数据集当前批次中有配对关系的所述动量文本嵌入向量，负样本为无配对关系的动量文本嵌入向量和所述文本嵌入队列中的全部嵌入向量。

S1023、根据所述动量代码嵌入向量，构建代码嵌入队列。

S1024、在查询文本-代码对比学习中，针对每个所述文本嵌入向量，正样本为该文本嵌入向量在数据集当前批次中有配对关系的所述动量代码嵌入向量，负样本为无配对关系的动量代码嵌入向量和所述代码嵌入队列中的全部嵌入向量。

在具体实施中，动量编码器缓慢更新的特性确保了邻近的批次之间编码出的嵌入向量拥有较高的一致性，本申请实施例使用队列保存每个批次的动量嵌入向量, 以便于后续批次能够在该队列中选择负样本, 从而实现了跨批次的对比学习。

经过代码编码器、动量代码编码器、查询文本编码器以及动量查询文本编码器四种编码器编码后，每一对代码-查询文本样本会被转换为代码嵌入向量和动量代码嵌入向量、文本嵌入向量和动量文本嵌入向量四种嵌入特征，假定训练的批次大小为，该批次中第/>对样本对应的代码嵌入向量和动量代码嵌入向量、文本嵌入向量和动量文本嵌入向量依次分别为/>，/>，/>和/>，其中/>。且这些嵌入经过了L2归一化，L2范数为1，于是可以利用向量点积计算余弦相似度。

这里，动量对比学习由两个对称的部分代码-查询文本对比学习（C2Q）和查询文本-代码对比学习（Q2C）组成，在本申请实施例中，同时维护文本嵌入队列以及代码嵌入队列。

具体的，代码-查询文本对比学习中，对于某个代码的代码嵌入向量，其正样本为与之匹配的查询文本对应的动量文本嵌入向量/>，而其余的动量文本嵌入向量/>(以及文本嵌入队列中所有的嵌入（表示为/>，/>为队列大小）都成为其负样本。

进一步的，可以使用InfoNCE计算代码-查询文本对比学习在动量对比学习中的第一对比损失函数，如下公式所示：

其中，代表代码-查询文本对比学习在动量对比学习中的第一对比损失；/>代表向量点积；/>代表第/>个代码嵌入向量，/>，B代表训练的批次大小；/>代表与代码嵌入向量/>相匹配的动量文本嵌入向量；/>代表当前批次中的第/>个动量文本嵌入向量，/>；/>代表文本嵌入队列中的第/>个动量文本嵌入向量，，/>为队列大小；/>代表温度超参数，用于控制相似度的平滑程度。

这里，第一对比损失函数旨在提高正样本对的相似度，拉近嵌入空间中的距离，同时降低负样本的相似度，推远嵌入空间中的距离。每个批次的损失计算完毕后，当前批次的全部动量文本嵌入向量加入到文本嵌入队列中，同时将文本嵌入队列中最先入队的B个嵌入出队并舍去，完成队列的更新。这种更新机制使得队列中存储的都是最近若干批次的嵌入，保证了一致性。

进一步的，对于查询文本-代码对比学习对应的第二对比损失函数，其计算方式与代码-查询文本对比学习类似，具体过程不再进行赘述。

作为一种可能的实施方式，对于模型训练时的代码而言，通常的做法是将与当前代码配对的查询作为唯一的正样本，而其他所有查询均作为需要被推远的负样本，实际上，这些查询负样本中与当前代码也存在不同程度的语义相关性，如果强迫模型在对比学习中去学习固定的标签，认为只有正样本是完全匹配的，而其他负样本是完全不匹配的，会导致所有负样本被无差别地推远，并不符合实际情况，故采用动量蒸馏方法加以解决。

具体的，针对动量对比学习单元输出的每个动量嵌入向量，确定该动量嵌入向量与对应的正样本与全部负样本之间的相似度；将所述相似度进行Softmax处理后确定为动量蒸馏的伪标签，作为参数更新编码器的辅助监督信号。

这里，默认的相似度标签中，正样本对的标签为1，负样本对的标签为0。而在动量蒸馏中，动量编码器被视为教师模型，参数更新编码器被视为学生模型，教师模型的输出作为指导学生模型学习的另一种标签。具体而言，动量编码器编码的嵌入与所有正样本和负样本的相似度经过Softmax处理后将被作为伪标签，其值介于0和1之间，作为参数更新编码器的辅助监督信号。

进一步的，动量对比学习单元对应的动量对比损失函数可以通过如下方式实现：针对代码-查询文本对比学习，确定代码嵌入向量与对应的正样本与全部负样本之间的第一相似度预测标签，以及代码-查询文本对比学习在动量对比学习中的第一对比损失函数；根据动量代码嵌入向量与对应的正样本与全部负样本之间的伪标签，以及第一相似度预测标签确定代码-查询文本对比学习对应的第一动量蒸馏损失函数；根据第一对比损失函数、第一动量蒸馏损失函数以及预设权重超参数，确定代码-查询文本对比损失函数。

针对查询文本-代码对比学习，确定文本嵌入向量与对应的正样本与全部负样本之间的第二相似度预测标签，以及查询文本-代码对比学习在动量对比学习中的第二对比损失函数；根据动量文本嵌入向量与对应的正样本与全部负样本之间的伪标签，以及第二相似度预测标签确定查询文本-代码对比学习对应的第二动量蒸馏损失函数；根据第二对比损失函数、第二动量蒸馏损失函数以及预设权重超参数，确定查询文本-代码对比损失函数；聚合代码-查询文本对比损失函数以及查询文本-代码对比损失函数，确定动量对比损失函数。

示例性的，以代码-查询文本对比学习为例，动量代码嵌入向量与对应的正样本与全部负样本之间的伪标签可以通过如下公式计算：

其中，代表当前批次中第/>个动量代码嵌入向量与第/>个正样本或负样本之间的伪标签, 此处/>，/>，B代表训练的批次大小；/>为队列大小；/>代表第/>个动量代码嵌入向量；/>代表当前批次中的第/>个动量文本嵌入向量，；/>代表文本嵌入队列中的第/>个动量文本嵌入向量，/>；代表/>和/>的并集中的第/>个向量。/>代表温度超参数。

进一步的，相似度预测标签可以根据如下公式进行计算：

其中，代表当前批次中第/>个代码嵌入向量与第/>个正样本或负样本之间的相似度预测标签, 此处/>，/>，B代表训练的批次大小；/>为队列大小；/>代表第/>个代码嵌入向量；/>代表当前批次中的第/>个动量文本嵌入向量，；/>代表文本嵌入队列中的第/>个动量文本嵌入向量，/>；代表/>和/>的并集中的第/>个向量。/>代表温度超参数。

进一步的，第一动量蒸馏损失函数可以通过如下公式进行计算：

其中，代表第一动量蒸馏损失；/>代表相似度预测标签；/>代表动量代码嵌入向量与对应的正样本或负样本之间的伪标签；B代表训练的批次大小；/>为队列大小。

进一步的，使用预设的权重超参数将第一动量蒸馏损失函数与第一对比损失函数组合在一起加权求和构成代码-查询文本对比损失函数。可以通过如下公式计算：

需要说明的是，查询文本-代码对比损失的计算过程中，伪标签、相似度预测标签、第二动量蒸馏损失函数、第二对比损失函数的计算方式以及加权求和方式与代码-查询文本对比学习类似，对此不再进行赘述。

这里，动量对比学习单元最终的动量对比损失函数可以通过如下公式计算：

其中，代表最终的动量对比损失函数；/>代表代码-查询文本对比损失函数；/>代表查询文本-代码对比损失函数。

S103、将所述文本嵌入向量和所述代码嵌入向量输入至所述跨模态匹配单元，确定对应的融合嵌入表征，并根据所述融合嵌入表征确定所述跨模态匹配单元对应的跨模态匹配损失函数。

在具体实施中，在对比学习中，代码和查询文本的嵌入向量分别通过独立的编码过程得到，然后能够计算余弦相似度以衡量两种嵌入之间的语义距离，这种架构称为双编码器架构，但是代码和查询文本的嵌入向量只在最后的相似度计算阶段显式地产生交互，仅依靠简单的余弦相似度不足以弥补两种模态之间的语义差距。

在本申请实施例中，采用跨模态匹配单元，利用单编码器架构进行补充，代码和文本嵌入向量被一同输入到联合编码器中, 在其内部进行特征融合与交互，并输出融合嵌入，该融合嵌入同时包含了两种输入的语义信息，能够被进行二分类，判断两种输入在语义上是否匹配。

这里，跨模态匹配单元采用Transformer编码器架构的代码预训练模型的最后6层；其中，在多头注意力子层和前馈层之间增加交叉注意力子层。

具体的，本申请实施例选取Transformer编码器架构的代码预训练模型（如UniXcoder）的最后6层。每一层中，在多头注意力子层和前馈层之间增加交叉注意力子层，构造出类似于Transformer解码器的结构，以此作为上述提到的联合编码器。

这样，引入交叉注意力子层能够同时输入代码和查询的信息，实现语义交互。此外，该子层的初始化参数与多头注意力子层的预训练参数相同，能够从预训练中受益。

在具体实施中，动量对比学习模块中由梯度更新编码器输出的隐状态和/>作为跨模态匹配模块的输入。/>(/>为该查询的词元数量，/>为Transformer隐状态维度) 被输入多头注意力子层，计算每一个注意力头：

其中，代表注意力头；/>标识第/>个注意力头，/>表示每个头的隐状态维度，假设头的数量为/>，则满足/>=d，设置/>，/>。而/>,/>和/>分别表示第/>个头的query，key和value(维度均为/>)，且/>,/>和/>(维度均为/>) 为其对应的映射矩阵，满足/>，/>以及/>。最后，将所有/>数量的头拼接在一起，并通过另一个线性映射/>得到多头注意力子层的输出/>:

这里，交叉注意力子层与多头注意力子层的结构相同，执行注意力计算时的区别在于输入来源不同，其中中的/>被替换为上一交叉注意力子层的输出/>，而/>和/>中的/>均被替换为/>。简而言之，query源于自然语言查询输入的隐状态，而key和value源于代码输入的隐状态。

当经过6层联合编码器的计算后得到最后的输出，选取首个词元[CLS]标记的嵌入/>作为该融合嵌入的最终表征。本申请实施例将融合嵌入表征/>经过线性变换，并通过Softmax转换为二分类预测标签/>。同时假设/>表示真实标签，其中/>表示代码-查询文本对的序号，/>表示类别，且0代表错误匹配，1代表正确匹配。然后使用交叉熵计算跨模态匹配的损失：

其中，代表跨模态匹配损失函数；/>代表训练的批次大小。

S104、根据所述跨模态匹配损失函数以及所述动量对比损失函数确定目标损失函数，根据所述目标损失函数训练所述代码搜索模型。

在具体实施中，目标损失函数通过如下公式计算：

其中，代表目标损失函数；/>代表最终的动量对比损失函数；/>代表跨模态匹配损失函数；/>代表控制对比损失和匹配损失相对权重的超参数。

S105、将目标查询文本以及待搜索源代码输入至训练好的所述代码搜索模型，确定与所述目标查询文本匹配的目标源代码。

本申请实施例提供的一种代码搜索方法，应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，根据获取的查询文本以及源代码，通过预训练编码器编码查询文本对应的文本嵌入向量和动量文本嵌入向量，以及源代码对应的代码嵌入向量和动量代码嵌入向量；采用动量对比学习单元，根据动量文本嵌入向量以及动量代码嵌入向量，分别构建代码嵌入向量与文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数；将文本嵌入向量和代码嵌入向量输入至跨模态匹配单元，确定对应的融合嵌入表征，并根据融合嵌入表征确定跨模态匹配单元对应的跨模态匹配损失函数；根据跨模态匹配损失函数以及动量对比学习单元对应的动量对比损失函数确定目标损失函数，根据目标损失函数训练代码搜索模型；将目标查询文本以及待搜索源代码输入至训练好的代码搜索模型，确定与目标查询文本匹配的目标源代码。可以在微调阶段通过动量对比学习改善普通对比学习的负样本选择范围局限于单一批次内的缺点, 通过维护队列以实现跨批次选择, 从而学习更好的嵌入空间，同时跨模态匹配辅助任务通过交叉注意力融合代码和查询的信息，在嵌入层面弥补两者的语义差距。

本领域技术人员可以理解，在具体实施方式的上述方法中，各步骤的撰写顺序并不意味着严格的执行顺序而对实施过程构成任何限定，各步骤的具体执行顺序应当以其功能和可能的内在逻辑确定。

基于同一发明构思，本申请实施例中还提供了与代码搜索方法对应的代码搜索装置，由于本申请实施例中的装置解决问题的原理与本申请实施例上述代码搜索方法相似，因此装置的实施可以参见方法的实施，重复之处不再赘述。

请参阅图3，图3为本申请实施例提供的一种代码搜索装置的示意图。如图3中所示，本申请实施例提供的代码搜索装置300应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，包括：

嵌入向量编码模块310，用于获取查询文本以及源代码，通过所述预训练编码器编码所述查询文本对应的文本嵌入向量和动量文本嵌入向量，以及所述源代码对应的代码嵌入向量和动量代码嵌入向量。

动量对比学习模块320，用于采用所述动量对比学习模块，根据所述动量文本嵌入向量以及所述动量代码嵌入向量，分别构建所述代码嵌入向量与所述文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数。

跨模态匹配模块330，用于将所述文本嵌入向量和所述代码嵌入向量输入至所述跨模态匹配模块，确定对应的融合嵌入表征，并根据所述融合嵌入表征确定所述跨模态匹配模块对应的跨模态匹配损失函数。

训练模块340，用于根据所述跨模态匹配损失函数以及所述动量对比学习模块对应的动量对比损失函数确定目标损失函数，根据所述目标损失函数训练所述代码搜索模型。

搜索模块350，用于将目标查询文本以及待搜索源代码输入至训练好的所述代码搜索模型，确定与所述目标查询文本匹配的目标源代码。

关于装置中的各模块的处理流程、以及各模块之间的交互流程的描述可以参照上述方法实施例中的相关说明，这里不再详述。

本申请实施例提供的一种代码搜索装置，应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，根据获取的查询文本以及源代码，通过预训练编码器编码查询文本对应的文本嵌入向量和动量文本嵌入向量，以及源代码对应的代码嵌入向量和动量代码嵌入向量；采用动量对比学习单元，根据动量文本嵌入向量以及动量代码嵌入向量，分别构建代码嵌入向量与文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数；将文本嵌入向量和代码嵌入向量输入至跨模态匹配单元，确定对应的融合嵌入表征，并根据融合嵌入表征确定跨模态匹配单元对应的跨模态匹配损失函数；根据跨模态匹配损失函数以及动量对比学习单元对应的动量对比损失函数确定目标损失函数，根据目标损失函数训练代码搜索模型；将目标查询文本以及待搜索源代码输入至训练好的代码搜索模型，确定与目标查询文本匹配的目标源代码。可以在微调阶段通过动量对比学习改善普通对比学习的负样本选择范围局限于单一批次内的缺点, 通过维护队列以实现跨批次选择, 从而学习更好的嵌入空间，同时跨模态匹配辅助任务通过交叉注意力融合代码和查询的信息，在嵌入层面弥补两者的语义差距。

对应于图1中的代码搜索方法，本申请实施例还提供了一种电子设备400，如图4所示，为本申请实施例提供的电子设备400结构示意图，包括：

处理器41、存储器42、和总线43；存储器42用于存储执行指令，包括内存421和外部存储器422；这里的内存421也称内存储器，用于暂时存放处理器41中的运算数据，以及与硬盘等外部存储器422交换的数据，处理器41通过内存421与外部存储器422进行数据交换，当所述电子设备400运行时，所述处理器41与所述存储器42之间通过总线43通信，使得所述处理器41执行图1中的代码搜索方法的步骤。

本申请实施例还提供一种计算机可读存储介质，该计算机可读存储介质上存储有计算机程序，该计算机程序被处理器运行时执行上述方法实施例中所述的代码搜索方法的步骤。其中，该存储介质可以是易失性或非易失的计算机可读取存储介质。

本申请实施例还提供一种计算机程序产品，该计算机程序产品包括有计算机指令，所述计算机指令被处理器执行时可以执行上述方法实施例中所述的代码搜索方法的步骤，具体可参见上述方法实施例，在此不再赘述。

其中，上述计算机程序产品可以具体通过硬件、软件或其结合的方式实现。在一个可选实施例中，所述计算机程序产品具体体现为计算机存储介质，在另一个可选实施例中，计算机程序产品具体体现为软件产品，例如软件开发包（Software Development Kit，SDK）等等。

所属领域的技术人员可以清楚地了解到，为描述的方便和简洁，上述描述的装置的具体工作过程，可以参考前述方法实施例中的对应过程，在此不再赘述。在本申请所提供的几个实施例中，应该理解到，所揭露的装置和方法，可以通过其它的方式实现。以上所描述的装置实施例仅仅是示意性的，例如，所述单元的划分，仅仅为一种逻辑功能划分，实际实现时可以有另外的划分方式，又例如，多个单元或组件可以结合或者可以集成到另一个系统，或一些特征可以忽略，或不执行。另一点，所显示或讨论的相互之间的耦合或直接耦合或通信连接可以是通过一些通信接口，装置或单元的间接耦合或通信连接，可以是电性，机械或其它的形式。

所述作为分离部件说明的单元可以是或者也可以不是物理上分开的，作为单元显示的部件可以是或者也可以不是物理单元，即可以位于一个地方，或者也可以分布到多个网络单元上。可以根据实际的需要选择其中的部分或者全部单元来实现本实施例方案的目的。

另外，在本申请各个实施例中的各功能单元可以集成在一个处理单元中，也可以是各个单元单独物理存在，也可以两个或两个以上单元集成在一个单元中。

所述功能如果以软件功能单元的形式实现并作为独立的产品销售或使用时，可以存储在一个处理器可执行的非易失的计算机可读取存储介质中。基于这样的理解，本申请的技术方案本质上或者说对现有技术做出贡献的部分或者该技术方案的部分可以以软件产品的形式体现出来，该计算机软件产品存储在一个存储介质中，包括若干指令用以使得一台计算机设备（可以是个人计算机，服务器，或者网络设备等）执行本申请各个实施例所述方法的全部或部分步骤。而前述的存储介质包括：U盘、移动硬盘、只读存储器（Read-OnlyMemory，ROM）、随机存取存储器（Random Access Memory，RAM）、磁碟或者光盘等各种可以存储程序代码的介质。

最后应说明的是：以上所述实施例，仅为本申请的具体实施方式，用以说明本申请的技术方案，而非对其限制，本申请的保护范围并不局限于此，尽管参照前述实施例对本申请进行了详细的说明，本领域的普通技术人员应当理解：任何熟悉本技术领域的技术人员在本申请揭露的技术范围内，其依然可以对前述实施例所记载的技术方案进行修改或可轻易想到变化，或者对其中部分技术特征进行等同替换；而这些修改、变化或者替换，并不使相应技术方案的本质脱离本申请实施例技术方案的精神和范围，都应涵盖在本申请的保护范围之内。因此，本申请的保护范围应所述以权利要求的保护范围为准。

Claims (9)

1.一种代码搜索方法，其特征在于，应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，所述方法包括：

获取查询文本以及源代码，通过所述预训练编码器编码所述查询文本对应的文本嵌入向量和动量文本嵌入向量，以及所述源代码对应的代码嵌入向量和动量代码嵌入向量；

采用所述动量对比学习单元，根据所述动量文本嵌入向量以及所述动量代码嵌入向量，分别构建所述代码嵌入向量与所述文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数；

将所述文本嵌入向量和所述代码嵌入向量输入至所述跨模态匹配单元，确定对应的融合嵌入表征，并根据所述融合嵌入表征确定所述跨模态匹配单元对应的跨模态匹配损失函数；

根据所述跨模态匹配损失函数以及所述动量对比损失函数确定目标损失函数，根据所述目标损失函数训练所述代码搜索模型；

将目标查询文本以及待搜索源代码输入至训练好的所述代码搜索模型，确定与所述目标查询文本匹配的目标源代码。

2.根据权利要求1所述的方法，其特征在于，所述预训练编码器包括代码编码器、查询文本编码器、动量代码编码器以及动量查询文本编码器；

所述代码编码器用于编码所述代码嵌入向量，所述动量代码编码器用于编码所述动量代码嵌入向量；

所述查询文本编码器用于编码所述文本嵌入向量，所述动量查询文本编码器用于编码所述动量文本嵌入向量；

其中，所述代码编码器与所述查询文本编码器为利用反向传播进行梯度更新的编码器；

所述动量代码编码器与动量查询文本编码器为利用预设动量更新规则进行参数更新的编码器。

3.根据权利要求1所述的方法，其特征在于，基于以下步骤构建正样本和负样本，具体包括：

根据所述动量文本嵌入向量，构建文本嵌入队列；

在代码-查询文本对比学习中，针对每个所述代码嵌入向量，正样本为该代码嵌入向量在数据集当前批次中有配对关系的所述动量文本嵌入向量，负样本为无配对关系的动量文本嵌入向量和所述文本嵌入队列中的全部嵌入向量；

根据所述动量代码嵌入向量，构建代码嵌入队列；

在查询文本-代码对比学习中，针对每个所述文本嵌入向量，正样本为该文本嵌入向量在数据集当前批次中有配对关系的所述动量代码嵌入向量，负样本为无配对关系的动量代码嵌入向量和所述代码嵌入队列中的全部嵌入向量。

4.根据权利要求3所述的方法，其特征在于，所述动量对比学习单元还用于：

针对所述动量对比学习单元输出的每个动量嵌入向量，确定该动量嵌入向量与对应的正样本与全部负样本之间的相似度；

将所述相似度进行Softmax处理后确定为动量蒸馏的伪标签，作为参数更新编码器的辅助监督信号。

5.根据权利要求4所述的方法，其特征在于，基于以下步骤确定所述动量对比损失函数：

确定所述代码-查询文本对比学习在动量对比学习中的第一对比损失函数，以及所述查询文本-代码对比学习在动量对比学习中的第二对比损失函数；

针对所述代码-查询文本对比学习，确定所述代码嵌入向量与对应的正样本与全部负样本之间的第一相似度预测标签；

根据所述动量代码嵌入向量与对应的正样本与全部负样本之间的所述伪标签，以及所述第一相似度预测标签确定所述代码-查询文本对比学习对应的第一动量蒸馏损失函数；

根据所述第一对比损失函数、所述第一动量蒸馏损失函数以及预设权重超参数，确定代码-查询文本对比损失函数；

针对所述查询文本-代码对比学习，确定所述文本嵌入向量与对应的正样本与全部负样本之间的第二相似度预测标签；

根据所述动量文本嵌入向量与对应的正样本与全部负样本之间的所述伪标签，以及所述第二相似度预测标签确定所述查询文本-代码对比学习对应的第二动量蒸馏损失函数；

根据所述第二对比损失函数、所述第二动量蒸馏损失函数以及预设权重超参数，确定查询文本-代码对比损失函数；

聚合所述代码-查询文本对比损失函数以及所述查询文本-代码对比损失函数，确定所述动量对比损失函数。

6.根据权利要求1所述的方法，其特征在于：

所述跨模态匹配单元采用Transformer编码器架构的代码预训练模型的最后6层；

其中，在多头注意力子层和前馈层之间增加交叉注意力子层。

7.一种代码搜索装置，其特征在于，应用于包括预训练编码器、动量对比学习单元以及跨模态匹配单元的代码搜索模型，所述装置包括：

嵌入向量编码模块，用于获取查询文本以及源代码的嵌入向量，通过所述预训练编码器编码所述查询文本对应的文本嵌入向量和动量文本嵌入向量，以及所述源代码对应的代码嵌入向量和动量代码嵌入向量；

动量对比学习模块，用于采用所述动量对比学习模块，根据所述动量文本嵌入向量以及所述动量代码嵌入向量，分别构建所述代码嵌入向量与所述文本嵌入向量对应的正样本及负样本进行动量对比学习，在最大化正样本之间的相似度的同时，最小化负样本之间的相似度，并确定对应的动量对比损失函数；

跨模态匹配模块，用于将所述文本嵌入向量和所述代码嵌入向量输入至所述跨模态匹配模块，确定对应的融合嵌入表征，并根据所述融合嵌入表征确定所述跨模态匹配模块对应的跨模态匹配损失函数；

训练模块，用于根据所述跨模态匹配损失函数以及所述动量对比损失函数确定目标损失函数，根据所述目标损失函数训练所述代码搜索模型；

搜索模块，用于将目标查询文本以及待搜索源代码输入至训练好的所述代码搜索模型，确定与所述目标查询文本匹配的目标源代码。

8.一种电子设备，其特征在于，包括：处理器、存储器和总线，所述存储器存储有所述处理器可执行的机器可读指令，当电子设备运行时，所述处理器与所述存储器之间通过总线通信，所述机器可读指令被所述处理器执行时执行如权利要求1至6中任一项所述的代码搜索方法的步骤。

9.一种计算机可读存储介质，其特征在于，该计算机可读存储介质上存储有计算机程序，该计算机程序被处理器运行时执行如权利要求1至6中任一项所述的代码搜索方法的步骤。