settings.py

import os


DATA_NAME = os.environ.get("DATA_NAME", "USPTO")
EXP_NAME = os.environ.get("EXP_NAME", "")

NUM_GPU = int(os.environ.get("NUM_GPUS_PER_NODE", 1))


TRAIN_BATCH_SIZE = int(os.environ.get("TRAIN_BATCH_SIZE", 4096))
VAL_BATCH_SIZE = int(os.environ.get("VAL_BATCH_SIZE", 4096))
TEST_BATCH_SIZE = int(os.environ.get("TEST_BATCH_SIZE", 2048 * NUM_GPU))

NUM_NODES = int(os.environ.get("NUM_NODES", 1))
ACCUMULATION_COUNT = int(os.environ.get("ACCUMULATION_COUNT", 1))
NUM_WORKERS = int(os.environ.get("NUM_WORKERS", 16))


class Args:
    
    # train #
    exp_name = EXP_NAME
    train_path = os.environ.get("TRAIN_FILE") 
    val_path = os.environ.get("VAL_FILE")
    test_path = os.environ.get("TEST_FILE")
    model_path = os.environ.get("MODEL_PATH")
    result_path = os.environ.get("RESULT_PATH")
    data_name = f"{DATA_NAME}"
    log_file = f"EZSolver"
    load_from = str(os.environ.get("LOAD_FROM", ''))
    resume = False

    #specific reaction for EZFlow
    exam_mode = os.environ.get("EXAM_MODE", "False").lower() == "true"
    target_index = int(os.environ.get("TARGET_INDEX", 0))

    backend = "nccl"
    local_rank = int(os.environ.get("LOCAL_RANK", -1))
    num_workers = NUM_WORKERS
    emb_dim = int(os.environ.get("EMB_DIM", 256))
    enc_num_layers = 12
    post_processing_layers = 1
    enc_heads = 32
    enc_filter_size = 2048
    dropout = float(os.environ.get("DROPOUT", 0.0))
    attn_dropout = float(os.environ.get("ATTN_DROPOUT", 0.0))
    rel_pos = "emb_only"
    shared_attention_layer = 0
    sigma = float(os.environ.get("SIGMA", 0.09))
    train_batch_size = (TRAIN_BATCH_SIZE / ACCUMULATION_COUNT / NUM_GPU / NUM_NODES)
    val_batch_size = (VAL_BATCH_SIZE / ACCUMULATION_COUNT / NUM_GPU / NUM_NODES)
    test_batch_size = TEST_BATCH_SIZE
    batch_type = "tokens_sum"
    lr = float(os.environ.get("LR", 0.0001))
    beta1 = 0.9
    beta2 = 0.998
    eps = 1e-9
    weight_decay = float(os.environ.get("WEIGHT_DECAY", 1e-2))
    warmup_steps = int(os.environ.get("WARMUP_STEPS", 30000))
    clip_norm = float(os.environ.get("CLIP_NORM", 20000.0))
    c_weight = float(os.environ.get("C_WEIGHT", 1.0))   
    num_train_samples = int(os.environ.get("NUM_TRAIN_SAMPLES", 5))
    num_val_samples = int(os.environ.get("NUM_VAL_SAMPLES", 5))  

    epoch = int(os.environ.get("EPOCH", 100))
    max_steps = int(os.environ.get("MAX_STEPS", 3000000))
    accumulation_count = ACCUMULATION_COUNT
    save_iter = int(os.environ.get("SAVE_ITER", 30000))
    log_iter = int(os.environ.get("LOG_ITER", 100))
    eval_iter = int(os.environ.get("EVAL_ITER", 30000))
    inf_iter = int(os.environ.get("INF_ITER", 50))

    sample_size = int(os.environ.get("SAMPLE_SIZE", 16))
    rbf_low = 0.0
    rbf_high = float(os.environ.get("RBF_HIGH", 12.0))
    rbf_gap = float(os.environ.get("RBF_GAP", 0.1))

    reactions = os.environ.get("REACTIONS")
    mechanism_results = os.environ.get("MECHANISM_RESULTS")
    beam_size = int(os.environ.get("BEAM_SIZE", 10))
    max_depth = int(os.environ.get("MAX_DEPTH", 4))
    bidirectional = os.environ.get("BIDIRECTIONAL", "True").lower() == "true"
    filter_2d = os.environ.get("FILTER_2D", "True").lower() == "true"
    filter_chem = os.environ.get("FILTER_CHEM", "True").lower() == "true"

    use_wandb = os.environ.get("USE_WANDB", "False").lower() == "true"

    seed = int(os.environ.get("SEED", 42))