A code retrieval model training method, code retrieval method and device

Abstract

The embodiment of the present invention provides a training method, a code retrieval method and a device for a code retrieval model, and relates to the field of data processing technology. The training method of the code retrieval model includes: obtaining a code index and a corresponding similarity index in a target code library; analyzing the code index to obtain at least one corresponding code retrieval query information; using a dual-tower model to be trained to encode the code retrieval query information, the code index and the similarity index respectively, to obtain a first feature vector, a second feature vector and a third feature vector; calculating a first similarity between the first feature vector and the second feature vector and a second similarity between the first feature vector and the third feature vector; calculating the contrast loss of the dual-tower model to be trained based on the first similarity and the second similarity, and adjusting the parameters of the dual-tower model to be trained according to the contrast loss to obtain a trained code retrieval model. According to the embodiment of the present invention, the accuracy of code retrieval is improved.

Description

Technical Field

[0001] The present invention relates to the field of data processing technology, and in particular to a code retrieval model training method, a code retrieval method and a device. Background Art

[0002] Retrieval-augmented Generation (RAG) is a model that combines retrieval and generation technologies. It generates answers or content by referencing information from an external knowledge base. It has strong interpretability and customization capabilities and is suitable for multiple natural language processing tasks such as question-answering systems, document generation, and intelligent assistants. RAG is based on a large language model (LLM) and combines a vector database with an embedding model. Specifically, RAG's embedding model encodes the content in the external knowledge base and stores the encoded feature vectors through a vector database. After the query to be retrieved is input into the LLM, the stored feature vectors are used as the context of the LLM to guide the LLM to generate the corresponding response information, and the output of the LLM is optimized so that it can reference a knowledge base outside the training data (i.e., an external knowledge base) to assist in generating responses.

[0003] Although the development of RAG enables LLM to connect to external knowledge bases to assist in generating responses, the retrieval accuracy of RAG in some areas of professional knowledge is not satisfactory. For example, when processing programming tasks, such as code retrieval, the query information entered by the user is a natural language description, and RAG's Embedding model (i.e., embedding model) needs to retrieve a piece of code after encoding the query information. Due to the difference in structure between code and natural language description, RAG cannot effectively understand the meaning of the natural language description when processing it, resulting in the generation of code retrieval results with different vector representations, which affects the accuracy of the code retrieval results. Summary of the invention

[0004] The purpose of the embodiments of the present invention is to provide a code retrieval model training method, a code retrieval method and a device to improve the accuracy of code retrieval. The specific technical solution is as follows:

[0005] In a first aspect, an embodiment of the present invention provides a method for training a code retrieval model, the method comprising:

[0006] Obtaining a code index in a target code library and a similarity index corresponding to the code index; the code index is information corresponding to a code block in the target code library, including at least a function name corresponding to the code block and comment information of the function; the similarity index is information corresponding to a code block whose similarity to the code index meets a preset condition;

[0007] Analyze the code index to obtain at least one code search query information corresponding to the code index; the number of the code search query information is the same as the number of the similarity indexes;

[0008] Using the dual-tower model to be trained, the code retrieval query information is encoded to obtain a first feature vector, the code index is encoded to obtain a second feature vector, and the similarity index is encoded to obtain a third feature vector;

[0009] Calculating a first similarity between the first feature vector and the second feature vector, and a second similarity between the first feature vector and the third feature vector;

[0010] Based on the first similarity and the second similarity, the contrast loss of the dual-tower model to be trained is calculated, and the parameters of the dual-tower model to be trained are adjusted according to the contrast loss to obtain a trained code retrieval model.

[0011] Optionally, analyzing the code index to obtain at least one code search query information corresponding to the code index includes:

[0012] The code index is input into a pre-trained large model, so that the pre-trained large model analyzes the code index based on a specified prompt word to obtain at least one code retrieval query information corresponding to the code index.

[0013] Optionally, the using the dual-tower model to be trained to encode the code retrieval query information to obtain a first feature vector, encoding the code index to obtain a second feature vector, and encoding the similarity index to obtain a third feature vector includes:

[0014] Encoding the code retrieval query information using a first encoding network of the dual-tower model to be trained to obtain a first feature vector;

[0015] Encoding the code index using a second encoding network of the dual-tower model to be trained to obtain a second feature vector;

[0016] Using a third encoding network of the dual-tower model to be trained to encode the similarity index to obtain a third feature vector;

[0017] The network parameters of the first coding network and the second coding network are different, and the network parameters of the second coding network and the third coding network are the same.

[0018] Optionally, the contrast loss of the dual-tower model to be trained is calculated using the following expression:

[0019] ;

[0020] in, represents the loss value of contrast loss, Indicates the number of code retrieval query information, To specify a constant value range of (0,2), Indicates The feature vector corresponding to the code retrieval query information, Indicates The feature vector of the code index corresponding to the code retrieval query information, Indicates The code retrieves the query information and The similarity between the code indexes corresponding to the code retrieval query information, Indicates The feature vector of similarity index corresponding to the code retrieval query information, Indicates The code retrieves the query information and The similarity between similar indexes corresponding to the code retrieval query information.

[0021] Optionally, the parameters of the dual-tower model to be trained are updated using the following expressions:

[0022] ;

[0023] in, represents the parameters of the dual-tower model to be trained, represents the learning rate, Represents the gradient of the loss function corresponding to the contrast loss with respect to the parameters of the dual-tower model to be trained.

[0024] In a second aspect, an embodiment of the present invention provides a code retrieval method, the method comprising:

[0025] Get the code retrieval query information to be retrieved;

[0026] Inputting the code search query information to be searched into a pre-trained code search model to perform code search, and obtaining a target code index; wherein the pre-trained code search model is trained using the code search model training method described in the first aspect above; the target code index is information corresponding to the retrieved code block, including at least a function name corresponding to the code block and comment information of the function;

[0027] The corresponding target code is obtained based on the target code index.

