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"
		# mpi_cmd="mpirun -n $mpi_num"
	else
		mpi_cmd=""
	fi

	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 ""
}

run_baichuan(){
	get_path

	engine_dir=${model_home}/engine/baichuan${model_version_num}-${model_size}/${precision}/${tp}
	tokenizer_dir=${model_home}/downloads/baichuan${model_version_num}-${model_size}

	env_cmd="PYTORCH_NO_XPU_MEMORY_CACHING=0 XMLIR_D_XPU_L3_SIZE=0 "
	required_cmd="--engine_dir=$engine_dir --tokenizer_dir=$tokenizer_dir"
	options="--max_output_len=128 --log_level=info"
	# inputs="--input_text='世界前五的高峰是？'"

	prof_cmd="--performance_test_scale=\
1x512x512E1x1024x1024E1x2048x64E1x2048x2048E\
2x512x512E2x1024x1024E2x2048x64E2x2048x2048E\
4x512x512E4x1024x1024E4x2048x64E4x2048x2048E\
8x512x512E8x1024x1024E8x2048x64E8x2048x2048E\
16x2048x2048E\
32x128x2048E32x2048x128E\
64x128x128E"
	prof_cmd="--performance_test_scale=20x1024x1024E32x1024x1024E48x1024x1024"
	cmd="${env_cmd} $mpi_cmd python3 run.py ${required_cmd} ${options} ${inputs} ${prof_cmd}"

	echo "==================== cmd ======================"
	echo $cmd
	eval $cmd

}

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

model_name="$1"
# device_id="$2"
log_file=./logs/${model_name}_"$(date '+%Y-%m-%d-%H:%M:%S')".log
export XPU_VISIBLE_DEVICES=$2
# export XPU_VISIBLE_DEVICES=0,1



case $model_name in
 "v1_13b_fp16_tp1")
	run_baichuan |& tee $log_file 2>&1 &
   ;;
 "v1_13b_int8_tp1")
	run_baichuan |& tee $log_file 2>&1 &
;;
 "v1_13b_fp16_tp2")
	run_baichuan |& tee $log_file 2>&1 &
   ;;
 "v1_13b_int8_tp2")
	run_baichuan |& tee $log_file 2>&1 &
   ;;
 "v1_7b_fp16_tp1")
	run_baichuan |& tee $log_file 2>&1 &
   ;;
 "v1_7b_int8_tp1")
	run_baichuan |& tee $log_file 2>&1 &
   ;;
 "v2_13b_fp16_tp1")
	run_baichuan |& tee $log_file 2>&1 &
   ;;
 "v2_13b_int8_tp1")
	run_baichuan |& tee $log_file 2>&1 &
   ;;
 "v2_13b_fp16_tp2")
	run_baichuan
   ;;
 "v2_13b_int8_tp2")
	run_baichuan |& tee $log_file 2>&1 &
   ;;
 *)
   echo "Unknown model name: $model_name"
   exit 1
   ;;
esac