#!/usr/bin/env python3 # Copyright 2025 Xiaomi Corp. (authors: Han Zhu) # # See ../../../../LICENSE for clarification regarding multiple authors # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ This script generates lhotse manifest files from TSV files for custom datasets. Each line of the TSV files should be in one of the following formats: 1. "{uniq_id}\t{text}\t{wav_path}" if the text corresponds to the full wav", 2. "{uniq_id}\t{text}\t{wav_path}\t{start_time}\t{end_time} if text corresponds to part of the wav. The start_time and end_time specify the start and end times of the text within the wav, which should be in seconds. Note: {uniq_id} must be unique for each line. Usage: Suppose you have two TSV files: "custom_train.tsv" and "custom_dev.tsv", where "custom" is your dataset name, "train"/"dev" are used for training and validation respectively. (1) Prepare the training data python3 -m zipvoice.bin.prepare_dataset \ --tsv-path data/raw/custom_train.tsv \ --prefix "custom" \ --subset "train" \ --num-jobs 20 \ --output-dir "data/manifests" The output file would be "data/manifests/custom_cuts_train.jsonl.gz". (2) Prepare the validation data python3 -m zipvoice.bin.prepare_dataset \ --tsv-path data/raw/custom_dev.tsv \ --prefix "custom" \ --subset "dev" \ --num-jobs 1 \ --output-dir "data/manifests" The output file would be "data/manifests/custom_cuts_dev.jsonl.gz". """ import argparse import logging import re from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import List, Optional, Tuple from lhotse import CutSet, validate_recordings_and_supervisions from lhotse.audio import Recording, RecordingSet from lhotse.qa import fix_manifests from lhotse.supervision import SupervisionSegment, SupervisionSet from lhotse.utils import Pathlike from tqdm.auto import tqdm def get_args(): parser = argparse.ArgumentParser() parser.add_argument( "--tsv-path", type=str, help="The path of the tsv file. Each line should be in the format: " "{uniq_id}\t{text}\t{wav_path}\t{start_time}\t{end_time} " "if text corresponds to part of the wav or {uniq_id}\t{text}\t{wav_path} " "if the text corresponds to the full wav", ) parser.add_argument( "--prefix", type=str, default="custom", help="Prefix of the output manifest file.", ) parser.add_argument( "--subset", type=str, default="train", help="Subset name manifest file, typically train or dev.", ) parser.add_argument( "--num-jobs", type=int, default=20, help="Number of jobs to processing.", ) parser.add_argument( "--output-dir", type=str, default="data/manifests", help="The destination directory of manifest files.", ) parser.add_argument( "--sampling-rate", type=int, default=24000, help="The target sampling rate.", ) return parser.parse_args() def _parse_recording( wav_path: str, ) -> Tuple[Recording, str]: """ :param wav_path: Path to the audio file :return: a tuple of "recording" and "recording_id" """ recording_id = wav_path.replace("/", "_").replace(".", "_") recording = Recording.from_file(path=wav_path, recording_id=recording_id) return recording, recording_id def _parse_supervision( supervision: List, recording_dict: dict ) -> Optional[SupervisionSegment]: """ :param line: A line from the TSV file :param recording_dict: Dictionary mapping recording IDs to Recording objects :return: A SupervisionSegment object """ uniq_id, text, wav_path, start, end = supervision try: recording_id = wav_path.replace("/", "_").replace(".", "_") recording = recording_dict[recording_id] duration = end - start if end is not None else recording.duration assert duration <= recording.duration, f"Duration {duration} is greater than " f"recording duration {recording.duration}" text = re.sub("_", " ", text) # "_" is treated as padding symbol text = re.sub(r"\s+", " ", text) # remove extra whitespace return SupervisionSegment( id=f"{uniq_id}", recording_id=recording.id, start=start, duration=duration, channel=recording.channel_ids, text=text.strip(), ) except Exception as e: logging.warning(f"Error processing line: {e}") return None def prepare_dataset( tsv_path: Pathlike, prefix: str, subset: str, sampling_rate: int, num_jobs: int, output_dir: Pathlike, ): """ Returns the manifests which consist of the Recordings and Supervisions :param tsv_path: Path to the TSV file :param output_dir: Path where to write the manifests :param num_jobs: Number of processes for parallel processing :return: The CutSet containing the data """ logging.info(f"Preparing {prefix} dataset {subset} subset.") output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) file_name = f"{prefix}_cuts_{subset}.jsonl.gz" if (output_dir / file_name).is_file(): logging.info(f"{file_name} exists, skipping.") return # Step 1: Read all unique recording paths recordings_path_set = set() supervision_list = list() with open(tsv_path, "r") as fr: for line in fr: items = line.strip().split("\t") if len(items) == 3: uniq_id, text, wav_path = items start, end = 0, None elif len(items) == 5: uniq_id, text, wav_path, start, end = items start, end = float(start), float(end) else: raise ValueError( f"Invalid line format: {line}," "requries to be 3 columns or 5 columns" ) recordings_path_set.add(wav_path) supervision_list.append((uniq_id, text, wav_path, start, end)) logging.info("Starting to process recordings...") # Step 2: Process recordings futures = [] recording_dict = {} with ThreadPoolExecutor(max_workers=num_jobs) as ex: for wav_path in tqdm(recordings_path_set, desc="Submitting jobs"): futures.append(ex.submit(_parse_recording, wav_path)) for future in tqdm(futures, desc="Processing recordings"): try: recording, recording_id = future.result() recording_dict[recording_id] = recording except Exception as e: logging.warning( f"Error processing recording {recording_id} with error: {e}" ) recording_set = RecordingSet.from_recordings(recording_dict.values()) logging.info("Starting to process supervisions...") # Step 3: Process supervisions supervisions = [] for supervision in tqdm(supervision_list, desc="Processing supervisions"): seg = _parse_supervision(supervision, recording_dict) if seg is not None: supervisions.append(seg) logging.info("Processing Cuts...") # Step 4: Create and validate manifests supervision_set = SupervisionSet.from_segments(supervisions) recording_set, supervision_set = fix_manifests(recording_set, supervision_set) validate_recordings_and_supervisions(recording_set, supervision_set) cut_set = CutSet.from_manifests( recordings=recording_set, supervisions=supervision_set ) cut_set = cut_set.sort_by_recording_id() cut_set = cut_set.resample(sampling_rate) cut_set = cut_set.trim_to_supervisions(keep_overlapping=False) logging.info(f"Saving file to {output_dir / file_name}") # Step 5: Write manifests to disk cut_set.to_file(output_dir / file_name) logging.info("Done!") if __name__ == "__main__": formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" logging.basicConfig(format=formatter, level=logging.INFO, force=True) args = get_args() prepare_dataset( tsv_path=args.tsv_path, prefix=args.prefix, subset=args.subset, sampling_rate=args.sampling_rate, num_jobs=args.num_jobs, output_dir=args.output_dir, )