build_baichuan() {
    get_path

    cmd="XTCL_BUILD_DEBUG=1 python3 build.py ${tp_cmd} --model_version $model_name \
                    --model_dir ${model_home}/downloads/baichuan${model_version_num}-${model_size} \
                    --dtype float16 ${int8_cmd} \
                    --use_gemm_plugin float16 \
                    --use_gpt_attention_plugin float16 \
                    --output_dir ${model_home}/engine/baichuan${model_version_num}-${model_size}/${precision}/${tp}"
	echo "******************** cmd *********************"
	echo $cmd
	eval ${cmd} |& tee ${log_file}
}

get_path(){
	model_home=/home/workspace
	model_version=$(echo $model_name | cut -d "_" -f 1)
	model_size=$(echo $model_name | cut -d "_" -f 2)
	precision=$(echo $model_name | cut -d "_" -f 3)
	tp=$(echo $model_name | cut -d "_" -f 4)
    model_name=${model_version}_${model_size}
	model_version_num=$(echo $model_version | grep -o '[0-9]\+')

	if [[ "$model_version_num" == "1" ]]; then
		model_version_num=""
	fi

	mpi_num=$(echo $tp | cut -d "p" -f 2)
	if (( $mpi_num > 1 )); then
		mpi_cmd="mpirun --allow-run-as-root -n $mpi_num"
		tp_cmd="--parallel_build --world_size $mpi_num"
	else
		mpi_cmd=""
		tp_cmd=""
	fi

    if [[ "$precision" == "int8" ]]; then
        int8_cmd="--use_weight_only"
    else
        int8_cmd=""
    fi

	echo "------------------------------------------------------"
	log_file=./logs/relay_${model_name}_"$(date '+%Y-%m-%d-%H:%M:%S')".log
	echo "log file -> ${log_file} "

	echo -e "\033[1;31m" # 设置红色字体
	echo "Model version      Model size      Precision      TP"
	echo -e "\033[0m" # 重置字体颜色
	echo "------------------------------------------------------"
	echo -e "\033[0;32m"
	echo "$model_version" "          " "$model_size" "          " "$precision" "        " "$tp"
	echo ""
}


if [ "$#" -ne 1 ]; then
 echo "Usage: $0 -m=<model_name>"
 exit 1
fi

model_name="$1"

case $model_name in
 "v1_13b_fp16_tp1")
	build_baichuan
   ;;
 "v1_13b_int8_tp1")
	build_baichuan
   ;;
 "v1_13b_fp16_tp2")
	build_baichuan
   ;;
 "v1_13b_int8_tp2")
	build_baichuan
   ;;
 "v1_7b_fp16_tp1")
	build_baichuan
   ;;
 "v1_7b_int8_tp1")
	build_baichuan
   ;;
 "v2_13b_fp16_tp1")
	build_baichuan
   ;;
 "v2_13b_int8_tp1")
	build_baichuan
   ;;
 "v2_13b_fp16_tp2")
	build_baichuan
   ;;
 "v2_13b_int8_tp2")
	build_baichuan
   ;;
 *)
   echo "Unknown model name: $model_name"
   exit 1
   ;;
esac