load_path='/data2/mnskim/checkpoints/tbg1/ckpts/models/worldformer/exp_bartbase_5_5_hf_new/model_epoch9.pt'

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

n_epochs=20

inferred_templates="/data2/mnskim/data/tbg/results/template_cls/exp_template_binary_cls_ddp4_model_epoch1.pt/test_cached.pkl"

CUDA_VISIBLE_DEVICES=7 python -m worldformer2.training.train_bart --train_file $train_file \
                                                                                   --test_file $test_file \
                                                                                   --cache_path $cache_path \
                                                                                   --load_path $load_path \
                                                                                   --n_epochs $n_epochs \
                                                                                   --train_batchsize 4 \
                                                                                   --n_accum 4 \
                                                                                   --mode 'predict' \
                                                                                   --lr 1e-4 \
                                                                                   --ckpt_policy 'best_only' \
                                                                                   --tokenizer_path $tokenizer_path \
                                                                                   --graph_encoder_tokenizer_path $graph_encoder_tokenizer_path \
                                                                                   --action_decoder_tokenizer_path $action_decoder_tokenizer_path \
                                                                                   --graph_decoder_tokenizer_path $graph_decoder_tokenizer_path \
                                                                                   --num_workers 1 \
                                                                                   --log_freq 32 \
                                                                                   --task_mode 'full_multitask' \
                                                                                   --add_atomic_tokens \
                                                                                   --inferred_templates $inferred_templates
                                                                                   #--save_gen
                                                                                   #--add_ckg
                                                                                   #--tie_embeddings
                                                                                   #--refresh_cache 