[0028] In a third aspect, an embodiment of the present invention provides a training device for a code retrieval model, the device comprising:

[0029] A first acquisition module is used to acquire a code index in a target code library and a similarity index corresponding to the code index; the code index is information corresponding to a code block in the target code library, including at least a function name corresponding to the code block and comment information of the function; the similarity index is information corresponding to a code block whose similarity to the code index meets a preset condition;

[0030] An analysis module, configured to analyze the code index to obtain at least one code search query information corresponding to the code index; the number of the code search query information is the same as the number of the similarity indexes;

[0031] An encoding module, used for encoding the code retrieval query information using the dual-tower model to be trained to obtain a first feature vector, encoding the code index to obtain a second feature vector, and encoding the similarity index to obtain a third feature vector;

[0032] A calculation module, used for calculating a first similarity between the first feature vector and the second feature vector, and a second similarity between the first feature vector and the third feature vector;

[0033] A training module is used to calculate the contrast loss of the dual-tower model to be trained based on the first similarity and the second similarity, and adjust the parameters of the dual-tower model to be trained according to the contrast loss to obtain a trained code retrieval model.

[0034] In a fourth aspect, an embodiment of the present invention provides a code retrieval device, the device comprising:

[0035] The second acquisition module is used to obtain the code search query information to be searched;

[0036] A retrieval module, used for inputting the code retrieval query information to be retrieved into a pre-trained code retrieval model to perform code retrieval and obtain a target code index; wherein the pre-trained code retrieval model is trained by the code retrieval model training method described in the first aspect above; the target code index is information corresponding to the retrieved code block, including at least a function name corresponding to the code block and comment information of the function;

[0037] The third acquisition module is used to acquire the corresponding target code based on the target code index.

[0038] In a fifth aspect, an embodiment of the present invention provides an electronic device, including a processor, a communication interface, a memory and a communication bus, wherein the processor, the communication interface and the memory communicate with each other through the communication bus;

[0039] Memory, used to store computer programs;

[0040] The processor is used to implement any of the above methods when executing the program stored in the memory.

[0041] In a sixth aspect, an embodiment of the present invention further provides a computer-readable storage medium, wherein the computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, any of the above-mentioned methods is implemented.

[0042] In a seventh aspect, an embodiment of the present invention further provides a computer program product comprising instructions, which, when executed on a computer, enables the computer to execute any of the above-described methods.

[0043] Beneficial effects of the embodiments of the present invention:

[0044] The embodiments of the present invention provide a method for training a code retrieval model, a method for code retrieval and a device. The code retrieval model is trained by contrastive learning. During the training process, the features of code retrieval query information and the code index and similarity index of the corresponding code block are extracted respectively, thereby enhancing the ability of the trained code retrieval model to grasp the features of natural language and code language. By contrastive learning, the natural language (code retrieval query information) and positive samples (code index) in the input training code retrieval model can be continuously approached in the vector space, thereby generating a connection. At the same time, the negative samples (similar index) and the natural language in the input training code retrieval model are gradually separated in the vector space, so that the positive samples and negative samples corresponding to the input can be distinguished, thereby improving the model's ability to distinguish and understand the semantics of professional domain knowledge and improving the accuracy of code retrieval.

[0045] Of course, it is not necessary to achieve all of the advantages described above at the same time to implement any product or method of the present invention. BRIEF DESCRIPTION OF THE DRAWINGS

[0046] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings in the following description are only some embodiments of the present invention. For ordinary technicians in this field, other embodiments can also be obtained based on these drawings.

[0047] Figure 1 A schematic diagram of a flow chart of a method for training a code retrieval model provided by an embodiment of the present invention;

[0048] Figure 2 A schematic diagram of the structure of a double-tower model to be trained provided in an embodiment of the present invention;

[0049] Figure 3 A schematic diagram of code retrieval model training provided by an embodiment of the present invention;

[0050] Figure 4 A schematic diagram of a flow chart of a code retrieval method provided by an embodiment of the present invention;

[0051] Figure 5 A schematic diagram of the structure of a training device for a code retrieval model provided by an embodiment of the present invention;

[0052] Figure 6 A schematic diagram of the structure of a code retrieval device provided by an embodiment of the present invention;

[0053] Figure 7 A schematic diagram of the structure of an electronic device provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0054] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments in the present invention, all other embodiments obtained by ordinary technicians in this field based on this application belong to the scope of protection of the present invention.

[0055] In order to improve the accuracy of code retrieval, an embodiment of the present invention provides a code retrieval model training method, a code retrieval method and a device. The execution subject of the code retrieval model training method and the code retrieval method provided by the embodiment of the present invention can be an electronic device, which can be a client device or a server device, etc.

[0056] The training method of the code retrieval model provided in the embodiment of the present invention may include the following steps:

[0057] Obtain a code index in the target code library and a similarity index corresponding to the code index; wherein the code index is information corresponding to a code block in the target code library, including at least a function name corresponding to the code block and comment information of the function; and the similarity index is information corresponding to a code block whose similarity to the code index meets a preset condition;

[0058] Analyze the code index to obtain at least one code search query information corresponding to the code index; the number of the code search query information is the same as the number of similar indexes;

[0059] Using the dual-tower model to be trained, the code retrieval query information is encoded to obtain a first feature vector, the code index is encoded to obtain a second feature vector, and the similarity index is encoded to obtain a third feature vector;

[0060] Calculating a first similarity between the first eigenvector and the second eigenvector, and a second similarity between the first eigenvector and the third eigenvector;

[0061] Based on the first similarity and the second similarity, the contrast loss of the dual-tower model to be trained is calculated, and the parameters of the dual-tower model to be trained are adjusted according to the contrast loss to obtain a trained code retrieval model.

