text: # File   : text.py
  # Author : Zhengkun Tian
  # Email  : zhengkun.tian@outlook.com
  
  import torch
  import logging
  from otrans.data import *
  from torch.utils.data import Dataset
  
  
  
  class TextDataset(Dataset):
      def __init__(self, params, datadict, is_eval=False):
  
          self.params = params
          self.is_eval = is_eval
          self.src_unit2idx = load_vocab(params['src_vocab'])
          self.tgt_unit2idx = load_vocab(params['tgt_vocab'])
          self.reverse = params['reverse'] if 'reverse' in params else False
  
          if self.reverse:
              logging.info('Reverse the src and tgt sequence!')
  
          self.src_list = []
          self.tgt_dict = {}
          for src_file in datadict['src']:
              with open(src_file, 'r', encoding='utf-8') as t:
                  for line in t:
                      parts = line.strip().split()
                      utt_id = parts[0]
                      label = []
                      for c in parts[1:]:
                          label.append(self.src_unit2idx[c] if c in self.src_unit2idx else self.src_unit2idx[UNK_TOKEN])
                      self.src_list.append((utt_id, label))
  
          for tgt_file in datadict['tgt']:
              with open(tgt_file, 'r', encoding='utf-8') as t:
                  for line in t:
                      parts = line.strip().split()
                      utt_id = parts[0]
                      label = []
                      for c in parts[1:]:
                          label.append(self.tgt_unit2idx[c] if c in self.tgt_unit2idx else self.tgt_unit2idx[UNK_TOKEN])
                      self.tgt_dict[utt_id] = label
  
          assert len(self.src_list) == len(self.tgt_dict)
  
          self.lengths = len(self.src_list)
  
      def __getitem__(self, index):
          idx, src_seq = self.src_list[index]
          tgt_seq = self.tgt_dict[idx]
  
          if self.reverse:
              src_seq.reverse()
              tgt_seq.reverse()
  
          return idx, src_seq, tgt_seq
  
      def __len__(self):
          return self.lengths
  
      @property
      def src_vocab_size(self):
          return len(self.src_unit2idx)
  
      @property
      def tgt_vocab_size(self):
          return len(self.tgt_unit2idx)
      
      @property
      def src_idx2unit(self):
          return {i: c for (c, i) in self.src_unit2idx.items()}
  
      @property
      def tgt_idx2unit(self):
          return {i: c for (c, i) in self.tgt_unit2idx.items()}
  
  