Code retrieval method and apparatus based on large language model

1 . A computer-implemented method for code retrieval based on a large language model, comprising:

acquiring an initial query corresponding to code continuous writing at an input end;

acquiring a code retrieval query according to the initial query to obtain a retrieval vector of the code retrieval query, the code retrieval query being a query for retrieving a corresponding code from a target code library;

acquiring a first index of the target code library, the first index comprising a plurality of code blocks and a plurality of code block vectors;

performing similarity calculation on the retrieval vector and the plurality of code block vectors in the first index; and obtaining at least one target code block corresponding to the code block vectors with top N similarity calculation results, N being a positive integer greater than or equal to 1;

acquiring a second index of the target code library, the second index being a code architecture knowledge graph, nodes in the code architecture knowledge graph representing code files, and edges in the code architecture knowledge graph representing dependency relationships between the code files;

acquiring a target code file corresponding to the target code block according to a source code file corresponding to the target code block and the second index;

acquiring a retrieval result corresponding to the code retrieval query which comprises the target code block and the target code file corresponding to the target code block,

constructing a prompt according to the target code block and the target code file;

obtaining an output code generation result according to the prompt by calling the large language model; and

returning the output code generation result output by the large language model to the input end to continue with the code continuous writing.

2 . The method according to claim 1 , wherein the acquiring a second index of the target code library comprises:

acquiring all code files in the target code library;

generating a directory tree according to all the code files and directory information of each code file in the target code library; and

inputting the directory tree into a code architecture extraction model to obtain the output second index,

the code architecture extraction model is obtained by pre-training.

3 . The method according to claim 2 , further comprising:

acquiring training data, the training data comprising a sample directory tree and a label code architecture knowledge graph;

inputting the sample directory tree into an initial neural network model to obtain an output prediction code architecture knowledge graph;

calculating a loss function value according to the prediction code architecture knowledge graph and the label code architecture knowledge graph; and

adjusting parameters of the initial neural network model according to the loss function value to obtain the code architecture extraction model.

4 . The method according to claim 1 , wherein the acquiring a target code file corresponding to the target code block according to a source code file corresponding to the target code block and the second index comprises:

determining a source node corresponding to the source code file in the second index according to the source code file corresponding to the target code block;

selecting a target node having an edge with the source node from the second index; and

acquiring the target code file corresponding to the target node from the target code library.

5 . An electronic device, comprising:

at least one processor; and

a memory communicatively connected with the at least one processor;

wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor to enable the at least one processor to perform a method for code retrieval based on a large language model, wherein the method for code retrieval based on a large language model comprises:

acquiring an initial query corresponding to code continuous writing at an input end;

acquiring a code retrieval query according to the initial query to obtain a retrieval vector of the code retrieval query, the code retrieval query being a query for retrieving a corresponding code from a target code library;

acquiring a first index of the target code library, the first index comprising a plurality of code blocks and a plurality of code block vectors;

performing similarity calculation on the retrieval vector and the plurality of code block vectors in the first index; and obtaining at least one target code block corresponding to the code block vectors with top N similarity calculation results, N being a positive integer greater than or equal to 1;

acquiring a second index of the target code library, the second index being a code architecture knowledge graph, nodes in the code architecture knowledge graph representing code files, and edges in the code architecture knowledge graph representing dependency relationships between the code files;

acquiring a target code file corresponding to the target code block according to a source code file corresponding to the target code block and the second index;

acquiring a retrieval result corresponding to the code retrieval query which comprises the target code block and the target code file corresponding to the target code block,

constructing a prompt according to the target code block and the target code file;

obtaining an output code generation result according to the prompt by calling the large language model; and

returning the output code generation result output by the large language model to the input end to continue with the code continuous writing.

6 . The electronic device according to claim 5 , wherein the acquiring a second index of the target code library comprises:

acquiring all code files in the target code library;

generating a directory tree according to all the code files and directory information of each code file in the target code library; and

inputting the directory tree into a code architecture extraction model to obtain the output second index,

the code architecture extraction model is obtained by pre-training.

7 . The electronic device according to claim 6 , further comprising a training unit configured to:

acquiring training data, the training data comprising a sample directory tree and a label code architecture knowledge graph;

inputting the sample directory tree into an initial neural network model to obtain an output prediction code architecture knowledge graph;

calculating a loss function value according to the prediction code architecture knowledge graph and the label code architecture knowledge graph; and

adjusting parameters of the initial neural network model according to the loss function value to obtain the code architecture extraction model.

8 . The electronic device according to claim 5 , wherein the acquiring a target code file corresponding to the target code block according to a source code file corresponding to the target code block and the second index comprises:

determining a source node corresponding to the source code file in the second index according to the source code file corresponding to the target code block;

selecting a target node having an edge with the source node from the second index; and

acquiring the target code file corresponding to the target node from the target code library.

9 . A non-transitory computer readable storage medium with computer instructions stored thereon, wherein the computer instructions are used for causing a method for code retrieval based on a large language model, wherein the method for code retrieval based on a large language model comprises:

acquiring an initial query corresponding to code continuous writing at an input end;

acquiring a code retrieval query according to the initial query to obtain a retrieval vector of the code retrieval query, the code retrieval query being a query for retrieving a corresponding code from a target code library;

acquiring a first index of the target code library, the first index comprising a plurality of code blocks and a plurality of code block vectors;

performing similarity calculation on the retrieval vector and the plurality of code block vectors in the first index; and obtaining at least one target code block corresponding to the code block vectors with top N similarity calculation results, N being a positive integer greater than or equal to 1;

acquiring a second index of the target code library, the second index being a code architecture knowledge graph, nodes in the code architecture knowledge graph representing code files, and edges in the code architecture knowledge graph representing dependency relationships between the code files;

acquiring a target code file corresponding to the target code block according to a source code file corresponding to the target code block and the second index;

acquiring a retrieval result corresponding to the code retrieval query which comprises the target code block and the target code file corresponding to the target code block,

constructing a prompt according to the target code block and the target code file;

obtaining an output code generation result according to the prompt by calling the large language model; and

returning the output code generation result output by the large language model to the input end to continue with the code continuous writing.

10 . The non-transitory computer readable storage medium according to claim 9 , wherein the acquiring a second index of the target code library comprises:

acquiring all code files in the target code library;

generating a directory tree according to all the code files and directory information of each code file in the target code library; and

inputting the directory tree into a code architecture extraction model to obtain the output second index,

the code architecture extraction model is obtained by pre-training.

11 . The non-transitory computer readable storage medium according to claim 10 , further comprising:

acquiring training data, the training data comprising a sample directory tree and a label code architecture knowledge graph;

inputting the sample directory tree into an initial neural network model to obtain an output prediction code architecture knowledge graph;

calculating a loss function value according to the prediction code architecture knowledge graph and the label code architecture knowledge graph; and

adjusting parameters of the initial neural network model according to the loss function value to obtain the code architecture extraction model.

12 . The non-transitory computer readable storage medium according to claim 9 , wherein the acquiring a target code file corresponding to the target code block according to a source code file corresponding to the target code block and the second index comprises:

determining a source node corresponding to the source code file in the second index according to the source code file corresponding to the target code block;

selecting a target node having an edge with the source node from the second index; and

acquiring the target code file corresponding to the target node from the target code library.