[0062] In the embodiment of the present invention, a code retrieval model is trained by contrastive learning. During the training process, features of code retrieval query information and code index and similarity index of the corresponding code block are extracted respectively, thereby enhancing the ability of the trained code retrieval model to grasp features of natural language and code language. Contrastive learning can be used to make the natural language (code retrieval query information) and positive samples (code index) in the input training code retrieval model continuously approach each other in the vector space, thereby generating a connection. Meanwhile, negative samples (similar index) and the natural language in the input training code retrieval model gradually move away from each other in the vector space, so as to be able to distinguish positive samples and negative samples corresponding to the input, thereby improving the model's ability to distinguish and understand the semantics of professional domain knowledge and improving the accuracy of code retrieval.

[0063] Among them, contrastive learning is a machine learning method, the core idea of ​​which is to learn the representation or features of data by comparing the similarities and differences between different samples. In the process of training a model using contrastive learning, a pair of positive and negative samples will be provided to the model, where the ground truth is called a positive sample and the other is a negative sample. By comparing positive and negative samples, the model will bring the input instance closer to the positive sample and move the input instance further away from the negative sample in the learned embedding space.

[0064] like Figure 1 As shown, the training method of the code retrieval model provided by the embodiment of the present invention may include the following steps:

[0065] S101, obtaining a code index in a target code library and a similarity index corresponding to the code index.

[0066] The target code library can be a local knowledge base of the electronic device, and the local knowledge base contains code bases in different technical fields. Specifically, the target code library can be a code base in any specified field, which contains code blocks or functions with different algorithm requirements. The specified field can be, for example, the financial field, the mathematical field, etc.

[0067] The code index is the information corresponding to the code block in the target code library, including at least the function name and the comment information of the function corresponding to the code block, and illustratively, it may also include the code corresponding to the function, etc. The similarity index is the information corresponding to the code block whose similarity with the code index meets the preset conditions, and also includes the function name and the comment information of the function corresponding to the code block, etc. Among them, the preset conditions may be the M code indexes with the highest similarity to the code index, or the code index with the similarity to the code index higher than the set threshold, etc. The set threshold may be set according to the actual situation.

[0068] The similarity index corresponding to the code index may belong to the same target code library as the code index, or may not belong to the target code library. The similarity index corresponding to the code index may be obtained through similarity calculation or may be manually screened.

[0069] For example, taking the code library in the financial field as the target code library, the code index of one of the code blocks or functions used for economic calculation can be the interface library of the function. The specific interface format of the interface library can be expressed as follows:

[0070] def stockTopNumDailyIncrease(payload)

[0071] """

[0072] Get the data of the top num stocks with the largest increase in the specified date range.

[0073] Applicable data formats: bar charts and tables

[0074] parameter:

[0075] - num (int): the number of stocks to get

[0076] - start_date (str): start date, format is 'YYYYMMDD'

[0077] - end_date (str): end date, format is 'YYYYMMDD'

[0078] - db_name (str): database name

[0079] return:

[0080] - list: List of the top num stocks with the largest increase

[0081] """

[0082] In the above, stockTopNumDailyIncrease(payload) indicates the function name, and the content between """"" and """"" indicates the comment information of the function.

[0083] S102: Analyze the code index to obtain at least one code search query information corresponding to the code index.

[0084] In an example, the comment information of the function in the code index may be parsed to obtain at least one code search query information corresponding to the code index.

[0085] Preferably, the number of code retrieval query information obtained is the same as the number of similarity indexes obtained above, so that the same code index corresponds to the same number of code retrieval query information and similarity indexes, achieving alignment of the number of training samples and facilitating subsequent training of the dual-tower model to be trained.

[0086] In an optional implementation, the code index may be input into a pre-trained large model, so that the pre-trained large model analyzes the code index based on a specified prompt word to obtain at least one code retrieval query information corresponding to the code index.

[0087] Among them, the pre-trained large model can be based on the GPT 4.0 model and trained according to sample code indexes, annotated retrieval results, and specified prompt words.

[0088] Exemplarily, the designated prompt words may be expressed as: You are a professional demand design artificial intelligence (AI), and your task is to parse the provided code interface document and generate multiple corresponding user requirements that can be solved using the code interface function. The generated requirements must be in Chinese, the requirements must use the first person, and there is no need to add additional requirements and purposes, and there is no need to respond to text other than the requirements.

[0089] Among them, the code interface document is the code index, and the generated user requirements are the code retrieval query information.

[0090] Exemplarily, the code index is taken as a positive sample, and the corresponding similar index is a negative sample. The positive sample and the negative sample include the function name and the function's annotation information. The positive sample (i.e., the code index) is represented by the following function name + function annotation information:

[0091] "positive_function": "stockTopNumDailyIncrease(payload)

[0092] Get the data of the top num stocks with the largest increase in the specified date range.

[0093] Applicable data formats: bar charts and tables parameter payload (object): request data.

[0094] num(number): The number of stocks to be obtained.

[0095] start_date (string): start date, for example ''YYYYMMDD''.

[0096] end_date (string): End date, for example ''YYYYMMDD''.

[0097] db_name (string): Database name.

[0098] return

[0099] Promise: A Promise object containing the response data.

[0100] Among them, "positive_function" represents a positive sample, followed by the function name stockTopNumDailyIncrease(payload) and the function's comment information.

[0101] The above code index is input into the pre-trained big model. Under the guidance of the above specified prompt words, the big model generates 8 corresponding code search query information, which can be specifically expressed as follows:

[0102] "queries": [

[0103] ""Query the top 15 stocks with the highest growth rate in mid-August of xxxx."",

[0104] ""Get the top ten stocks with the highest growth rate."",

[0105] ""Get the top 20 stock data with the highest increase between March 1st and March 15th, XXXX."",

[0106] ""Give me the top five stocks with the largest gains."",

[0107] ""Give me the top five stocks with the largest gains in the fourth quarter of 2017."",

[0108] ""Get the stock data with the highest growth rate in XXXX."",

[0109] ""Query the top ten stocks with the highest growth rate from February 5 to February 20, XXXX."",

[0110] ""Get the top 25 stock data with the highest growth rate in the first half of XXXX (January 1 to June 30). "" ]

