model

CoCoSoDa

@@ -1 +0,0 @@
-Subproject commit 2f2bf8e7994acef846ede7c1078a0b18bc4154d9

README.md

@@ -0,0 +1,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.
+
+![1](Figure/CoCoSoDa.png)
+
+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.

dataset/get_data.sh

@@ -0,0 +1,26 @@
+# wget https://s3.amazonaws.com/code-search-net/CodeSearchNet/v2/python.zip
+# wget https://s3.amazonaws.com/code-search-net/CodeSearchNet/v2/java.zip
+# wget https://s3.amazonaws.com/code-search-net/CodeSearchNet/v2/ruby.zip
+# wget https://s3.amazonaws.com/code-search-net/CodeSearchNet/v2/javascript.zip
+# wget https://s3.amazonaws.com/code-search-net/CodeSearchNet/v2/go.zip
+# wget https://s3.amazonaws.com/code-search-net/CodeSearchNet/v2/php.zip
+
+wget https://huggingface.co/datasets/code_search_net/resolve/main/data/python.zip
+wget https://huggingface.co/datasets/code_search_net/resolve/main/data/java.zip
+wget https://huggingface.co/datasets/code_search_net/resolve/main/data/ruby.zip
+wget https://huggingface.co/datasets/code_search_net/resolve/main/data/javascript.zip
+wget https://huggingface.co/datasets/code_search_net/resolve/main/data/go.zip
+wget https://huggingface.co/datasets/code_search_net/resolve/main/data/php.zip
+
+unzip python.zip
+unzip java.zip
+unzip ruby.zip
+unzip javascript.zip
+unzip go.zip
+unzip php.zip
+rm *.zip
+rm *.pkl
+
+python preprocess.py
+rm -r */final
+cd ..

dataset/preprocess.py

@@ -0,0 +1,47 @@
+import json
+import os
+
+for language in ['ruby','go','java','javascript','php','python']:
+    print(language)
+    train,valid,test,codebase=[],[],[], []
+    for root, dirs, files in os.walk(language+'/final'):
+        for file in files:
+            temp=os.path.join(root,file)
+            if '.jsonl' in temp:
+                if 'train' in temp:
+                    train.append(temp)
+                elif 'valid' in temp:
+                    valid.append(temp)
+                    codebase.append(temp)
+                elif 'test' in temp:
+                    test.append(temp) 
+                    codebase.append(temp)
+                    
+    train_data,valid_data,test_data,codebase_data={},{},{},{}
+    for files,data in [[train,train_data],[valid,valid_data],[test,test_data],[codebase,codebase_data]]:
+            for file in files:
+                if '.gz' in file:
+                    os.system("gzip -d {}".format(file))
+                    file=file.replace('.gz','')
+                with open(file) as f:
+                    for line in f:
+                        line=line.strip()
+                        js=json.loads(line)
+                        data[js['url']]=js
+                        
+    with open('{}/codebase.jsonl'.format(language),'w') as f3:
+        for tag,data in [['train',train_data],['valid',valid_data],['test',test_data],['test',test_data],['codebase',codebase_data]]:
+            with open('{}/{}.jsonl'.format(language,tag),'w') as f1, open("{}/{}.txt".format(language,tag)) as f2:
+                for line in f2:
+                    line=line.strip()
+                    if line in data:
+                        js=data[line]
+                        if tag in ['valid','test']:
+                            js['original_string']=''
+                            js['code']=''
+                            js['code_tokens']=[]
+                        if tag=='codebase':
+                            js['docstring']=''
+                            js['docstring_tokens']=[]
+                        f1.write(json.dumps(js)+'\n')
+                    

dataset/utils.py

@@ -0,0 +1,46 @@
+import pickle
+import os
+import json
+import prettytable as pt
+import numpy as np
+import math
+import logging
+logger = logging.getLogger(__name__)
+
+def read_json_file(filename):
+    with open(filename, 'r') as fp:
+        data = fp.readlines()
+    if len(data) == 1:
+        data = json.loads(data[0])
+    else:
+        data = [json.loads(line) for line in data]
+    return data 
+
+
+def save_json_data(data_dir, filename, data):
+    os.makedirs(data_dir, exist_ok=True)
+    file_name = os.path.join(data_dir, filename)
+    with open(file_name, 'w') as output:
+        if type(data) == list:
+            if type(data[0]) in [str, list,dict]:
+                for item in data:
+                    output.write(json.dumps(item))
+                    output.write('\n')
+
+            else:
+                json.dump(data, output)
+        elif type(data) == dict:
+            json.dump(data, output)
+        else:
+            raise RuntimeError('Unsupported type: %s' % type(data))
+    logger.info("saved dataset in " + file_name)
+
+def save_pickle_data(path_dir, filename, data):
+    full_path = path_dir + '/' + filename
+    print("Save dataset to: %s" % full_path)
+    if not os.path.exists(path_dir):
+        os.makedirs(path_dir)
+
+    with open(full_path, 'wb') as output:
+        pickle.dump(data, output,protocol=4)
+

model.py

