template.conf

# SLURM settings
export ACCOUNT=${ACCOUNT:-"YOUR_ACCOUNT"}
export PARTITION=${PARTITION:-"YOUR_PARTITION"}
export RUN_NAME=${RUN_NAME:-"YOUR_RUN_NAME"}
export CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-"YOUR_CONTAINER_MOUNTS"}

# Output path
export OUTPUT_PATH=${OUTPUT_PATH:-"YOUR_OUTPUT_DIR/${MODEL}-TP${TP}PP${PP}EP${EP}VPP${VPP}-MBS${MBS}GBS${GBS}"}

# Data paths
declare -A DATA_PATHS=(
    # DATASET:MODEL:                DATA_PATH
    ["{YOUR_DATASET}:Mixtral-8x2B"]="YOUR_DATA_PATH"
    ["{YOUR_DATASET}:Mixtral-8x7B"]="YOUR_DATA_PATH"
    ["{YOUR_DATASET}:Mixtral-8x22B"]="YOUR_DATA_PATH"
    ["{YOUR_DATASET}:DeepSeek-V2"]="YOUR_DATA_PATH"
    ["{YOUR_DATASET}:DeepSeek-V2-Lite"]="YOUR_DATA_PATH"
    ["{YOUR_DATASET}:DeepSeek-V3"]="YOUR_DATA_PATH"
    ["{YOUR_DATASET}:Qwen2-57B-A14B"]="YOUR_DATA_PATH"
)

# Tokenizer models
declare -A TOKENIZER_MODELS=(
    # DATASET:MODEL:  TOKENIZER_MODEL
    ["Mixtral-8x2B"]="YOUR_TOKENIZER_MODEL"
    ["Mixtral-8x7B"]="YOUR_TOKENIZER_MODEL"
    ["Mixtral-8x22B"]="YOUR_TOKENIZER_MODEL"
)

# Checkpoint paths
declare -A CHECKPOINT_PATHS=(
    # MODEL:         CHECKPOINT_PATH
    ["Mixtral-8x7B"]="YOUR_CHECKPOINT_PATH"
    ["Mixtral-8x22B"]="YOUR_CHECKPOINT_PATH"
    ["DeepSeek-V2"]="YOUR_CHECKPOINT_PATH"
    ["DeepSeek-V2-Lite"]="YOUR_CHECKPOINT_PATH"
    ["DeepSeek-V3"]="YOUR_CHECKPOINT_PATH"
    ["Qwen2-57B-A14B"]="YOUR_CHECKPOINT_PATH"
)

# Function to get path based on MODEL
get_path_based_on_model() {
    local -n paths=$1
    local model_name=$2
    # Verify if model_name exists in paths array
    if [[ -v paths[$model_name] ]]; then
        echo "${paths[$model_name]}"
    else
        echo ""
    fi
}

# Set DATA_PATH based on MODEL
export DATA_PATH=${DATA_PATH:-${DATA_PATHS["${DATASET}:${MODEL}"]:-""}}

# Set TOKENIZER_MODEL based on MODEL
export TOKENIZER_MODEL=${TOKENIZER_MODEL:-${TOKENIZER_MODELS["${MODEL}"]:-""}}

# Set LOAD_PATH based on MODEL
if [[ ${PRETRAIN} == 0 ]]; then
    export LOAD_PATH=${LOAD_PATH:-${CHECKPOINT_PATHS["${MODEL}"]:-""}}
fi