[0112] queries represent code retrieval query information. The number of code retrieval query information is the same as the number of similarity indexes. Correspondingly, the similarity indexes mentioned above are also 8. Specifically, negative samples (i.e. similarity indexes) are represented by the following function name + function annotation information:

[0113] "negative_functions": [

[0114] "stockFundPerformance(payload)

[0115] Get performance data for a specified stock over a specified time period.

[0116] Applicable data forms: candle charts, mixed charts and tables

[0117] parameter

[0118] payload (object): request data.

[0119] fund_name: (string): stock name.

[0120] start_date (string): start date, for example ''YYYYMMDD''.

[0121] end_date (string): End date, for example ''YYYYMMDD''.

[0122] db_name (string): Database name.

[0123] return

[0124] Promise: A Promise object containing the response data. ",

[0125] "stockMaxDrawdown(payload)

[0126] Get the maximum drawdown of a specified stock in a specified time period.

[0127] Applicable data form: table

[0128] parameter

[0129] payload (object): request data.

[0130] fund_name (string or array): stock name or list of stock names.

[0131] start_date (string): start date, for example ''YYYYMMDD''.

[0132] end_date (string): End date, for example ''YYYYMMDD''.

[0133] db_name (string): Database name.

[0134] return

[0135] Promise: A Promise object containing the response data. ",

[0136] "stockHistoricalVolatility(payload)

[0137] Get the historical volatility of a specified stock over a specified time period.

[0138] Applicable data form: table

[0139] parameter

[0140] payload (object): request data.

[0141] fund_name (string or array): stock name or list of stock names.

[0142] start_date (string): start date, for example ''YYYYMMDD''.

[0143] end_date (string): End date, for example ''YYYYMMDD''.

[0144] db_name (string): Database name.

[0145] return

[0146] Promise: A Promise object containing the response data. ",

[0147] "stockDailyInfo(payload)

[0148] Get daily important indicators for a specified stock within a specified date range.

[0149] Applicable data form: table

[0150] parameter

[0151] payload (object): request data.

[0152] fund_name (string or array): stock name or list of stock names.

[0153] start_date (string): start date, for example ''YYYYMMDD''.

[0154] end_date (string): End date, for example ''YYYYMMDD''.

[0155] db_name (string): Database name.

[0156] return

[0157] Promise: A Promise object containing the response data. ",

[0158] "fundTopNumDailyIncrease(payload)

[0159] Get the top num funds with the highest growth rate within the specified date range.

[0160] Applicable data formats: bar charts and tables

[0161] parameter

[0162] payload (object): request data.

[0163] num (number): The number of funds to obtain.

[0164] start_date (string): start date, for example ''YYYYMMDD''.

[0165] end_date (string): End date, for example ''YYYYMMDD''.

[0166] db_name (string): Database name.

[0167] return

[0168] Promise: A Promise object containing the response data. ",

[0169] "fundPortfolio(payload)

[0170] Get the holdings data of the specified fund within the specified date range.

[0171] Applicable data formats: tables and pie charts

[0172] parameter

[0173] payload (object): request data.

[0174] fund_name (string or array): fund name or list of fund names.

[0175] start_date (string): start date, for example ''YYYYMMDD''.

[0176] end_date (string): End date, for example ''YYYYMMDD''.

[0177] db_name (string): Database name.

[0178] return

[0179] Promise: A Promise object containing the response data. ",

[0180] "stockMonthlyCloseChange(payload)

[0181] Get the monthly closing price changes of the specified stock in the specified time period.

[0182] Applicable data forms: line chart and table

[0183] parameter

[0184] payload (object): request data.

[0185] fund_name (string or array): stock name or list of stock names.

[0186] start_date (string): start date, for example ''YYYYMMDD''.

[0187] end_date (string): End date, for example ''YYYYMMDD''.

[0188] db_name (string): Database name.

[0189] return

[0190] Promise: A Promise object containing the response data. ",

[0191] "stockWeeklyCloseChange(payload)

[0192] Get the weekly closing price changes of the specified stock in the specified time period.

[0193] Applicable data forms: line chart and table

[0194] parameter

[0195] payload (object): request data.

[0196] fund_name (string or array): stock name or list of stock names.

[0197] start_date (string): start date, for example ''YYYYMMDD''.

[0198] end_date (string): End date, for example ''YYYYMMDD''.

[0199] db_name (string): Database name.

[0200] return

[0201] Promise: A Promise object containing the response data. ]

[0203] Among them, "negative_functions" represents negative samples, and the following [ ] contains 8 function names and function annotation information. The 8 function names are: stockFundPerformance(payload), stockMaxDrawdown(payload), stockHistoricalVolatility(payload), stockDailyInfo(payload), fundTopNumDailyIncrease(payload), fundPortfolio(payload), stockMonthlyCloseChange(payload), stockWeeklyCloseChange(payload).

[0204] The positive samples (code index), negative samples (similar index) and code retrieval query information obtained above are combined into triplets as training samples, and then the dual-tower model to be trained is trained.

[0205] like Figure 1 As shown, S103, using the dual-tower model to be trained, the code retrieval query information is encoded to obtain a first feature vector, the code index is encoded to obtain a second feature vector, and the similarity index is encoded to obtain a third feature vector.

[0206] In an embodiment of the present invention, the dual-tower model to be trained is fine-tuned using the BERT-base-Chinese model as a baseline model, and the dual-tower model to be trained includes a first encoding network, a second encoding network, and a third encoding network.

[0207] In an optional implementation, the first encoding network of the twin-tower model to be trained can be used to encode the code retrieval query information to obtain a first feature vector; the second encoding network of the twin-tower model to be trained can be used to encode the code index to obtain a second feature vector; the third encoding network of the twin-tower model to be trained can be used to encode the similarity index to obtain a third feature vector; wherein the network parameters of the first encoding network and the second encoding network are different, and the network parameters of the second encoding network and the third encoding network are the same.

