shard=3

exp_name='debug_rm_rm'

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/rafbart_v2_fulltest'
cache_path='/data2/mnskim/data/tbg/caches/datasets/debug_rafbart_v2_fulltest'

n_epochs=10

load_ckpt="/data2/mnskim/checkpoints/tbg1/ckpts/models/worldformer/exp_rafbart_6_finetune_epoch0/model_epoch3.pt"

gpu=$((shard + 4))
#gpu=2
CUDA_VISIBLE_DEVICES=$gpu 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 4 \
                                                                                       --test_batchsize 4 \
                                                                                       --n_accum 4 \
                                                                                       --mode 'predict' \
                                                                                       --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 \
                                                                                       --add_atomic_tokens \
                                                                                       --load_path $load_ckpt \
                                                                                       --train_negative_mode 'sample' \
                                                                                       --val_negative_mode 'sample' \
                                                                                       --test_negative_mode 'all' \
                                                                                       --use_state_vec \
                                                                                       --use_graph_vec \
                                                                                       --task_mode 'raf_multitask' \
                                                                                       --test_roms_shard $shard \
                                                                                       #--use_oracle_template

