Systems and methods for semantic code search
Embodiments described herein provides a contrastive learning framework that leverages hard negative examples, that are mined globally from the entire training corpus for a given query to improve the quality of code and natural language representations. Specifically, similar examples from the training corpus are extracted and used as hard negatives in an online manner during training while keeping the minibatch construction random.
1. A method for a natural language code search system, the method comprising:
receiving, via a communication interface, a training corpus of bimodal pairs, wherein at least one training pair from the training corpus includes a natural language description and a corresponding programming language snippet;
encoding, at a beginning of a first training epoch for training a code search neural network model using the training corpus, by an encoder associated with a first set of model parameters of the code search neural network model implemented on one or more processors, the training corpus of natural language descriptions and programming language snippet into a set of representations in a feature space, wherein the natural language description is encoded into a natural language representation and the corresponding programming language snippet is encoded into a programming language representation, wherein the natural language representation and the programming language representation form a positive pair;
determining, by the one or more processors, a set of nearest neighbor representations to at least one of the natural language representation or the programming language representation among the set of representations in the feature space of the training corpus;
forming a set of negative pairs, including forming each negative pair in the set of negative pairs to include the natural language representation or the programming language representation, and one nearest neighbor representation in the set of nearest neighbor representations;
training the code search neural network model including updating the first set of model parameters at an end of the first training epoch based at least in part on a contrastive learning loss comparing the positive pair and the set of negative pairs;
dynamically re-encoding, at a next training epoch for training the code search neural network model using the training corpus, by the encoder associated with the updated first set of model parameters, the training corpus of natural language descriptions and programming language snippet into an updated set of representations in the feature space;
training the code search neural network model at the next training epoch using the contrastive learning loss of the positive pair and the set of negative pairs selected based on the updated set of representations;
receiving, by the communication interface, a natural language search query for a programming language snippet; and
outputting, by the trained code search neural network model, the programming language snippet based on an input of the natural language search query.
2. The method of claim 1 , wherein the set of nearest neighbor representations are selected from a plurality of programming language representations that are generated from the training corpus, as having closest distances to the natural language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a selected programming language representation and the natural language representation.
3. The method of claim 1 , wherein the set of nearest neighbor representations are selected from a plurality of natural language representations that are generated from the training corpus, as having closest distances to the natural language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a corresponding programming language representation of a selected natural language representation and the natural language representation.
4. The method of claim 1 , wherein the set of nearest neighbor representations are selected from a plurality of programming language representations that are generated from the training corpus, as having closest distances to the programming language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a corresponding natural language representation of a selected programming language representation and the programming language representation.
5. The method of claim 1 , wherein the set of nearest neighbor representations are selected from a plurality of natural language representations that are generated from the training corpus, as having closest distances to the programming language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a selected natural language representation and the programming language representation.
6. The method of claim 1 , wherein the encoder is a Transformer encoder.
7. The method of claim 1 , wherein the natural language representation and the programming language representation are obtained by constructing a FAISS index of the at least one training pair at a beginning of training with initial model parameters for the encoder.
8. The method of claim 7 , wherein the FAISS index is periodically re-constructed using updated model parameters for the encoder after a training epoch.
9. The method of claim 1 , wherein the contrastive learning loss corresponds to the formed set of negative pairs, and the method further comprises:
determining a specific set of negative pairs that maximizes the contrastive learning loss during the first training epoch; and
updating parameters of the encoder by minimizing the contrastive learning loss computed from the specific set of negative pairs during the next training epoch.
10. A system for natural language code search, the system comprising:
a communication interface receiving a training corpus of bimodal pairs and a natural language search query for a predicted programming language snippet, wherein at least one training pair from the training corpus includes a natural language description and a corresponding programming language snippet;
a memory storing an encoder associated with a first set of model parameters of a code search neural network model and a plurality of processor-executable instructions; and
a processor executing the instructions to perform operations comprising:
encoding, at a beginning of a first training epoch for training the code search neural network model using the training corpus, by the encoder, the training corpus of natural language descriptions and programming language snippet into a set of representations in a feature space, wherein the natural language description is encoded into a natural language representation and the corresponding programming language snippet is encoded into a programming language representation, wherein the natural language representation and the programming language representation form a positive pair;
determining, by the processor, a set of nearest neighbor representations to at least one of the natural language representation or the programming language representation among the set of representations in the feature space of the training corpus;
forming a set of negative pairs, including forming each negative pair in the set of negative pairs to include the natural language representation or the programming language representation, and one nearest neighbor representation in the set of nearest neighbor representations;
training the code search neural network model including updating the first set of model parameters at an end of the first training epoch based at least in part on a contrastive learning loss comparing the positive pair and the set of negative pairs;
dynamically re-encoding, at a next training epoch for training the code search neural network model using the training corpus, by the encoder associated with the updated first set of model parameters, the training corpus of natural language descriptions and programming language snippet into an updated set of representations in the feature space;
training the code search neural network model at the next training epoch using the contrastive learning loss of the positive pair and the set of negative pairs selected based on the updated set of representations; and
outputting, by the trained code search neural network model, the predicted programming language snippet based on an input of the natural language search query.
11. The system of claim 10 , wherein the set of nearest neighbor representations are selected from a plurality of programming language representations that are generated from the training corpus, as having closest distances to the natural language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a selected programming language representation and the natural language representation.
12. The system of claim 10 , wherein the set of nearest neighbor representations are selected from a plurality of natural language representations that are generated from the training corpus, as having closest distances to the natural language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a corresponding programming language representation of a selected natural language representation and the natural language representation.
13. The system of claim 10 , wherein the set of nearest neighbor representations are selected from a plurality of programming language representations that are generated from the training corpus, as having closest distances to the programming language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a corresponding natural language representation of a selected programming language representation and the programming language representation.
14. The system of claim 10 , wherein the set of nearest neighbor representations are selected from a plurality of natural language representations that are generated from the training corpus, as having closest distances to the programming language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a selected natural language representation and the programming language representation.
15. The system of claim 10 , wherein the natural language representation and the programming language representation are obtained by constructing a FAISS index of the at least one training pair at a beginning of training with initial model parameters for the encoder.
16. The system of claim 15 , wherein the FAISS index is periodically re-constructed using updated model parameters for the encoder after a training epoch.
17. The system of claim 10 , wherein the contrastive learning loss corresponds to the formed set of negative pairs, and the processor executing the instructions to perform operations comprising:
determining a specific set of negative pairs that maximizes the contrastive learning loss during the first training epoch; and
updating parameters of the encoder by minimizing the contrastive learning loss computed from the specific set of negative pairs during the next training epoch.
18. A non-transitory processor-readable storage medium storing a plurality of processor-executable instructions for natural language code search, the instructions being executed by a processor to perform operations comprising:
receiving, via a communication interface, a training corpus of bimodal pairs, wherein at least one training pair from the training corpus includes a natural language description and a corresponding programming language snippet;
encoding, at a beginning of a first training epoch for training a code search neural network model using the training corpus, by an encoder associated with a first set of model parameters of the code search neural network model implemented on one or more processors, the training corpus of natural language descriptions and programming language snippet into a set of representations in a feature space, wherein the natural language description is encoded into a natural language representation and the corresponding programming language snippet is encoded into a programming language representation, wherein the natural language representation and the programming language representation form a positive pair;
determining, by the one or more processors, a set of nearest neighbor representations to at least one of the natural language representation or the programming language representation among the set of representations in the feature space of the training corpus;
forming a set of negative pairs, wherein at least one including forming each negative pair in the set of negative pairs to include includes the natural language representation or the programming language representation, and one nearest neighbor representation in the set of nearest neighbor representations;
training the code search neural network model including updating the first set of model parameters at an end of the first training epoch based at least in part on a contrastive learning loss comparing the positive pair and the set of negative pairs;
dynamically re-encoding, at a next training epoch for training the code search neural network model using the training corpus, by the encoder associated with the updated first set of model parameters, the training corpus of natural language descriptions and programming language snippet into an updated set of representations in the feature space;
training the code search neural network model at the next training epoch using the contrastive learning loss of the positive pair and the set of negative pairs selected based on the updated set of representations;
receiving, by the communication interface, a natural language search query for a predicted programming language snippet; and
outputting, by the trained code search neural network model, the predicted programming language snippet based on an input of the natural language search query.
19. The non-transitory processor-readable storage medium of claim 18 , wherein the set of nearest neighbor representations are selected from a plurality of programming language representations that are generated from the training corpus, as having closest distances to the natural language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a selected programming language representation and the natural language representation.
20. The non-transitory processor-readable storage medium of claim 18 , wherein the set of nearest neighbor representations are selected from a plurality of natural language representations that are generated from the training corpus, as having closest distances to the natural language representation in the feature space, and
wherein the set of negative pairs contains a negative pair formed by a corresponding programming language representation of a selected natural language representation and the natural language representation.