modules/loader/default_loader.py

import io, os
import dill as pickle
import torch
from torch.utils.data import DataLoader
from torchtext.data import BucketIterator, Dataset, Example, Field
from torchtext.datasets import TranslationDataset, Multi30k, IWSLT, WMT14
from collections import Counter

import modules.constants as const
from utils.save import load_vocab_from_path
import laonlp

class DefaultLoader:
  def __init__(self, train_path_or_name, language_tuple=None, valid_path=None, eval_path=None, option=None):
    """Load training/eval data file pairing, process and create data iterator for training """
    self._language_tuple = language_tuple
    self._train_path = train_path_or_name
    self._eval_path = eval_path
    self._option = option
  
  @property
  def language_tuple(self):
    """DefaultLoader will use the default lang option @bleu_batch_iter <sos>, hence, None"""
    return None, None

  def tokenize(self, sentence):
    return sentence.strip().split()

  def detokenize(self, list_of_tokens):
    """Differentiate between [batch, len] and [len]; joining tokens back to strings"""
    if( len(list_of_tokens) == 0 or isinstance(list_of_tokens[0], str)):
      # [len], single sentence version
      return " ".join(list_of_tokens)
    else:
      # [batch, len], batch sentence version
      return [" ".join(tokens) for tokens in list_of_tokens]

  def _train_path_is_name(self):
    return os.path.isfile(self._train_path + self._language_tuple[0]) and os.path.isfile(self._train_path + self._language_tuple[1])

  def create_length_constraint(self, token_limit):
    """Filter an iterator if it pass a token limit"""
    return lambda x: len(x.src) <= token_limit and len(x.trg) <= token_limit

  def build_field(self, **kwargs):
    """Build fields that will handle the conversion from token->idx and vice versa. To be overriden by MultiLoader."""
    return Field(lower=False, tokenize=laonlp.tokenize.word_tokenize), Field(lower=False, tokenize=self.tokenize, init_token=const.DEFAULT_SOS, eos_token=const.DEFAULT_EOS, is_target=True)

  def build_vocab(self, fields, model_path=None, data=None, **kwargs):
    """Build the vocabulary object for torchtext Field. There are three flows:
      - if the model path is present, it will first try to load the pickled/dilled vocab object from path. This is accessed on continued training & standalone inference
      - if that failed and data is available, try to build the vocab from that data. This is accessed on first time training
      - if data is not available, search for set of two vocab files and read them into the fields. This is accessed on first time training
    TODO: expand on the vocab file option (loading pretrained vectors as well)
    """
    src_field, trg_field = fields
    if(model_path is None or not load_vocab_from_path(model_path, self._language_tuple, fields)):
      # the condition will try to load vocab pickled to model path.
      if(data is not None):
        print("Building vocab from received data.")
        # build the vocab using formatted data.
        src_field.build_vocab(data, **kwargs)
        trg_field.build_vocab(data, **kwargs)
      else:
        print("Building vocab from preloaded text file.")
        # load the vocab values from external location (a formatted text file). Initialize values as random
        external_vocab_location = self._option.get("external_vocab_location", None)
        src_ext, trg_ext = self._language_tuple
        # read the files and create a mock Counter object; which then is passed to vocab's class
        # see Field.build_vocab for the options used with vocab_cls
        vocab_src = external_vocab_location + src_ext
        with io.open(vocab_src, "r", encoding="utf-8") as svf:
          mock_counter = Counter({w.strip():1 for w in svf.readlines()})
          special_tokens = [src_field.unk_token, src_field.pad_token, src_field.init_token, src_field.eos_token]
          src_field.vocab = src_field.vocab_cls(mock_counter, specials=special_tokens, min_freq=5, **kwargs)
        vocab_trg = external_vocab_location + trg_ext
        with io.open(vocab_trg, "r", encoding="utf-8") as tvf:
          mock_counter = Counter({w.strip():1 for w in tvf.readlines()})
          special_tokens = [trg_field.unk_token, trg_field.pad_token, trg_field.init_token, trg_field.eos_token]
          trg_field.vocab = trg_field.vocab_cls(mock_counter, specials=special_tokens, min_freq=5, **kwargs)
    else:
      print("Load vocab from path successful.")

  def create_iterator(self, fields, model_path=None):
    """Create the iterator needed to load batches of data and bind them to existing fields
    NOTE: unlike the previous loader, this one inputs list of tokens instead of a string, which necessitate redefinining of translate_sentence pipe"""
    if(not self._train_path_is_name()):
      # load the default torchtext dataset by name
      # TODO load additional arguments in the config
      dataset_cls = next( (s for s in [Multi30k, IWSLT, WMT14] if s.__name__ == self._train_path), None )
      if(dataset_cls is None):
        raise ValueError("The specified train path {:s}(+{:s}/{:s}) does neither point to a valid files path nor is a name of torchtext dataset class.".format(self._train_path, *self._language_tuple))
      src_suffix, trg_suffix = ext = self._language_tuple
#      print(ext, fields)
      self._train_data, self._valid_data, self._eval_data = dataset_cls.splits(exts=ext, fields=fields) #, split_ratio=self._option.get("train_test_split", const.DEFAULT_TRAIN_TEST_SPLIT)
    else:
      # create dataset from path. Also add all necessary constraints (e.g lengths trimming/excluding)
      src_suffix, trg_suffix = ext = self._language_tuple
      filter_fn = self.create_length_constraint(self._option.get("train_max_length", const.DEFAULT_TRAIN_MAX_LENGTH))
      self._train_data = TranslationDataset(self._train_path, ext, fields, filter_pred=filter_fn)
      self._valid_data = self._eval_data = TranslationDataset(self._eval_path, ext, fields)
#    first_sample = self._train_data[0]; raise Exception("{} {}".format(first_sample.src, first_sample.trg))
    # whatever created, we now have the two set of data ready. add the necessary constraints/filtering/etc.
    train_data = self._train_data
    eval_data = self._eval_data
    # now we can execute build_vocab. This function will try to load vocab from model_path, and if fail, build the vocab from train_data
    build_vocab_kwargs = self._option.get("build_vocab_kwargs", {})
    self.build_vocab(fields, data=train_data, model_path=model_path, **build_vocab_kwargs)
#    raise Exception("{}".format(len(src_field.vocab)))
    # crafting iterators
    train_iter = BucketIterator(train_data, batch_size=self._option.get("batch_size", const.DEFAULT_BATCH_SIZE), device=self._option.get("device", const.DEFAULT_DEVICE) )
    eval_iter = BucketIterator(eval_data, batch_size=self._option.get("eval_batch_size", const.DEFAULT_EVAL_BATCH_SIZE), device=self._option.get("device", const.DEFAULT_DEVICE), train=False )
    return train_iter, eval_iter