#!/bin/bash
# Activate conda
source ~/miniconda3/etc/profile.d/conda.sh
conda activate bpe

# select GPU
export CUDA_VISIBLE_DEVICES=6

work_dir=/home/n5huang/dna_token
data_dir=${work_dir}/output_tokens
cache_dir=${data_dir}/cache

runnum=len2_5120

# -p to use existing directory
mkdir -p ${work_dir}/pretrain/models/${runnum}
output_dir=${work_dir}/pretrain/models/${runnum}

# CHANGED: 
# --config_name ${work_dir}/BPE_merge/config_5120.json -> --config_name ${work_dir}/pretrain/models/len2_5120/config.json 
# --tokenizer_name ${work_dir}/merge_bpe/... -> --tokenizer_name ${work_dir}/tokenizer_evaluation/merge_bpe/...
# all tf-idf -> len2_5120 (model, tokenizer, data)
# ADDED:
# --resume_from_checkpoint /home/n5huang/dna_token/pretrain/models/len2_5120/checkpoint-100000
# REMOVED:
# overwrite_output_dir True

cp /home/n5huang/dna-tokenizer/BPE_merge/04.pretrain_200k.sh ${output_dir}
python /home/n5huang/dna-tokenizer/BPE_merge/run_mlm_200k.py \
					--output_dir ${output_dir} \
					--model_type bert \
					--tokenizer_name ${work_dir}/tokenizer_evaluation/merge_bpe/vocab_5120/merge_tokenizer_unigram_len2.json \
					--config_name ${work_dir}/pretrain/models/len2_5120/config.json \
					--project_name token_eval_len2_5120 \
					--do_train True \
					--model_max_length 512 \
					--max_seq_length 512 \
					--line_by_line True \
					--pad_to_max_length True \
					--train_file ${data_dir}/merge_bpe_5120_allchr_all_tokenized_train_chrOnly.tsv \
					--validation_file ${data_dir}/merge_bpe_5120_allchr_all_tokenized_val_chrOnly.tsv \
					--cache_dir ${cache_dir} \
					--use_fast_tokenizer True \
					--do_eval True \
					--gradient_accumulation_steps 1 \
					--per_device_train_batch_size 96 \
					--per_device_eval_batch_size 96 \
					--save_steps 1000 \
					--save_total_limit 10 \
					--max_steps 200000 \
					--logging_steps 1000 \
					--learning_rate 4e-5 \
					--adam_epsilon 1e-6 \
					--weight_decay 0.01 \
					--adam_beta1 0.9 \
					--adam_beta2 0.98 \
					--mlm_probability 0.15 \
					--warmup_steps 10000 \
					--seed 42 \
					--preprocessing_num_workers 8 \
					--resume_from_checkpoint /home/n5huang/dna_token/pretrain/models/len2_5120/checkpoint-100000