text: # -*- coding: utf-8 -*-
  # Author: XuMing <xuming624@qq.com>
  # Brief: Date reader for training seq2seq model
  from collections import Counter
  
  # Define constants associated with the usual special tokens.
  PAD_ID = 0
  GO_ID = 1
  EOS_ID = 2
  
  PAD_TOKEN = 'PAD'
  EOS_TOKEN = 'EOS'
  GO_TOKEN = 'GO'
  
  
  class Reader:
      def __init__(self, config, train_path=None, token_2_id=None,
                   special_tokens=(), dataset_copies=1):
          self.config = config
          self.dataset_copies = dataset_copies
          # Vocabulary
          max_vocab_size = config.max_vocab_size
          if train_path is None:
              self.token_2_id = token_2_id
          else:
              token_counts = Counter()
              for tokens in self.read_tokens(train_path):
                  token_counts.update(tokens)
  
              self.token_counts = token_counts
              # Get max_vocabulary size words
              count_pairs = sorted(token_counts.items(), key=lambda k: (-k[1], k[0]))
              vocab, _ = list(zip(*count_pairs))
              vocab = list(vocab)
              # Insert the special tokens to the beginning
              vocab[0:0] = special_tokens
              full_token_id = list(zip(vocab, range(len(vocab))))
              self.full_token_2_id = dict(full_token_id)
              self.token_2_id = dict(full_token_id[:max_vocab_size])
          self.id_2_token = {v: k for k, v in self.token_2_id.items()}
  
      def read_tokens(self, path):
          """
          Reads the given file line by line and yields the list of tokens present
          in each line.
  
          :param path:
          :return:
          """
          raise NotImplementedError("Must implement read_tokens")
  
      def unknown_token(self):
          raise NotImplementedError("Must implement unknow_tokens")
  
      def read_samples_by_string(self, path):
          """
          Reads the given file line by line and yields the word-form of each
          derived sample.
  
          :param path:
          :return:
          """
          raise NotImplementedError("Must implement read_samples")
  
      def convert_token_2_id(self, token):
          """
          Token to id
          :param token:
          :return:
          """
          token_id = token if token in self.token_2_id else self.unknown_token()
          return self.token_2_id[token_id]
  
      def convert_id_2_token(self, id):
          """
          Word id to token
          :param id:
          :return:
          """
          return self.id_2_token[id]
  
      def is_unknown_token(self, token):
          """
          True if the given token is out of vocabulary
          :param token:
          :return:
          """
          return token not in self.token_2_id or token == self.unknown_token()
  
      def sentence_2_token_ids(self, sentence):
          """
          Convert a sentence to word ids
          :param sentence:
          :return:
          """
          return [self.convert_token_2_id(w) for w in sentence.split()]
  
      def token_ids_2_tokens(self, word_ids):
          """
          Convert a list of word ids to words
          :param word_ids:
          :return:
          """
          return [self.convert_id_2_token(w) for w in word_ids]
  
      def read_samples(self, path):
          """
          Read sample of path's data
          :param path:
          :return: generate list
          """
          for source_words, target_words in self.read_samples_by_string(path):
              source = [self.convert_token_2_id(w) for w in source_words]
              target = [self.convert_token_2_id(w) for w in target_words]
              target.append(EOS_ID)
              yield source, target
  
      def build_dataset(self, path):
          dataset = [[] for _ in self.config.buckets]
          # Copy the data set for different dropouts
          for _ in range(self.dataset_copies):
              for source, target in self.read_samples(path):
                  for bucket_id, (source_size, target_size) in enumerate(self.config.buckets):
                      if len(source) < source_size and len(target) < target_size:
                          dataset[bucket_id].append([source, target])
                          break
          return dataset
  