egs/zipvoice_dialog/local/prepare_opendialog.py

#!/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 prepares lhotse manifest files from the raw OpenDialog datasets.

We assume that you have downloaded the OpenDialog dataset and untarred the
tar files in audio/en and audio/zh so that the mp3 files are placed under
these two directories.

Download OpenDialog at https://huggingface.co/datasets/k2-fsa/OpenDialog
or https://www.modelscope.cn/datasets/k2-fsa/OpenDialog

"""

import argparse
import json
import logging
import math
import re
from concurrent.futures import ThreadPoolExecutor
from functools import partial
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.cut import Cut
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(
        "--dataset-path",
        type=str,
        help="The path of OpenDialog dataset.",
    )

    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 = Path(wav_path).stem
    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
    """

    def _round_down(num, ndigits=0):
        factor = 10**ndigits
        return math.floor(num * factor) / factor

    uniq_id, text, wav_path, start, end = supervision
    try:
        recording_id = Path(wav_path).stem

        recording = recording_dict[recording_id]
        duration = (
            _round_down(end - start, ndigits=8)
            if end is not None
            else _round_down(recording.duration, ndigits=8)
        )
        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.info(f"Error processing line: {e}")
        return None


def prepare_subset(
    jsonl_path: Pathlike,
    lang: str,
    sampling_rate: int,
    num_jobs: int,
    output_dir: Pathlike,
):
    """
    Returns the manifests which consist of the Recordings and Supervisions

    :param jsonl_path: Path to the jsonl file
    :param lang: Language of the subset
    :param sampling_rate: Target sampling rate of the audio
    :param num_jobs: Number of processes for parallel processing
    :param output_dir: Path where to write the manifests
    """
    logging.info(f"Preparing {lang} subset")

    # Step 1: Read all unique recording paths
    logging.info(f"Reading {jsonl_path}")
    recordings_path_set = set()
    supervision_list = list()
    with open(jsonl_path, "r") as fr:
        for line in fr:
            try:
                items = json.loads(line)
                uniq_id, text, wav_path = items["id"], items["text"], items["path"]
                start, end = 0, None
                recordings_path_set.add(jsonl_path.parent / wav_path)
                supervision_list.append((uniq_id, text, wav_path, start, end))
            except Exception as e:
                logging.warning(f"Error {e} when decoding JSON line: {line}")
                continue
    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()
    if sampling_rate != 24000:
        # All OpenDialog audios are 24kHz
        cut_set = cut_set.resample(sampling_rate)
    cut_set = cut_set.trim_to_supervisions(keep_overlapping=False)

    logging.info("Saving cuts to disk...")
    # Step 5: Write manifests to disk
    cut_set.to_file(output_dir / f"opendialog_cuts_raw_{lang.upper()}-all.jsonl.gz")
    dev_cut_set = cut_set.subset(first=1000)
    dev_cut_set.to_file(output_dir / f"opendialog_cuts_raw_{lang.upper()}-dev.jsonl.gz")

    def remove_dev(c: Cut, set: set):
        if c.id in set:
            return False
        return True

    _remove_dev = partial(remove_dev, set=set(dev_cut_set.ids))
    train_cut_set = cut_set.filter(_remove_dev)
    train_cut_set.to_file(
        output_dir / f"opendialog_cuts_raw_{lang.upper()}-train.jsonl.gz"
    )


def prepare_dataset(
    dataset_path: Pathlike,
    sampling_rate: int,
    num_jobs: int,
    output_dir: Pathlike,
):
    for lang in ["en", "zh"]:
        jsonl_path = dataset_path / f"manifest.{lang}.jsonl"
        prepare_subset(
            jsonl_path=jsonl_path,
            lang=lang,
            sampling_rate=sampling_rate,
            num_jobs=num_jobs,
            output_dir=output_dir,
        )


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()
    dataset_path = Path(args.dataset_path)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    prepare_dataset(
        dataset_path=dataset_path,
        sampling_rate=args.sampling_rate,
        num_jobs=args.num_jobs,
        output_dir=output_dir,
    )