text: from collections import Counter
  import numpy as np
  from torch.utils.data import Sampler
  
  
  class RepeatsBatchSampler(Sampler):
      def __init__(self, keys, subset_index=None):
          """
          Batch sampler where each batch consists of all samples with identical keys value.
  
          Args:
              keys (Any): list of keys used to group indicies.
              subset_index (list of indices, optional): List of indices to subselect entries from keys. Defaults to None, where all indices from keys are used.
          """
  
          if subset_index is None:
              subset_index = np.arange(len(keys))
          _, inv = np.unique(keys[subset_index], return_inverse=True)
          self.repeat_index = np.unique(inv)
          self.repeat_sets = inv
          self.subset_index = subset_index
  
      def __iter__(self):
          for u in self.repeat_index:
              yield list(self.subset_index[self.repeat_sets == u])
  
      def __len__(self):
          return len(self.repeat_index)
  
  
  class SubsetSequentialSampler(Sampler):
      def __init__(self, indices):
          """
          Samples elements sequentially from a given list of indices, without replacement.
          Arguments:
              indices (list): a list of indices
          """
          self.indices = indices
  
      def __iter__(self):
          return (self.indices[i] for i in range(len(self.indices)))
  
      def __len__(self):
          return len(self.indices)
  
  
  class SampledSubsetRandomSampler(Sampler):
      """"
      Samples elements randomly from sampled subset of indices.
      Arguments:
          indices (sequence): a sequence of indices
          num_samples (int): number of samples to draw
      """
  
      def __init__(self, indices, num_samples):
          self.indices = indices
          self.num_samples = num_samples
          self.replace = num_samples > len(indices)
  
      def __iter__(self):
          indices = np.random.choice(
              self.indices, size=self.num_samples, replace=self.replace
          )
          return iter(indices.tolist())
  
      def __len__(self):
          return self.num_samples
  
  
  class SampledSubsetSequentialSampler(Sampler):
      r"""Samples elements sequentially from sampled subset of indices.
      Arguments:
          indices (sequence): a sequence of indices
          num_samples (int): number of samples to draw
      """
  
      def __init__(self, indices, num_samples):
          self.indices = indices
          self.num_samples = num_samples
          self.replace = num_samples > len(indices)
  
      def __iter__(self):
          indices = np.random.choice(
              self.indices, size=self.num_samples, replace=self.replace
          )
          sorted_indices = np.sort(indices)
          return iter(sorted_indices.tolist())
  
      def __len__(self):
          return self.num_samples
  
  
  class SubSubsetRandomSequentialSampler(Sampler):
      """Samples a subset from a dataset randomly.
      Then iterates through those indices sequentially.
      Arguments:
          indices (list): a list of indices
          subsubset_size (int): size of subset of subset
      """
  
      def __init__(self, indices, subsubset_size):
          self.indices = indices
          self.subsubset_size = subsubset_size
  
      def __iter__(self):
          subsubset_indices = np.random.choice(
              self.indices, size=self.subsubset_size, replace=False
          )
          return (subsubset_indices[i] for i in range(len(subsubset_indices)))
  
      def __len__(self):
          return self.subsubset_size
  
  
  class BalancedSubsetSampler(Sampler):
      """Samples elements randomly from a given list of indices, without replacement, balanced by occurence of types.
      Arguments:
          indices (list): a list of indices
      """
  
      def configure_sampler(self, indices, types, mode="shortest"):
          self.indices = indices
          c = Counter(types[indices])
          if mode == "longest":
              self.num_samples = max(c.values())
              self.replacement = True
          elif mode == "shortest":
              self.num_samples = min(c.values())
              self.replacement = False
  
          for e, n in c.items():
              c[e] = 1 / n
          self.types = types
          self.weights = torch.DoubleTensor([c[types[i]] for i in indices])
  
      def __init__(self, indices, types, mode="shortest"):
          self.configure_sampler(indices, types, mode)
  
      def __iter__(self):
          selection = torch.multinomial(self.weights, self.num_samples, self.replacement)
          # print(Counter(self.types[self.indices[i]] for i in selection), len(selection)))
          return (self.indices[i] for i in selection)
  
      def __len__(self):
          return self.num_samples
  