text: # -*- coding: utf-8 -*-
  
  from collections import Counter
  from itertools import chain
  import io
  import codecs
  import sys
  
  import torch
  import torchtext
  
  from onmt.Utils import aeq
  from onmt.io.DatasetBase import ONMTDatasetBase, PAD_WORD, BOS_WORD, EOS_WORD
  
  
  class TextDataset(ONMTDatasetBase):
      """ Dataset for data_type=='text'
  
          Build `Example` objects, `Field` objects, and filter_pred function
          from text corpus.
  
          Args:
              fields (dict): a dictionary of `torchtext.data.Field`.
                  Keys are like 'src', 'tgt', 'src_map', and 'alignment'.
              src_examples_iter (dict iter): preprocessed source example
                  dictionary iterator.
              tgt_examples_iter (dict iter): preprocessed target example
                  dictionary iterator.
              num_src_feats (int): number of source side features.
              num_tgt_feats (int): number of target side features. 
              src_seq_length (int): maximum source sequence length.
              tgt_seq_length (int): maximum target sequence length.
              dynamic_dict (bool): create dynamic dictionaries?
              use_filter_pred (bool): use a custom filter predicate to filter
                  out examples?
      """
      def __init__(self, fields, src_examples_iter, tgt_examples_iter, weights_iter,
                   num_src_feats=0, num_tgt_feats=0,
                   src_seq_length=0, tgt_seq_length=0,
                   dynamic_dict=True, use_filter_pred=True):
          self.data_type = 'text'
  
          # self.src_vocabs: mutated in dynamic_dict, used in
          # collapse_copy_scores and in Translator.py
          self.src_vocabs = []
  
          self.n_src_feats = num_src_feats
          self.n_tgt_feats = num_tgt_feats
  
          # Each element of an example is a dictionary whose keys represents
          # at minimum the src tokens and their indices and potentially also
          # the src and tgt features and alignment information.
          if weights_iter is not None:
              examples_iter = (self._join_dicts(src, tgt, weights) for src, tgt, weights in
                               zip(src_examples_iter, tgt_examples_iter, weights_iter))
          elif tgt_examples_iter is not None:
              examples_iter = (self._join_dicts(src, tgt) for src, tgt in
                               zip(src_examples_iter, tgt_examples_iter))
          else:
              examples_iter = src_examples_iter
  
          if dynamic_dict:
              examples_iter = self._dynamic_dict(examples_iter)
  
          # Peek at the first to see which fields are used.
          ex, examples_iter = self._peek(examples_iter)
          keys = ex.keys()
  
          out_fields = [(k, fields[k]) if k in fields else (k, None)
                        for k in keys]
          example_values = ([float(ex[k][0]) if k == 'weights' else ex[k] for k in keys] for ex in examples_iter)
          out_examples = (self._construct_example_fromlist(
                              ex_values, out_fields)
                          for ex_values in example_values)
          # If out_examples is a generator, we need to save the filter_pred
          # function in serialization too, which would cause a problem when
          # `torch.save()`. Thus we materialize it as a list.
          out_examples = list(out_examples)
  
          def filter_pred(example):
              return 0 < len(example.src) <= src_seq_length \
                 and 0 < len(example.tgt) <= tgt_seq_length
  
          filter_pred = filter_pred if use_filter_pred else lambda x: True
  
          super(TextDataset, self).__init__(
              out_examples, out_fields, filter_pred
          )
  
      def sort_key(self, ex):
          """ Sort using length of source sentences. """
          return len(ex.src)
  
      @staticmethod
      def collapse_copy_scores(scores, batch, tgt_vocab, src_vocabs):
          """
          Given scores from an expanded dictionary
          corresponeding to a batch, sums together copies,
          with a dictionary word when it is ambigious.
          """
          offset = len(tgt_vocab)
          for b in range(batch.batch_size):
              blank = []
              fill = []
              index = batch.indices.data[b]
              src_vocab = src_vocabs[index]
              for i in range(1, len(src_vocab)):
                  sw = src_vocab.itos[i]
                  ti = tgt_vocab.stoi[sw]
                  if ti != 0:
                      blank.append(offset + i)
                      fill.append(ti)
              if blank:
                  blank = torch.Tensor(blank).type_as(batch.indices.data)
                  fill = torch.Tensor(fill).type_as(batch.indices.data)
                  scores[:, b].index_add_(1, fill,
                                          scores[:, b].index_select(1, blank))
                  scores[:, b].index_fill_(1, blank, 1e-10)
          return scores
  
      @staticmethod
      def make_text_examples_nfeats_tpl(path, truncate, side):
          """
          Args:
              path (str): location of a src or tgt file.
              truncate (int): maximum sequence length (0 for unlimited).
              side (str): "src" or "tgt".
  
          Returns:
              (example_dict iterator, num_feats) tuple.
          """
          assert side in ['src', 'tgt']
  
          if path is None:
              return (None, 0)
  
          # All examples have same number of features, so we peek first one
          # to get the num_feats.
          examples_nfeats_iter = \
              TextDataset.read_text_file(path, truncate, side)
  
          first_ex = next(examples_nfeats_iter)
          num_feats = first_ex[1]
  
          # Chain back the first element - we only want to peek it.
          examples_nfeats_iter = chain([first_ex], examples_nfeats_iter)
          examples_iter = (ex for ex, nfeats in examples_nfeats_iter)
  
          return (examples_iter, num_feats)
  
      @staticmethod
      def read_text_file(path, truncate, side):
          """
          Args:
              path (str): location of a src or tgt file.
              truncate (int): maximum sequence length (0 for unlimited).
              side (str): "src" or "tgt".
  
          Yields:
              (word, features, nfeat) triples for each line.
          """
          with codecs.open(path, "r", "utf-8") as corpus_file:
              for i, line in enumerate(corpus_file):
                  line = line.strip().split()
                  if truncate:
                      line = line[:truncate]
  
                  words, feats, n_feats = \
                      TextDataset.extract_text_features(line)
  
                  example_dict = {side: words, "indices": i}
                  if feats:
                      prefix = side + "_feat_"
                      example_dict.update((prefix + str(j), f)
                                          for j, f in enumerate(feats))
                  yield example_dict, n_feats
  
      @staticmethod
      def get_fields(n_src_features, n_tgt_features):
          """
          Args:
              n_src_features (int): the number of source features to
                  create `torchtext.data.Field` for.
              n_tgt_features (int): the number of target features to
                  create `torchtext.data.Field` for.
  
          Returns:
              A dictionary whose keys are strings and whose values
              are the corresponding Field objects.
          """
          fields = {}
  
          fields["src"] = torchtext.data.Field(
              pad_token=PAD_WORD,
              include_lengths=True)
  
          for j in range(n_src_features):
              fields["src_feat_"+str(j)] = \
                  torchtext.data.Field(pad_token=PAD_WORD)
  
          fields["tgt"] = torchtext.data.Field(
              init_token=BOS_WORD, eos_token=EOS_WORD,
              pad_token=PAD_WORD)
  
          for j in range(n_tgt_features):
              fields["tgt_feat_"+str(j)] = \
                  torchtext.data.Field(init_token=BOS_WORD, eos_token=EOS_WORD,
                                       pad_token=PAD_WORD)
  
          def make_src(data, _):
              src_size = max([t.size(0) for t in data])
              src_vocab_size = max([t.max() for t in data]) + 1
              alignment = torch.zeros(src_size, len(data), src_vocab_size)
              for i, sent in enumerate(data):
                  for j, t in enumerate(sent):
                      alignment[j, i, t] = 1
              return alignment
  
          fields["src_map"] = torchtext.data.Field(
              use_vocab=False, dtype=torch.long,
              postprocessing=make_src, sequential=False)
  
          def make_tgt(data, _):
              tgt_size = max([t.size(0) for t in data])
              alignment = torch.zeros(tgt_size, len(data)).long()
              for i, sent in enumerate(data):
                  alignment[:sent.size(0), i] = sent
              return alignment
  
          fields["alignment"] = torchtext.data.Field(
              use_vocab=False, dtype=torch.long,
              postprocessing=make_tgt, sequential=False)
  
          fields["indices"] = torchtext.data.Field(
              use_vocab=False, dtype=torch.long,
              sequential=False)
  
          fields["weights"] = torchtext.data.Field(use_vocab=False, dtype=torch.float, sequential=False)
          
          return fields
  
      @staticmethod
      def get_num_features(corpus_file, side):
          """
          Peek one line and get number of features of it.
          (All lines must have same number of features).
          For text corpus, both sides are in text form, thus
          it works the same.
  
          Args:
              corpus_file (str): file path to get the features.
              side (str): 'src' or 'tgt'.
  
          Returns:
              number of features on `side`.
          """
          with codecs.open(corpus_file, "r", "utf-8") as cf:
              f_line = cf.readline().strip().split()
              _, _, num_feats = TextDataset.extract_text_features(f_line)
  
          return num_feats
  
      # Below are helper functions for intra-class use only.
      def _dynamic_dict(self, examples_iter):
          for example in examples_iter:
              src = example["src"]
              src_vocab = torchtext.vocab.Vocab(Counter(src))
              self.src_vocabs.append(src_vocab)
              # Mapping source tokens to indices in the dynamic dict.
              src_map = torch.LongTensor([src_vocab.stoi[w] for w in src])
              example["src_map"] = src_map
  
              if "tgt" in example:
                  tgt = example["tgt"]
                  mask = torch.LongTensor(
                          [0] + [src_vocab.stoi[w] for w in tgt] + [0])
                  example["alignment"] = mask
              yield example
  
  
  class ShardedTextCorpusIterator(object):
      """
      This is the iterator for text corpus, used for sharding large text
      corpus into small shards, to avoid hogging memory.
  
      Inside this iterator, it automatically divides the corpus file into
      shards of size `shard_size`. Then, for each shard, it processes
      into (example_dict, n_features) tuples when iterates.
      """
      def __init__(self, corpus_path, line_truncate, side, shard_size,
                   assoc_iter=None):
          """
          Args:
              corpus_path: the corpus file path.
              line_truncate: the maximum length of a line to read.
                              0 for unlimited.
              side: "src" or "tgt".
              shard_size: the shard size, 0 means not sharding the file.
              assoc_iter: if not None, it is the associate iterator that
                          this iterator should align its step with.
          """
          try:
              # The codecs module seems to have bugs with seek()/tell(),
              # so we use io.open().
              self.corpus = io.open(corpus_path, "r", encoding="utf-8")
          except IOError:
              sys.stderr.write("Failed to open corpus file: %s" % corpus_path)
              sys.exit(1)
  
          self.line_truncate = line_truncate
          self.side = side
          self.shard_size = shard_size
          self.assoc_iter = assoc_iter
          self.last_pos = 0
          self.line_index = -1
          self.eof = False
  
      def __iter__(self):
          """
          Iterator of (example_dict, nfeats).
          On each call, it iterates over as many (example_dict, nfeats) tuples
          until this shard's size equals to or approximates `self.shard_size`.
          """
          iteration_index = -1
          if self.assoc_iter is not None:
              # We have associate iterator, just yields tuples
              # util we run parallel with it.
              while self.line_index < self.assoc_iter.line_index:
                  line = self.corpus.readline()
                  if line == '':
                      raise AssertionError(
                          "Two corpuses must have same number of lines!")
  
                  self.line_index += 1
                  iteration_index += 1
                  yield self._example_dict_iter(line, iteration_index)
  
              if self.assoc_iter.eof:
                  self.eof = True
                  self.corpus.close()
          else:
              # Yield tuples util this shard's size reaches the threshold.
              self.corpus.seek(self.last_pos)
              while True:
                  if self.shard_size != 0 and self.line_index % 64 == 0:
                      # This part of check is time consuming on Py2 (but
                      # it is quite fast on Py3, weird!). So we don't bother
                      # to check for very line. Instead we chekc every 64
                      # lines. Thus we are not dividing exactly per
                      # `shard_size`, but it is not too much difference.
                      cur_pos = self.corpus.tell()
                      if cur_pos >= self.last_pos + self.shard_size:
                          self.last_pos = cur_pos
                          raise StopIteration
  
                  line = self.corpus.readline()
                  if line == '':
                      self.eof = True
                      self.corpus.close()
                      raise StopIteration
  
                  self.line_index += 1
                  iteration_index += 1
                  yield self._example_dict_iter(line, iteration_index)
  
      def hit_end(self):
          return self.eof
  
      @property
      def num_feats(self):
          # We peek the first line and seek back to
          # the beginning of the file.
          saved_pos = self.corpus.tell()
  
          line = self.corpus.readline().split()
          if self.line_truncate:
              line = line[:self.line_truncate]
          _, _, self.n_feats = TextDataset.extract_text_features(line)
  
          self.corpus.seek(saved_pos)
  
          return self.n_feats
  
      def _example_dict_iter(self, line, index):
          line = line.split()
          if self.line_truncate:
              line = line[:self.line_truncate]
          words, feats, n_feats = TextDataset.extract_text_features(line)
          example_dict = {self.side: words, "indices": index}
          if feats:
              # All examples must have same number of features.
              aeq(self.n_feats, n_feats)
  
              prefix = self.side + "_feat_"
              example_dict.update((prefix + str(j), f)
                                  for j, f in enumerate(feats))
  
          return example_dict
  