matching_parallel.py

#!/usr/bin/env python3
# Copyright    2023  Xiaomi Corp.        (authors: Wei Kang)
#
# 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.


import argparse
import logging
import os
from datetime import datetime
from multiprocessing.pool import Pool
from multiprocessing.pool import ThreadPool
from pathlib import Path
from queue import Queue
from threading import Barrier, Thread

from lhotse import CutSet, MonoCut, load_manifest_lazy
from lhotse.serialization import SequentialJsonlWriter
from textsearch import AttributeDict

from matching import align, load_data, get_params, split, write


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--manifest-in",
        type=str,
        help="""The manifest generated by transcript stage containing book path,
        recordings path and recognition results.
        """,
    )
    parser.add_argument(
        "--manifest-out",
        type=str,
        help="""The file name of the new manifests to write to. 
        """,
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=50,
        help="""The number of cuts in a batch.
        """,
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=6,
        help="""The number of workers run in parallel, each worker will process
        a batch of cuts.
        """,
    )
    return parser.parse_args()


def dataloader(
    worker_index: int,
    params: AttributeDict,
    cuts_queue: Queue,
    data_queue: Queue,
    barrier: Barrier,
):
    """
    Load texts (the references) from disk, and construct the sourced_text object.

    The item in cuts_queue is a list of MonoCut red from manifests.

    The item in data_queue is :

      {
          "num_query_tokens": int, # the total number of tokens of queries
          "cuts": List[MonoCut],   # The cuts from cuts_queue
          "cut_indexes": Tuple[int, int], # A tuple of (cut index, supervision index)
          "sourced_text": SourcedText, # The sourced_text constructed from batch_cuts.
      }
    """
    while True:
        try:
            batch_cuts = cuts_queue.get()
            if batch_cuts is None:
                # We put None to cuts_queue here, because we only put one None
                # at the end of cuts_queue, but there might be more than one
                # dataloader, putting more None to cuts_queue to make all
                # dataloader exit normally.
                cuts_queue.put(None)
                break

            data = load_data(
                params=params, batch_cuts=batch_cuts, worker_index=worker_index
            )
            data["cuts"] = batch_cuts
            data_queue.put(data)

        except Exception as e:
            logging.error(f"Worker[{worker_index}] caught {type(e)}: e")
            continue
    barrier.wait()
    if worker_index == 0:
        # Only need to put one None to data_queue, we will handle the existing
        # of multiple aligner in aligner thread.
        data_queue.put(None)
    logging.info(f"Worker[{worker_index}] dataloader done.")


def aligner(
    worker_index: int,
    params: AttributeDict,
    data_queue: Queue,
    align_queue: Queue,
    barrier: Barrier,
    thread_pool: ThreadPool,
):
    """
    Align each query in the sourced_text to its corresponding segment in references

    The item in data_queue is:
      {
          "num_query_tokens": int,        # the total number of tokens of queries
          "cuts": List[MonoCut],          # A list of MonoCut from cuts_queue
          "cut_indexes": List[Tuple[int, int]], # A tuple of (cut index, supervision index)
          "sourced_text": SourcedText,    # The sourced_text constructed from batch_cuts.
      }

    The item in align_queue is:
      {
          "cuts": List[MonoCut],  # A list of MonoCut inherited from data_queue.
          "cut_indexes": List[Tuple[int, int]], # Inherit from data_queue.
          "alignments": alignments,  # The alignment results.
          "sourced_text": SourcedText, # Inherit from data_queue.
      }

    """
    while True:
        try:
            item = data_queue.get()
            if item is None:
                # We put None to data_queue here, because we only put one None
                # at the end of data_queue in dataloader thread, but there might
                # be more than one aligner, putting more None to data_queue to
                # make all aligner exit normally.
                data_queue.put(None)
                break
            sourced_text = item["sourced_text"]
            num_query_tokens = item["num_query_tokens"]

            alignments = align(
                params=params,
                num_query_tokens=num_query_tokens,
                sourced_text=sourced_text,
                thread_pool=thread_pool,
                worker_index=worker_index,
            )

            align_queue.put(
                {
                    "cuts": item["cuts"],
                    "cut_indexes": item["cut_indexes"],
                    "alignments": alignments,
                    "sourced_text": sourced_text,
                }
            )
        except Exception as e:
            logging.error(f"Worker[{worker_index}] caught {type(e)}: e")
            continue
    barrier.wait()
    if worker_index == 0:
        # Only need to put one None to align_queue, we will handle the existing
        # of multiple splitter in splitter thread.
        align_queue.put(None)
    logging.info(f"Worker[{worker_index}] aligner done.")


