icefall/egs/librilight/SSL/prepare.sh
2024-10-24 00:43:16 -07:00

101 lines
3.1 KiB
Bash
Executable File

#!/usr/bin/env bash
# fix segmentation fault reported in https://github.com/k2-fsa/icefall/issues/674
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
set -eou pipefail
nj=32
# run step 0 to step 4 by default
stage=0
stop_stage=4
# We assume dl_dir (download dir) contains the following
# directories and files. If not, they will be downloaded
# by this script automatically.
#
# - $dl_dir/LibriLight
# - small
# - medium
# - large
#
# You can download them from
# - https://dl.fbaipublicfiles.com/librilight/data/small.tar
# - https://dl.fbaipublicfiles.com/librilight/data/medium.tar
# - https://dl.fbaipublicfiles.com/librilight/data/large.tar
dl_dir=$PWD/download
. shared/parse_options.sh || exit 1
# All files generated by this script are saved in "data".
# You can safely remove "data" and rerun this script to regenerate it.
mkdir -p data
log() {
# This function is from espnet
local fname=${BASH_SOURCE[1]##*/}
echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
}
log "Running prepare.sh"
log "dl_dir: $dl_dir"
if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
log "Stage 1: Prepare Libri-Light manifest"
# We assume that you have downloaded the Libri-Light corpus
# to $dl_dir/LibriLight
mkdir -p data/manifests
if [ ! -e data/manifests/.librilight.done ]; then
lhotse prepare librilight -j $nj $dl_dir/LibriLight data/manifests
touch data/manifests/.librilight.done
fi
fi
if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
log "Stage 2: Preprocess Libri-Light manifest"
mkdir -p data/kmeans
if [ ! -f data/kmeans/.preprocess_complete ]; then
python3 ./local/preprocess_librilight.py
touch data/kmeans/.preprocess_complete
fi
fi
if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
log "Stage 3: Split medium and large subset into pieces"
num_per_split=10000
split_dir=data/kmeans/medium_split
if [ ! -f $split_dir/.split_completed ]; then
lhotse split-lazy ./data/kmeans/librilight_cuts_medium_raw.jsonl.gz $split_dir $num_per_split
touch $split_dir/.split_completed
fi
split_dir=data/kmeans/large_split
if [ ! -f $split_dir/.split_completed ]; then
lhotse split-lazy ./data/kmeans/librilight_cuts_large_raw.jsonl.gz $split_dir $num_per_split
touch $split_dir/.split_completed
fi
fi
if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
log "Stage 4: Extract SSL target for librilight"
if [ ! -e download/hubert_base_ls960.pt ]; then
wget https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960.pt -P download
fi
if [ ! -e download/hubert_base_ls960_L9_km500.bin ]; then
wget https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960_L9_km500.bin -P download
fi
if [ ! -e data/kmeans/.extract_small.done ]; then
./local/extract_kmeans.py --subset small
touch data/kmeans/.extract_small.done
fi
if [ ! -e data/kmeans/.extract_medium.done ]; then
./local/extract_kmeans.py --subset medium
touch data/kmeans/.extract_medium.done
fi
if [ ! -e data/kmeans/.extract_large.done ]; then
./local/extract_kmeans.py --subset large
touch data/kmeans/.extract_large.done
fi
fi