pipeline/clean/merge-parallel.py

"""
Merges multiple parallel corpora into a single "source" language file, and a single
"target" language file, each.

For instance:

  dataset1.en.zst dataset1.ru.zst
  dataset2.en.zst dataset2.ru.zst
  dataset3.en.zst dataset3.ru.zst

  Gets merged into:

  corpus.en.zst
  corpus.ru.zst
"""

import argparse
from contextlib import ExitStack
from glob import glob
from pathlib import Path
from typing import Generator, Optional
from pipeline.common.datasets import (
    FilteringStep,
    Statistics,
    WeakStringSet,
    shuffle_with_max_lines,
)
from pipeline.common.downloads import get_human_readable_file_size, read_lines, write_lines
from pipeline.common.logging import get_logger

logger = get_logger(__file__)


class FilteringStatistics(Statistics):
    """
    Gather statistics about the filtering process.
    """

    def __init__(self, dataset_path: Path) -> None:
        super().__init__(dataset_path)
        self.parallel_corpus = FilteringStep(
            "The parallel corpora are merged and deduplicated",
        )
        self.final_truncated = FilteringStep("The final result can be truncated by max_lines")
        self.datasets = []

    def add_parallel_dataset(self, location: str):
        # e.g. /path/to/ada83_v1.en.zst
        path = Path(location)
        # e.g. ada83_v1
        dataset_stem = Path(path.stem).stem
        step = FilteringStep(dataset_stem)
        self.datasets.append(step)
        return step


def log_dataset(location: str):
    logger.info(f"Reading dataset {location}")


class DeduplicateCorpus:
    def __init__(
        self,
        datasets_src: list[Path],
        datasets_trg: list[Path],
        src_outpath: Path,
        trg_outpath: Path,
        stats: FilteringStatistics,
    ) -> None:
        self.datasets_src: list[Path] = datasets_src
        self.datasets_trg: list[Path] = datasets_trg
        self.src_outpath: Path = src_outpath
        self.trg_outpath: Path = trg_outpath
        self.stats: FilteringStatistics = stats
        self.dataset_stats: FilteringStep = None

    def run(
        self,
        total_corpus_bytes: int,
        max_lines: Optional[int],
    ):
        stats = self.stats
        with ExitStack() as stack:
            src_outfile = stack.enter_context(write_lines(self.src_outpath))
            trg_outfile = stack.enter_context(write_lines(self.trg_outpath))

            if max_lines:
                for line in shuffle_with_max_lines(
                    line_stream=self.yield_lines_string(stack),
                    seed=38540735095,
                    max_lines=max_lines,
                    total_byte_size=total_corpus_bytes,
                ):
                    src_line, trg_line = line.split("\t")
                    src_outfile.write(src_line)
                    trg_outfile.write(trg_line)

                stats.final_truncated.visited = stats.parallel_corpus.kept
                stats.final_truncated.kept = min(max_lines, stats.parallel_corpus.kept)
            else:
                for src_line, trg_line in self.yield_lines_tuple(stack):
                    src_outfile.write(src_line)
                    trg_outfile.write(trg_line)

                stats.final_truncated.kept = stats.parallel_corpus.kept
                stats.final_truncated.visited = stats.parallel_corpus.kept

    def yield_lines_tuple(self, stack: ExitStack) -> Generator[tuple[str, str], None, None]:
        strings_seen = WeakStringSet()
        stats = self.stats
        src_lines: Generator[str, None, None] = stack.enter_context(
            read_lines(self.datasets_src, on_enter_location=self.on_enter_location)
        )
        trg_lines: Generator[str, None, None] = stack.enter_context(
            read_lines(self.datasets_trg, on_enter_location=log_dataset)
        )

        for src_line, trg_line in zip(src_lines, trg_lines):
            # No separator is needed as the newline is included.
            line = src_line + trg_line

            if line in strings_seen:
                stats.parallel_corpus.filtered += 1
                self.dataset_stats.filtered += 1
            else:
                stats.parallel_corpus.kept += 1
                self.dataset_stats.kept += 1

                strings_seen.add(line)

                yield src_line, trg_line

    def yield_lines_string(self, stack: ExitStack) -> Generator[str, None, None]:
        for src_line, trg_line in self.yield_lines_tuple(stack):
            if "\t" in src_line or "\t" in trg_line:
                logger.error("A line contained a tab character, skipping:")
                logger.error(f" src: {src_line}")
                logger.error(f" trg: {src_line}")
            else:
                yield f"{src_line}\t{trg_line}"

    def on_enter_location(self, location):
        log_dataset(location)
        self.dataset_stats = self.stats.add_parallel_dataset(location)


