import os

# GPT chain search
seed_list = [1]
target = 'single_chain_search'
dir_suffix = '2_step_reasoning'
lr = 2e-5
gpu_id = 0
batch_size = 100
scheduler = 'GradualWarmupScheduler_CosineAnnealingLR'
model = 'GPT'
N_train = 200000
N_test = 1000

dname = ['train', 'test']
dmode = ['train', 'test']
dmask = [0, 1]
dshow = [1, 1]
dpercent = [1, 1]

dn = ' '.join(map(str, dname))
dp = ' '.join(map(str, dpercent))
dmode = ' '.join(map(str, dmode))
dmask = ' '.join(map(str, dmask))
dshow = ' '.join(map(str, dshow))

L, H = 3, 1
suffix = f'{L}L{H}H'

os.system(f'CUDA_VISIBLE_DEVICES={gpu_id} python3 -m main -N_train {N_train} -N_test {N_test} -seed 1 -func {target} -lr {lr} -m {model}\
                -scheduler {scheduler} -ne 200 -nl {L} -nh {H} -bs {batch_size} -dir_suffix {dir_suffix} \
                -dmode {dmode} -dp {dp} -dn {dn} -dmask {dmask} -dshow {dshow} -suffix {suffix}\
                --use_mod_list_specific True -ple 1 -pae 5 -plae 10 -sme 5 -sl 13\
                --optim_T_max 200 --optim_eta_min 1e-5 --optim_multiplier 3 --optim_total_epoch 200')