text: # packages
  import sys
  sys.path.append("C:/BERTVision/code/torch")
  import torch, h5py
  import numpy as np
  
  # prepare torch data set
  class COLAH5Processor(torch.utils.data.Dataset):
      NAME = 'COLAH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
      def __init__(self, type, args):
          self.type = type
          self.args = args
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          # todo: generalize locations
          self.train_embed_path = self.file_path + 'cola_bert_embeds.h5'
          self.train_label_path = self.file_path + 'cola_labels.h5'
          self.train_idx_path = self.file_path + 'cola_idx.h5'
  
          self.val_embed_path = self.file_path + 'cola_dev_bert_embeds.h5'
          self.val_label_path = self.file_path + 'cola_dev_labels.h5'
          self.val_idx_path = self.file_path + 'cola_dev_idx.h5'
  
          # if train, initialize the train data
          if self.type == 'train':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
          # if train, initialize the dev data
          if self.type == 'dev':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.val_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.val_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.val_idx_path, 'r')['idx']
  
              with h5py.File(self.val_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
      def __len__(self):
          return self.dataset_len
  
      def __getitem__(self, idx):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(idx):
              idx = idx.tolist()
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'labels': self.labels[idx],
                      'idx': self.idx[idx]}
              return sample
  
          if self.type == 'dev':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'labels': self.labels[idx],
                      'idx': self.idx[idx]}
              return sample
  
  #
  
  
  # prepare torch data set
  class MNLIH5Processor(torch.utils.data.Dataset):
      NAME = 'MNLIH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
      def __init__(self, type, args, shard=False, **kwargs):
          # init configurable string
          self.type = type
          # receive arg parser
          self.args = args
          # init shard for sampling large ds if specified
          self.shard = shard
          # set seed if given
          self.seed = kwargs
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          # todo: generalize locations
          self.train_embed_path = self.file_path + 'mnli_bert_embeds.h5'
          self.train_label_path = self.file_path + 'mnli_labels.h5'
          self.train_idx_path = self.file_path + 'mnli_idx.h5'
  
          self.dev_matched_embed_path = self.file_path + 'mnli_dev_matched_bert_embeds.h5'
          self.dev_matched_label_path = self.file_path + 'mnli_dev_matched_labels.h5'
          self.dev_matched_idx_path = self.file_path + 'mnli_dev_matched_idx.h5'
  
          self.dev_mismatched_embed_path = self.file_path + 'mnli_dev_mismatched_bert_embeds.h5'
          self.dev_mismatched_label_path = self.file_path + 'mnli_dev_mismatched_labels.h5'
          self.dev_mismatched_idx_path = self.file_path + 'mnli_dev_mismatched_idx.h5'
  
          # if train, initialize the train data
          if self.type == 'train' and self.shard == True:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # np seed
              np.random.seed(self.seed['seed'])
              # 15% of data set
              self.random_indices_len = int(self.args.shard*self.dataset_len)
              # select randomly from indices, without replacement
              self.random_indices = np.random.choice(np.arange(self.dataset_len), size=self.random_indices_len, replace=False)
              # sort for now
              self.random_indices = np.sort(self.random_indices)
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          elif self.type == 'train' and self.shard == False:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          # if matched, initialize the dev data
          if self.type == 'dev_matched':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.dev_matched_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.dev_matched_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.dev_matched_idx_path, 'r')['idx']
  
              with h5py.File(self.dev_matched_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              self.random_indices = np.arange(self.dataset_len)
  
          # if matched, initialize the dev data
          if self.type == 'dev_mismatched':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.dev_mismatched_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.dev_mismatched_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.dev_mismatched_idx_path, 'r')['idx']
  
              with h5py.File(self.dev_mismatched_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              self.random_indices = np.arange(self.dataset_len)
  
      def __len__(self):
          if self.type == 'train' and self.shard == True:
              return len(self.random_indices)
          else:
              return self.dataset_len
  
      def __getitem__(self, random_indices):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(random_indices):
              random_indices = random_indices.tolist()
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
          if self.type == 'dev_matched':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
          if self.type == 'dev_mismatched':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
  #
  
  
  # prepare torch data set
  class MSRH5Processor(torch.utils.data.Dataset):
      NAME = 'MSRH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
      def __init__(self, type, args, shard=False, **kwargs):
          # init configurable string
          self.type = type
          # receive arg parser
          self.args = args
          # init shard for sampling large ds if specified
          self.shard = shard
          # set seed if given
          self.seed = kwargs
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          # todo: generalize locations
          self.train_embed_path = self.file_path + 'msr_bert_embeds.h5'
          self.train_label_path = self.file_path + 'msr_labels.h5'
          self.train_idx_path = self.file_path + 'msr_idx.h5'
  
          self.val_embed_path = self.file_path + 'msr_dev_bert_embeds.h5'
          self.val_label_path = self.file_path + 'msr_dev_labels.h5'
          self.val_idx_path = self.file_path + 'msr_dev_idx.h5'
  
          # if train, initialize the train data
          if self.type == 'train' and self.shard == True:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # np seed
              np.random.seed(self.seed['seed'])
              # 15% of data set
              self.random_indices_len = int(self.args.shard*self.dataset_len)
              # select randomly from indices, without replacement
              self.random_indices = np.random.choice(np.arange(self.dataset_len), size=self.random_indices_len, replace=False)
              # sort for now
              self.random_indices = np.sort(self.random_indices)
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          elif self.type == 'train' and self.shard == False:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          # if train, initialize the dev data
          if self.type == 'dev':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.val_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.val_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.val_idx_path, 'r')['idx']
  
              with h5py.File(self.val_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              self.random_indices = np.arange(self.dataset_len)
  
      def __len__(self):
          if self.type == 'train' and self.shard == True:
              return len(self.random_indices)
          else:
              return self.dataset_len
  
      def __getitem__(self, random_indices):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(random_indices):
              random_indices = random_indices.tolist()
  
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
          if self.type == 'dev':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
  #
  
  
  # prepare torch data set
  class QNLIH5Processor(torch.utils.data.Dataset):
      NAME = 'QNLIH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
      def __init__(self, type, args, shard=False, **kwargs):
          # init configurable string
          self.type = type
          # receive arg parser
          self.args = args
          # init shard for sampling large ds if specified
          self.shard = shard
          # set seed if given
          self.seed = kwargs
  
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          # h5 locs
          self.train_embed_path = self.file_path + 'qnli_bert_embeds.h5'
          self.train_label_path = self.file_path + 'qnli_labels.h5'
          self.train_idx_path = self.file_path + 'qnli_idx.h5'
  
          self.val_embed_path = self.file_path + 'qnli_dev_bert_embeds.h5'
          self.val_label_path = self.file_path + 'qnli_dev_labels.h5'
          self.val_idx_path = self.file_path + 'qnli_dev_idx.h5'
  
          # if train, initialize the train data
          if self.type == 'train' and self.shard == True:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # np seed
              np.random.seed(self.seed['seed'])
              # 15% of data set
              self.random_indices_len = int(self.args.shard*self.dataset_len)
              # select randomly from indices, without replacement
              self.random_indices = np.random.choice(np.arange(self.dataset_len), size=self.random_indices_len, replace=False)
              # sort for now
              self.random_indices = np.sort(self.random_indices)
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          elif self.type == 'train' and self.shard == False:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          # if train, initialize the dev data
          if self.type == 'dev':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.val_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.val_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.val_idx_path, 'r')['idx']
  
              with h5py.File(self.val_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              self.random_indices = np.arange(self.dataset_len)
  
      def __len__(self):
          if self.type == 'train' and self.shard == True:
              return len(self.random_indices)
          else:
              return self.dataset_len
  
      def __getitem__(self, random_indices):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(random_indices):
              random_indices = random_indices.tolist()
  
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
          if self.type == 'dev':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
  
  #
  
  
  # prepare torch data set
  class QQPH5Processor(torch.utils.data.Dataset):
      NAME = 'QQPH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
      def __init__(self, type, args, shard=False, **kwargs):
          # init configurable string
          self.type = type
          # receive arg parser
          self.args = args
          # init shard for sampling large ds if specified
          self.shard = shard
          # set seed if given
          self.seed = kwargs
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          # h5 locs
          self.train_embed_path = self.file_path + 'qqpairs_bert_embeds.h5'
          self.train_label_path = self.file_path + 'qqpairs_labels.h5'
          self.train_idx_path = self.file_path + 'qqpairs_idx.h5'
  
          self.val_embed_path = self.file_path + 'qqpairs_dev_bert_embeds.h5'
          self.val_label_path = self.file_path + 'qqpairs_dev_labels.h5'
          self.val_idx_path = self.file_path + 'qqpairs_dev_idx.h5'
  
          # if train, initialize the train data
          if self.type == 'train' and self.shard == True:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # np seed
              np.random.seed(self.seed['seed'])
              # 15% of data set
              self.random_indices_len = int(self.args.shard*self.dataset_len)
              # select randomly from indices, without replacement
              self.random_indices = np.random.choice(np.arange(self.dataset_len), size=self.random_indices_len, replace=False)
              # sort for now
              self.random_indices = np.sort(self.random_indices)
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          elif self.type == 'train' and self.shard == False:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          # if train, initialize the dev data
          if self.type == 'dev':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.val_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.val_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.val_idx_path, 'r')['idx']
  
              with h5py.File(self.val_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              self.random_indices = np.arange(self.dataset_len)
  
      def __len__(self):
          if self.type == 'train' and self.shard == True:
              return len(self.random_indices)
          else:
              return self.dataset_len
  
      def __getitem__(self, random_indices):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(random_indices):
              random_indices = random_indices.tolist()
  
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
          if self.type == 'dev':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
  
  #
  
  
  # prepare torch data set
  class RTEH5Processor(torch.utils.data.Dataset):
      NAME = 'RTEH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
      def __init__(self, type, args):
          self.type = type
          self.args = args
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          # todo: generalize locations
          self.train_embed_path = self.file_path + 'rte_bert_embeds.h5'
          self.train_label_path = self.file_path + 'rte_labels.h5'
          self.train_idx_path = self.file_path + 'rte_idx.h5'
  
          self.val_embed_path = self.file_path + 'rte_dev_bert_embeds.h5'
          self.val_label_path = self.file_path + 'rte_dev_labels.h5'
          self.val_idx_path = self.file_path + 'rte_dev_idx.h5'
  
          # if train, initialize the train data
          if self.type == 'train':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
          # if train, initialize the dev data
          if self.type == 'dev':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.val_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.val_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.val_idx_path, 'r')['idx']
  
              with h5py.File(self.val_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
      def __len__(self):
          return self.dataset_len
  
      def __getitem__(self, idx):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(idx):
              idx = idx.tolist()
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'labels': self.labels[idx],
                      'idx': self.idx[idx]}
              return sample
  
          if self.type == 'dev':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'labels': self.labels[idx],
                      'idx': self.idx[idx]}
              return sample
  
  
  # prepare torch data set
  class SQuADH5Processor(torch.utils.data.Dataset):
      NAME = 'SQuADH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
  
      def __init__(self, type, args):
          self.type = type
          self.args = args
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          self.train_embed_path = self.file_path + 'squad_train_embeds.h5'
          self.train_start_path = self.file_path + 'squad_train_start_labels.h5'
          self.train_end_path = self.file_path + 'squad_train_end_labels.h5'
          self.train_idx_path = self.file_path + 'squad_train_indices.h5'
  
          self.val_embed_path = self.file_path + 'squad_dev_embeds.h5'
          self.val_start_path = self.file_path + 'squad_dev_start_labels.h5'
          self.val_end_path = self.file_path + 'squad_dev_end_labels.h5'
          self.val_idx_path = self.file_path + 'squad_dev_indices.h5'
  
          # if train, initialize the train data
          if self.type == 'train':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # start ids are shaped: [batch_sz, ]
              self.start = h5py.File(self.train_start_path, 'r')["start_ids"]
              # end ids are shaped: [batch_sz, ]
              self.end = h5py.File(self.train_end_path, 'r')["end_ids"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['indices']
  
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
          # if train, initialize the dev data
          if self.type == 'dev':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.val_embed_path, 'r')["embeds"]
              # start ids are shaped: [batch_sz, ]
              self.start = h5py.File(self.val_start_path, 'r')["start_ids"]
              # end ids are shaped: [batch_sz, ]
              self.end = h5py.File(self.val_end_path, 'r')["end_ids"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.val_idx_path, 'r')['indices']
  
              with h5py.File(self.val_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
      def __len__(self):
          return self.dataset_len
  
      def __getitem__(self, idx):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(idx):
              idx = idx.tolist()
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'start_ids': self.start[idx],
                      'end_ids': self.end[idx],
                      'idx': self.idx[idx]}
              return sample
  
          if self.type == 'dev':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'start_ids': self.start[idx],
                      'end_ids': self.end[idx],
                      'idx': self.idx[idx]}
              return sample
  
  #
  
  # prepare torch data set
  class SSTH5Processor(torch.utils.data.Dataset):
      NAME = 'SSTH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
      def __init__(self, type, args, shard=False, **kwargs):
          # init configurable string
          self.type = type
          # receive arg parser
          self.args = args
          # init shard for sampling large ds if specified
          self.shard = shard
          # set seed if given
          self.seed = kwargs
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          # h5 locs
          self.train_embed_path = self.file_path + 'sst_bert_embeds.h5'
          self.train_label_path = self.file_path + 'sst_labels.h5'
          self.train_idx_path = self.file_path + 'sst_idx.h5'
  
          self.val_embed_path = self.file_path + 'sst_dev_bert_embeds.h5'
          self.val_label_path = self.file_path + 'sst_dev_labels.h5'
          self.val_idx_path = self.file_path + 'sst_dev_idx.h5'
  
          # if train, initialize the train data
          if self.type == 'train' and self.shard == True:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # np seed
              np.random.seed(self.seed['seed'])
              # 15% of data set
              self.random_indices_len = int(self.args.shard*self.dataset_len)
              # select randomly from indices, without replacement
              self.random_indices = np.random.choice(np.arange(self.dataset_len), size=self.random_indices_len, replace=False)
              # sort for now
              self.random_indices = np.sort(self.random_indices)
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          elif self.type == 'train' and self.shard == False:
              # initialize ds length
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
          # if train, initialize the dev data
          if self.type == 'dev':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.val_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.val_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.val_idx_path, 'r')['idx']
  
              with h5py.File(self.val_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
              self.random_indices = np.arange(self.dataset_len)
  
      def __len__(self):
          if self.type == 'train' and self.shard == True:
              return len(self.random_indices)
          else:
              return self.dataset_len
  
      def __getitem__(self, random_indices):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(random_indices):
              random_indices = random_indices.tolist()
  
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
          if self.type == 'dev':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[random_indices],
                      'labels': self.labels[random_indices],
                      'idx': self.idx[random_indices]}
              return sample
  
  
  
  # prepare torch data set
  class STSBH5Processor(torch.utils.data.Dataset):
      NAME = 'STSBH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
      def __init__(self, type, args):
          self.type = type
          self.args = args
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          # todo: generalize locations
          self.train_embed_path = self.file_path + 'stsb_bert_embeds.h5'
          self.train_label_path = self.file_path + 'stsb_labels.h5'
          self.train_idx_path = self.file_path + 'stsb_idx.h5'
  
          self.val_embed_path = self.file_path + 'stsb_dev_bert_embeds.h5'
          self.val_label_path = self.file_path + 'stsb_dev_labels.h5'
          self.val_idx_path = self.file_path + 'stsb_dev_idx.h5'
  
          # if train, initialize the train data
          if self.type == 'train':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
          # if train, initialize the dev data
          if self.type == 'dev':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.val_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.val_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.val_idx_path, 'r')['idx']
  
              with h5py.File(self.val_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
      def __len__(self):
          return self.dataset_len
  
      def __getitem__(self, idx):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(idx):
              idx = idx.tolist()
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'labels': self.labels[idx],
                      'idx': self.idx[idx]}
              return sample
  
          if self.type == 'dev':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'labels': self.labels[idx],
                      'idx': self.idx[idx]}
              return sample
  
  
  #
  # prepare torch data set
  class WNLIH5Processor(torch.utils.data.Dataset):
      NAME = 'WNLIH5'
      '''
      This class lazily emits batches from H5 files for deep learning.
  
      Parameters
      ----------
      type : string
          A string used to flag conditional statements in order to use/retrieve
          the right data set. [train or dev]
  
      Returns
      -------
      sample : tensor [layers, tokens, features]
          A single sample of data indexed by the torch data set.
      '''
      def __init__(self, type, args):
          self.type = type
          self.args = args
  
          if self.args.checkpoint == 'bert-base-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\'
          elif self.args.checkpoint == 'bert-large-uncased':
              self.file_path = 'C:\\w266\\data\\h5py_embeds\\bert_large\\'
  
          # todo: generalize locations
          self.train_embed_path = self.file_path + 'wnli_bert_embeds.h5'
          self.train_label_path = self.file_path + 'wnli_labels.h5'
          self.train_idx_path = self.file_path + 'wnli_idx.h5'
  
          self.val_embed_path = self.file_path + 'wnli_dev_bert_embeds.h5'
          self.val_label_path = self.file_path + 'wnli_dev_labels.h5'
          self.val_idx_path = self.file_path + 'wnli_dev_idx.h5'
  
          # if train, initialize the train data
          if self.type == 'train':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.train_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.train_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.train_idx_path, 'r')['idx']
  
              with h5py.File(self.train_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
          # if train, initialize the dev data
          if self.type == 'dev':
              # embeds are shaped: [batch_sz, layers, tokens, features]
              self.embeddings = h5py.File(self.val_embed_path, 'r')["embeds"]
              # labels are shaped: [batch_sz, ]
              self.labels = h5py.File(self.val_label_path, 'r')["labels"]
              # idx ids are shaped: [batch_sz, ]
              self.idx = h5py.File(self.val_idx_path, 'r')['idx']
  
              with h5py.File(self.val_embed_path, 'r') as file:
                  self.dataset_len = len(file["embeds"])
  
      def __len__(self):
          return self.dataset_len
  
      def __getitem__(self, idx):
          '''
          Torch's lazy emission system
          '''
          if torch.is_tensor(idx):
              idx = idx.tolist()
  
          if self.type == 'train':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'labels': self.labels[idx],
                      'idx': self.idx[idx]}
              return sample
  
          if self.type == 'dev':
              # emits [layers, tokens, features]
              sample = {'embeddings': self.embeddings[idx],
                      'labels': self.labels[idx],
                      'idx': self.idx[idx]}
              return sample
  
  #
  