# -*- coding: utf-8 -*- import os import math import codecs import torch import pyonmttok from onmt.constants import DefaultTokens from onmt.transforms import TransformPipe class IterOnDevice(torch.utils.data.IterableDataset): """Sent items from `iterable` on `device_id` and yield.""" def __init__(self, iterable, device_id): super(IterOnDevice).__init__() self.iterable = iterable self.device_id = device_id # temporary as long as translation_server and scoring_preparator still use lists if hasattr(iterable, "transforms"): self.transform = TransformPipe.build_from( [iterable.transforms[name] for name in iterable.transforms] ) @staticmethod def batch_to_device(tensor_batch, device_id): """Move `batch` to `device_id`, cpu if `device_id` < 0.""" device = torch.device(device_id) if device_id >= 0 else torch.device("cpu") for key in tensor_batch.keys(): if key != "src_ex_vocab": tensor_batch[key] = tensor_batch[key].to(device) def __iter__(self): for tensor_batch in self.iterable: self.batch_to_device(tensor_batch, self.device_id) yield tensor_batch def build_vocab(opt, specials): """Build vocabs dict to be stored in the checkpoint based on vocab files having each line [token, count] Args: opt: src_vocab, tgt_vocab, n_src_feats Return: vocabs: {'src': pyonmttok.Vocab, 'tgt': pyonmttok.Vocab, 'src_feats' : [pyonmttok.Vocab, ...]}, 'data_task': seq2seq or lm 'decoder_start_token': DefaultTokens.BOS } """ def _pad_vocab_to_multiple(vocab, multiple): vocab_size = len(vocab) if vocab_size % multiple == 0: return vocab target_size = int(math.ceil(vocab_size / multiple)) * multiple for i in range(target_size - vocab_size): vocab.add_token(DefaultTokens.VOCAB_PAD + str(i)) return vocab default_specials = opt.default_specials vocabs = {} src_vocab = _read_vocab_file(opt.src_vocab, opt.src_words_min_frequency) src_specials = [ item for item in (default_specials + specials["src"]) if item not in src_vocab ] if DefaultTokens.SEP in src_specials and ( "<0x0A>" in src_vocab or "Ċ" in src_vocab ): # this is hack: if the special separator ｟newline｠is returned because of the # "docify" transform.get_specials we don't add it if the corresponding newline code # is already included in the sentencepiece or BPE-with-gpt2-pretok. src_specials.remove(DefaultTokens.SEP) src_vocab = pyonmttok.build_vocab_from_tokens( src_vocab, maximum_size=opt.src_vocab_size, special_tokens=src_specials ) src_vocab.default_id = src_vocab[DefaultTokens.UNK] if opt.vocab_size_multiple > 1: src_vocab = _pad_vocab_to_multiple(src_vocab, opt.vocab_size_multiple) vocabs["src"] = src_vocab if opt.share_vocab: vocabs["tgt"] = src_vocab else: tgt_vocab = _read_vocab_file(opt.tgt_vocab, opt.tgt_words_min_frequency) tgt_specials = [ item for item in (default_specials + specials["tgt"]) if item not in tgt_vocab ] if DefaultTokens.SEP in tgt_specials and ( "<0x0A>" in tgt_vocab or "Ċ" in src_vocab ): tgt_specials.remove(DefaultTokens.SEP) tgt_vocab = pyonmttok.build_vocab_from_tokens( tgt_vocab, maximum_size=opt.tgt_vocab_size, special_tokens=tgt_specials ) tgt_vocab.default_id = tgt_vocab[DefaultTokens.UNK] if opt.vocab_size_multiple > 1: tgt_vocab = _pad_vocab_to_multiple(tgt_vocab, opt.vocab_size_multiple) vocabs["tgt"] = tgt_vocab if opt.n_src_feats > 0: src_feats_vocabs = [] for i in range(opt.n_src_feats): src_f_vocab = _read_vocab_file(f"{opt.src_vocab}_feat{i}", 1) src_f_vocab = pyonmttok.build_vocab_from_tokens( src_f_vocab, maximum_size=0, minimum_frequency=1, special_tokens=default_specials, ) src_f_vocab.default_id = src_f_vocab[DefaultTokens.UNK] if opt.vocab_size_multiple > 1: src_f_vocab = _pad_vocab_to_multiple( src_f_vocab, opt.vocab_size_multiple ) src_feats_vocabs.append(src_f_vocab) vocabs["src_feats"] = src_feats_vocabs vocabs["data_task"] = opt.data_task vocabs["decoder_start_token"] = opt.decoder_start_token return vocabs def _read_vocab_file(vocab_path, min_count): """Loads a vocabulary from the given path. Args: vocab_path (str): Path to utf-8 text file containing vocabulary. Each token should be on a line, may followed with a count number seperate by space if `with_count`. No extra whitespace is allowed. min_count (int): retains only tokens with min_count frequency. """ if not os.path.exists(vocab_path): raise RuntimeError("Vocabulary not found at {}".format(vocab_path)) else: with codecs.open(vocab_path, "rb", "utf-8") as f: lines = [line.strip("\n") for line in f if line.strip("\n")] first_line = lines[0].split(None, 1) has_count = len(first_line) == 2 and first_line[-1].isdigit() if has_count: vocab = [] for line in lines: if int(line.split(None, 1)[1]) >= min_count: vocab.append(line.split(None, 1)[0]) else: vocab = lines return vocab def vocabs_to_dict(vocabs): """ Convert a dict of pyonmttok vocabs into a plain text dict to be saved in the checkpoint """ vocabs_dict = {} vocabs_dict["src"] = vocabs["src"].ids_to_tokens vocabs_dict["tgt"] = vocabs["tgt"].ids_to_tokens if "src_feats" in vocabs.keys(): vocabs_dict["src_feats"] = [ feat_vocab.ids_to_tokens for feat_vocab in vocabs["src_feats"] ] vocabs_dict["data_task"] = vocabs["data_task"] if "decoder_start_token" in vocabs.keys(): vocabs_dict["decoder_start_token"] = vocabs["decoder_start_token"] else: vocabs_dict["decoder_start_token"] = DefaultTokens.BOS return vocabs_dict def dict_to_vocabs(vocabs_dict): """ Convert a dict formatted vocabs (as stored in a checkpoint) into a dict of pyonmttok vocabs objects. """ vocabs = {} vocabs["data_task"] = vocabs_dict["data_task"] if "decoder_start_token" in vocabs_dict.keys(): vocabs["decoder_start_token"] = vocabs_dict["decoder_start_token"] else: vocabs["decoder_start_token"] = DefaultTokens.BOS vocabs["src"] = pyonmttok.build_vocab_from_tokens(vocabs_dict["src"]) if vocabs_dict["src"] == vocabs_dict["tgt"]: vocabs["tgt"] = vocabs["src"] else: vocabs["tgt"] = pyonmttok.build_vocab_from_tokens(vocabs_dict["tgt"]) if "src_feats" in vocabs_dict.keys(): vocabs["src_feats"] = [] for feat_vocab in vocabs_dict["src_feats"]: vocabs["src_feats"].append(pyonmttok.build_vocab_from_tokens(feat_vocab)) return vocabs