def splitter(
    worker_index: int,
    params: AttributeDict,
    align_queue: Queue,
    write_queue: Queue,
    barrier: Barrier,
    process_pool: Pool,
):
    """
    Split the query into smaller segments.

    The item in align_queue is:
      {
          "cuts": List[MonoCut],  # A list of MonoCut inherited from data_queue.
          "cut_indexes": List[Tuple[int, int]], # Inherit from data_queue.
          "alignments": alignments,  # The alignment results.
          "sourced_text": SourcedText, # Inherit from data_queue.
      }

    The item in write_queue is:
      {
          "cuts": List[MonoCut],  # A list of MonoCut inherited from align_queue.
          "segments": segments,  # The segmented results (contains cut_indexes).
      }
    """
    while True:
        try:
            item = align_queue.get()
            if item is None:
                # We put None to align_queue here, because we only put one None
                # at the end of align_queue in aligner thread, but there might
                # be more than one splitter, putting more None to align_queue to
                # make all splitter exit normally.
                align_queue.put(None)
                break

            alignments = item["alignments"]
            sourced_text = item["sourced_text"]
            cut_indexes = item["cut_indexes"]

            results = split(
                params=params,
                sourced_text=sourced_text,
                alignments=alignments,
                cut_indexes=cut_indexes,
                process_pool=process_pool,
                worker_index=worker_index,
            )
            write_queue.put({"cuts": item["cuts"], "segments": results})

        except Exception as e:
            logging.error(
                f"Worker[{worker_index}] Splitter caught {type(e)}: e"
            )
            continue
    barrier.wait()
    if worker_index == 0:
        write_queue.put(None)
    logging.info(f"Worker[{worker_index}] splitter done.")


def writer(
    params: AttributeDict,
    write_queue: Queue,
    cuts_writer: SequentialJsonlWriter,
):
    """
    Write the segmented results to disk as new manifests.

    The item in write_queue is:
      {
          "cuts": List[MonoCut],  # A list of MonoCut inherited from align_queue.
          "segments": segments,  # The segmented results (contains cut_indexes).
      }
    """
    while True:
        try:
            item = write_queue.get()
            if item is None:
                break

            results = item["segments"]
            batch_cuts = item["cuts"]
            write(
                params=params,
                batch_cuts=batch_cuts,
                results=results,
                cuts_writer=cuts_writer,
            )
        except Exception as e:
            logging.error(f"Writer caught {type(e)}: e")
            continue
    logging.info(f"Writer done.")


def main():
    args = get_args()
    params = get_params()
    params.update(vars(args))

    logging.info(f"params : {params}")

    raw_cuts = load_manifest_lazy(params.manifest_in)
    cuts_writer = CutSet.open_writer(params.manifest_out, overwrite=True)

    # thread_pool to run the levenshtein alignment.
    # we use thread_pool here because the levenshtein run on C++ with GIL released.
    thread_pool = ThreadPool()

    # process_pool to split the query into segments.
    # we use process_pool here because the splitting algorithm run on Python
    # (we can not get accelerating with thread_pool because of the GIL)
    process_pool = Pool()

    # We create several producer and consumer patterns to make each part of the
    # whole pipeline run in parallel.
    # 1. Read cuts to cuts_queue;
    # 2. Construct sourced_text from cuts_queue and put it into data_queue;
    # 3. Align each query in sourced_text from data_queue, put the result to align_queue;
    # 4. Read alignments from align_queue then split into smaller segments, and
    #    put the results to write_queue.
    # 5. Writer write the item from write_queue to disk.
    cuts_queue = Queue(params.num_workers * 4)
    data_queue = Queue(params.num_workers * 4)
    align_queue = Queue(params.num_workers * 2)
    write_queue = Queue(params.num_workers * 8)

    dataloader_threads = []
    num_dataloaders = params.num_workers * 2
    dataloader_barrier = Barrier(num_dataloaders)
    for i in range(num_dataloaders):
        dataloader_threads.append(
            Thread(
                target=dataloader,
                args=(
                    i,
                    params,
                    cuts_queue,
                    data_queue,
                    dataloader_barrier,
                ),
            )
        )
        dataloader_threads[-1].start()

    aligner_threads = []
    num_aligners = params.num_workers
    aligner_barrier = Barrier(num_aligners)
    for i in range(num_aligners):
        aligner_threads.append(
            Thread(
                target=aligner,
                args=(
                    i,
                    params,
                    data_queue,
                    align_queue,
                    aligner_barrier,
                    thread_pool,
                ),
            )
        )
        aligner_threads[-1].start()

    # splitting is fast, so we only need one splitter.
    splitter_threads = []
    num_splitters = max(1, params.num_workers // 4)
    splitter_barrier = Barrier(num_splitters)
    for i in range(num_splitters):
        splitter_threads.append(
            Thread(
                target=splitter,
                args=(
                    i,
                    params,
                    align_queue,
                    write_queue,
                    splitter_barrier,
                    process_pool,
                ),
            )
        )
        splitter_threads[-1].start()

    # only one thread to write results
    writer_thread = Thread(
        target=writer,
        args=(
            params,
            write_queue,
            cuts_writer,
        ),
    )
    writer_thread.start()

    batch_cuts = []
    logging.info(f"Start processing...")
    for i, cut in enumerate(raw_cuts):
        if len(batch_cuts) >= params.batch_size:
            cuts_queue.put(batch_cuts)
            batch_cuts = []
            logging.info(f"Number of cuts have been loaded is {i}")
        batch_cuts.append(cut)
    if len(batch_cuts):
        cuts_queue.put(batch_cuts)
    cuts_queue.put(None)

    for t in dataloader_threads + aligner_threads + splitter_threads:
        t.join()
    writer_thread.join()
    cuts_writer.close()


if __name__ == "__main__":
    formatter = (
        "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
    )
    now = datetime.now()
    data_time = now.strftime("%Y-%m-%d-%H-%M-%S")
    os.makedirs("logs", exist_ok=True)
    log_file_name = f"logs/matching_{data_time}"
    logging.basicConfig(
        level=logging.INFO,
        format=formatter,
        handlers=[logging.FileHandler(log_file_name), logging.StreamHandler()],
    )

    main()