#!/bin/bash
source /usr/local/Ascend/nnal/atb/set_env.sh --cxx_abi=0
source /usr/local/Ascend/ascend-toolkit/set_env.sh

script_path=$(realpath "${BASH_SOURCE[0]}")
script_dir=$(dirname "$script_path")
parent_dir=$(dirname "$script_dir")
MindSpeed_Core_MS_PATH=$(dirname "$parent_dir")

switch=${1:-msa}
if [ "$switch" != "nomsa" ]; then
    export PYTHONPATH=${MindSpeed_Core_MS_PATH}/msadapter/:${MindSpeed_Core_MS_PATH}/MSAdapter/:$PYTHONPATH
    export PYTHONPATH=${MindSpeed_Core_MS_PATH}/msadapter/msa_thirdparty/:${MindSpeed_Core_MS_PATH}/MSAdapter/msa_thirdparty:$PYTHONPATH
fi

export PYTHONPATH=${MindSpeed_Core_MS_PATH}/MindSpeed-LLM/:${MindSpeed_Core_MS_PATH}/Megatron-LM/:${MindSpeed_Core_MS_PATH}/MindSpeed/:$PYTHONPATH
echo "..............................................done set PYTHONPATH"
echo $PYTHONPATH