exp_name='exp_bartbase_5_5_hf_new'
#exp_name='exp_basemodel_2'
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'

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/v7_bart_baseline'

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=4 python -m worldformer2.training.train_bart --train_file $train_file \
                                                                           --test_file $test_file \
                                                                           --cache_path $cache_path \
                                                                           --save_path $save_path \
                                                                           --n_epochs $n_epochs \
                                                                           --train_batchsize 2 \
                                                                           --n_accum 8 \
                                                                           --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 \
                                                                           --task_mode 'full_multitask' \
                                                                           --add_atomic_tokens \
                                                                           --wandb_project 'tbg' \
                                                                           --wandb_name $exp_name \
                                                                           #--refresh_cache \
                                                                           #--pretrained_ckpt $pretrained_ckpt \
                                                                           #--init_mode 'graph'
                                                                           #--add_ckg \
                                                                           #--refresh_cache 