def sample_corpus(
    artifacts: Path, name: str, sample_size: int, src_outpath: Path, trg_outpath: Path
):
    """
    Generate a sample of the corpus data with the following format:

    e.g.
    > cat artifacts/corpus.sample.txt
    Sentence 1 in source language
    Sentence 1 in target language

    Sentence 2 in source language
    Sentence 2 in target language

    Sentence 3 in source language
    Sentence 3 in target language
    ...
    """
    total_byte_size = src_outpath.stat().st_size + trg_outpath.stat().st_size

    with ExitStack() as stack:
        sample_path = artifacts / f"{name}.sample.txt"

        src_lines = stack.enter_context(read_lines(src_outpath))
        trg_lines = stack.enter_context(read_lines(trg_outpath))
        sample_outfile = stack.enter_context(
            write_lines(
                sample_path,
                # The browser won't know the encoding when viewing this sample without including
                # a "byte order mark", which python can do via this encoding.
                encoding="utf-8-sig",
            )
        )

        def join_src_trg():
            for src_line, trg_line in zip(src_lines, trg_lines):
                # The src and trg line each have a newline at the end. This means that
                # each sentence pair will be separate by a blank line to make for easy
                # scanning of datasets.
                yield f"{src_line}{trg_line}\n"

        logger.info("Stream in:")
        logger.info(f" - {src_outpath}")
        logger.info(f" - {trg_outpath}")
        logger.info(f"Write a {sample_size:,} line sample of the merged corpus:")
        logger.info(f" - {sample_path}")

        for line in shuffle_with_max_lines(
            line_stream=join_src_trg(),
            seed=9834523434,
            max_lines=sample_size,
            total_byte_size=total_byte_size,
        ):
            sample_outfile.write(line)


def get_datasets(src: str, trg: str, datasets_glob: str):
    dataset_paths: list[str] = glob(datasets_glob)
    datasets_src: list[Path] = []
    datasets_trg: list[Path] = []
    dataset_paths.sort()

    total_corpus_bytes = 0

    for dataset in dataset_paths:
        path = Path(dataset)
        if dataset.endswith(f"{src}.zst"):
            datasets_src.append(path)
        elif dataset.endswith(f"{trg}.zst"):
            datasets_trg.append(path)
        else:
            raise Exception(f"Dataset does not match naming scheme: {dataset}")

        formatted_size, bytes = get_human_readable_file_size(path)
        logger.info(f" - {path} ({formatted_size})")
        total_corpus_bytes += bytes

    return datasets_src, datasets_trg, total_corpus_bytes


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        # Preserves whitespace in the help text.
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument(
        "--src",
        type=str,
        help="The source locale",
    )

    parser.add_argument(
        "--trg",
        type=str,
        help="The target locale",
    )

    parser.add_argument(
        "--datasets_glob",
        type=str,
        help="A glob-style path to the mono datasets, e.g. /path/to/*.zst",
    )

    parser.add_argument(
        "--max_lines",
        type=str,
        default="None",
        help="The (optionally) maximum number of sentences that will be merged.",
    )

    parser.add_argument(
        "--sample_size", type=int, default=10_000, help="Generate a random sample of sentences."
    )

    parser.add_argument(
        "--artifacts",
        type=Path,
        help="The path to the artifacts directory.",
    )

    parser.add_argument(
        "--name",
        type=str,
        help='The final corpus name, e.g. "corpus" will output a "corpus.en.zst" file.',
    )

    args = parser.parse_args()

    datasets_src, datasets_trg, total_corpus_bytes = get_datasets(
        args.src, args.trg, args.datasets_glob
    )

    logger.info("Parallel datasets:")

    src_outpath = args.artifacts / f"{args.name}.{args.src}.zst"
    trg_outpath = args.artifacts / f"{args.name}.{args.trg}.zst"

    stats = FilteringStatistics(args.artifacts / args.name)

    max_lines: Optional[int] = None
    if args.max_lines != "None":
        max_lines = int(args.max_lines)

    deduplicate_corpus = DeduplicateCorpus(
        datasets_src,
        datasets_trg,
        src_outpath,
        trg_outpath,
        stats,
    )

    deduplicate_corpus.run(total_corpus_bytes, max_lines)

    sample_corpus(
        artifacts=args.artifacts,
        name=args.name,
        sample_size=args.sample_size,
        src_outpath=src_outpath,
        trg_outpath=trg_outpath,
    )

    stats.save_json()


if __name__ == "__main__":
    main()