File size: 3,176 Bytes
51c57f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
# CoCoSoDa: Effective Contrastive Learning for Code Search
Our approach adopts the pre-trained model as the base code/query encoder and optimizes it using multimodal contrastive learning and soft data augmentation.
CoCoSoDa is comprised of the following four components:
* **Pre-trained code/query encoder** captures the semantic information of a code snippet or a natural language query and maps it into a high-dimensional embedding space.
as the code/query encoder.
* **Momentum code/query encoder** encodes the samples (code snippets or queries) of current and previous mini-batches to enrich the negative samples.
* **Soft data augmentation** is to dynamically mask or replace some tokens in a sample (code/query) to generate a similar sample as a form of data augmentation.
* **Multimodal contrastive learning loss function** is used as the optimization objective and consists of inter-modal and intra-modal contrastive learning loss. They are used to minimize the distance of the representations of similar samples and maximize the distance of different samples in the embedding space.
## Source code
### Environment
```
conda create -n CoCoSoDa python=3.6 -y
conda activate CoCoSoDa
pip install torch==1.10 transformers==4.12.5 seaborn==0.11.2 fast-histogram nltk==3.6.5 networkx==2.5.1 tree_sitter tqdm prettytable gdown more-itertools tensorboardX sklearn
```
### Data
```
cd dataset
bash get_data.sh
```
Data statistic is shown in this Table.
| PL | Training | Validation | Test | Candidate Codes|
| :--------- | :------: | :----: | :----: |:----: |
| Ruby | 24,927 | 1,400 | 1,261 |4,360|
| JavaScript | 58,025 | 3,885 | 3,291 |13,981|
| Java | 164,923 | 5,183 | 10,955 |40,347|
| Go | 167,288 | 7,325 | 8,122 |28,120|
| PHP | 241,241 | 12,982 | 14,014 |52,660|
| Python | 251,820 | 13,914 | 14,918 |43,827|
It will take about 10min.
### Training and Evaualtion
We have uploaded the pre-trained model to [huggingface](https://huggingface.co/). You can directly download [DeepSoftwareAnalytics/CoCoSoDa](https://huggingface.co/DeepSoftwareAnalytics/CoCoSoDa) and fine-tune it.
#### Pre-training (Optional)
```
bash run_cocosoda.sh $lang
```
The optimized model is saved in `./saved_models/cocosoda/`. You can upload them to [huggingface](https://huggingface.co/).
It will take about 3 days.
#### Fine-tuning
```
lang=java
bash run_fine_tune.sh $lang
```
#### Zero-shot running
```
lang=python
bash run_zero-shot.sh $lang
```
### Results
#### The Model Evaluated with MRR
| Model | Ruby | Javascript | Go | Python | Java | PHP | Avg. |
| -------------- | :-------: | :--------: | :-------: | :-------: | :-------: | :-------: | :-------: |
| CoCoSoDa | **0.818**| **0.764**| **0.921** |**0.757**| **0.763**| **0.703** |**0.788**|
## Appendix
The description of baselines, addtional experimetal results and discussion are shown in `Appendix/Appendix.pdf`.
## Contact
Feel free to contact Ensheng Shi (enshengshi@qq.com) if you have any further questions or no response to github issue for more than 1 day. |