run.py

# 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()