Systems and methods for natural language code search
Embodiments are directed to translating a natural language query into a code snippet in a programing language that semantically represents the query. The embodiments include a cascading neural network that includes an encoder network and a classifier network. The encoder network being faster but less accurate than the classifier network. The encoder network is trained using a contrastive learning framework to identify code candidates from a large set of code snippets. The classifier network is trained using a binary classifier to identify the code snippet that semantically represents the query from the code candidates.
1. A method for translating a natural language query into a code snippet in a programming language, the method comprising:
generating, at an encoder network, a plurality of code snippet indices and a natural language query representation from a plurality of code snippets and the natural language query, respectively;
determining one or more code candidates based on distance metrics between the natural language query representations and the plurality of code snippet indices;
generating text-code pairs from the natural language query and the determined one or more code candidates based on the distance metrics, a text-code pair including the natural language query and a code candidate from the determined one or more code candidates;
jointly encoding, by a single encoder of a classifier network, an input sequence concatenating the natural language query and the code candidate from the text-code pair into one joint representation, wherein the single encoder and the encoder network share one or more encoder layers;
predicting, by a classification head of the classifier network, a matching probability that the natural language query and the code candidate are a match using an input of the one joint representation outputted from the single encoder; and
outputting the code candidate in the programming language in response to the natural language query based on the matching probability.
2. The method of claim 1 , further comprising:
training, the encoder network to determine the code candidates on a contrastive loss function.
3. The method of claim 1 , further comprising:
training, the classifier network to determine the code snippet from the pairs using a cross-entropy objective function.
4. The method of claim 1 , wherein the classifier network sequentially follows the encoder network, and the encoder network is an order of magnitude faster and an order of magnitude less accurate than the classifier network.
5. The method of claim 1 , wherein the encoder network is trained on a different loss function from the classifier network.
6. The method of claim 1 , wherein the encoder network shares a portion of a neural network structure with the classifier network.
7. The method of claim 1 , wherein generating the code candidates further comprises:
generating encodings from the natural language query; and
determining, using the code snippet index, the encodings of the code candidates that are within a distance determined by a distance function from the encodings of the natural language query.
8. The method of claim 1 , wherein determining the code snippet further comprises:
determining a confidence score that a code candidate of each pair is the semantic representation of the natural language query;
ranking confidence scores of the pairs; and
selecting a code candidate of a pair corresponding to a highest confidence score as the code snippet that is the semantic representation of the natural language query.
9. A system for translating a natural language query into a code snippet in a programming language, the system comprising:
a memory configured to store a cascading neural network;
a processor coupled to the memory and configured to execute instructions for causing the cascading neural network to:
generate, at an encoder network of the cascading neural network, a plurality of code snippet indices and a natural language query representation from a plurality of code snippets and the natural language query, respectively;
determine one or more code candidates based on distance metrics between the natural language query representation and the plurality of code snippet indices;
generate text-code pairs from the natural language query and the determined one or more code candidates based on the distance metrics, a text-code pair including the natural language query and a code candidate from the determined one or more code candidates;
jointly encode, by a single encoder of a classifier network, an input sequence concatenating the natural language query and the code candidate from the text-code pair into one joint representation, wherein the single encoder and the encoder network share one or more encoder layers;
predict, by a classification head of the classifier network, a matching probability that the natural language query and the code candidate are a match using an input of the one joint representation outputted from the single encoder; and
output the code candidate in the programming language in response to the natural language query, based on the matching probability.
10. The system of claim 9 , wherein the processor is further configured to:
train, the encoder network to determine the code candidates on a contrastive loss function; and
train, the classifier network to determine the code snippet from the pairs using a cross-entropy objective function.
11. The system of claim 9 , wherein the classifier network sequentially follows the encoder network, and the encoder network is an order of magnitude faster and an order of magnitude less accurate than the classifier network.
12. The system of claim 9 , wherein the encoder network shares a portion of a neural network structure with the classifier network.
13. The system of claim 9 , wherein to generate the code candidates the processor is further configured to:
generate encodings from the natural language query; and
determine, using the code snippet index, encodings of the code candidates that are within a distance determined by a distance function from the encodings of the natural language query.
14. The system of claim 9 , wherein to determine the code snippet the processor is further configured to:
determine a confidence score that a code candidate of each pair is a semantic representation of the natural language query;
rank confidence scores of the pairs; and
select a code candidate of a pair corresponding to a highest confidence score as the code snippet that is the semantic representation of the natural language query.
15. A non-transitory computer readable medium having instructions stored thereon, that when executed by a processor cause the processor to perform operations for translating a natural language query into a code snippet in a programming language, the operations comprising:
generating, at an encoder network, a plurality of code snippet indices and a natural language query representation from a plurality of code snippets and the natural language query, respectively;
determining one or more code candidates based on distance metrics between the natural language query representations and the plurality of code snippet indices;
generating text-code pairs from the natural language query and the determined one or more code candidates based on the distance metrics, a text-code pair including the natural language query and a code candidate from the determined one or more code candidates;
jointly encoding, by a single encoder of a classifier network, an input sequence concatenating the natural language query and the code candidate from the text-code pair into one joint representation, wherein the single encoder and the encoder network share one or more encoder layers;
predicting, by a classification head of the classifier network, a matching probability that the natural language query and the code candidate are a match using an input of the one joint representation outputted from the single encoder; and
outputting the code candidate in the programming language in response to the natural language query based on the matching probability.
16. The non-transitory computer readable medium of claim 15 , further comprising:
training, the encoder network to determine the code candidates on a contrastive loss function; and
training, the classifier network to determine the code snippet from the pairs using a cross-entropy objective function.
17. The non-transitory computer readable medium of claim 15 , wherein the classifier network sequentially follows the encoder network, and the encoder network is an order of magnitude faster and an order of magnitude less accurate than the classifier network.
18. The non-transitory computer readable medium of claim 15 , wherein the encoder network shares a portion of a neural network structure with the classifier network.
19. The non-transitory computer readable medium of claim 15 , wherein generating the code candidates further comprises:
generating encodings from the natural language query; and
determining, using the code snippet index, encodings of the code candidates that are within a distance determined by a distance function from the encodings of the natural language query.
20. The non-transitory computer readable medium of claim 15 , wherein determining the code snippet further comprises:
determining a confidence score that a code candidate of each pair is a semantic representation of the natural language query;
ranking confidence scores of the pairs; and
selecting a code candidate of a pair corresponding to a highest confidence score as the code snippet that is the semantic representation of the natural language query.