onmt/inputters/text_corpus.py

"""Module that contain shard utils for dynamic data."""
import os
from onmt.utils.logging import logger
from onmt.constants import CorpusName, CorpusTask
from onmt.transforms import TransformPipe
from onmt.inputters.text_utils import process, parse_features, append_features_to_text
from contextlib import contextmanager
import itertools


@contextmanager
def exfile_open(filename, *args, **kwargs):
    """Extended file opener enables open(filename=None).

    This context manager enables open(filename=None) as well as regular file.
    filename None will produce endlessly None for each iterate,
    while filename with valid path will produce lines as usual.

    Args:
        filename (str|None): a valid file path or None;
        *args: args relate to open file using codecs;
        **kwargs: kwargs relate to open file using codecs.

    Yields:
        `None` repeatly if filename==None,
        else yield from file specified in `filename`.
    """
    if filename is None:
        from itertools import repeat

        _file = repeat(None)
    else:
        import codecs

        _file = codecs.open(filename, *args, **kwargs)
    yield _file
    if filename is not None and _file:
        _file.close()


class ParallelCorpus(object):
    """A parallel corpus file pair that can be loaded to iterate."""

    def __init__(
        self, name, src, tgt, align=None, n_src_feats=0, src_feats_defaults=None
    ):
        """Initialize src & tgt side file path."""
        self.id = name
        self.src = src
        self.tgt = tgt
        self.align = align
        self.n_src_feats = n_src_feats
        self.src_feats_defaults = src_feats_defaults

    def load(self, offset=0, stride=1):
        """
        Load file and iterate by lines.
        `offset` and `stride` allow to iterate only on every
        `stride` example, starting from `offset`.
        """

        def make_ex(sline, tline, align):
            sline, sfeats = parse_features(
                sline,
                n_feats=self.n_src_feats,
                defaults=self.src_feats_defaults,
            )
            # 'src_original' and 'tgt_original' store the
            # original line before tokenization. These
            # fields are used later on in the feature
            # transforms.
            example = {
                "src": sline,
                "tgt": tline,
                "src_original": sline,
                "tgt_original": tline,
            }
            if align is not None:
                example["align"] = align
            if sfeats is not None:
                example["src_feats"] = [f for f in sfeats]
            return example

        if isinstance(self.src, list):
            fs = self.src
            ft = [] if self.tgt is None else self.tgt
            fa = [] if self.align is None else self.align
            for i, (sline, tline, align) in enumerate(
                itertools.zip_longest(fs, ft, fa)
            ):
                if (i // stride) % stride == offset:
                    yield make_ex(sline, tline, align)
        else:
            with exfile_open(self.src, mode="rb") as fs, exfile_open(
                self.tgt, mode="rb"
            ) as ft, exfile_open(self.align, mode="rb") as fa:
                for i, (sline, tline, align) in enumerate(zip(fs, ft, fa)):
                    if (i // stride) % stride == offset:
                        if tline is not None:
                            tline = tline.decode("utf-8")
                        if align is not None:
                            align = align.decode("utf-8")
                        yield make_ex(sline.decode("utf-8"), tline, align)

    def __str__(self):
        cls_name = type(self).__name__
        return (
            f"{cls_name}({self.id}, {self.src}, {self.tgt}, "
            f"align={self.align}, "
            f"n_src_feats={self.n_src_feats}, "
            f'src_feats_defaults="{self.src_feats_defaults}")'
        )


def get_corpora(opts, task=CorpusTask.TRAIN, src=None, tgt=None, align=None):
    corpora_dict = {}
    if task == CorpusTask.TRAIN:
        for corpus_id, corpus_dict in opts.data.items():
            if corpus_id != CorpusName.VALID:
                corpora_dict[corpus_id] = ParallelCorpus(
                    corpus_id,
                    corpus_dict["path_src"],
                    corpus_dict["path_tgt"],
                    # corpus_dict["path_align"],            ### new adding 
                    n_src_feats=opts.n_src_feats,
                    src_feats_defaults=opts.src_feats_defaults,
                )
    elif task == CorpusTask.VALID:
        if CorpusName.VALID in opts.data.keys():
            corpora_dict[CorpusName.VALID] = ParallelCorpus(
                CorpusName.VALID,
                opts.data[CorpusName.VALID]["path_src"],
                opts.data[CorpusName.VALID]["path_tgt"],
                # opts.data[CorpusName.VALID]["path_align"], ### new adding 
                n_src_feats=opts.n_src_feats,
                src_feats_defaults=opts.src_feats_defaults,
            )
        else:
            return None
    else:
        corpora_dict[CorpusName.INFER] = ParallelCorpus(
            CorpusName.INFER,
            src if src else opts.src,
            tgt if tgt else opts.tgt,
            align if align else None,
            n_src_feats=opts.n_src_feats,
            src_feats_defaults=opts.src_feats_defaults,
        )
    return corpora_dict


class ParallelCorpusIterator(object):
    """An iterator dedicated to ParallelCorpus.

    Args:
        corpus (ParallelCorpus): corpus to iterate;
        transform (TransformPipe): transforms to be applied to corpus;
        skip_empty_level (str): security level when encouter empty line;
        stride (int): iterate corpus with this line stride;
        offset (int): iterate corpus with this line offset.
    """

    def __init__(
        self, corpus, transform, skip_empty_level="warning", stride=1, offset=0
    ):
        self.cid = corpus.id
        self.corpus = corpus
        self.transform = transform
        if skip_empty_level not in ["silent", "warning", "error"]:
            raise ValueError(f"Invalid argument skip_empty_level={skip_empty_level}")
        self.skip_empty_level = skip_empty_level
        self.stride = stride
        self.offset = offset

    def _tokenize(self, stream):
        for example in stream:
            example["src"] = example["src"].strip("\n").split()
            example["src_original"] = example["src_original"].strip("\n").split()
            if "src_feats" in example:
                example["src_feats"] = [
                    feat.strip("\n").split() for feat in example["src_feats"]
                ]
            if example["tgt"] is not None:
                example["tgt"] = example["tgt"].strip("\n").split()
                example["tgt_original"] = example["tgt_original"].strip("\n").split()
            if "align" in example:
                example["align"] = example["align"].strip("\n").split()
            yield example

    def _transform(self, stream):
        for example in stream:
            # NOTE: moved to dynamic_iterator.py cf process()
            # item = self.transform.apply(
            # example, is_train=self.infinitely, corpus_name=self.cid)
            item = (example, self.transform, self.cid)
            if item is not None:
                yield item
        report_msg = self.transform.stats()
        if report_msg != "":
            logger.info(
                "* Transform statistics for {}({:.2f}%):\n{}\n".format(
                    self.cid, 100 / self.stride, report_msg
                )
            )

    def _add_index(self, stream):
        for i, item in enumerate(stream):
            example = item[0]
            line_number = i * self.stride + self.offset
            example["indices"] = line_number
            if example["tgt"] is not None:
                if (
                    len(example["src"]) == 0
                    or len(example["tgt"]) == 0
                    or ("align" in example and example["align"] == 0)
                ):
                    # empty example: skip
                    empty_msg = f"Empty line  in {self.cid}#{line_number}."
                    if self.skip_empty_level == "error":
                        raise IOError(empty_msg)
                    elif self.skip_empty_level == "warning":
                        logger.warning(empty_msg)
                    if len(example["src"]) == 0 and len(example["tgt"]) == 0:
                        yield item
                    continue
            yield item

    def __iter__(self):
        corpus_stream = self.corpus.load(stride=self.stride, offset=self.offset)
        tokenized_corpus = self._tokenize(corpus_stream)
        transformed_corpus = self._transform(tokenized_corpus)
        indexed_corpus = self._add_index(transformed_corpus)
        yield from indexed_corpus


def build_corpora_iters(
    corpora, transforms, corpora_info, skip_empty_level="warning", stride=1, offset=0
):
    """Return `ParallelCorpusIterator` for all corpora defined in opts."""
    corpora_iters = dict()
    for c_id, corpus in corpora.items():
        transform_names = corpora_info[c_id].get("transforms", [])
        corpus_transform = [
            transforms[name] for name in transform_names if name in transforms
        ]
        transform_pipe = TransformPipe.build_from(corpus_transform)
        corpus_iter = ParallelCorpusIterator(
            corpus,
            transform_pipe,
            skip_empty_level=skip_empty_level,
            stride=stride,
            offset=offset,
        )
        corpora_iters[c_id] = corpus_iter
    return corpora_iters


def save_transformed_sample(opts, transforms, n_sample=3):
    """Save transformed data sample as specified in opts."""

    if n_sample == -1:
        logger.info(f"n_sample={n_sample}: Save full transformed corpus.")
    elif n_sample == 0:
        logger.info(f"n_sample={n_sample}: no sample will be saved.")
        return
    elif n_sample > 0:
        logger.info(f"Save {n_sample} transformed example/corpus.")
    else:
        raise ValueError(f"n_sample should >= -1, get {n_sample}.")

    corpora = get_corpora(opts, CorpusTask.TRAIN)
    datasets_iterables = build_corpora_iters(
        corpora, transforms, opts.data, skip_empty_level=opts.skip_empty_level
    )
    sample_path = os.path.join(os.path.dirname(opts.save_data), CorpusName.SAMPLE)
    os.makedirs(sample_path, exist_ok=True)
    for c_name, c_iter in datasets_iterables.items():
        dest_base = os.path.join(sample_path, "{}.{}".format(c_name, CorpusName.SAMPLE))
        with open(dest_base + ".src", "w", encoding="utf-8") as f_src, open(
            dest_base + ".tgt", "w", encoding="utf-8"
        ) as f_tgt:
            bucket = []
            for i, ex in enumerate(c_iter):
                if i > n_sample:
                    break
                else:
                    bucket.append(ex)
            pro_bucket = process(CorpusTask.TRAIN, bucket)
            if pro_bucket is not None:
                for maybe_example in pro_bucket:
                    if maybe_example is not None:
                        src_line, tgt_line = (
                            maybe_example["src"]["src"],
                            maybe_example["tgt"]["tgt"],
                        )

                        if "feats" in maybe_example["src"]:
                            src_feats_lines = maybe_example["src"]["feats"]
                        else:
                            src_feats_lines = []

                        src_pretty_line = append_features_to_text(
                            src_line, src_feats_lines
                        )

                        f_src.write(src_pretty_line + "\n")
                        f_tgt.write(tgt_line + "\n")