[0208] For example, the dual-tower model to be trained is as follows: Figure 2As shown, the code index as a positive sample, the similarity index as a negative sample, and the code retrieval query information are input into the dual-tower model to be trained in the form of a triple. The first encoding network of the dual-tower model to be trained encodes the code retrieval query information to obtain a first feature vector, the second encoding network encodes the code index to obtain a second feature vector, and the third encoding network encodes the similarity index to obtain a third feature vector. Among them, the first feature vector, the second feature vector, and the third feature vector can be PyTorch tensors. The network structure and type of the first encoding network, the second encoding network, and the third encoding network are the same, for example, they can all be multi-layer convolutional neural networks.

[0209] In the embodiment of the present invention, the dual-tower model to be trained encodes the code retrieval query information and the positive samples and the negative samples respectively, and the obtained first feature vector, second feature vector and third feature vector have no intersection, so that the features of the code retrieval query information and the positive samples and the negative samples can be extracted respectively, thereby enhancing the ability of the dual-tower model to be trained to grasp the features of natural language and code language, so that the dual-tower model to be trained has a stronger semantic understanding ability.

[0210] S104, calculating a first similarity between the first feature vector and the second feature vector, and a second similarity between the first feature vector and the third feature vector.

[0211] For example, Figure 2 As shown, the cosine similarity between the first feature vector and the second feature vector, and the cosine similarity between the first feature vector and the third feature vector can be calculated. Of course, the similarity can also be calculated using the Euclidean distance, which is not limited in the embodiment of the present invention.

[0212] S105, based on the first similarity and the second similarity, calculating the contrast loss of the dual-tower model to be trained, and adjusting the parameters of the dual-tower model to be trained according to the contrast loss to obtain a trained code retrieval model.

[0213] In an optional implementation, the contrast loss of the dual-tower model to be trained can be calculated using the following expression (ie, loss function):

[0214]

[0215] in, represents the loss value of contrast loss, Indicates the number of code retrieval query information, To specify a constant value range of (0,2), Indicates The feature vector corresponding to the code retrieval query information, Indicates The feature vector of the code index corresponding to the code retrieval query information, Indicates The code retrieves the query information and The similarity between the code indexes corresponding to the code retrieval query information, Indicates The feature vector of similarity index corresponding to the code retrieval query information, Indicates The code retrieves the query information and The similarity between similar indexes corresponding to the code retrieval query information.

[0216] Specifically, For ensuring that the similarity between the feature vector corresponding to the code retrieval query information and the feature vector of the code index corresponding to the code retrieval query information (i.e., the positive similarity) is higher than the similarity between the feature vector corresponding to the code retrieval query information and the feature vector of the similar index corresponding to the code retrieval query information (i.e., the negative similarity) by a certain degree, exemplary, The value can be 0.2. Can indicate the The feature vector corresponding to the code retrieval query information and the The cosine similarity between the feature vectors of the code index corresponding to the code retrieval query information. The larger the cosine similarity, the higher the cosine similarity. The feature vector corresponding to the code retrieval query information and the The closer the feature vectors of the code indexes corresponding to the code retrieval query information are, the calculation method can be:

[0217]

[0218] Similarly, That means The feature vector corresponding to the code retrieval query information and the The cosine similarity between the feature vectors of similar indices corresponding to the code retrieval query information.

[0219] When the loss value of the contrast loss of the dual-tower model to be trained meets the preset threshold, that is, when the contrast loss of the dual-tower model to be trained converges, the similarity between the feature vector corresponding to the code retrieval query information and the feature vector of the code index corresponding to the code retrieval query information is as high as possible. At the same time, the similarity between the feature vector corresponding to the code retrieval query information and the feature vector of the similar index corresponding to the code retrieval query information is as low as possible. At this time, the dual-tower model to be trained has a good ability to distinguish between positive samples and negative samples, and has a high retrieval or matching performance. The training of the dual-tower model to be trained is completed, and a trained code retrieval model is obtained.

[0220] Alternatively, the number of iterations is set, and when the set number of iterations is met, the training of the dual-tower model to be trained is completed to obtain a trained code retrieval model. The preset threshold and the set number of iterations can be set according to requirements.

[0221] In an optional implementation, the following expression can be used to update the parameters of the dual-tower model to be trained:

[0222]

[0223] in, represents the parameters of the dual-tower model to be trained, represents the learning rate, Represents the gradient of the loss function corresponding to the contrast loss with respect to the parameters of the dual-tower model to be trained.

[0224] Exemplarily, the number of iterations is set to 20 times. In each iteration, the triples consisting of code index, similarity index and code retrieval query information are input into the dual-tower model to be trained in batches. The batch size can be set as needed, such as 32, etc. The dual-tower model to be trained encodes the code index, similarity index and code retrieval query information respectively to obtain the corresponding feature vector, and then optimizes the parameters of the dual-tower model to be trained by calculating the contrast loss through the loss function of contrast learning. Specifically, during the training process, the code index, similarity index and code retrieval query information are first converted into the input format required by the dual-tower model to be trained, such as adding token ids and attention masks, etc., to ensure the consistency of the data processed in the same batch. Afterwards, the first encoding network of the dual-tower model to be trained is used to encode the code retrieval query information to obtain the first feature vector, the second encoding network is used to encode the code index to obtain the second feature vector, and the third encoding network is used to encode the similarity index to obtain the third feature vector. Then, the above loss function is used to calculate the cosine similarity of each feature vector obtained to obtain the contrast loss of the dual-tower model to be trained. Then, the gradient of the loss function with respect to the parameters of the dual-tower model to be trained is transferred to each network of the dual-tower model to be trained, and the optimizer is used to use the above expression Update the parameters of the dual-tower model to be trained. After the parameters are back-propagated, the scheduler is used to dynamically adjust the learning rate according to the average loss of the current iteration cycle. To speed up the convergence of the dual-tower model to be trained. After completing the training of the dual-tower model to be trained and obtaining the trained code retrieval model, the parameters of the trained code retrieval model, the specific parameters in each network, and the tokenizer (the tokenizer of the dual-tower model to be trained) parameters can be stored to ensure that the trained code retrieval model can be reproduced or further fine-tuned. For example, the parameters of the encoder in each network can be stored in JSON format, and the tokenizer can be stored in a reusable pre-trained format.