@@ -0,0 +1,396 @@
+import torch
+import torch.nn as nn
+from prettytable import PrettyTable
+from torch.nn.modules.activation import Tanh
+import copy
+import logging
+logger = logging.getLogger(__name__)
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup,
+                  RobertaConfig, RobertaModel, RobertaTokenizer)
+def whitening_torch_final(embeddings):
+    mu = torch.mean(embeddings, dim=0, keepdim=True)
+    cov = torch.mm((embeddings - mu).t(), embeddings - mu)
+    u, s, vt = torch.svd(cov)
+    W = torch.mm(u, torch.diag(1/torch.sqrt(s)))
+    embeddings = torch.mm(embeddings - mu, W)
+    return embeddings
+
+class BaseModel(nn.Module): 
+    def __init__(self, ):
+        super().__init__()
+        
+    def model_parameters(self):
+        table = PrettyTable()
+        table.field_names = ["Layer Name", "Output Shape", "Param #"]
+        table.align["Layer Name"] = "l"
+        table.align["Output Shape"] = "r"
+        table.align["Param #"] = "r"
+        for name, parameters in self.named_parameters():
+            if parameters.requires_grad:
+                table.add_row([name, str(list(parameters.shape)), parameters.numel()])
+        return table
+class Model(BaseModel):   
+    def __init__(self, encoder):
+        super(Model, self).__init__()
+        self.encoder = encoder
+      
+    def forward(self, code_inputs=None, nl_inputs=None): 
+        # code_inputs [bs, seq] 
+        if code_inputs is not None:
+            outputs = self.encoder(code_inputs,attention_mask=code_inputs.ne(1))[0] #[bs, seq_len, dim]
+            outputs = (outputs*code_inputs.ne(1)[:,:,None]).sum(1)/code_inputs.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度，
+            return torch.nn.functional.normalize(outputs, p=2, dim=1)
+        else:
+            outputs = self.encoder(nl_inputs,attention_mask=nl_inputs.ne(1))[0]
+            outputs = (outputs*nl_inputs.ne(1)[:,:,None]).sum(1)/nl_inputs.ne(1).sum(-1)[:,None]
+            return torch.nn.functional.normalize(outputs, p=2, dim=1)
+        
+ 
+class Multi_Loss_CoCoSoDa( BaseModel):
+
+    def __init__(self, base_encoder, args, mlp=False):
+        super(Multi_Loss_CoCoSoDa, self).__init__()
+
+        self.K = args.moco_k
+        self.m = args.moco_m
+        self.T = args.moco_t
+        dim= args.moco_dim
+
+        # create the encoders
+        # num_classes is the output fc dimension
+        self.code_encoder_q = base_encoder
+        self.code_encoder_k = copy.deepcopy(base_encoder)
+        self.nl_encoder_q = base_encoder
+        # self.nl_encoder_q = RobertaModel.from_pretrained("roberta-base")   
+        self.nl_encoder_k = copy.deepcopy(self.nl_encoder_q)
+        self.mlp = mlp
+        self.time_score= args.time_score
+        self.do_whitening = args.do_whitening
+        self.do_ineer_loss = args.do_ineer_loss 
+        self.agg_way = args.agg_way
+        self.args = args
+
+        for param_q, param_k in zip(self.code_encoder_q.parameters(), self.code_encoder_k.parameters()):
+            param_k.data.copy_(param_q.data)  # initialize
+            param_k.requires_grad = False  # not update by gradient
+
+        for param_q, param_k in zip(self.nl_encoder_q.parameters(), self.nl_encoder_k.parameters()):
+            param_k.data.copy_(param_q.data)  # initialize
+            param_k.requires_grad = False  # not update by gradient
+
+        # create the code queue
+        torch.manual_seed(3047)
+        torch.cuda.manual_seed(3047)
+        self.register_buffer("code_queue", torch.randn(dim,self.K ))
+        self.code_queue = nn.functional.normalize(self.code_queue, dim=0)
+        self.register_buffer("code_queue_ptr", torch.zeros(1, dtype=torch.long))
+        # create the masked code queue
+        self.register_buffer("masked_code_queue", torch.randn(dim, self.K ))
+        self.masked_code_queue = nn.functional.normalize(self.masked_code_queue, dim=0)
+        self.register_buffer("masked_code_queue_ptr", torch.zeros(1, dtype=torch.long))
+
+
+        # create the nl queue
+        self.register_buffer("nl_queue", torch.randn(dim, self.K ))
+        self.nl_queue = nn.functional.normalize(self.nl_queue, dim=0)
+        self.register_buffer("nl_queue_ptr", torch.zeros(1, dtype=torch.long))
+        # create the masked nl  queue
+        self.register_buffer("masked_nl_queue", torch.randn(dim, self.K ))
+        self.masked_nl_queue= nn.functional.normalize(self.masked_nl_queue, dim=0)
+        self.register_buffer("masked_nl_queue_ptr", torch.zeros(1, dtype=torch.long))
+
+
+
+
+    @torch.no_grad()
+    def _momentum_update_key_encoder(self):
+        """
+        Momentum update of the key encoder
+        % key encoder的Momentum update
+        """
+        for param_q, param_k in zip(self.code_encoder_q.parameters(), self.code_encoder_k.parameters()):
+            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)
+        for param_q, param_k in zip(self.nl_encoder_q.parameters(), self.nl_encoder_k.parameters()):
+            param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)
+        if self.mlp:
+            for param_q, param_k in zip(self.code_encoder_q_fc.parameters(), self.code_encoder_k_fc.parameters()):
+                param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)
+            for param_q, param_k in zip(self.nl_encoder_q_fc.parameters(), self.nl_encoder_k_fc.parameters()):
+                param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self, keys, option='code'):
+        # gather keys before updating queue
+        # keys = concat_all_gather(keys)
+
+        batch_size = keys.shape[0]
+        if option == 'code':
+            code_ptr = int(self.code_queue_ptr)
+            assert self.K % batch_size == 0  # for simplicity
+
+            # replace the keys at ptr (dequeue and enqueue)
+            try:
+                self.code_queue[:, code_ptr:code_ptr + batch_size] = keys.T
+            except:
+                print(code_ptr)
+                print(batch_size)
+                print(keys.shape)
+                exit(111)
+            code_ptr = (code_ptr + batch_size) % self.K  # move pointer  ptr->pointer
+
+            self.code_queue_ptr[0] = code_ptr
+        
+        elif option == 'masked_code':
+            masked_code_ptr = int(self.masked_code_queue_ptr)
+            assert self.K % batch_size == 0  # for simplicity
+
+            # replace the keys at ptr (dequeue and enqueue)
+            try:
+                self.masked_code_queue[:, masked_code_ptr:masked_code_ptr + batch_size] = keys.T
+            except:
+                print(masked_code_ptr)
+                print(batch_size)
+                print(keys.shape)
+                exit(111)
+            masked_code_ptr = (masked_code_ptr + batch_size) % self.K  # move pointer  ptr->pointer
+
+            self.masked_code_queue_ptr[0] = masked_code_ptr
+        
+        elif option == 'nl':
+
+            nl_ptr = int(self.nl_queue_ptr)
+            assert self.K % batch_size == 0  # for simplicity
+
+            # replace the keys at ptr (dequeue and enqueue)
+            self.nl_queue[:, nl_ptr:nl_ptr + batch_size] = keys.T
+            nl_ptr = (nl_ptr + batch_size) % self.K  # move pointer  ptr->pointer
+
+            self.nl_queue_ptr[0] = nl_ptr
+        elif option == 'masked_nl':
+
+            masked_nl_ptr = int(self.masked_nl_queue_ptr)
+            assert self.K % batch_size == 0  # for simplicity
+
+            # replace the keys at ptr (dequeue and enqueue)
+            self.masked_nl_queue[:, masked_nl_ptr:masked_nl_ptr + batch_size] = keys.T
+            masked_nl_ptr = (masked_nl_ptr + batch_size) % self.K  # move pointer  ptr->pointer
+
+            self.masked_nl_queue_ptr[0] = masked_nl_ptr
+
+    
+
+    def forward(self, source_code_q, source_code_k, nl_q,nl_k):
+        """
+        Input:
+            im_q: a batch of query images
+            im_k: a batch of key images
+        Output:
+            logits, targets
+        """
+        if not self.args.do_multi_lang_continue_pre_train:
+            # logger.info(".do_multi_lang_continue_pre_train")
+            outputs = self.code_encoder_q(source_code_q, attention_mask=source_code_q.ne(1))[0]
+            code_q  = (outputs*source_code_q.ne(1)[:,:,None]).sum(1)/source_code_q.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度，
+            code_q  = torch.nn.functional.normalize(code_q, p=2, dim=1)
+            # compute query features for nl
+            outputs= self.nl_encoder_q(nl_q, attention_mask=nl_q.ne(1))[0]  # queries: NxC   bs*feature_dim
+            nl_q = (outputs*nl_q.ne(1)[:,:,None]).sum(1)/nl_q.ne(1).sum(-1)[:,None]
+            nl_q = torch.nn.functional.normalize(nl_q, p=2, dim=1)
+            code2nl_logits = torch.einsum("ab,cb->ac", code_q,nl_q )
+            # loss = self.loss_fct(scores*20, torch.arange(code_inputs.size(0), device=scores.device))
+            code2nl_logits /= self.T
+            # label
+            code2nl_label = torch.arange(code2nl_logits.size(0), device=code2nl_logits.device)
+            return code2nl_logits,code2nl_label, None, None
+        if self.agg_way == "avg":
+            # compute query features for source code
+            outputs = self.code_encoder_q(source_code_q, attention_mask=source_code_q.ne(1))[0]
+            code_q  = (outputs*source_code_q.ne(1)[:,:,None]).sum(1)/source_code_q.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度，
+            code_q  = torch.nn.functional.normalize(code_q, p=2, dim=1)
+            # compute query features for nl
+            outputs= self.nl_encoder_q(nl_q, attention_mask=nl_q.ne(1))[0]  # queries: NxC   bs*feature_dim
+            nl_q = (outputs*nl_q.ne(1)[:,:,None]).sum(1)/nl_q.ne(1).sum(-1)[:,None]
+            nl_q = torch.nn.functional.normalize(nl_q, p=2, dim=1)
+
+            # compute key features
+            with torch.no_grad():  # no gradient to keys
+                self._momentum_update_key_encoder()  # update the key encoder
+
+                # shuffle for making use of BN
+                # im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)
+
+                # masked code
+                outputs = self.code_encoder_k(source_code_k, attention_mask=source_code_k.ne(1))[0]  # keys: NxC
+                code_k  = (outputs*source_code_k.ne(1)[:,:,None]).sum(1)/source_code_k.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度，
+                code_k  = torch.nn.functional.normalize( code_k, p=2, dim=1)
+                # masked nl
+                outputs = self.nl_encoder_k(nl_k, attention_mask=nl_k.ne(1))[0]   # keys: bs*dim
+                nl_k = (outputs*nl_k.ne(1)[:,:,None]).sum(1)/nl_k.ne(1).sum(-1)[:,None]
+                nl_k = torch.nn.functional.normalize(nl_k, p=2, dim=1)
+
+        elif self.agg_way == "cls_pooler":
+            # logger.info(self.agg_way )
+            # compute query features for source code
+            outputs = self.code_encoder_q(source_code_q, attention_mask=source_code_q.ne(1))[1]
+            code_q  = torch.nn.functional.normalize(code_q, p=2, dim=1)
+            # compute query features for nl
+            outputs= self.nl_encoder_q(nl_q, attention_mask=nl_q.ne(1))[1]  # queries: NxC   bs*feature_dim
+            nl_q = torch.nn.functional.normalize(nl_q, p=2, dim=1)
+
+            # compute key features
+            with torch.no_grad():  # no gradient to keys
+                self._momentum_update_key_encoder()  # update the key encoder
+
+                # shuffle for making use of BN
+                # im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)
+
+                # masked code
+                outputs = self.code_encoder_k(source_code_k, attention_mask=source_code_k.ne(1))[1]  # keys: NxC
+                code_k  = torch.nn.functional.normalize( code_k, p=2, dim=1)
+                # masked nl
+                outputs = self.nl_encoder_k(nl_k, attention_mask=nl_k.ne(1))[1]   # keys: bs*dim
+                nl_k = torch.nn.functional.normalize(nl_k, p=2, dim=1)
+
+        elif self.agg_way == "avg_cls_pooler":
+            # logger.info(self.agg_way )
+            outputs = self.code_encoder_q(source_code_q, attention_mask=source_code_q.ne(1))
+            code_q_cls = outputs[1]
+            outputs = outputs[0]
+            code_q_avg  = (outputs*source_code_q.ne(1)[:,:,None]).sum(1)/source_code_q.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度，
+            code_q = code_q_cls + code_q_avg 
+            code_q  = torch.nn.functional.normalize(code_q, p=2, dim=1)
+            # compute query features for nl
+            outputs= self.nl_encoder_q(nl_q, attention_mask=nl_q.ne(1))
+            nl_q_cls = outputs[1]
+            outputs= outputs[0]  # queries: NxC   bs*feature_dim
+            nl_q_avg = (outputs*nl_q.ne(1)[:,:,None]).sum(1)/nl_q.ne(1).sum(-1)[:,None]
+            nl_q = nl_q_avg  + nl_q_cls 
+            nl_q = torch.nn.functional.normalize(nl_q, p=2, dim=1)
+
+            # compute key features
+            with torch.no_grad():  # no gradient to keys
+                self._momentum_update_key_encoder()  # update the key encoder
+
+                # shuffle for making use of BN
+                # im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)
+
+                # masked code
+
+                outputs = self.code_encoder_k(source_code_k, attention_mask=source_code_k.ne(1))
+                code_k_cls = outputs[1]  # keys: NxC
+                outputs = outputs[0]
+                code_k_avg  = (outputs*source_code_k.ne(1)[:,:,None]).sum(1)/source_code_k.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度，
+                code_k =  code_k_cls + code_k_avg
+                code_k  = torch.nn.functional.normalize( code_k, p=2, dim=1)
+                # masked nl
+                outputs = self.nl_encoder_k(nl_k, attention_mask=nl_k.ne(1))
+                nl_k_cls = outputs[1]   # keys: bs*dim
+                outputs = outputs[0]
+                nl_k_avg = (outputs*nl_k.ne(1)[:,:,None]).sum(1)/nl_k.ne(1).sum(-1)[:,None]
+                nl_k =  nl_k_cls + nl_k_avg
+                nl_k = torch.nn.functional.normalize(nl_k, p=2, dim=1)
+            
+        # ## do_whitening
+        # if self.do_whitening:
+        #     code_q = whitening_torch_final(code_q)
+        #     code_k = whitening_torch_final(code_k)
+        #     nl_q = whitening_torch_final(nl_q)
+        #     nl_k = whitening_torch_final(nl_k)
+
+
+        ## code vs nl
+        code2nl_pos = torch.einsum('nc,bc->nb', [code_q, nl_q])
+        # negative logits: NxK
+        code2nl_neg = torch.einsum('nc,ck->nk', [code_q, self.nl_queue.clone().detach()])
+        # logits: Nx(n+K)
+        code2nl_logits = torch.cat([self.time_score*code2nl_pos, code2nl_neg], dim=1)
+        # apply temperature
+        code2nl_logits /= self.T
+        # label
+        code2nl_label = torch.arange(code2nl_logits.size(0), device=code2nl_logits.device)
+
+        ## code vs masked nl
+        code2maskednl_pos = torch.einsum('nc,bc->nb', [code_q, nl_k])
+        # negative logits: NxK
+        code2maskednl_neg = torch.einsum('nc,ck->nk', [code_q, self.masked_nl_queue.clone().detach()])
+        # logits: Nx(n+K)
+        code2maskednl_logits = torch.cat([self.time_score*code2maskednl_pos, code2maskednl_neg], dim=1)
+        # apply temperature
+        code2maskednl_logits /= self.T
+        # label
+        code2maskednl_label = torch.arange(code2maskednl_logits.size(0), device=code2maskednl_logits.device)
+
+        ## nl vs code
+        # nl2code_pos = torch.einsum('nc,nc->n', [nl_q, code_k]).unsqueeze(-1)
+        nl2code_pos = torch.einsum('nc,bc->nb', [nl_q, code_q])
+        # negative logits: bsxK
+        nl2code_neg = torch.einsum('nc,ck->nk', [nl_q, self.code_queue.clone().detach()])
+        # nl2code_logits: bsx(n+K)
+        nl2code_logits = torch.cat([self.time_score*nl2code_pos, nl2code_neg], dim=1)
+        # apply temperature
+        nl2code_logits /= self.T
+        # label
+        nl2code_label = torch.arange(nl2code_logits.size(0), device=nl2code_logits.device)
+
+        ## nl vs masked code
+        # nl2code_pos = torch.einsum('nc,nc->n', [nl_q, code_k]).unsqueeze(-1)
+        nl2maskedcode_pos = torch.einsum('nc,bc->nb', [nl_q, code_k])
+        # negative logits: bsxK
+        nl2maskedcode_neg = torch.einsum('nc,ck->nk', [nl_q, self.masked_code_queue.clone().detach()])
+        # nl2code_logits: bsx(n+K)
+        nl2maskedcode_logits = torch.cat([self.time_score*nl2maskedcode_pos, nl2maskedcode_neg], dim=1)
+        # apply temperature
+        nl2maskedcode_logits /= self.T
+        # label
+        nl2maskedcode_label = torch.arange(nl2maskedcode_logits.size(0), device=nl2maskedcode_logits.device)
+        
+        #logit 4*bsx(1+K)
+        inter_logits = torch.cat((code2nl_logits, code2maskednl_logits, nl2code_logits ,nl2maskedcode_logits ), dim=0)
+
+        # labels: positive key indicators
+        # inter_labels = torch.zeros(inter_logits.shape[0], dtype=torch.long).cuda()
+        inter_labels =  torch.cat((code2nl_label, code2maskednl_label, nl2code_label, nl2maskedcode_label), dim=0)
+
+        if self.do_ineer_loss:
+            # logger.info("do_ineer_loss")
+            ## code vs masked code
+            code2maskedcode_pos = torch.einsum('nc,bc->nb', [code_q, code_k])
+            # negative logits: NxK
+            code2maskedcode_neg = torch.einsum('nc,ck->nk', [code_q, self.masked_code_queue.clone().detach()])
+            # logits: Nx(n+K)
+            code2maskedcode_logits = torch.cat([self.time_score*code2maskedcode_pos, code2maskedcode_neg], dim=1)
+            # apply temperature
+            code2maskedcode_logits /= self.T
+            # label
+            code2maskedcode_label = torch.arange(code2maskedcode_logits.size(0), device=code2maskedcode_logits.device)
+
+
+            ## nl vs masked nl
+            # nl2code_pos = torch.einsum('nc,nc->n', [nl_q, code_k]).unsqueeze(-1)
+            nl2maskednl_pos = torch.einsum('nc,bc->nb', [nl_q, nl_k])
+            # negative logits: bsxK
+            nl2maskednl_neg = torch.einsum('nc,ck->nk', [nl_q, self.masked_nl_queue.clone().detach()])
+            # nl2code_logits: bsx(n+K)
+            nl2maskednl_logits = torch.cat([self.time_score*nl2maskednl_pos, nl2maskednl_neg], dim=1)
+            # apply temperature
+            nl2maskednl_logits /= self.T
+            # label
+            nl2maskednl_label = torch.arange(nl2maskednl_logits.size(0), device=nl2maskednl_logits.device)
+            
+
+            #logit 6*bsx(1+K)
+            inter_logits = torch.cat((inter_logits, code2maskedcode_logits, nl2maskednl_logits), dim=0)
+
+            # labels: positive key indicators
+            # inter_labels = torch.zeros(inter_logits.shape[0], dtype=torch.long).cuda()
+            inter_labels =  torch.cat(( inter_labels, code2maskedcode_label, nl2maskednl_label ), dim=0)
+
+
+        # dequeue and enqueue
+        self._dequeue_and_enqueue(code_q, option='code')
+        self._dequeue_and_enqueue(nl_q, option='nl')
+        self._dequeue_and_enqueue(code_k, option='masked_code')
+        self._dequeue_and_enqueue(nl_k, option='masked_nl')
+
+        return inter_logits, inter_labels, code_q, nl_q 
+

parser/DFG.py

