File size: 2,776 Bytes
b410583 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
WORKDIR="your_CodeT5_path/CodeT5"
export PYTHONPATH=$WORKDIR
TASK=${1}
SUB_TASK=${2}
MODEL_TAG=${3}
GPU=${4}
DATA_NUM=${5}
BS=${6}
LR=${7}
SRC_LEN=${8}
TRG_LEN=${9}
PATIENCE=${10}
EPOCH=${11}
WARMUP=${12}
MODEL_DIR=${13}
SUMMARY_DIR=${14}
RES_FN=${15}
if [[ $DATA_NUM == -1 ]]; then
DATA_TAG='all'
else
DATA_TAG=$DATA_NUM
EPOCH=1
fi
if [[ ${TASK} == 'multi_task' ]]; then
FULL_MODEL_TAG=${MODEL_TAG}_${DATA_TAG}_lr${LR}_s${16}
else
FULL_MODEL_TAG=${MODEL_TAG}_${DATA_TAG}_lr${LR}_bs${BS}_src${SRC_LEN}_trg${TRG_LEN}_pat${PATIENCE}_e${EPOCH}
fi
if [[ ${SUB_TASK} == none ]]; then
OUTPUT_DIR=${MODEL_DIR}/${TASK}/${FULL_MODEL_TAG}
else
OUTPUT_DIR=${MODEL_DIR}/${TASK}/${SUB_TASK}/${FULL_MODEL_TAG}
fi
CACHE_DIR=${OUTPUT_DIR}/cache_data
RES_DIR=${OUTPUT_DIR}/prediction
LOG=${OUTPUT_DIR}/train.log
mkdir -p ${OUTPUT_DIR}
mkdir -p ${CACHE_DIR}
mkdir -p ${RES_DIR}
if [[ $MODEL_TAG == roberta ]]; then
MODEL_TYPE=roberta
TOKENIZER=roberta-base
MODEL_PATH=roberta-base
elif [[ $MODEL_TAG == codebert ]]; then
MODEL_TYPE=roberta
TOKENIZER=roberta-base
MODEL_PATH=microsoft/codebert-base
elif [[ $MODEL_TAG == bart_base ]]; then
MODEL_TYPE=bart
TOKENIZER=facebook/bart-base
MODEL_PATH=facebook/bart-base
elif [[ $MODEL_TAG == codet5_small ]]; then
MODEL_TYPE=codet5
TOKENIZER=Salesforce/codet5-small
MODEL_PATH=Salesforce/codet5-small
elif [[ $MODEL_TAG == codet5_base ]]; then
MODEL_TYPE=codet5
TOKENIZER=Salesforce/codet5-base
MODEL_PATH=Salesforce/codet5-base
elif [[ $MODEL_TAG == codet5_large ]]; then
MODEL_TYPE=codet5
TOKENIZER=Salesforce/codet5-large
MODEL_PATH=Salesforce/codet5-large
fi
if [[ ${TASK} == 'multi_task' ]]; then
RUN_FN=${WORKDIR}/run_multi_gen.py
MULTI_TASK_AUG='--max_steps '${16}' --save_steps '${17}' --log_steps '${18}
elif [[ ${TASK} == 'clone' ]]; then
RUN_FN=${WORKDIR}/run_clone.py
elif [[ ${TASK} == 'defect' ]] && [[ ${MODEL_TYPE} == 'roberta' || ${MODEL_TYPE} == 'bart' ]]; then
RUN_FN=${WORKDIR}/run_defect.py
else
RUN_FN=${WORKDIR}/run_gen.py
fi
CUDA_VISIBLE_DEVICES=${GPU} \
python ${RUN_FN} ${MULTI_TASK_AUG} \
--do_train --do_eval --do_eval_bleu --do_test \
--task ${TASK} --sub_task ${SUB_TASK} --model_type ${MODEL_TYPE} --data_num ${DATA_NUM} \
--num_train_epochs ${EPOCH} --warmup_steps ${WARMUP} --learning_rate ${LR}e-5 --patience ${PATIENCE} \
--tokenizer_name=${TOKENIZER} --model_name_or_path=${MODEL_PATH} --data_dir ${WORKDIR}/data \
--cache_path ${CACHE_DIR} --output_dir ${OUTPUT_DIR} --summary_dir ${SUMMARY_DIR} \
--save_last_checkpoints --always_save_model --res_dir ${RES_DIR} --res_fn ${RES_FN} \
--train_batch_size ${BS} --eval_batch_size ${BS} --max_source_length ${SRC_LEN} --max_target_length ${TRG_LEN} \
2>&1 | tee ${LOG}
|