[0225] By applying the embodiment of the present invention, a code retrieval model is trained by contrastive learning. During the training process, features of code retrieval query information and the code index and similarity index of the corresponding code block are extracted respectively, thereby enhancing the ability of the trained code retrieval model to grasp the features of natural language and code language. By contrastive learning, the natural language (code retrieval query information) and positive samples (code index) in the input training code retrieval model can be brought closer in the vector space, thereby generating a connection. Meanwhile, negative samples (similar index) and the natural language in the input training code retrieval model are gradually moved away in the vector space, so that the positive samples and negative samples corresponding to the input can be distinguished, thereby improving the model's ability to distinguish and understand the semantics of professional domain knowledge, improving the accuracy of code retrieval, and effectively solving the problem of poor retrieval and matching effects of retrieval models for texts in different formats in related technologies.

[0226] For example, Figure 3 As shown, the training of the code retrieval model of the embodiment of the present invention includes three parts: data generation, model fine-tuning and model evaluation.

[0227] In the data generation phase, obtain the target code library ( Figure 3 The code index in the function library ( Figure 3 positive samples) and similarity indexes corresponding to code indices ( Figure 3 negative samples), input the code index into the pre-trained large model, so that the pre-trained large model analyzes the code index based on the specified prompt word, and obtains the code retrieval query information corresponding to the code index ( Figure 3 of user queries).

[0228] In the model fine-tuning stage, the ternary data set consisting of code retrieval query information, code index and similarity index is used as the training data set, and the code retrieval query information, code index and similarity index are input into the dual-tower model to be trained ( Figure 3 In a baseline model), the code retrieval query information is encoded to obtain a first feature vector, the code index is encoded to obtain a second feature vector, and the similarity index is encoded to obtain a third feature vector, a first similarity between the first feature vector and the second feature vector, and a second similarity between the first feature vector and the third feature vector are calculated, and based on the first similarity and the second similarity, a contrast loss of the dual-tower model to be trained is calculated, and the parameters of the dual-tower model to be trained are adjusted according to the contrast loss ( Figure 3 Fine-tuning training), and obtain the trained code retrieval model ( Figure 3 optimization model).

[0229] In the model evaluation phase, user queries are input into the trained code retrieval model for retrieval, and the retrieval results are compared with the true value retrieval results to evaluate whether the model meets expectations.

[0230] Exemplarily, in the model evaluation stage, multiple groups of verification samples, i.e., multiple groups of user queries and true value retrieval results, can be used to input user queries into the baseline model and the trained code retrieval model respectively, and compare the model retrieval results with the true value retrieval results, so as to evaluate the trained code retrieval model while verifying whether the retrieval accuracy of the trained code retrieval model is significantly higher than the retrieval accuracy of the baseline model.

[0231] In the process of verification, multiple groups of verification samples and model output results are compared as shown in Table 1 below, where the baseline model uses the bert-base-chinese model:

[0232] Table 1 Comparison of model evaluation results

[0233] ;

[0234]

[0235] Table 1 only exemplarily shows the comparison of the retrieval results of the baseline model and the trained code retrieval model using 4 groups of verification samples. It can be seen from Table 1 that the retrieval accuracy of the trained code retrieval model trained by the code retrieval model training method provided by the embodiment of the present invention is significantly better than that of the baseline model.

[0236] like Figure 4 As shown, the code retrieval method provided by the embodiment of the present invention includes:

[0237] S401, obtaining search query information of the code to be searched.

[0238] The code to be retrieved retrieves the query information, that is, the user query described above.

[0239] S402: Input the code search query information to be searched into a pre-trained code search model to perform code search and obtain a target code index.

[0240] Among them, the pre-trained code retrieval model is trained using the above-mentioned code retrieval model training method; the target code index is the information corresponding to the retrieved code block, at least including the function name corresponding to the code block and the function's comment information.

[0241] S403, obtaining a corresponding target code based on the target code index.

[0242] When the function name and the function comment information corresponding to the code block are retrieved, the target code corresponding to the code block is obtained.

[0243] In the embodiment of the present invention, since the pre-trained code retrieval model is obtained by training in a comparative learning manner, the features of the code retrieval query information and the code index and similarity index of the corresponding code block are extracted respectively during the training process, thereby enhancing the ability of the trained code retrieval model to grasp the features of natural language and code language. By using comparative learning, the natural language (code retrieval query information) and positive samples (code index) in the input trained code retrieval model can be continuously approached in the vector space, thereby generating a connection. At the same time, the negative samples (similar index) and the natural language in the input trained code retrieval model are gradually separated in the vector space, so that the positive samples and negative samples corresponding to the input can be distinguished, thereby improving the model's ability to distinguish and understand the semantics of professional domain knowledge, and improving the accuracy of code retrieval. Furthermore, the trained code retrieval model can more accurately retrieve retrieval results matching the input.

[0244] Corresponding to the above method embodiments, the embodiments of the present invention also provide corresponding device embodiments.

[0245] like Figure 5 As shown, an embodiment of the present invention provides a training device for a code retrieval model, the device comprising:

[0246] The first acquisition module 501 is used to acquire a code index in the target code library and a similarity index corresponding to the code index; the code index is information corresponding to a code block in the target code library, including at least a function name and function annotation information corresponding to the code block; the similarity index is information corresponding to a code block whose similarity to the code index meets a preset condition;

[0247] The analysis module 502 is used to analyze the code index to obtain at least one code search query information corresponding to the code index; the number of the code search query information is the same as the number of similar indexes;

[0248] An encoding module 503, used to encode the code retrieval query information using the dual-tower model to be trained to obtain a first feature vector, encode the code index to obtain a second feature vector, and encode the similarity index to obtain a third feature vector;

[0249] A calculation module 504, configured to calculate a first similarity between the first feature vector and the second feature vector, and a second similarity between the first feature vector and the third feature vector;

[0250] The training module 505 is used to calculate the contrast loss of the dual-tower model to be trained based on the first similarity and the second similarity, and adjust the parameters of the dual-tower model to be trained according to the contrast loss to obtain a trained code retrieval model.

[0251] Optionally, the above-mentioned analysis module 502 is specifically used to input the code index into a pre-trained large model, so that the pre-trained large model analyzes the code index based on a specified prompt word to obtain at least one code retrieval query information corresponding to the code index.

[0252] Optionally, the encoding module 503 is specifically used to:

[0253] Encode the code retrieval query information using the first encoding network of the dual-tower model to be trained to obtain a first feature vector;

[0254] Encoding the code index using the second encoding network of the dual-tower model to be trained to obtain a second feature vector;

[0255] Using the third encoding network of the dual-tower model to be trained to encode the similarity index to obtain a third feature vector;

[0256] The network parameters of the first coding network and the second coding network are different, and the network parameters of the second coding network and the third coding network are the same.

[0257] Optionally, use the following expression to calculate the contrast loss of the two-tower model to be trained:

[0258] ;

[0259] in, represents the loss value of contrast loss, Indicates the number of code retrieval query information, To specify a constant value range of (0,2), Indicates The feature vector corresponding to the code retrieval query information, Indicates The feature vector of the code index corresponding to the code retrieval query information, Indicates The code retrieves the query information and The similarity between the code indexes corresponding to the code retrieval query information, Indicates The feature vector of similarity index corresponding to the code retrieval query information, Indicates The code retrieves the query information and The similarity between similar indexes corresponding to the code retrieval query information.

[0260] Optionally, the parameters of the dual-tower model to be trained are updated using the following expressions:

[0261] ;

[0262] in, represents the parameters of the dual-tower model to be trained, represents the learning rate, Represents the gradient of the loss function corresponding to the contrast loss with respect to the parameters of the dual-tower model to be trained.

[0263] like Figure 6 As shown, an embodiment of the present invention provides a code retrieval device, the device comprising:

[0264] The second acquisition module 601 is used to acquire query information of the code to be retrieved;

[0265] The retrieval module 602 is used to input the code retrieval query information to be retrieved into the pre-trained code retrieval model to perform code retrieval and obtain a target code index; wherein the pre-trained code retrieval model is trained by the training method of the above-mentioned code retrieval model; the target code index is the information corresponding to the retrieved code block, including at least the function name corresponding to the code block and the comment information of the function;

[0266] The third acquisition module 603 is used to acquire the corresponding target code based on the target code index.

[0267] The embodiment of the present invention further provides an electronic device, such as Figure 7 As shown, it includes a processor 701 , a communication interface 702 , a memory 703 and a communication bus 704 , wherein the processor 701 , the communication interface 702 , and the memory 703 communicate with each other via the communication bus 704 .

[0268] Memory 703, used for storing computer programs;

[0269] The processor 701 is used to implement any of the above methods when executing the program stored in the memory 703 to achieve the same technical effect.

[0270] The communication bus mentioned in the above electronic device can be a Peripheral Component Interconnect (PCI) bus or an Extended Industry Standard Architecture (EISA) bus, etc. The communication bus can be divided into an address bus, a data bus, a control bus, etc. For ease of representation, only one thick line is used in the figure, but it does not mean that there is only one bus or one type of bus.

[0271] The communication interface is used for communication between the above electronic device and other devices.

[0272] The memory may include a random access memory (RAM) or a non-volatile memory (NVM), such as at least one disk memory. Optionally, the memory may also be at least one storage device located away from the aforementioned processor.

[0273] The above-mentioned processor can be a general-purpose processor, including a central processing unit (CPU), a network processor (NP), etc.; it can also be a digital signal processor (DSP), an application specific integrated circuit (ASIC), a field programmable gate array (FPGA) or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components.

[0274] In another embodiment of the present invention, a computer-readable storage medium is provided, in which a computer program is stored. When the computer program is executed by a processor, any of the above methods is implemented to achieve the same technical effect.

[0275] In another embodiment of the present invention, a computer program product including instructions is provided. When the computer program product is run on a computer, the computer executes any one of the methods in the above embodiments to achieve the same technical effect.

[0276] In the above embodiments, it can be implemented in whole or in part by software, hardware, firmware or any combination thereof. When implemented by software, it can be implemented in whole or in part in the form of a computer program product. The computer program product includes one or more computer instructions. When the computer program instructions are loaded and executed on a computer, the process or function described in the embodiment of the present invention is generated in whole or in part. The computer may be a general-purpose computer, a special-purpose computer, a computer network, or other programmable device. The computer instructions may be stored in a computer-readable storage medium, or transmitted from one computer-readable storage medium to another computer-readable storage medium. For example, the computer instructions may be transmitted from one website site, computer, server or data center to another website site, computer, server or data center by wired (e.g., coaxial cable, optical fiber, digital subscriber line (DSL)) or wireless (e.g., infrared, wireless, microwave, etc.) means. The computer-readable storage medium may be any available medium that a computer can access or a data storage device such as a server or data center that includes one or more available media integrated. The available medium may be a magnetic medium (e.g., a floppy disk, a hard disk, a tape), an optical medium (e.g., a DVD), or a semiconductor medium (e.g., a solid-state drive Solid State Disk (SSD)), etc.

[0277] It should be noted that, in this article, relational terms such as first and second, etc. are only used to distinguish one entity or operation from another entity or operation, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Moreover, the terms "include", "comprise" or any other variants thereof are intended to cover non-exclusive inclusion, so that a process, method, article or device including a series of elements includes not only those elements, but also other elements not explicitly listed, or also includes elements inherent to such process, method, article or device. In the absence of further restrictions, the elements defined by the sentence "comprise a ..." do not exclude the existence of other identical elements in the process, method, article or device including the elements.

[0278] Each embodiment in this specification is described in a related manner, and the same or similar parts between the embodiments can be referred to each other, and each embodiment focuses on the differences from other embodiments. In particular, for embodiments such as devices / electronic devices, since they are basically similar to method embodiments, the description is relatively simple, and the relevant parts can be referred to the partial description of the method embodiments.

[0279] The above description is only a preferred embodiment of the present invention and is not intended to limit the protection scope of the present invention. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present invention are included in the protection scope of the present invention.

Claims

1. A method for training a code retrieval model, characterized in that: The method comprises: Obtain a code index in a target code library and a similarity index corresponding to the code index; the code index is information corresponding to a code block in the target code library, including at least a function name and function annotation information corresponding to the code block; the similarity index is information corresponding to a code block whose similarity to the code index meets a preset condition; the code index and the similarity index are expressed in the form of function name + function annotation information; the preset condition includes: a set number of code indexes with the highest similarity to the code index, or a code index with a similarity to the code index higher than a set threshold; Inputting the code index into a pre-trained large model, allowing the pre-trained large model to analyze the code index based on a specified prompt word to obtain a plurality of code retrieval query information corresponding to the code index; the number of the code retrieval query information is the same as the number of the similarity indexes; Using the dual-tower model to be trained, the code retrieval query information is encoded to obtain a first feature vector, the code index is encoded to obtain a second feature vector, and the similarity index is encoded to obtain a third feature vector; Calculating a first similarity between the first feature vector and the second feature vector, and a second similarity between the first feature vector and the third feature vector; Based on the first similarity and the second similarity, the contrast loss of the dual-tower model to be trained is calculated, and the parameters of the dual-tower model to be trained are adjusted according to the contrast loss to obtain a trained code retrieval model.

2. The method according to claim 1, characterized in that The method of using the dual-tower model to be trained to encode the code retrieval query information to obtain a first feature vector, encoding the code index to obtain a second feature vector, and encoding the similarity index to obtain a third feature vector includes: Encoding the code retrieval query information using a first encoding network of the dual-tower model to be trained to obtain a first feature vector; Encoding the code index using a second encoding network of the dual-tower model to be trained to obtain a second feature vector; Using a third encoding network of the dual-tower model to be trained to encode the similarity index to obtain a third feature vector; The network parameters of the first coding network and the second coding network are different, and the network parameters of the second coding network and the third coding network are the same.

3. The method according to claim 1, characterized in that The contrast loss of the dual-tower model to be trained is calculated using the following expression: ； in, represents the loss value of contrast loss, Indicates the number of code retrieval query information, To specify a constant value range of (0,2), Indicates The feature vector corresponding to the code retrieval query information, Indicates The feature vector of the code index corresponding to the code retrieval query information, Indicates The code retrieves the query information and The similarity between the code indexes corresponding to the code retrieval query information, Indicates The feature vector of similarity index corresponding to the code retrieval query information, Indicates The code retrieves the query information and The similarity between similar indexes corresponding to the code retrieval query information.

4. The method according to claim 1, characterized in that: Use the following expression to update the parameters of the dual-tower model to be trained: ； in, represents the parameters of the dual-tower model to be trained, represents the learning rate, Represents the gradient of the loss function corresponding to the contrast loss with respect to the parameters of the dual-tower model to be trained.

5. A code retrieval method, characterized in that: The method comprises: Get the code retrieval query information to be retrieved; Inputting the code search query information to be searched into a pre-trained code search model to perform code search, and obtaining a target code index; wherein the pre-trained code search model is trained by any method of claims 1-4; the target code index is information corresponding to the retrieved code block, including at least a function name corresponding to the code block and comment information of the function; The corresponding target code is obtained based on the target code index.

6. A training device for a code retrieval model, characterized in that: The device comprises: A first acquisition module is used to acquire a code index in a target code library and a similarity index corresponding to the code index; the code index is information corresponding to a code block in the target code library, including at least a function name and function annotation information corresponding to the code block; the similarity index is information corresponding to a code block whose similarity to the code index meets a preset condition; the code index and the similarity index are expressed in the form of function name + function annotation information; the preset condition includes: a set number of code indexes with the highest similarity to the code index, or a code index with a similarity to the code index higher than a set threshold; An analysis module, used for inputting the code index into a pre-trained large model, so that the pre-trained large model analyzes the code index based on a specified prompt word to obtain a plurality of code search query information corresponding to the code index; the number of the code search query information is the same as the number of the similarity indexes; An encoding module, used for encoding the code retrieval query information using the dual-tower model to be trained to obtain a first feature vector, encoding the code index to obtain a second feature vector, and encoding the similarity index to obtain a third feature vector; A calculation module, used for calculating a first similarity between the first feature vector and the second feature vector, and a second similarity between the first feature vector and the third feature vector; A training module is used to calculate the contrast loss of the dual-tower model to be trained based on the first similarity and the second similarity, and adjust the parameters of the dual-tower model to be trained according to the contrast loss to obtain a trained code retrieval model.

7. A code retrieval device, characterized in that: The device comprises: The second acquisition module is used to obtain the code search query information to be searched; A retrieval module, used for inputting the code retrieval query information to be retrieved into a pre-trained code retrieval model to perform code retrieval and obtain a target code index; wherein the pre-trained code retrieval model is trained by any method of claims 1-4; the target code index is information corresponding to the retrieved code block, including at least a function name corresponding to the code block and comment information of the function; The third acquisition module is used to acquire the corresponding target code based on the target code index.

8. An electronic device, characterized in that: It includes a processor, a communication interface, a memory and a communication bus, wherein the processor, the communication interface and the memory communicate with each other through the communication bus; Memory, used to store computer programs; A processor, for implementing any of the methods described in claims 1-5 when executing a program stored in a memory.

9. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, the method according to any one of claims 1 to 5 is implemented.

Images