#!/usr/bin/env bash
|
|
. ./path.sh || exit 1;
|
|
# machines configuration
|
CUDA_VISIBLE_DEVICES="0,1"
|
gpu_num=2
|
count=1
|
train_cmd=utils/run.pl
|
infer_cmd=utils/run.pl
|
|
# general configuration
|
lang=zh
|
nlsyms_txt=none # Non-linguistic symbol list if existing.
|
cleaner=none # Text cleaner.
|
g2p=none # g2p method (needed if token_type=phn).
|
lm_fold_length=150 # fold_length for LM training.
|
word_vocab_size=10000 # Size of word vocabulary.
|
token_type=char
|
lm_token_list=
|
|
nj=10
|
## path to AISHELL2 trans
|
lm_train_text=
|
lm_dev_text=
|
lm_test_text=
|
|
train_data_path_and_name_and_type=${lm_train_text},text,text
|
train_shape_file=
|
valid_data_path_and_name_and_type=${lm_dev_text},text,text
|
valid_shape_file=
|
lm_config=conf/train_lm_transformer.yaml
|
exp_dir=./data
|
tag=exp1
|
model_dir="baseline_$(basename "${lm_config}" .yaml)_${lang}_${token_type}_${tag}"
|
lm_exp=${exp_dir}/exp/${model_dir}
|
inference_lm=valid.loss.ave.pb # Language model path for decoding.
|
|
stage=0
|
stop_stage=3
|
|
. utils/parse_options.sh || exit 1;
|
|
# Set bash to 'debug' mode, it will exit on :
|
# -e 'error', -u 'undefined variable', -o ... 'error in pipeline', -x 'print commands',
|
set -e
|
set -u
|
set -o pipefail
|
|
min() {
|
local a b
|
a=$1
|
for b in "$@"; do
|
if [ "${b}" -le "${a}" ]; then
|
a="${b}"
|
fi
|
done
|
echo "${a}"
|
}
|
|
# you can set gpu num for decoding here
|
gpuid_list=$CUDA_VISIBLE_DEVICES # set gpus for decoding, e.g., gpuid_list=2,3, the same as training stage by default
|
ngpu=$(echo $gpuid_list | awk -F "," '{print NF}')
|
|
mkdir -p ${exp_dir}/exp/${model_dir}
|
token_list=${exp_dir}/exp/${model_dir}/vocab.txt
|
blank="<blank>" # CTC blank symbole
|
sos="<s>" # sos symbole
|
eos="</s>" # eos symbole
|
oov="<unk>" # Out of vocabulary symbol.
|
if [ ${stage} -le 0 ] && [ ${stop_stage} -ge 0 ]; then
|
if [ "${token_type}" = char ] || [ "${token_type}" = word ]; then
|
echo "Stage 0: Generate character level token_list from ${lm_train_text}"
|
|
# The first symbol in token_list must be "<blank>":
|
# 0 is reserved for CTC-blank for ASR and also used as ignore-index in the other task
|
python -m funasr.bin.tokenize_text \
|
--token_type "${token_type}" \
|
--input "${lm_train_text}" \
|
--output "${token_list}" \
|
--non_linguistic_symbols "${nlsyms_txt}" \
|
--field 2- \
|
--cleaner "${cleaner}" \
|
--g2p "${g2p}" \
|
--write_vocabulary true \
|
--add_symbol "${blank}:0" \
|
--add_symbol "${sos}:1" \
|
--add_symbol "${eos}:2" \
|
--add_symbol "${oov}:-1"
|
|
else
|
echo "Error: not supported --token_type '${token_type}'"
|
exit 2
|
fi
|
|
## use_word_lm=false
|
## # Create word-list for word-LM training
|
## if ${use_word_lm} && [ "${token_type}" != word ]; then
|
## echo "Generate word level token_list from ${lm_train_text}"
|
## python -m funasr.bin.tokenize_text \
|
## --token_type word \
|
## --input "${lm_train_text}" \
|
## --output "${token_list}" \
|
## --field 2- \
|
## --cleaner "${cleaner}" \
|
## --g2p "${g2p}" \
|
## --write_vocabulary true \
|
## --vocabulary_size "${word_vocab_size}" \
|
## --add_symbol "${blank}:0" \
|
## --add_symbol "${sos}:1" \
|
## --add_symbol "${eos}:2" \
|
## --add_symbol "${oov}:-1"
|
## fi
|
|
lm_token_list="${token_list}"
|
|
fi
|
|
|
if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
|
echo "stage 1: Data preparation"
|
|
# 1. Split the key file
|
_logdir="${exp_dir}/exp/${model_dir}/log"
|
mkdir -p "${_logdir}"
|
# Get the minimum number among ${nj} and the number lines of input files
|
_nj=$(min "${nj}" "$(<${lm_train_text} wc -l)" "$(<${lm_dev_text} wc -l)")
|
|
key_file="${lm_train_text}"
|
split_scps=""
|
for n in $(seq ${_nj}); do
|
split_scps+=" ${_logdir}/train.${n}.scp"
|
done
|
# shellcheck disable=SC2086
|
utils/split_scp.pl "${key_file}" ${split_scps}
|
|
key_file="${lm_dev_text}"
|
split_scps=""
|
for n in $(seq ${_nj}); do
|
split_scps+=" ${_logdir}/dev.${n}.scp"
|
done
|
# shellcheck disable=SC2086
|
utils/split_scp.pl "${key_file}" ${split_scps}
|
|
# 2. Submit jobs
|
## python ../../funasr/bin/lm_train.py \
|
${train_cmd} JOB=1:"${_nj}" "${_logdir}"/stats.JOB.log \
|
python -m funasr.bin.lm_train \
|
--collect_stats true \
|
--use_preprocessor true \
|
--token_type "${token_type}" \
|
--token_list "${lm_token_list}" \
|
--non_linguistic_symbols "${nlsyms_txt}" \
|
--cleaner "${cleaner}" \
|
--g2p "${g2p}" \
|
--train_data_path_and_name_and_type "${lm_train_text},text,text" \
|
--valid_data_path_and_name_and_type "${lm_dev_text},text,text" \
|
--train_shape_file "${_logdir}/train.JOB.scp" \
|
--valid_shape_file "${_logdir}/dev.JOB.scp" \
|
--output_dir "${_logdir}/stats.JOB" \
|
--config ${lm_config} || { cat "${_logdir}"/stats.*.log; exit 1; }
|
|
# 3. Aggregate shape files
|
_opts=
|
for i in $(seq "${_nj}"); do
|
_opts+="--input_dir ${_logdir}/stats.${i} "
|
done
|
lm_stats_dir=${exp_dir}/exp/${model_dir}
|
# shellcheck disable=SC2086
|
python -m funasr.bin.aggregate_stats_dirs ${_opts} --output_dir "${lm_stats_dir}"
|
|
# Append the num-tokens at the last dimensions. This is used for batch-bins count
|
<"${lm_stats_dir}/train/text_shape" \
|
awk -v N="$(<${lm_token_list} wc -l)" '{ print $0 "," N }' \
|
>"${lm_stats_dir}/train/text_shape.${token_type}"
|
|
<"${lm_stats_dir}/valid/text_shape" \
|
awk -v N="$(<${lm_token_list} wc -l)" '{ print $0 "," N }' \
|
>"${lm_stats_dir}/valid/text_shape.${token_type}"
|
|
train_shape_file=${lm_stats_dir}/train/text_shape.${token_type}
|
valid_shape_file=${lm_stats_dir}/valid/text_shape.${token_type}
|
|
fi
|
|
# Training Stage
|
world_size=$gpu_num # run on one machine
|
if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
|
echo "stage 2: Training"
|
mkdir -p ${lm_exp}
|
mkdir -p ${lm_exp}/log
|
INIT_FILE=${lm_exp}/ddp_init
|
if [ -f $INIT_FILE ];then
|
rm -f $INIT_FILE
|
fi
|
init_method=file://$(readlink -f $INIT_FILE)
|
echo "$0: init method is $init_method"
|
for ((i = 0; i < $gpu_num; ++i)); do
|
{
|
rank=$i
|
local_rank=$i
|
gpu_id=$(echo $CUDA_VISIBLE_DEVICES | cut -d',' -f$[$i+1])
|
python ../../../funasr/bin/lm_train.py \
|
--gpu_id ${gpu_id} \
|
--use_preprocessor true \
|
--token_type "${token_type}" \
|
--token_list "${lm_token_list}" \
|
--non_linguistic_symbols "${nlsyms_txt}" \
|
--cleaner "${cleaner}" \
|
--train_data_path_and_name_and_type "${train_data_path_and_name_and_type}" \
|
--train_shape_file "${train_shape_file}" \
|
--valid_data_path_and_name_and_type "${valid_data_path_and_name_and_type}" \
|
--valid_shape_file "${valid_shape_file}" \
|
--fold_length "${lm_fold_length}" \
|
--resume true \
|
--output_dir "${lm_exp}" \
|
--config ${lm_config} \
|
--ngpu ${gpu_num} \
|
--num_worker_count ${count} \
|
--multiprocessing_distributed true \
|
--dist_init_method ${init_method} \
|
--dist_world_size ${world_size} \
|
--dist_rank ${rank} \
|
--local_rank ${local_rank} 1> ${exp_dir}/exp/${model_dir}/log/train.log.$i 2>&1
|
} &
|
done
|
wait
|
fi
|
|
# Testing Stage
|
gpu_num=1
|
if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
|
echo "Stage 3: Calc perplexity: ${lm_test_text}"
|
|
python ../../../funasr/bin/lm_inference.py \
|
--output_dir "${lm_exp}/perplexity_test" \
|
--ngpu "${gpu_num}" \
|
--batch_size 1 \
|
--train_config "${lm_exp}"/config.yaml \
|
--model_file "${lm_exp}/${inference_lm}" \
|
--data_path_and_name_and_type "${lm_test_text},text,text" \
|
--num_workers 1 \
|
--split_with_space false
|
fi
|
