exp_name='exp_rafbart_6_pretrain'

save_path="/data2/mnskim/checkpoints/tbg1/ckpts/models/worldformer/${exp_name}"
#lr=1e-4
lr=3e-5
#mkdir -p $save_path

# Tokenizer
tokenizer_path='facebook/bart-base'
graph_encoder_tokenizer_path='facebook/bart-base'
action_decoder_tokenizer_path='facebook/bart-base'
graph_decoder_tokenizer_path='facebook/bart-base'

# new data
train_file='/data2/mnskim/data/tbg/JerichoWorld/datasets/v7/train.json'
test_file='/data2/mnskim/data/tbg/JerichoWorld/datasets/v7/test.json'

cache_path='/data2/mnskim/data/tbg/caches/datasets/rafbart_pretrain'

n_epochs=10

pretrained_ckpt='/data2/mnskim/checkpoints/tbg1/ckpts/models/worldformer/exp_newdata_v3_bart_pretrain_atomic2020_b256_epochfix/model_epoch0.pt'

CUDA_VISIBLE_DEVICES=7 python /home/mnskim/workspace/tbg/tbg1/worldformer2/training/train_rafbart_v2.py --train_file $train_file \
                                                                                       --test_file $test_file \
                                                                                       --cache_path $cache_path \
                                                                                       --save_path $save_path \
                                                                                       --n_epochs $n_epochs \
                                                                                       --train_batchsize 4 \
                                                                                       --val_batchsize 16 \
                                                                                       --test_batchsize 8 \
                                                                                       --n_accum 4 \
                                                                                       --mode 'train' \
                                                                                       --lr $lr \
                                                                                       --ckpt_policy 'all' \
                                                                                       --tokenizer_path $tokenizer_path \
                                                                                       --graph_encoder_tokenizer_path $tokenizer_path \
                                                                                       --action_decoder_tokenizer_path $tokenizer_path \
                                                                                       --graph_decoder_tokenizer_path $tokenizer_path \
                                                                                       --num_workers 8 \
                                                                                       --log_freq 32 \
                                                                                       --use_state_vec \
                                                                                       --use_graph_vec \
                                                                                       --task_mode 'full_multitask' \
                                                                                       --add_atomic_tokens \
                                                                                       --skip_nohole \
                                                                                       --wandb_project 'tbg' \
                                                                                       --wandb_name $exp_name \
                                                                                       #--refresh_cache \
                                                                                       #--dist-backend 'nccl' \
                                                                                       #--multiprocessing-distributed \
                                                                                       #--world-size 1 \
                                                                                       #--local_rank 0 \
                                                                                       #--dist-url 'tcp://127.0.0.1:3456' \
                                                                                       #--refresh_cache \
                                                                                       #--wandb_project 'tbg' \
                                                                                       #--wandb_name $exp_name \
                                                                                       #--pretrained_ckpt $pretrained_ckpt \
                                                                                       #--init_mode 'graph'
                                                                                       #--add_ckg \
                                                                                       