@@ -0,0 +1,1184 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+from tree_sitter import Language, Parser
+from .utils import (remove_comments_and_docstrings,
+                   tree_to_token_index,
+                   index_to_code_token,
+                   tree_to_variable_index)
+
+
+def DFG_python(root_node,index_to_code,states):
+    assignment=['assignment','augmented_assignment','for_in_clause']
+    if_statement=['if_statement']
+    for_statement=['for_statement']
+    while_statement=['while_statement']
+    do_first_statement=['for_in_clause'] 
+    def_statement=['default_parameter']
+    states=states.copy() 
+    if (len(root_node.children)==0 or root_node.type=='string') and root_node.type!='comment':        
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_python(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states        
+    elif root_node.type in assignment:
+        if root_node.type=='for_in_clause':
+            right_nodes=[root_node.children[-1]]
+            left_nodes=[root_node.child_by_field_name('left')]
+        else:
+            if root_node.child_by_field_name('right') is None:
+                return [],states
+            left_nodes=[x for x in root_node.child_by_field_name('left').children if x.type!=',']
+            right_nodes=[x for x in root_node.child_by_field_name('right').children if x.type!=',']
+            if len(right_nodes)!=len(left_nodes):
+                left_nodes=[root_node.child_by_field_name('left')]
+                right_nodes=[root_node.child_by_field_name('right')]
+            if len(left_nodes)==0:
+                left_nodes=[root_node.child_by_field_name('left')]
+            if len(right_nodes)==0:
+                right_nodes=[root_node.child_by_field_name('right')]
+        DFG=[]
+        for node in right_nodes:
+            temp,states=DFG_python(node,index_to_code,states)
+            DFG+=temp
+            
+        for left_node,right_node in zip(left_nodes,right_nodes):
+            left_tokens_index=tree_to_variable_index(left_node,index_to_code)
+            right_tokens_index=tree_to_variable_index(right_node,index_to_code)
+            temp=[]
+            for token1_index in left_tokens_index:
+                idx1,code1=index_to_code[token1_index]
+                temp.append((code1,idx1,'computedFrom',[index_to_code[x][1] for x in right_tokens_index],
+                             [index_to_code[x][0] for x in right_tokens_index]))
+                states[code1]=[idx1]
+            DFG+=temp        
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in ['elif_clause','else_clause']:
+                temp,current_states=DFG_python(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                temp,new_states=DFG_python(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for i in range(2):
+            right_nodes=[x for x in root_node.child_by_field_name('right').children if x.type!=',']
+            left_nodes=[x for x in root_node.child_by_field_name('left').children if x.type!=',']
+            if len(right_nodes)!=len(left_nodes):
+                left_nodes=[root_node.child_by_field_name('left')]
+                right_nodes=[root_node.child_by_field_name('right')]
+            if len(left_nodes)==0:
+                left_nodes=[root_node.child_by_field_name('left')]
+            if len(right_nodes)==0:
+                right_nodes=[root_node.child_by_field_name('right')]
+            for node in right_nodes:
+                temp,states=DFG_python(node,index_to_code,states)
+                DFG+=temp
+            for left_node,right_node in zip(left_nodes,right_nodes):
+                left_tokens_index=tree_to_variable_index(left_node,index_to_code)
+                right_tokens_index=tree_to_variable_index(right_node,index_to_code)
+                temp=[]
+                for token1_index in left_tokens_index:
+                    idx1,code1=index_to_code[token1_index]
+                    temp.append((code1,idx1,'computedFrom',[index_to_code[x][1] for x in right_tokens_index],
+                                 [index_to_code[x][0] for x in right_tokens_index]))
+                    states[code1]=[idx1]
+                DFG+=temp   
+            if  root_node.children[-1].type=="block":
+                temp,states=DFG_python(root_node.children[-1],index_to_code,states)
+                DFG+=temp 
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_python(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_python(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_python(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+        
+
+def DFG_java(root_node,index_to_code,states):
+    assignment=['assignment_expression']
+    def_statement=['variable_declarator']
+    increment_statement=['update_expression']
+    if_statement=['if_statement','else']
+    for_statement=['for_statement']
+    enhanced_for_statement=['enhanced_for_statement']
+    while_statement=['while_statement']
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type=='string') and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_java(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_java(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_java(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_java(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_java(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_java(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="local_variable_declaration":
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in enhanced_for_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        body=root_node.child_by_field_name('body')
+        DFG=[]
+        for i in range(2):
+            temp,states=DFG_java(value,index_to_code,states)
+            DFG+=temp       
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)        
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            temp,states=DFG_java(body,index_to_code,states)
+            DFG+=temp                       
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_java(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_java(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_java(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+def DFG_csharp(root_node,index_to_code,states):
+    assignment=['assignment_expression']
+    def_statement=['variable_declarator']
+    increment_statement=['postfix_unary_expression']
+    if_statement=['if_statement','else']
+    for_statement=['for_statement']
+    enhanced_for_statement=['for_each_statement']
+    while_statement=['while_statement']
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type=='string') and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        if len(root_node.children)==2:
+            name=root_node.children[0]
+            value=root_node.children[1]
+        else:
+            name=root_node.children[0]
+            value=None
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_csharp(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_csharp(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_csharp(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_csharp(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="local_variable_declaration":
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in enhanced_for_statement:
+        name=root_node.child_by_field_name('left')
+        value=root_node.child_by_field_name('right')
+        body=root_node.child_by_field_name('body')
+        DFG=[]
+        for i in range(2):
+            temp,states=DFG_csharp(value,index_to_code,states)
+            DFG+=temp       
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)        
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            temp,states=DFG_csharp(body,index_to_code,states)
+            DFG+=temp                       
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_csharp(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+
+    
+    
+def DFG_ruby(root_node,index_to_code,states):
+    assignment=['assignment','operator_assignment']
+    if_statement=['if','elsif','else','unless','when']
+    for_statement=['for']
+    while_statement=['while_modifier','until']
+    do_first_statement=[] 
+    def_statement=['keyword_parameter']
+    if (len(root_node.children)==0 or root_node.type=='string') and root_node.type!='comment':
+        states=states.copy()
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_ruby(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states        
+    elif root_node.type in assignment:
+        left_nodes=[x for x in root_node.child_by_field_name('left').children if x.type!=',']
+        right_nodes=[x for x in root_node.child_by_field_name('right').children if x.type!=',']
+        if len(right_nodes)!=len(left_nodes):
+            left_nodes=[root_node.child_by_field_name('left')]
+            right_nodes=[root_node.child_by_field_name('right')]
+        if len(left_nodes)==0:
+            left_nodes=[root_node.child_by_field_name('left')]
+        if len(right_nodes)==0:
+            right_nodes=[root_node.child_by_field_name('right')]
+        if root_node.type=="operator_assignment":
+            left_nodes=[root_node.children[0]]
+            right_nodes=[root_node.children[-1]]
+
+        DFG=[]
+        for node in right_nodes:
+            temp,states=DFG_ruby(node,index_to_code,states)
+            DFG+=temp
+            
+        for left_node,right_node in zip(left_nodes,right_nodes):
+            left_tokens_index=tree_to_variable_index(left_node,index_to_code)
+            right_tokens_index=tree_to_variable_index(right_node,index_to_code)
+            temp=[]
+            for token1_index in left_tokens_index:
+                idx1,code1=index_to_code[token1_index]
+                temp.append((code1,idx1,'computedFrom',[index_to_code[x][1] for x in right_tokens_index],
+                             [index_to_code[x][0] for x in right_tokens_index]))
+                states[code1]=[idx1]
+            DFG+=temp        
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement:
+                temp,current_states=DFG_ruby(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                temp,new_states=DFG_ruby(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for i in range(2):
+            left_nodes=[root_node.child_by_field_name('pattern')]
+            right_nodes=[root_node.child_by_field_name('value')]
+            assert len(right_nodes)==len(left_nodes)
+            for node in right_nodes:
+                temp,states=DFG_ruby(node,index_to_code,states)
+                DFG+=temp
+            for left_node,right_node in zip(left_nodes,right_nodes):
+                left_tokens_index=tree_to_variable_index(left_node,index_to_code)
+                right_tokens_index=tree_to_variable_index(right_node,index_to_code)
+                temp=[]
+                for token1_index in left_tokens_index:
+                    idx1,code1=index_to_code[token1_index]
+                    temp.append((code1,idx1,'computedFrom',[index_to_code[x][1] for x in right_tokens_index],
+                                 [index_to_code[x][0] for x in right_tokens_index]))
+                    states[code1]=[idx1]
+                DFG+=temp 
+            temp,states=DFG_ruby(root_node.child_by_field_name('body'),index_to_code,states)
+            DFG+=temp 
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_ruby(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_ruby(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_ruby(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+def DFG_go(root_node,index_to_code,states):
+    assignment=['assignment_statement',]
+    def_statement=['var_spec']
+    increment_statement=['inc_statement']
+    if_statement=['if_statement','else']
+    for_statement=['for_statement']
+    enhanced_for_statement=[]
+    while_statement=[]
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type=='string') and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_go(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_go(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_go(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_go(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in states:
+            if key not in new_states:
+                new_states[key]=states[key]
+            else:
+                new_states[key]+=states[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_go(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_go(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="for_clause":
+                if child.child_by_field_name('update') is not None:
+                    temp,states=DFG_go(child.child_by_field_name('update'),index_to_code,states)
+                    DFG+=temp                 
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_go(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_go(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+    
+    
+
+def DFG_php(root_node,index_to_code,states):
+    assignment=['assignment_expression','augmented_assignment_expression']
+    def_statement=['simple_parameter']
+    increment_statement=['update_expression']
+    if_statement=['if_statement','else_clause']
+    for_statement=['for_statement']
+    enhanced_for_statement=['foreach_statement']
+    while_statement=['while_statement']
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type=='string') and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('default_value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_php(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_php(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_php(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_php(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in states:
+            if key not in new_states:
+                new_states[key]=states[key]
+            else:
+                new_states[key]+=states[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_php(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_php(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="assignment_expression":               
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in enhanced_for_statement:
+        name=None
+        value=None
+        for child in root_node.children:
+            if child.type=='variable_name' and value is None:
+                value=child
+            elif child.type=='variable_name' and name is None:
+                name=child
+                break
+        body=root_node.child_by_field_name('body')
+        DFG=[]
+        for i in range(2):
+            temp,states=DFG_php(value,index_to_code,states)
+            DFG+=temp       
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)        
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            temp,states=DFG_php(body,index_to_code,states)
+            DFG+=temp                       
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_php(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states        
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_php(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_php(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+
+def DFG_javascript(root_node,index_to_code,states):
+    assignment=['assignment_pattern','augmented_assignment_expression']
+    def_statement=['variable_declarator']
+    increment_statement=['update_expression']
+    if_statement=['if_statement','else']
+    for_statement=['for_statement']
+    enhanced_for_statement=[]
+    while_statement=['while_statement']
+    do_first_statement=[]    
+    states=states.copy()
+    if (len(root_node.children)==0 or root_node.type=='string') and root_node.type!='comment':
+        idx,code=index_to_code[(root_node.start_point,root_node.end_point)]
+        if root_node.type==code:
+            return [],states
+        elif code in states:
+            return [(code,idx,'comesFrom',[code],states[code].copy())],states
+        else:
+            if root_node.type=='identifier':
+                states[code]=[idx]
+            return [(code,idx,'comesFrom',[],[])],states
+    elif root_node.type in def_statement:
+        name=root_node.child_by_field_name('name')
+        value=root_node.child_by_field_name('value')
+        DFG=[]
+        if value is None:
+            indexs=tree_to_variable_index(name,index_to_code)
+            for index in indexs:
+                idx,code=index_to_code[index]
+                DFG.append((code,idx,'comesFrom',[],[]))
+                states[code]=[idx]
+            return sorted(DFG,key=lambda x:x[1]),states
+        else:
+            name_indexs=tree_to_variable_index(name,index_to_code)
+            value_indexs=tree_to_variable_index(value,index_to_code)
+            temp,states=DFG_javascript(value,index_to_code,states)
+            DFG+=temp            
+            for index1 in name_indexs:
+                idx1,code1=index_to_code[index1]
+                for index2 in value_indexs:
+                    idx2,code2=index_to_code[index2]
+                    DFG.append((code1,idx1,'comesFrom',[code2],[idx2]))
+                states[code1]=[idx1]   
+            return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in assignment:
+        left_nodes=root_node.child_by_field_name('left')
+        right_nodes=root_node.child_by_field_name('right')
+        DFG=[]
+        temp,states=DFG_javascript(right_nodes,index_to_code,states)
+        DFG+=temp            
+        name_indexs=tree_to_variable_index(left_nodes,index_to_code)
+        value_indexs=tree_to_variable_index(right_nodes,index_to_code)        
+        for index1 in name_indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in value_indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]   
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in increment_statement:
+        DFG=[]
+        indexs=tree_to_variable_index(root_node,index_to_code)
+        for index1 in indexs:
+            idx1,code1=index_to_code[index1]
+            for index2 in indexs:
+                idx2,code2=index_to_code[index2]
+                DFG.append((code1,idx1,'computedFrom',[code2],[idx2]))
+            states[code1]=[idx1]
+        return sorted(DFG,key=lambda x:x[1]),states   
+    elif root_node.type in if_statement:
+        DFG=[]
+        current_states=states.copy()
+        others_states=[]
+        flag=False
+        tag=False
+        if 'else' in root_node.type:
+            tag=True
+        for child in root_node.children:
+            if 'else' in child.type:
+                tag=True
+            if child.type not in if_statement and flag is False:
+                temp,current_states=DFG_javascript(child,index_to_code,current_states)
+                DFG+=temp
+            else:
+                flag=True
+                temp,new_states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp
+                others_states.append(new_states)
+        others_states.append(current_states)
+        if tag is False:
+            others_states.append(states)        
+        new_states={}
+        for dic in others_states:
+            for key in dic:
+                if key not in new_states:
+                    new_states[key]=dic[key].copy()
+                else:
+                    new_states[key]+=dic[key]
+        for key in states:
+            if key not in new_states:
+                new_states[key]=states[key]
+            else:
+                new_states[key]+=states[key]
+        for key in new_states:
+            new_states[key]=sorted(list(set(new_states[key])))
+        return sorted(DFG,key=lambda x:x[1]),new_states
+    elif root_node.type in for_statement:
+        DFG=[]
+        for child in root_node.children:
+            temp,states=DFG_javascript(child,index_to_code,states)
+            DFG+=temp
+        flag=False
+        for child in root_node.children:
+            if flag:
+                temp,states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp                
+            elif child.type=="variable_declaration":               
+                flag=True
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states
+    elif root_node.type in while_statement:  
+        DFG=[]
+        for i in range(2):
+            for child in root_node.children:
+                temp,states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp    
+        dic={}
+        for x in DFG:
+            if (x[0],x[1],x[2]) not in dic:
+                dic[(x[0],x[1],x[2])]=[x[3],x[4]]
+            else:
+                dic[(x[0],x[1],x[2])][0]=list(set(dic[(x[0],x[1],x[2])][0]+x[3]))
+                dic[(x[0],x[1],x[2])][1]=sorted(list(set(dic[(x[0],x[1],x[2])][1]+x[4])))
+        DFG=[(x[0],x[1],x[2],y[0],y[1]) for x,y in sorted(dic.items(),key=lambda t:t[0][1])]
+        return sorted(DFG,key=lambda x:x[1]),states    
+    else:
+        DFG=[]
+        for child in root_node.children:
+            if child.type in do_first_statement:
+                temp,states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp
+        for child in root_node.children:
+            if child.type not in do_first_statement:
+                temp,states=DFG_javascript(child,index_to_code,states)
+                DFG+=temp
+        
+        return sorted(DFG,key=lambda x:x[1]),states
+
+
+     

parser/__init__.py

@@ -0,0 +1,5 @@
+from .utils import (remove_comments_and_docstrings,
+                   tree_to_token_index,
+                   index_to_code_token,
+                   tree_to_variable_index)
+from .DFG import DFG_python,DFG_java,DFG_ruby,DFG_go,DFG_php,DFG_javascript,DFG_csharp

parser/build.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. 
+# Licensed under the MIT license.
+
+from tree_sitter import Language, Parser
+
+Language.build_library(
+  # Store the library in the `build` directory
+  'my-languages.so',
+
+  # Include one or more languages
+  [
+    'tree-sitter-go',
+    'tree-sitter-javascript',
+    'tree-sitter-python',
+    'tree-sitter-php',
+    'tree-sitter-java',
+    'tree-sitter-ruby',
+    'tree-sitter-c-sharp',
+  ]
+)
+

parser/build.sh

@@ -0,0 +1,8 @@
+git clone https://github.com/tree-sitter/tree-sitter-go
+git clone https://github.com/tree-sitter/tree-sitter-javascript
+git clone https://github.com/tree-sitter/tree-sitter-python
+git clone https://github.com/tree-sitter/tree-sitter-ruby
+git clone https://github.com/tree-sitter/tree-sitter-php
+git clone https://github.com/tree-sitter/tree-sitter-java
+git clone https://github.com/tree-sitter/tree-sitter-c-sharp
+python build.py

parser/utils.py

@@ -0,0 +1,98 @@
+import re
+from io import StringIO
+import  tokenize
+def remove_comments_and_docstrings(source,lang):
+    if lang in ['python']:
+        """
+        Returns 'source' minus comments and docstrings.
+        """
+        io_obj = StringIO(source)
+        out = ""
+        prev_toktype = tokenize.INDENT
+        last_lineno = -1
+        last_col = 0
+        for tok in tokenize.generate_tokens(io_obj.readline):
+            token_type = tok[0]
+            token_string = tok[1]
+            start_line, start_col = tok[2]
+            end_line, end_col = tok[3]
+            ltext = tok[4]
+            if start_line > last_lineno:
+                last_col = 0
+            if start_col > last_col:
+                out += (" " * (start_col - last_col))
+            # Remove comments:
+            if token_type == tokenize.COMMENT:
+                pass
+            # This series of conditionals removes docstrings:
+            elif token_type == tokenize.STRING:
+                if prev_toktype != tokenize.INDENT:
+            # This is likely a docstring; double-check we're not inside an operator:
+                    if prev_toktype != tokenize.NEWLINE:
+                        if start_col > 0:
+                            out += token_string
+            else:
+                out += token_string
+            prev_toktype = token_type
+            last_col = end_col
+            last_lineno = end_line
+        temp=[]
+        for x in out.split('\n'):
+            if x.strip()!="":
+                temp.append(x)
+        return '\n'.join(temp)
+    elif lang in ['ruby']:
+        return source
+    else:
+        def replacer(match):
+            s = match.group(0)
+            if s.startswith('/'):
+                return " " # note: a space and not an empty string
+            else:
+                return s
+        pattern = re.compile(
+            r'//.*?$|/\*.*?\*/|\'(?:\\.|[^\\\'])*\'|"(?:\\.|[^\\"])*"',
+            re.DOTALL | re.MULTILINE
+        )
+        temp=[]
+        for x in re.sub(pattern, replacer, source).split('\n'):
+            if x.strip()!="":
+                temp.append(x)
+        return '\n'.join(temp)
+
+def tree_to_token_index(root_node):
+    if (len(root_node.children)==0 or root_node.type=='string') and root_node.type!='comment':
+        return [(root_node.start_point,root_node.end_point)]
+    else:
+        code_tokens=[]
+        for child in root_node.children:
+            code_tokens+=tree_to_token_index(child)
+        return code_tokens
+    
+def tree_to_variable_index(root_node,index_to_code):
+    if (len(root_node.children)==0 or root_node.type=='string') and root_node.type!='comment':
+        index=(root_node.start_point,root_node.end_point)
+        _,code=index_to_code[index]
+        if root_node.type!=code:
+            return [(root_node.start_point,root_node.end_point)]
+        else:
+            return []
+    else:
+        code_tokens=[]
+        for child in root_node.children:
+            code_tokens+=tree_to_variable_index(child,index_to_code)
+        return code_tokens    
+
+def index_to_code_token(index,code):
+    start_point=index[0]
+    end_point=index[1]
+    if start_point[0]==end_point[0]:
+        s=code[start_point[0]][start_point[1]:end_point[1]]
+    else:
+        s=""
+        s+=code[start_point[0]][start_point[1]:]
+        for i in range(start_point[0]+1,end_point[0]):
+            s+=code[i]
+        s+=code[end_point[0]][:end_point[1]]   
+    return s
+   

run.py

@@ -0,0 +1,1420 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Fine-tuning the library models for language modeling on a text file (GPT, GPT-2, BERT, RoBERTa).
+GPT and GPT-2 are fine-tuned using a causal language modeling (CLM) loss while BERT and RoBERTa are fine-tuned
+using a masked language modeling (MLM) loss.
+"""
+
+from unittest import removeResult
+import torch.nn.functional as F
+import argparse
+import logging
+import os
+# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:24"
+import pickle
+import random
+import torch
+import json
+from random import choice
+import numpy as np
+from itertools import cycle
+from model import Model,Multi_Loss_CoCoSoDa
+from torch.nn import CrossEntropyLoss
+from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler
+from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup,
+                  RobertaConfig, RobertaModel, RobertaTokenizer)
+
+logger = logging.getLogger(__name__)
+from tqdm import tqdm
+import multiprocessing
+cpu_cont = 16
+
+from parser import DFG_python,DFG_java,DFG_ruby,DFG_go,DFG_php,DFG_javascript
+from parser import (remove_comments_and_docstrings,
+                   tree_to_token_index,
+                   index_to_code_token,
+                   tree_to_variable_index)
+from tree_sitter import Language, Parser
+import sys
+sys.path.append("dataset")
+torch.cuda.set_per_process_memory_fraction(0.8)
+from utils import save_json_data, save_pickle_data
+dfg_function={
+    'python':DFG_python,
+    'java':DFG_java,
+    'ruby':DFG_ruby,
+    'go':DFG_go,
+    'php':DFG_php,
+    'javascript':DFG_javascript
+}
+
+parsers={}        
+for lang in dfg_function:
+    LANGUAGE = Language('parser/my-languages.so', lang)
+    parser = Parser()
+    parser.set_language(LANGUAGE) 
+    parser = [parser,dfg_function[lang]]    
+    parsers[lang]= parser
+    
+
+ruby_special_token = ['keyword', 'identifier', 'separators', 'simple_symbol', 'constant', 'instance_variable',
+ 'operator', 'string_content', 'integer', 'escape_sequence', 'comment', 'hash_key_symbol',
+  'global_variable', 'heredoc_beginning', 'heredoc_content', 'heredoc_end', 'class_variable',]
+
+java_special_token = ['keyword', 'identifier', 'type_identifier',  'separators', 'operator', 'decimal_integer_literal',
+ 'void_type', 'string_literal', 'decimal_floating_point_literal', 
+ 'boolean_type', 'null_literal', 'comment', 'hex_integer_literal', 'character_literal']
+
+go_special_token = ['keyword', 'identifier', 'separators', 'type_identifier', 'int_literal', 'operator', 
+'field_identifier', 'package_identifier', 'comment',  'escape_sequence', 'raw_string_literal',
+'rune_literal', 'label_name', 'float_literal']
+
+javascript_special_token =['keyword', 'separators', 'identifier', 'property_identifier', 'operator', 
+'number', 'string_fragment', 'comment', 'regex_pattern', 'shorthand_property_identifier_pattern', 
+'shorthand_property_identifier', 'regex_flags', 'escape_sequence', 'statement_identifier']
+
+php_special_token =['text', 'php_tag', 'name', 'operator', 'keyword', 'string', 'integer', 'separators', 'comment', 
+'escape_sequence', 'ERROR',  'boolean', 'namespace', 'class', 'extends']
+
+python_special_token =['keyword', 'identifier', 'separators', 'operator', '"', 'integer', 
+'comment', 'none', 'escape_sequence']
+
+
+special_token={
+    'python':python_special_token,
+    'java':java_special_token,
+    'ruby':ruby_special_token,
+    'go':go_special_token,
+    'php':php_special_token,
+    'javascript':javascript_special_token
+}
+
+all_special_token = []
+for key, value in special_token.items():
+    all_special_token = list(set(all_special_token ).union(set(value)))
+
+def lalign(x, y, alpha=2):
+    x = torch.tensor(x)
+    y= torch.tensor(y)
+    return (x - y).norm(dim=1).pow(alpha).mean()
+    # code2nl_pos = torch.einsum('nc,nc->n', [x, y]).unsqueeze(-1)
+
+    # return code2nl_pos.mean()
+
+def lunif(x, t=2):
+    x = torch.tensor(x)
+    sq_pdist = torch.pdist(x, p=2).pow(2)
+    return sq_pdist.mul(-t).exp().mean().log()
+
+
+
+def cal_r1_r5_r10(ranks):
+    r1,r5,r10= 0,0,0
+    data_len= len(ranks)
+    for item in ranks:
+        if item >=1:
+            r1 +=1
+            r5 += 1 
+            r10 += 1
+        elif item >=0.2:
+            r5+= 1
+            r10+=1
+        elif item >=0.1:
+            r10 +=1
+    result = {"R@1":round(r1/data_len,3), "R@5": round(r5/data_len,3),  "R@10": round(r10/data_len,3)}
+    return result
+
+#remove comments, tokenize code and extract dataflow                                        
+def extract_dataflow(code, parser,lang):
+    #remove comments
+    try:
+        code=remove_comments_and_docstrings(code,lang)
+    except:
+        pass    
+    #obtain dataflow
+    if lang=="php":
+        code="<?php"+code+"?>"    
+    try:
+        tree = parser[0].parse(bytes(code,'utf8'))    
+        root_node = tree.root_node  
+        tokens_index=tree_to_token_index(root_node)     
+        code=code.split('\n')
+        code_tokens=[index_to_code_token(x,code) for x in tokens_index]  
+        index_to_code={}
+        for idx,(index,code) in enumerate(zip(tokens_index,code_tokens)):
+            index_to_code[index]=(idx,code)  
+        try:
+            DFG,_=parser[1](root_node,index_to_code,{}) 
+        except:
+            DFG=[]
+        DFG=sorted(DFG,key=lambda x:x[1])
+        indexs=set()
+        for d in DFG:
+            if len(d[-1])!=0:
+                indexs.add(d[1])
+            for x in d[-1]:
+                indexs.add(x)
+        new_DFG=[]
+        for d in DFG:
+            if d[1] in indexs:
+                new_DFG.append(d)
+        dfg=new_DFG
+    except:
+        dfg=[]
+    return code_tokens,dfg
+
+#remove comments, tokenize code and extract dataflow                                        
+def tokenizer_source_code(code, parser,lang):
+    #remove comments
+    try:
+        code=remove_comments_and_docstrings(code,lang)
+    except:
+        pass    
+    #obtain dataflow
+    if lang=="php":
+        code="<?php"+code+"?>"    
+    try:
+        tree = parser[0].parse(bytes(code,'utf8'))    
+        root_node = tree.root_node  
+        tokens_index=tree_to_token_index(root_node)     
+        code=code.split('\n')
+        code_tokens=[index_to_code_token(x,code) for x in tokens_index]  
+    except:
+        dfg=[]
+    return code_tokens
+
+class InputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 code_tokens,
+                 code_ids,
+                #  position_idx,
+                #  dfg_to_code,
+                #  dfg_to_dfg,                 
+                 nl_tokens,
+                 nl_ids,
+                 url,
+
+    ):
+        self.code_tokens = code_tokens
+        self.code_ids = code_ids
+        # self.position_idx=position_idx
+        # self.dfg_to_code=dfg_to_code
+        # self.dfg_to_dfg=dfg_to_dfg        
+        self.nl_tokens = nl_tokens
+        self.nl_ids = nl_ids
+        self.url=url
+
+
+class TypeAugInputFeatures(object):
+    """A single training/test features for a example."""
+    def __init__(self,
+                 code_tokens,
+                 code_ids,
+                #  position_idx,
+                 code_type,
+                 code_type_ids,                 
+                 nl_tokens,
+                 nl_ids,
+                 url,
+
+    ):
+        self.code_tokens = code_tokens
+        self.code_ids = code_ids
+        # self.position_idx=position_idx
+        self.code_type=code_type
+        self.code_type_ids=code_type_ids        
+        self.nl_tokens = nl_tokens
+        self.nl_ids = nl_ids
+        self.url=url
+
+def convert_examples_to_features(js):
+    js,tokenizer,args=js
+    #code
+    if args.lang == "java_mini":
+        parser=parsers["java"]
+    else:
+        parser=parsers[js["language"]]
+    # code
+    code_tokens=tokenizer_source_code(js['original_string'],parser,args.lang)
+    code_tokens=" ".join(code_tokens[:args.code_length-2])
+    code_tokens=tokenizer.tokenize(code_tokens)[:args.code_length-2]
+    code_tokens =[tokenizer.cls_token]+code_tokens+[tokenizer.sep_token]
+    code_ids =  tokenizer.convert_tokens_to_ids(code_tokens)
+    padding_length = args.code_length - len(code_ids)
+    code_ids+=[tokenizer.pad_token_id]*padding_length   
+
+    #nl
+    nl=' '.join(js['docstring_tokens'])
+    nl_tokens=tokenizer.tokenize(nl)[:args.nl_length-2]
+    nl_tokens =[tokenizer.cls_token]+nl_tokens+[tokenizer.sep_token]
+    nl_ids =  tokenizer.convert_tokens_to_ids(nl_tokens)
+    padding_length = args.nl_length - len(nl_ids)
+    nl_ids+=[tokenizer.pad_token_id]*padding_length  
+
+    return InputFeatures(code_tokens,code_ids,nl_tokens,nl_ids,js['url'])
+
+
+def convert_examples_to_features_aug_type(js):
+    js,tokenizer,args=js
+    #code
+    if args.lang == "java_mini":
+        parser=parsers["java"]
+    else:
+        parser=parsers[js["language"]]
+    # code
+    token_type_role = js[ 'bpe_token_type_role']
+    code_token = [item[0] for item in token_type_role]
+    # code = ' '.join(code_token[:args.code_length-4])
+    # code_tokens = tokenizer.tokenize(code)[:args.code_length-4]
+    code_tokens = code_token[:args.code_length-4]
+    code_tokens =[tokenizer.cls_token,"<encoder-only>",tokenizer.sep_token]+code_tokens+[tokenizer.sep_token]
+    code_ids = tokenizer.convert_tokens_to_ids(code_tokens)
+    padding_length = args.code_length - len(code_ids)
+    code_ids += [tokenizer.pad_token_id]*padding_length
+
+    # code type
+    code_type_token = [item[-1] for item in token_type_role]
+    # code_type= ' '.join(code_type_token[:args.code_length-4])
+    # code_type_tokens = tokenizer.tokenize(code_type)[:args.code_length-4]
+    code_type_tokens = code_type_token[:args.code_length-4]
+    code_type_tokens =[tokenizer.cls_token,"<encoder-only>",tokenizer.sep_token]+code_type_tokens+[tokenizer.sep_token]
+    code_type_ids = tokenizer.convert_tokens_to_ids(code_type_tokens)
+    padding_length = args.code_length - len(code_type_ids)
+    code_type_ids += [tokenizer.pad_token_id]*padding_length
+
+    #nl
+    nl=' '.join(js['docstring_tokens'])
+    nl_tokens = tokenizer.tokenize(nl)[:args.nl_length-4]
+    nl_tokens = [tokenizer.cls_token,"<encoder-only>",tokenizer.sep_token]+nl_tokens+[tokenizer.sep_token]
+    nl_ids = tokenizer.convert_tokens_to_ids(nl_tokens)
+    padding_length = args.nl_length - len(nl_ids)
+    nl_ids += [tokenizer.pad_token_id]*padding_length 
+
+    return TypeAugInputFeatures(code_tokens,code_ids,code_type_tokens,code_type_ids,nl_tokens,nl_ids,js['url'])
+  
+
+class TextDataset(Dataset):
+    def __init__(self, tokenizer, args, file_path=None,pool=None):
+        self.args=args
+        prefix=file_path.split('/')[-1][:-6]
+        cache_file=args.output_dir+'/'+prefix+'.pkl'
+        n_debug_samples = args.n_debug_samples
+        # if 'codebase' in file_path:
+        #     n_debug_samples = 100000
+        if 'train' in file_path:
+            self.split = "train"
+        else:
+            self.split = "other"
+        if os.path.exists(cache_file):
+            self.examples=pickle.load(open(cache_file,'rb'))
+            if args.debug:
+                self.examples= self.examples[:n_debug_samples]
+        else:
+            self.examples = []
+            data=[]
+            if args.debug:
+                with open(file_path, encoding="utf-8") as f:
+                    for line in f:
+                        line=line.strip()
+                        js=json.loads(line)
+                        data.append((js,tokenizer,args))
+                        if len(data) >= n_debug_samples:
+                            break
+            else:
+                with open(file_path, encoding="utf-8") as f:
+                    for line in f:
+                        line=line.strip()
+                        js=json.loads(line)
+                        data.append((js,tokenizer,args))
+            
+            if self.args.data_aug_type == "replace_type":
+                self.examples=pool.map(convert_examples_to_features_aug_type, tqdm(data,total=len(data)))
+            else:
+                self.examples=pool.map(convert_examples_to_features, tqdm(data,total=len(data)))
+            
+        if 'train' in file_path:
+            for idx, example in enumerate(self.examples[:3]):
+                logger.info("*** Example ***")
+                logger.info("idx: {}".format(idx))
+                logger.info("code_tokens: {}".format([x.replace('\u0120','_') for x in example.code_tokens]))
+                logger.info("code_ids: {}".format(' '.join(map(str, example.code_ids))))             
+                logger.info("nl_tokens: {}".format([x.replace('\u0120','_') for x in example.nl_tokens]))
+                logger.info("nl_ids: {}".format(' '.join(map(str, example.nl_ids))))          
+                
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, item): 
+        if self.args.data_aug_type == "replace_type":
+            return (torch.tensor(self.examples[item].code_ids),
+                    torch.tensor(self.examples[item].code_type_ids),
+                    torch.tensor(self.examples[item].nl_ids))
+        else:
+            return (torch.tensor(self.examples[item].code_ids),
+                    torch.tensor(self.examples[item].nl_ids),
+                    torch.tensor(self.examples[item].code_tokens),
+                    torch.tensor(self.examples[item].nl_tokens))
+
+        
+def convert_examples_to_features_unixcoder(js,tokenizer,args):
+    """convert examples to token ids"""
+    code = ' '.join(js['code_tokens']) if type(js['code_tokens']) is list else ' '.join(js['code_tokens'].split())
+    code_tokens = tokenizer.tokenize(code)[:args.code_length-4]
+    code_tokens =[tokenizer.cls_token,"<encoder-only>",tokenizer.sep_token]+code_tokens+[tokenizer.sep_token]
+    code_ids = tokenizer.convert_tokens_to_ids(code_tokens)
+    padding_length = args.code_length - len(code_ids)
+    code_ids += [tokenizer.pad_token_id]*padding_length
+    
+    nl = ' '.join(js['docstring_tokens']) if type(js['docstring_tokens']) is list else ' '.join(js['doc'].split())
+    nl_tokens = tokenizer.tokenize(nl)[:args.nl_length-4]
+    nl_tokens = [tokenizer.cls_token,"<encoder-only>",tokenizer.sep_token]+nl_tokens+[tokenizer.sep_token]
+    nl_ids = tokenizer.convert_tokens_to_ids(nl_tokens)
+    padding_length = args.nl_length - len(nl_ids)
+    nl_ids += [tokenizer.pad_token_id]*padding_length    
+    
+    return InputFeatures(code_tokens,code_ids,nl_tokens,nl_ids,js['url'] if "url" in js else js["retrieval_idx"])
+
+class TextDataset_unixcoder(Dataset):
+    def __init__(self, tokenizer, args, file_path=None, pooler=None):
+        self.examples = []
+        data = []
+        n_debug_samples = args.n_debug_samples
+        with open(file_path) as f:
+            if "jsonl" in file_path:
+                for line in f:
+                    line = line.strip()
+                    js = json.loads(line)
+                    if 'function_tokens' in js:
+                        js['code_tokens'] = js['function_tokens']
+                    data.append(js)
+                    if args.debug  and len(data) >= n_debug_samples:
+                            break
+            elif "codebase"in file_path or "code_idx_map" in file_path:
+                js = json.load(f)
+                for key in js:
+                    temp = {}
+                    temp['code_tokens'] = key.split()
+                    temp["retrieval_idx"] = js[key]
+                    temp['doc'] = ""
+                    temp['docstring_tokens'] = ""
+                    data.append(temp)
+                    if  args.debug  and len(data) >= n_debug_samples:
+                            break
+            elif "json" in file_path:
+                for js in json.load(f):
+                    data.append(js)
+                    if args.debug and len(data) >= n_debug_samples:
+                            break 
+        # if "test" in file_path:
+        #     data = data[-200:]
+        for js in data:
+            self.examples.append(convert_examples_to_features_unixcoder(js,tokenizer,args))
+                
+        if "train" in file_path:
+            # self.examples = self.examples[:128]
+            for idx, example in enumerate(self.examples[:3]):
+                logger.info("*** Example ***")
+                logger.info("idx: {}".format(idx))
+                logger.info("code_tokens: {}".format([x.replace('\u0120','_') for x in example.code_tokens]))
+                logger.info("code_ids: {}".format(' '.join(map(str, example.code_ids))))
+                logger.info("nl_tokens: {}".format([x.replace('\u0120','_') for x in example.nl_tokens]))
+                logger.info("nl_ids: {}".format(' '.join(map(str, example.nl_ids))))                             
+        
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, i):   
+        return (torch.tensor(self.examples[i].code_ids),torch.tensor(self.examples[i].nl_ids))
+        # return (torch.tensor(self.examples[i].code_ids),
+        #             torch.tensor(self.examples[i].nl_ids),
+        #             [self.examples[i].code_tokens],
+        #             [self.examples[i].nl_tokens])
+            
+def set_seed(seed=42):
+    random.seed(seed)
+    os.environ['PYHTONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # all gpus
+    torch.backends.cudnn.deterministic = True
+
+
+def mask_tokens(inputs,tokenizer,mlm_probability):
+    """ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """
+    labels = inputs.clone()
+    # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)
+    probability_matrix = torch.full(labels.shape, mlm_probability).to(inputs.device)
+    special_tokens_mask = [tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in
+                           labels.tolist()] # for masking special token
+    probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool).to(inputs.device), value=0.0)
+    if tokenizer._pad_token is not None:
+        padding_mask = labels.eq(tokenizer.pad_token_id)
+        probability_matrix.masked_fill_(padding_mask, value=0.0) # masked padding
+        
+    masked_indices = torch.bernoulli(probability_matrix).bool() # will decide who will be masked
+    labels[~masked_indices] = -100  # We only compute loss on masked tokens
+
+    # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
+    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool().to(inputs.device) & masked_indices
+    inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
+
+    # 10% of the time, we replace masked input tokens with random word
+    indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool().to(inputs.device) & masked_indices & ~indices_replaced
+    random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long).to(inputs.device)
+    inputs[indices_random] = random_words[indices_random]
+
+    # The rest of the time (10% of the time) we keep the masked input tokens unchanged
+    return inputs, labels
+
+
+def replace_with_type_tokens(inputs,replaces,tokenizer,mlm_probability):
+    """ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """
+    labels = inputs.clone()
+    # We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)
+    probability_matrix = torch.full(labels.shape, mlm_probability).to(inputs.device)
+    special_tokens_mask = [tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in
+                           labels.tolist()] # for masking special token
+    probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool).to(inputs.device), value=0.0)
+    if tokenizer._pad_token is not None:
+        padding_mask = labels.eq(tokenizer.pad_token_id)
+        probability_matrix.masked_fill_(padding_mask, value=0.0) # masked padding
+        
+    masked_indices = torch.bernoulli(probability_matrix).bool() # will decide who will be masked
+    labels[~masked_indices] = -100  # We only compute loss on masked tokens
+
+    # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
+    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool().to(inputs.device) & masked_indices
+    inputs[indices_replaced] = replaces[indices_replaced] 
+
+    return inputs, labels
+
+def replace_special_token_with_type_tokens(inputs, speical_token_ids, tokenizer, mlm_probability):
+    """ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """
+    labels = inputs.clone()
+    probability_matrix = torch.full(labels.shape,0.0).to(inputs.device)   
+    probability_matrix.masked_fill_(labels.eq(speical_token_ids).to(inputs.device), value=mlm_probability)
+    masked_indices = torch.bernoulli(probability_matrix).bool() # will decide who will be masked
+    labels[~masked_indices] = -100  # We only compute loss on masked tokens
+
+    # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
+    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool().to(inputs.device) & masked_indices
+    inputs[indices_replaced] =  speical_token_ids
+
+    return inputs, labels
+
+def replace_special_token_with_mask(inputs, speical_token_ids, tokenizer, mlm_probability):
+    """ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """
+    labels = inputs.clone()
+    probability_matrix = torch.full(labels.shape,0.0).to(inputs.device)   
+    probability_matrix.masked_fill_(labels.eq(speical_token_ids).to(inputs.device), value=mlm_probability)
+    masked_indices = torch.bernoulli(probability_matrix).bool() # will decide who will be masked
+    labels[~masked_indices] = -100  # We only compute loss on masked tokens
+
+    # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
+    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool().to(inputs.device) & masked_indices
+    inputs[indices_replaced] =tokenizer.convert_tokens_to_ids(tokenizer.mask_token) 
+
+    return inputs, labels
+
+def train(args, model, tokenizer,pool):
+
+    """ Train the model """
+    if args.data_aug_type ==  "replace_type" :
+        train_dataset=TextDataset(tokenizer, args, args.train_data_file, pool)
+    else:
+        # if "unixcoder" in args.model_name_or_path or "coco" in args.model_name_or_path :
+        train_dataset=TextDataset_unixcoder(tokenizer, args, args.train_data_file, pool)
+        # else:
+        #     train_dataset=TextDataset(tokenizer, args, args.train_data_file, pool)
+    train_sampler = RandomSampler(train_dataset)
+    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size,num_workers=4,drop_last=True)
+
+    model.to(args.device)
+    if args.local_rank not in [-1, 0]:
+        torch.distributed.barrier()  
+    no_decay = ['bias', 'LayerNorm.weight']
+    optimizer_grouped_parameters = [
+        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+         'weight_decay': args.weight_decay},
+        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+    ]
+    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=1e-8)
+    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=len(train_dataloader)*args.num_train_epochs)
+
+    # multi-gpu training (should be after apex fp16 initialization)
+    if args.n_gpu > 1:
+        model = torch.nn.DataParallel(model)
+
+    # Train!
+    logger.info("***** Running training *****")
+    logger.info("  Num examples = %d", len(train_dataset))
+    logger.info("  Num Epochs = %d", args.num_train_epochs)
+    logger.info("  Num quene = %d", args.moco_k)
+    logger.info("  Instantaneous batch size per GPU = %d", args.train_batch_size//args.n_gpu)
+    logger.info("  Total train batch size  = %d", args.train_batch_size)
+    logger.info("  Total optimization steps = %d", len(train_dataloader)*args.num_train_epochs)
+    
+    model.zero_grad()
+    model.train()
+    tr_num,tr_loss,best_mrr=0,0,-1
+    loss_fct = CrossEntropyLoss()
+    # if args.model_type ==  "multi-loss-cocosoda" :
+    if args.model_type in  ["no_aug_cocosoda", "multi-loss-cocosoda"]  :
+        if args.do_continue_pre_trained:
+            logger.info("do_continue_pre_trained")
+        elif args.do_fine_tune:
+            logger.info("do_fine_tune")
+    special_token_list = special_token[args.lang]
+    special_token_id_list = tokenizer.convert_tokens_to_ids(special_token_list)
+    model_eval = model.module if hasattr(model,'module') else model
+    for idx in range(args.num_train_epochs): 
+        print(idx)
+        for step,batch in enumerate(train_dataloader):
+           
+                          #get inputs
+            code_inputs = batch[0].to(args.device)    
+            nl_inputs = batch[1].to(args.device)
+            #get code and nl vectors
+            nl_outputs = model_eval.nl_encoder_q(nl_inputs, attention_mask=nl_inputs.ne(1))
+            nl_vec =nl_outputs [1]
+            code_outputs = model_eval.code_encoder_q(code_inputs, attention_mask=code_inputs.ne(1))
+            code_vec =code_outputs [1]
+            # code_vec = model(code_inputs=code_inputs)
+            # nl_vec = model(nl_inputs=nl_inputs)
+            torch.cuda.empty_cache()
+            tr_num+=1  
+            #calculate scores and loss
+            scores = torch.einsum("ab,cb->ac",nl_vec,code_vec)
+            
+            loss = loss_fct(scores*20, torch.arange(code_inputs.size(0), device=scores.device))
+
+            tr_loss += loss.item()
+           
+            if (step+1)% args.eval_frequency==0:
+                logger.info("epoch {} step {} loss {}".format(idx,step+1,round(tr_loss/tr_num,5)))
+                tr_loss=0
+                tr_num=0
+            
+            #backward
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+            optimizer.step()
+            optimizer.zero_grad()
+            scheduler.step() 
+            torch.cuda.empty_cache()
+ 
+        results = evaluate(args, model, tokenizer,args.eval_data_file, pool, eval_when_training=True)
+        for key, value in results.items():
+            logger.info("  %s = %s", key, round(value,4))    
+            
+        #save best model
+        if results['eval_mrr']>best_mrr:
+            best_mrr=results['eval_mrr']
+            logger.info("  "+"*"*20)  
+            logger.info("  Best mrr:%s",round(best_mrr,4))
+            logger.info("  "+"*"*20)                          
+
+            checkpoint_prefix = 'checkpoint-best-mrr'
+            output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix))                        
+            if not os.path.exists(output_dir):
+                os.makedirs(output_dir)                        
+            model_to_save = model.module if hasattr(model,'module') else model
+            output_dir = os.path.join(output_dir, '{}'.format('model.bin')) 
+            torch.save(model_to_save.state_dict(), output_dir)
+            logger.info("Saving model checkpoint to %s", output_dir)
+
+        output_dir_epoch = os.path.join(args.output_dir, '{}'.format(idx))                        
+        if not os.path.exists(output_dir_epoch):
+            os.makedirs(output_dir_epoch)                        
+        model_to_save = model.module if hasattr(model,'module') else model
+        output_dir_epoch = os.path.join(output_dir_epoch, '{}'.format('model.bin')) 
+        torch.save(model_to_save.state_dict(), output_dir_epoch)
+        logger.info("Saving model checkpoint to %s", output_dir_epoch)
+
+def  multi_lang_continue_pre_train(args, model, tokenizer,pool):
+    """ Train the model """
+    #get training dataset
+    if "unixcoder" in args.model_name_or_path:
+        train_datasets = []
+        for train_data_file in args.couninue_pre_train_data_files:
+            train_dataset=TextDataset_unixcoder(tokenizer, args, train_data_file, pool)
+            train_datasets.append(train_dataset)
+    else:
+        train_datasets = []
+        for train_data_file in args.couninue_pre_train_data_files:
+            train_dataset=TextDataset(tokenizer, args, train_data_file, pool)
+            train_datasets.append(train_dataset)
+
+    train_samplers = [RandomSampler(train_dataset) for train_dataset in train_datasets]
+    # https://blog.csdn.net/weixin_44966641/article/details/124878064
+    train_dataloaders = [cycle(DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size,drop_last=True)) for train_dataset,train_sampler in zip(train_datasets,train_samplers)]
+    t_total = args.max_steps
+
+    #get optimizer and scheduler
+    # Prepare optimizer and schedule (linear warmup and decay)https://huggingface.co/transformers/v3.3.1/training.html
+    model.to(args.device)
+    if args.local_rank not in [-1, 0]:
+        torch.distributed.barrier()  
+    no_decay = ['bias', 'LayerNorm.weight']
+    optimizer_grouped_parameters = [
+        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+         'weight_decay': 0.01},
+        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+    ]
+    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=1e-8)
+    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.num_warmup_steps,num_training_steps=t_total)
+
+    # Train!
+    training_data_length = sum ([len(item) for item in train_datasets])
+    logger.info("***** Running training *****")
+    logger.info("  Num examples = %d", training_data_length)
+    logger.info("  Num Epochs = %d", args.num_train_epochs)
+    logger.info("  Num quene = %d", args.moco_k)
+    logger.info("  Instantaneous batch size per GPU = %d", args.train_batch_size//args.n_gpu)
+    logger.info("  Total train batch size  = %d", args.train_batch_size)
+
+    checkpoint_last = os.path.join(args.output_dir, 'checkpoint-last')
+    scheduler_last = os.path.join(checkpoint_last, 'scheduler.pt')
+    optimizer_last = os.path.join(checkpoint_last, 'optimizer.pt')
+    if os.path.exists(scheduler_last):
+        scheduler.load_state_dict(torch.load(scheduler_last, map_location="cpu"))
+    if os.path.exists(optimizer_last):
+        optimizer.load_state_dict(torch.load(optimizer_last, map_location="cpu"))    
+    if args.local_rank == 0:
+        torch.distributed.barrier()         
+    if args.fp16:
+        try:
+            from apex import amp
+        except ImportError:
+            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
+        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)
+
+    # multi-gpu training (should be after apex fp16 initialization)
+    if args.n_gpu > 1:
+        model = torch.nn.DataParallel(model)
+
+    # Distributed training (should be after apex fp16 initialization)
+    if args.local_rank != -1:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank%args.gpu_per_node],
+                                                          output_device=args.local_rank%args.gpu_per_node,
+                                                          find_unused_parameters=True)
+ 
+    loss_fct = CrossEntropyLoss()
+    set_seed(args.seed)  # Added here for reproducibility (even between python 2 and 3)
+    probs=[len(x) for x in train_datasets]
+    probs=[x/sum(probs) for x in probs]
+    probs=[x**0.7 for x in probs]
+    probs=[x/sum(probs) for x in probs]
+    # global_step = args.start_step
+    model.zero_grad()
+    model.train()
+
+    global_step = args.start_step
+    step=0
+    tr_loss, logging_loss,avg_loss,tr_nb, best_mrr = 0.0, 0.0,0.0,0,-1
+    tr_num=0
+    special_token_list = all_special_token 
+    special_token_id_list = tokenizer.convert_tokens_to_ids(special_token_list)
+    while True: 
+        
+        train_dataloader=np.random.choice(train_dataloaders, 1, p=probs)[0]
+        # train_dataloader=train_dataloader[0]
+        step+=1
+        batch=next(train_dataloader)
+        # source_ids= batch.to(args.device)
+        model.train()
+        # loss = model(source_ids)
+        code_inputs = batch[0].to(args.device)  
+        code_transformations_ids = code_inputs.clone()
+        nl_inputs = batch[1].to(args.device)
+        nl_transformations_ids= nl_inputs.clone()
+        
+        if step%4 == 0:
+            code_transformations_ids[:, 3:], _ = mask_tokens(code_inputs.clone()[:, 3:] ,tokenizer,args.mlm_probability)
+            nl_transformations_ids[:, 3:], _ = mask_tokens(nl_inputs.clone()[:, 3:] ,tokenizer,args.mlm_probability)
+        elif step%4 == 1:
+            code_types = code_inputs.clone()
+            code_transformations_ids[:, 3:], _ = replace_with_type_tokens(code_inputs.clone()[:, 3:] ,code_types.clone()[:, 3:],tokenizer,args.mlm_probability)
+        elif step%4 == 2:
+            random.seed( step)
+            choice_token_id  = choice(special_token_id_list)
+            code_transformations_ids[:, 3:], _ = replace_special_token_with_type_tokens(code_inputs.clone()[:, 3:], choice_token_id, tokenizer,args.mlm_probability)
+        elif step%4 == 3:
+            random.seed( step)
+            choice_token_id  = choice(special_token_id_list)
+            code_transformations_ids[:, 3:], _ = replace_special_token_with_mask(code_inputs.clone()[:, 3:], choice_token_id, tokenizer,args.mlm_probability)
+        
+
+        tr_num+=1   
+        inter_output, inter_target, _, _= model(source_code_q=code_inputs, source_code_k=code_transformations_ids, 
+                                    nl_q=nl_inputs , nl_k=nl_transformations_ids )
+        
+        
+        
+        # loss_fct = CrossEntropyLoss()
+        loss = loss_fct(20*inter_output, inter_target)
+
+        if args.n_gpu > 1:
+            loss = loss.mean() # mean() to average on multi-gpu parallel training
+
+            
+        if args.gradient_accumulation_steps > 1:
+            loss = loss / args.gradient_accumulation_steps
+
+        if args.fp16:
+            with amp.scale_loss(loss, optimizer) as scaled_loss:
+                scaled_loss.backward()
+            torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
+        else:
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+        tr_loss += loss.item()
+        if (step+1)% args.eval_frequency==0:
+            logger.info("step {} loss {}".format(step+1,round(tr_loss/tr_num,5)))
+            tr_loss=0
+            tr_num=0
+
+        if (step + 1) % args.gradient_accumulation_steps == 0:
+            optimizer.step()
+            optimizer.zero_grad()
+            scheduler.step()  
+            global_step += 1
+            output_flag=True
+            avg_loss=round((tr_loss - logging_loss) /(global_step- tr_nb),6)
+
+            if global_step %100 == 0:
+                logger.info(" global steps (step*gradient_accumulation_steps ): %s loss: %s", global_step, round(avg_loss,6))
+            if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
+                logging_loss = tr_loss
+                tr_nb=global_step
+
+            if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
+                checkpoint_prefix = 'checkpoint-mrr'
+                # results = evaluate(args, model, tokenizer,pool=pool,eval_when_training=True)
+                results = evaluate(args, model, tokenizer,args.eval_data_file, pool, eval_when_training=True)
+
+                # for key, value in results.items():
+                #     logger.info("  %s = %s", key, round(value,6))
+                logger.info("  %s = %s", 'eval_mrr', round(results['eval_mrr'],6))
+
+                if results['eval_mrr']>best_mrr:
+                    best_mrr=results['eval_mrr']
+                    logger.info("  "+"*"*20)  
+                    logger.info("  Best mrr:%s",round(best_mrr,4))
+                    logger.info("  "+"*"*20)                          
+
+                    output_dir = os.path.join(args.output_dir, '{}'.format('checkpoint-best-mrr'))                        
+                    if not os.path.exists(output_dir):
+                        os.makedirs(output_dir)                        
+                    model_to_save = model.module if hasattr(model,'module') else model
+                    output_dir = os.path.join(output_dir, '{}'.format('model.bin')) 
+                    torch.save(model_to_save.state_dict(), output_dir)
+                    logger.info("Saving model checkpoint to %s", output_dir)
+
+
+
+                # Save model checkpoint
+                output_dir = os.path.join(args.output_dir, '{}-{}-{}'.format(checkpoint_prefix, global_step,round(results['eval_mrr'],6)))
+                if not os.path.exists(output_dir):
+                    os.makedirs(output_dir)
+                model_to_save = model.module.code_encoder_q  if hasattr(model,'module') else model.code_encoder_q   # Take care of distributed/parallel training
+                model_to_save.save_pretrained(output_dir)
+                torch.save(args, os.path.join(output_dir, 'training_args.bin'))
+                logger.info("Saving model checkpoint to %s", output_dir)
+
+                # _rotate_checkpoints(args, checkpoint_prefix)
+
+                last_output_dir = os.path.join(args.output_dir, 'checkpoint-last')
+                if not os.path.exists(last_output_dir):
+                    os.makedirs(last_output_dir)
+                model_to_save.save_pretrained(last_output_dir)
+                idx_file = os.path.join(last_output_dir, 'idx_file.txt')
+                with open(idx_file, 'w', encoding='utf-8') as idxf:
+                    idxf.write(str(0) + '\n')
+
+                torch.save(optimizer.state_dict(), os.path.join(last_output_dir, "optimizer.pt"))
+                torch.save(scheduler.state_dict(), os.path.join(last_output_dir, "scheduler.pt"))
+                logger.info("Saving optimizer and scheduler states to %s", last_output_dir)
+
+                step_file = os.path.join(last_output_dir, 'step_file.txt')
+                with open(step_file, 'w', encoding='utf-8') as stepf:
+                    stepf.write(str(global_step) + '\n')
+
+            if args.max_steps > 0 and global_step > args.max_steps:
+                break
+
+
+def evaluate(args, model, tokenizer,file_name,pool, eval_when_training=False):
+    # if "unixcoder" in args.model_name_or_path or "coco" in args.model_name_or_path :
+    dataset_class = TextDataset_unixcoder
+    # else:
+        # dataset_class = TextDataset
+    query_dataset = dataset_class(tokenizer, args, file_name, pool)
+    query_sampler = SequentialSampler(query_dataset)
+    query_dataloader = DataLoader(query_dataset, sampler=query_sampler, batch_size=args.eval_batch_size,num_workers=4)
+    
+    code_dataset = dataset_class(tokenizer, args, args.codebase_file, pool)
+    code_sampler = SequentialSampler(code_dataset)
+    code_dataloader = DataLoader(code_dataset, sampler=code_sampler, batch_size=args.eval_batch_size,num_workers=4)    
+
+    # multi-gpu evaluate
+    if args.n_gpu > 1 and eval_when_training is False:
+        model = torch.nn.DataParallel(model)
+
+    # Eval!
+    logger.info("***** Running evaluation on %s *****"%args.lang)
+    logger.info("  Num queries = %d", len(query_dataset))
+    logger.info("  Num codes = %d", len(code_dataset))
+    logger.info("  Batch size = %d", args.eval_batch_size)
+
+    
+    model.eval()
+    model_eval = model.module if hasattr(model,'module') else model
+    code_vecs=[] 
+    nl_vecs=[]
+    for batch in query_dataloader:  
+        nl_inputs = batch[-1].to(args.device)
+        with torch.no_grad():
+            if args.model_type ==  "base" :
+                nl_vec = model(nl_inputs=nl_inputs) 
+
+            elif args.model_type in  ["cocosoda" ,"no_aug_cocosoda", "multi-loss-cocosoda"]:
+                outputs = model_eval.nl_encoder_q(nl_inputs, attention_mask=nl_inputs.ne(1))
+                if args.agg_way == "avg":
+                    outputs = outputs [0]
+                    nl_vec = (outputs*nl_inputs.ne(1)[:,:,None]).sum(1)/nl_inputs.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度，
+                elif args.agg_way == "cls_pooler":
+                    nl_vec =outputs [1]
+                elif args.agg_way == "avg_cls_pooler":
+                     nl_vec =outputs [1] +  (outputs[0]*nl_inputs.ne(1)[:,:,None]).sum(1)/nl_inputs.ne(1).sum(-1)[:,None] 
+                nl_vec  = torch.nn.functional.normalize( nl_vec, p=2, dim=1)
+                if args.do_whitening:
+                    nl_vec=whitening_torch_final(nl_vec)
+
+
+            
+            nl_vecs.append(nl_vec.cpu().numpy()) 
+
+    for batch in code_dataloader:
+        with torch.no_grad():
+            code_inputs = batch[0].to(args.device)
+            if args.model_type ==  "base" :
+                code_vec = model(code_inputs=code_inputs)
+            elif args.model_type in  ["cocosoda" ,"no_aug_cocosoda", "multi-loss-cocosoda"]:
+                # code_vec =  model_eval.code_encoder_q(code_inputs, attention_mask=code_inputs.ne(1))[1]
+                outputs = model_eval.code_encoder_q(code_inputs, attention_mask=code_inputs.ne(1))
+                if args.agg_way == "avg":
+                    outputs = outputs [0]
+                    code_vec  = (outputs*code_inputs.ne(1)[:,:,None]).sum(1)/code_inputs.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度，
+                elif args.agg_way == "cls_pooler":
+                    code_vec=outputs [1]
+                elif args.agg_way == "avg_cls_pooler":
+                     code_vec=outputs [1] +  (outputs[0]*code_inputs.ne(1)[:,:,None]).sum(1)/code_inputs.ne(1).sum(-1)[:,None] 
+                code_vec  = torch.nn.functional.normalize(code_vec, p=2, dim=1)
+                if args.do_whitening:
+                    code_vec=whitening_torch_final(code_vec)
+        
+            
+            
+            code_vecs.append(code_vec.cpu().numpy())  
+
+    model.train()    
+    code_vecs=np.concatenate(code_vecs,0)
+    nl_vecs=np.concatenate(nl_vecs,0)
+
+    scores=np.matmul(nl_vecs,code_vecs.T)
+    
+    sort_ids=np.argsort(scores, axis=-1, kind='quicksort', order=None)[:,::-1]    
+    
+    nl_urls=[]
+    code_urls=[]
+    for example in query_dataset.examples:
+        nl_urls.append(example.url)
+        
+    for example in code_dataset.examples:
+        code_urls.append(example.url)
+        
+    ranks=[]
+    for url, sort_id in zip(nl_urls,sort_ids):
+        rank=0
+        find=False
+        for idx in sort_id[:1000]:
+            if find is False:
+                rank+=1
+            if code_urls[idx]==url:
+                find=True
+        if find:
+            ranks.append(1/rank)
+        else:
+            ranks.append(0)
+    if args.save_evaluation_reuslt:
+        evaluation_result = {"nl_urls":nl_urls, "code_urls":code_urls,"sort_ids":sort_ids[:,:10],"ranks":ranks}
+        save_pickle_data(args.save_evaluation_reuslt_dir, "evaluation_result.pkl",evaluation_result)
+    result = cal_r1_r5_r10(ranks)
+    result["eval_mrr"]  = round(float(np.mean(ranks)),3)
+    return result
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    # soda
+    parser.add_argument('--data_aug_type',default="replace_type",choices=["replace_type", "random_mask" ,"other"], help="the ways of soda",required=False)
+    parser.add_argument('--aug_type_way',default="random_replace_type",choices=["random_replace_type", "replace_special_type" ,"replace_special_type_with_mask"], help="the ways of soda",required=False)
+    parser.add_argument('--print_align_unif_loss', action='store_true', help='print_align_unif_loss', required=False)
+    parser.add_argument('--do_ineer_loss', action='store_true', help='print_align_unif_loss', required=False)
+    parser.add_argument('--only_save_the_nl_code_vec', action='store_true', help='print_align_unif_loss', required=False)
+    parser.add_argument('--do_zero_short', action='store_true', help='print_align_unif_loss', required=False)
+    parser.add_argument('--agg_way',default="cls_pooler",choices=["avg", "cls_pooler","avg_cls_pooler" ], help="base is codebert/graphcoder/unixcoder",required=False)
+    parser.add_argument('--weight_decay',default=0.01, type=float,required=False)
+    parser.add_argument('--do_single_lang_continue_pre_train', action='store_true', help='do_single_lang_continue_pre_train', required=False)
+    parser.add_argument('--save_evaluation_reuslt', action='store_true', help='save_evaluation_reuslt', required=False)
+    parser.add_argument('--save_evaluation_reuslt_dir', type=str, help='save_evaluation_reuslt', required=False)
+    parser.add_argument('--epoch', type=int, default=50,
+                        help="random seed for initialization")
+    # new continue pre-training
+    parser.add_argument('--fp16', action='store_true',
+                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
+    parser.add_argument("--local_rank", type=int, default=-1,
+                        help="For distributed training: local_rank")
+    parser.add_argument("--loaded_model_filename", type=str, required=False,
+                        help="loaded_model_filename")
+    parser.add_argument("--loaded_codebert_model_filename", type=str, required=False,
+                        help="loaded_model_filename")
+    parser.add_argument('--do_multi_lang_continue_pre_train', action='store_true', help='do_multi_lang_continue_pre_train', required=False)
+    parser.add_argument("--couninue_pre_train_data_files", default=["dataset/ruby/train.jsonl",  "dataset/java/train.jsonl",], type=str, nargs='+', required=False,
+                        help="The input training data files (some json files).")
+    # parser.add_argument("--couninue_pre_train_data_files", default=["dataset/go/train.jsonl",  "dataset/java/train.jsonl",
+    # "dataset/javascript/train.jsonl",  "dataset/php/train.jsonl",  "dataset/python/train.jsonl",  "dataset/ruby/train.jsonl",], type=list, required=False,
+    #                     help="The input training data files (some json files).")
+    parser.add_argument('--do_continue_pre_trained', action='store_true', help='debug mode', required=False)
+    parser.add_argument('--do_fine_tune', action='store_true', help='debug mode', required=False)
+    parser.add_argument('--do_whitening', action='store_true', help='do_whitening https://github.com/Jun-jie-Huang/WhiteningBERT', required=False)
+    parser.add_argument("--time_score", default=1, type=int,help="cosine value * time_score")   
+    parser.add_argument("--max_steps", default=100, type=int, help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
+    parser.add_argument("--num_warmup_steps", default=0, type=int, help="num_warmup_steps")
+    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
+                        help="Number of updates steps to accumulate before performing a backward/update pass.")    
+    parser.add_argument('--logging_steps', type=int, default=50,
+                        help="Log every X updates steps.")
+    parser.add_argument('--save_steps', type=int, default=50,
+                        help="Save checkpoint every X updates steps.")
+    # new moco
+    parser.add_argument('--moco_type',default="encoder_queue",choices=["encoder_queue","encoder_momentum_encoder_queue" ], help="base is codebert/graphcoder/unixcoder",required=False)
+
+    
+    # debug
+    parser.add_argument('--use_best_mrr_model', action='store_true', help='cosine_space', required=False)
+    parser.add_argument('--debug', action='store_true', help='debug mode', required=False)
+    parser.add_argument('--n_debug_samples', type=int, default=100, required=False)
+    parser.add_argument("--max_codeblock_num", default=10, type=int,
+                        help="Optional NL input sequence length after tokenization.")    
+    parser.add_argument('--hidden_size', type=int, default=768, required=False)
+    parser.add_argument("--eval_frequency", default=1, type=int, required=False)
+    parser.add_argument("--mlm_probability", default=0.1, type=float, required=False)
+
+    # model type
+    parser.add_argument('--do_avg', action='store_true', help='avrage hidden status', required=False)
+    parser.add_argument('--model_type',default="base",choices=["base", "cocosoda","multi-loss-cocosoda","no_aug_cocosoda"], help="base is codebert/graphcoder/unixcoder",required=False)
+    # moco
+    # moco specific configs:
+    parser.add_argument('--moco_dim', default=768, type=int,
+                        help='feature dimension (default: 768)')
+    parser.add_argument('--moco_k', default=32, type=int,
+                        help='queue size; number of negative keys (default: 65536), which is divided by 32, etc.')
+    parser.add_argument('--moco_m', default=0.999, type=float,
+                        help='moco momentum of updating key encoder (default: 0.999)')
+    parser.add_argument('--moco_t', default=0.07, type=float,
+                        help='softmax temperature (default: 0.07)')
+
+    # options for moco v2
+    parser.add_argument('--mlp', action='store_true',help='use mlp head')
+
+    ## Required parameters
+    parser.add_argument("--train_data_file", default="dataset/java/train.jsonl", type=str, required=False,
+                        help="The input training data file (a json file).")
+    parser.add_argument("--output_dir", default="saved_models/pre-train", type=str, required=False,
+                        help="The output directory where the model predictions and checkpoints will be written.")
+    parser.add_argument("--eval_data_file", default="dataset/java/valid.jsonl", type=str,
+                        help="An optional input evaluation data file to evaluate the MRR(a jsonl file).")
+    parser.add_argument("--test_data_file", default="dataset/java/test.jsonl", type=str,
+                        help="An optional input test data file to test the MRR(a josnl file).")
+    parser.add_argument("--codebase_file", default="dataset/java/codebase.jsonl", type=str,
+                        help="An optional input test data file to codebase (a jsonl file).")  
+    
+    parser.add_argument("--lang", default="java", type=str,
+                        help="language.")  
+    
+    parser.add_argument("--model_name_or_path", default="DeepSoftwareAnalytics/CoCoSoDa", type=str,
+                        help="The model checkpoint for weights initialization.")
+    parser.add_argument("--config_name", default="DeepSoftwareAnalytics/CoCoSoDa", type=str,
+                        help="Optional pretrained config name or path if not the same as model_name_or_path")
+    parser.add_argument("--tokenizer_name", default="DeepSoftwareAnalytics/CoCoSoDa", type=str,
+                        help="Optional pretrained tokenizer name or path if not the same as model_name_or_path")
+    
+    parser.add_argument("--nl_length", default=50, type=int,
+                        help="Optional NL input sequence length after tokenization.")    
+    parser.add_argument("--code_length", default=100, type=int,
+                        help="Optional Code input sequence length after tokenization.") 
+    parser.add_argument("--data_flow_length", default=0, type=int,
+                        help="Optional Data Flow input sequence length after tokenization.",required=False) 
+    
+    parser.add_argument("--do_train", action='store_true',
+                        help="Whether to run training.")
+    parser.add_argument("--do_eval", action='store_true',
+                        help="Whether to run eval on the dev set.")
+    parser.add_argument("--do_test", action='store_true',
+                        help="Whether to run eval on the test set.")  
+    
+    parser.add_argument("--train_batch_size", default=4, type=int,
+                        help="Batch size for training.")
+    parser.add_argument("--eval_batch_size", default=4, type=int,
+                        help="Batch size for evaluation.")
+    parser.add_argument("--learning_rate", default=2e-5, type=float,
+                        help="The initial learning rate for Adam.")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float,
+                        help="Max gradient norm.")
+    parser.add_argument("--num_train_epochs", default=4, type=int,
+                        help="Total number of training epochs to perform.")
+
+    parser.add_argument('--seed', type=int, default=3407,
+                        help="random seed for initialization")  
+        
+    #print arguments
+    args = parser.parse_args()
+    return  args                     
+
+def create_model(args,model,tokenizer, config=None):
+    # logger.info("args.data_aug_type %s"%args.data_aug_type)
+    # replace token with type
+    if args.data_aug_type in ["replace_type" , "other"] and not args.only_save_the_nl_code_vec:
+        special_tokens_dict = {'additional_special_tokens': all_special_token}
+        logger.info(" new token %s"%(str(special_tokens_dict)))
+        num_added_toks = tokenizer.add_special_tokens(special_tokens_dict)
+        model.resize_token_embeddings(len(tokenizer))
+  
+    if (args.loaded_model_filename) and ("pytorch_model.bin" in args.loaded_model_filename):
+        logger.info("reload pytorch model from {}".format(args.loaded_model_filename))
+        model.load_state_dict(torch.load(args.loaded_model_filename),strict=False) 
+        # model.from_pretrain
+    if args.model_type ==  "base" :
+        model= Model(model)
+    elif args.model_type ==  "multi-loss-cocosoda":
+        model= Multi_Loss_CoCoSoDa(model,args, args.mlp)
+    if (args.loaded_model_filename) and ("pytorch_model.bin" not in args.loaded_model_filename) :
+        logger.info("reload model from {}".format(args.loaded_model_filename))
+        model.load_state_dict(torch.load(args.loaded_model_filename)) 
+        # model.load_state_dict(torch.load(args.loaded_model_filename,strict=False)) 
+        # model.from_pretrained(args.loaded_model_filename)  
+    if (args.loaded_codebert_model_filename) :
+        logger.info("reload pytorch model from {}".format(args.loaded_codebert_model_filename))
+        model.load_state_dict(torch.load(args.loaded_codebert_model_filename),strict=False)   
+    logger.info(model.model_parameters())
+
+
+    return model
+
+def main():
+    
+    args = parse_args()
+    #set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    #set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    pool = multiprocessing.Pool(cpu_cont)
+
+    # Set seed
+    set_seed(args.seed)
+
+    #build model
+
+    if "codet5" in   args.model_name_or_path:
+        config = T5Config.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
+        tokenizer =  RobertaTokenizer.from_pretrained(args.tokenizer_name)
+        model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path)
+        model = model.encoder
+    else:
+        config = RobertaConfig.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
+        tokenizer = RobertaTokenizer.from_pretrained(args.tokenizer_name)
+        model = RobertaModel.from_pretrained(args.model_name_or_path) 
+    model=create_model(args,model,tokenizer,config)
+
+    logger.info("Training/evaluation parameters %s", args)
+    args.start_step = 0
+
+    model.to(args.device)
+    
+    # Training
+    if args.do_multi_lang_continue_pre_train:
+        multi_lang_continue_pre_train(args, model, tokenizer, pool)
+        output_tokenizer_dir = os.path.join(args.output_dir,"tokenzier")                      
+        if not os.path.exists(output_tokenizer_dir):
+            os.makedirs( output_tokenizer_dir)    
+        tokenizer.save_pretrained( output_tokenizer_dir)
+    if args.do_train:
+        train(args, model, tokenizer, pool)
+     
+    
+    # Evaluation
+    results = {}
+
+    if args.do_eval:
+        checkpoint_prefix = 'checkpoint-best-mrr/model.bin'
+        output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix))  
+        if (not args.only_save_the_nl_code_vec) and (not args.do_zero_short) :
+            model.load_state_dict(torch.load(output_dir),strict=False)      
+        model.to(args.device)
+        result=evaluate(args, model, tokenizer,args.eval_data_file, pool)
+        logger.info("***** Eval valid results *****")
+        for key in sorted(result.keys()):
+            logger.info("  %s = %s", key, str(round(result[key],4)))
+            
+    if args.do_test:
+
+        logger.info("runnning test")
+        checkpoint_prefix = 'checkpoint-best-mrr/model.bin'
+        output_dir = os.path.join(args.output_dir, '{}'.format(checkpoint_prefix))  
+        if (not args.only_save_the_nl_code_vec) and (not args.do_zero_short) :
+            model.load_state_dict(torch.load(output_dir),strict=False)      
+        model.to(args.device)
+        result=evaluate(args, model, tokenizer,args.test_data_file, pool)
+        logger.info("***** Eval test results *****")
+        for key in sorted(result.keys()):
+            logger.info("  %s = %s", key, str(round(result[key],4)))
+        save_json_data(args.output_dir, "result.jsonl", result)
+    return results
+
+
+def gen_vector():
+    
+    args = parse_args()
+    #set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    #set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    pool = multiprocessing.Pool(cpu_cont)
+
+    # Set seed
+    set_seed(args.seed)
+    if "codet5" in   args.model_name_or_path:
+        config = T5Config.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
+        tokenizer =  RobertaTokenizer.from_pretrained(args.tokenizer_name)
+        model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path)
+        model = model.encoder
+    else:
+        config = RobertaConfig.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
+        tokenizer = RobertaTokenizer.from_pretrained(args.tokenizer_name)
+        model = RobertaModel.from_pretrained(args.model_name_or_path) 
+    model=create_model(args,model,tokenizer,config)
+
+    if args.data_aug_type ==  "replace_type" :
+        train_dataset=TextDataset(tokenizer, args, args.train_data_file, pool)
+    else:
+        # if "unixcoder" in args.model_name_or_path or "coco" in args.model_name_or_path :
+        train_dataset=TextDataset_unixcoder(tokenizer, args, args.train_data_file, pool)
+        # else:
+        #     train_dataset=TextDataset(tokenizer, args, args.train_data_file, pool)
+    train_sampler = SequentialSampler(train_dataset)
+    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size,num_workers=4,drop_last=False)
+
+    for idx in range(args.num_train_epochs):
+        output_dir_epoch = os.path.join(args.output_dir, '{}'.format(idx))
+        output_dir_epoch = os.path.join(output_dir_epoch, '{}'.format('model.bin'))
+
+        model.load_state_dict(torch.load(output_dir_epoch),strict=False)      
+        model.to(args.device)
+
+        model_eval = model.module if hasattr(model,'module') else model
+
+        all_nl_vec = []
+        all_code_vec = []
+
+        for step,batch in enumerate(train_dataloader):
+            code_inputs = batch[0].to(args.device)    
+            nl_inputs = batch[1].to(args.device)
+            #get code and nl vectors
+            nl_outputs = model_eval.nl_encoder_q(nl_inputs, attention_mask=nl_inputs.ne(1))
+            nl_vec =nl_outputs [1]
+            code_outputs = model_eval.code_encoder_q(code_inputs, attention_mask=code_inputs.ne(1))
+            code_vec =code_outputs [1]
+            all_nl_vec.append(nl_vec.detach().cpu().numpy())
+            all_code_vec.append(code_vec.detach().cpu().numpy())
+        all_nl_vec = np.concatenate(all_nl_vec, axis=0)
+        all_code_vec = np.concatenate(all_code_vec, axis=0)
+        print(all_nl_vec.shape, all_code_vec.shape)
+        np.save("/home/yiming/cocosoda/CoCoSoDa/saved_models/fine_tune/ruby/" + str(idx) + "/all_nl_vec.npy", all_nl_vec)
+        np.save("/home/yiming/cocosoda/CoCoSoDa/saved_models/fine_tune/ruby/" + str(idx) + "/all_code_vec.npy", all_code_vec)
+        idxs = [i for i in range(len(all_nl_vec))]
+        for epoch in range(1,2):
+            idxs_dir_path = "/home/yiming/cocosoda/CoCoSoDa/saved_models/codesearch_contrastive_learning/Model/Epoch_" + str(epoch)
+            if os.path.exists(idxs_dir_path):
+                pass
+            else:
+                os.mkdir(idxs_dir_path)
+            idxs_path = idxs_dir_path + "/index.json"
+            json_file = open(idxs_path, mode='w')
+            json.dump(idxs, json_file, indent=4)
+
+    if args.data_aug_type ==  "replace_type" :
+        test_dataset=TextDataset(tokenizer, args, args.test_data_file, pool)
+    else:
+        # if "unixcoder" in args.model_name_or_path or "coco" in args.model_name_or_path :
+        test_dataset=TextDataset_unixcoder(tokenizer, args, args.test_data_file, pool)
+        # else:
+        #     test_dataset=TextDataset(tokenizer, args, args.test_data_file, pool)
+    test_sampler = SequentialSampler(test_dataset)
+    test_dataloader = DataLoader(test_dataset, sampler=test_sampler, batch_size=args.train_batch_size,num_workers=4,drop_last=False)
+
+    for idx in range(args.num_train_epochs):
+        output_dir_epoch = os.path.join(args.output_dir, '{}'.format(idx))
+        output_dir_epoch = os.path.join(output_dir_epoch, '{}'.format('model.bin'))
+
+        model.load_state_dict(torch.load(output_dir_epoch),strict=False)      
+        model.to(args.device)
+
+        model_eval = model.module if hasattr(model,'module') else model
+
+        all_nl_vec = []
+        all_code_vec = []
+
+        for step,batch in enumerate(test_dataloader):
+            code_inputs = batch[0].to(args.device)    
+            nl_inputs = batch[1].to(args.device)
+            #get code and nl vectors
+            nl_outputs = model_eval.nl_encoder_q(nl_inputs, attention_mask=nl_inputs.ne(1))
+            nl_vec =nl_outputs [1]
+            code_outputs = model_eval.code_encoder_q(code_inputs, attention_mask=code_inputs.ne(1))
+            code_vec =code_outputs [1]
+            all_nl_vec.append(nl_vec.detach().cpu().numpy())
+            all_code_vec.append(code_vec.detach().cpu().numpy())
+        all_nl_vec = np.concatenate(all_nl_vec, axis=0)
+        all_code_vec = np.concatenate(all_code_vec, axis=0)
+        print(all_nl_vec.shape, all_code_vec.shape)
+        np.save("/home/yiming/cocosoda/CoCoSoDa/saved_models/fine_tune/ruby/" + str(idx) + "/test_all_nl_vec.npy", all_nl_vec)
+        np.save("/home/yiming/cocosoda/CoCoSoDa/saved_models/fine_tune/ruby/" + str(idx) + "/test_all_code_vec.npy", all_code_vec)
+
+def gen_label():
+    
+    args = parse_args()
+    #set log
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',level=logging.INFO )
+    #set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.n_gpu = torch.cuda.device_count()
+    args.device = device
+    logger.info("device: %s, n_gpu: %s",device, args.n_gpu)
+    
+    pool = multiprocessing.Pool(cpu_cont)
+
+    # # Set seed
+    # set_seed(args.seed)
+    # if "codet5" in   args.model_name_or_path:
+    #     config = T5Config.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
+    #     tokenizer =  RobertaTokenizer.from_pretrained(args.tokenizer_name)
+    #     model = T5ForConditionalGeneration.from_pretrained(args.model_name_or_path)
+    #     model = model.encoder
+    # else:
+    #     config = RobertaConfig.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
+    #     tokenizer = RobertaTokenizer.from_pretrained(args.tokenizer_name)
+    #     model = RobertaModel.from_pretrained(args.model_name_or_path) 
+    # model=create_model(args,model,tokenizer,config)
+
+    # if args.data_aug_type ==  "replace_type" :
+    #     train_dataset=TextDataset(tokenizer, args, args.train_data_file, pool)
+    # else:
+    #     # if "unixcoder" in args.model_name_or_path or "coco" in args.model_name_or_path :
+    #     train_dataset=TextDataset_unixcoder(tokenizer, args, args.train_data_file, pool)
+    #     # else:
+    #     #     train_dataset=TextDataset(tokenizer, args, args.train_data_file, pool)
+    # train_sampler = SequentialSampler(train_dataset)
+    # train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size,num_workers=4,drop_last=False)
+
+    code_list = []
+    docstring_list = []
+
+    with open(args.train_data_file, 'rt') as gz_file:
+        for line in gz_file:
+            data = json.loads(line)
+            code = data['code']
+            docstring = data['docstring']
+            
+            # 将 code 和 docstring 添加到列表中
+            code_list.append(code)
+            docstring_list.append(docstring)
+
+    print(len(code_list))
+    print(len(docstring_list))
+
+    # print(code_list[0])
+    # print(docstring_list[0])
+
+    code_output_file = '/home/yiming/cocosoda/CoCoSoDa/saved_models/fine_tune/ruby/code_list.json'
+    docstring_output_file = '/home/yiming/cocosoda/CoCoSoDa/saved_models/fine_tune/ruby/docstring_list.json'
+
+    # 存储 code_list 到 JSON 文件
+    with open(code_output_file, 'w') as file:
+        json.dump(code_list, file)
+
+    # 存储 docstring_list 到 JSON 文件
+    with open(docstring_output_file, 'w') as file:
+        json.dump(docstring_list, file)
+
+    code_list = []
+    docstring_list = []
+
+    with open(args.test_data_file, 'rt') as gz_file:
+        for line in gz_file:
+            data = json.loads(line)
+            code = data['code']
+            docstring = data['docstring']
+            
+            # 将 code 和 docstring 添加到列表中
+            code_list.append(code)
+            docstring_list.append(docstring)
+
+    print(len(code_list))
+    print(len(docstring_list))
+
+    code_output_file = '/home/yiming/cocosoda/CoCoSoDa/saved_models/fine_tune/ruby/test_code_list.json'
+    docstring_output_file = '/home/yiming/cocosoda/CoCoSoDa/saved_models/fine_tune/ruby/test_docstring_list.json'
+
+    # 存储 code_list 到 JSON 文件
+    with open(code_output_file, 'w') as file:
+        json.dump(code_list, file)
+
+    # 存储 docstring_list 到 JSON 文件
+    with open(docstring_output_file, 'w') as file:
+        json.dump(docstring_list, file)
+
+if __name__ == "__main__":
+    # main()
+    # gen_vector()
+    gen_label()
+

run_cocosoda.sh

@@ -0,0 +1,59 @@
+lang=ruby
+current_time=$(date "+%Y%m%d%H%M%S")
+# current_time=tmp
+
+code_length=64
+nl_length=64
+
+model_type=multi-loss-cocosoda #"base", "cocosoda" 
+moco_k=1024
+moco_m=0.999
+lr=2e-5
+moco_t=0.07
+
+epoch=10
+batch_size=128
+max_steps=100000
+save_steps=1000
+data_aug_type="replace_type"
+couninue_pre_train_data_files='dataset/java/train.jsonl dataset/javascript/train.jsonl  dataset/python/train.jsonl  dataset/php/train.jsonl  dataset/go/train.jsonl dataset/ruby/train.jsonl'
+CUDA_VISIBLE_DEVICES="0,1"
+base_model=unixcoder
+
+function continue_pre_train () {
+output_dir=./saved_models/cocosoda/
+mkdir -p $output_dir
+echo ${output_dir}
+CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}  python run.py  --eval_frequency  100 \
+    --moco_m ${moco_m} --moco_t ${moco_t}  \
+    --output_dir ${output_dir}  \
+    --moco_k ${moco_k} \
+    --model_type ${model_type} \
+    --data_aug_type other \
+    --config_name=microsoft/${base_model}-base  \
+    --model_name_or_path=microsoft/${base_model}-base \
+    --tokenizer_name=microsoft/${base_model}-base \
+    --lang=$lang \
+    --do_test \
+    --time_score 1 \
+    --do_multi_lang_continue_pre_train \
+    --max_steps ${max_steps} \
+    --save_steps ${save_steps} \
+    --gradient_accumulation_steps 1 \
+    --logging_steps 50 \
+    --couninue_pre_train_data_files  ${couninue_pre_train_data_files} \
+    --train_data_file=dataset/$lang/train.jsonl \
+    --eval_data_file=dataset/$lang/valid.jsonl  \
+    --test_data_file=dataset/$lang/test.jsonl \
+    --codebase_file=dataset/$lang/codebase.jsonl \
+    --num_train_epochs ${epoch} \
+    --code_length ${code_length} \
+    --nl_length ${nl_length} \
+    --train_batch_size ${batch_size} \
+    --eval_batch_size 64 \
+    --learning_rate ${lr} \
+    --seed 123456 2>&1| tee ${output_dir}/save_tokenizer.log
+}
+
+
+continue_pre_train 

run_fine_tune.sh

@@ -0,0 +1,55 @@
+lang=$1
+current_time=$(date "+%Y%m%d%H%M%S")
+
+code_length=64
+nl_length=64
+
+model_type=multi-loss-cocosoda #"base", "cocosoda" 
+moco_k=1024
+moco_m=0.999
+lr=2e-5
+moco_t=0.07
+
+batch_size=64
+max_steps=1000
+save_steps=100
+aug_type_way=random_replace_type 
+data_aug_type=random_mask
+
+base_model=DeepSoftwareAnalytics/CoCoSoDa
+epoch=5
+# echo ${base_model}
+CUDA_VISIBLE_DEVICES="0,1"
+# exit 111
+
+
+function fine-tune () {
+output_dir=./saved_models/fine_tune/${lang}
+mkdir -p $output_dir
+echo ${output_dir}
+ CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} python run.py   --eval_frequency  100 \
+    --moco_m ${moco_m} --moco_t ${moco_t}  \
+    --model_type ${model_type} \
+    --output_dir ${output_dir}  \
+    --data_aug_type ${data_aug_type} \
+    --moco_k ${moco_k} \
+    --config_name=${base_model}  \
+    --model_name_or_path=${base_model} \
+    --tokenizer_name=${base_model} \
+    --lang=$lang \
+    --do_train \
+    --do_test \
+    --train_data_file=dataset/$lang/train.jsonl \
+    --eval_data_file=dataset/$lang/valid.jsonl  \
+    --test_data_file=dataset/$lang/test.jsonl \
+    --codebase_file=dataset/$lang/codebase.jsonl \
+    --num_train_epochs ${epoch} \
+    --code_length ${code_length} \
+    --nl_length ${nl_length} \
+    --train_batch_size ${batch_size} \
+    --eval_batch_size 64 \
+    --learning_rate ${lr} \
+    --seed 123456 2>&1| tee ${output_dir}/running.log
+}
+
+fine-tune 

run_zero-shot.sh

@@ -0,0 +1,40 @@
+lang=$1
+current_time=$(date "+%Y%m%d%H%M%S")
+code_length=256
+nl_length=128
+model_type=cocosoda #"base", "cocosoda" ,"multi-loss-cocosoda"
+moco_k=1024
+moco_m=0.999
+lr=2e-5
+moco_t=0.07
+max_steps=1000
+aug_type_way=random_replace_type 
+data_aug_type=random_mask
+base_model=DeepSoftwareAnalytics/CoCoSoDa
+CUDA_VISIBLE_DEVICES=0
+
+function zero-shot () {
+output_dir=./saved_models/zero-shot/${lang}
+mkdir -p $output_dir
+
+ CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} python run.py   --eval_frequency  100 \
+    --do_zero_short \
+    --moco_m ${moco_m} --moco_t ${moco_t}  \
+    --model_type ${model_type} \
+    --output_dir ${output_dir}  \
+    --data_aug_type ${data_aug_type} \
+    --moco_k ${moco_k} \
+    --config_name=${base_model}  \
+    --model_name_or_path=${base_model} \
+    --tokenizer_name=${base_model} \
+    --lang=$lang \
+    --do_test \
+    --test_data_file=dataset/$lang/test.jsonl \
+    --codebase_file=dataset/$lang/codebase.jsonl \
+    --code_length ${code_length} \
+    --nl_length ${nl_length} \
+    --eval_batch_size 128 \
+    --learning_rate ${lr} \
+    --seed 123456 2>&1| tee ${output_dir}/running.log
+}
+zero-shot