File size: 11,891 Bytes
c668e80 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 |
"""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")
|