text: import numpy as np
  import torch
  from torch.utils.data import DataLoader
  from torch.utils.data.sampler import WeightedRandomSampler, SubsetRandomSampler
  from wilds.common.utils import get_counts, split_into_groups
  
  def get_train_loader(loader, dataset, batch_size,
          uniform_over_groups=None, grouper=None, distinct_groups=True, n_groups_per_batch=None, **loader_kwargs):
      """
      Constructs and returns the data loader for training.
      Args:
          - loader (str): Loader type. 'standard' for standard loaders and 'group' for group loaders,
                          which first samples groups and then samples a fixed number of examples belonging
                          to each group.
          - dataset (WILDSDataset or WILDSSubset): Data
          - batch_size (int): Batch size
          - uniform_over_groups (None or bool): Whether to sample the groups uniformly or according
                                                to the natural data distribution.
                                                Setting to None applies the defaults for each type of loaders.
                                                For standard loaders, the default is False. For group loaders,
                                                the default is True.
          - grouper (Grouper): Grouper used for group loaders or for uniform_over_groups=True
          - distinct_groups (bool): Whether to sample distinct_groups within each minibatch for group loaders.
          - n_groups_per_batch (int): Number of groups to sample in each minibatch for group loaders.
          - loader_kwargs: kwargs passed into torch DataLoader initialization.
      Output:
          - data loader (DataLoader): Data loader.
      """
      if loader == 'standard':
          if uniform_over_groups is None or not uniform_over_groups:
              return DataLoader(
                  dataset,
                  shuffle=True, # Shuffle training dataset
                  sampler=None,
                  collate_fn=dataset.collate,
                  batch_size=batch_size,
                  **loader_kwargs)
          else:
              assert grouper is not None
              groups, group_counts = grouper.metadata_to_group(
                  dataset.metadata_array,
                  return_counts=True)
              group_weights = 1 / group_counts
              weights = group_weights[groups]
  
              # Replacement needs to be set to True, otherwise we'll run out of minority samples
              sampler = WeightedRandomSampler(weights, len(dataset), replacement=True)
              return DataLoader(
                  dataset,
                  shuffle=False, # The WeightedRandomSampler already shuffles
                  sampler=sampler,
                  collate_fn=dataset.collate,
                  batch_size=batch_size,
                  **loader_kwargs)
  
      elif loader == 'group':
          if uniform_over_groups is None:
              uniform_over_groups = True
          assert grouper is not None
          assert n_groups_per_batch is not None
          if n_groups_per_batch > grouper.n_groups:
              raise ValueError(f'n_groups_per_batch was set to {n_groups_per_batch} but there are only {grouper.n_groups} groups specified.')
  
          group_ids = grouper.metadata_to_group(dataset.metadata_array)
          batch_sampler = GroupSampler(
              group_ids=group_ids,
              batch_size=batch_size,
              n_groups_per_batch=n_groups_per_batch,
              uniform_over_groups=uniform_over_groups,
              distinct_groups=distinct_groups)
  
          return DataLoader(dataset,
                shuffle=None,
                sampler=None,
                collate_fn=dataset.collate,
                batch_sampler=batch_sampler,
                drop_last=False,
                **loader_kwargs)
  
  def get_eval_loader(loader, dataset, batch_size, grouper=None, **loader_kwargs):
      """
      Constructs and returns the data loader for evaluation.
      Args:
          - loader (str): Loader type. 'standard' for standard loaders.
          - dataset (WILDSDataset or WILDSSubset): Data
          - batch_size (int): Batch size
          - loader_kwargs: kwargs passed into torch DataLoader initialization.
      Output:
          - data loader (DataLoader): Data loader.
      """
      if loader == 'standard':
          return DataLoader(
              dataset,
              shuffle=False, # Do not shuffle eval datasets
              sampler=None,
              collate_fn=dataset.collate,
              batch_size=batch_size,
              **loader_kwargs)
  
  class GroupSampler:
      """
          Constructs batches by first sampling groups,
          then sampling data from those groups.
          It drops the last batch if it's incomplete.
      """
      def __init__(self, group_ids, batch_size, n_groups_per_batch,
                   uniform_over_groups, distinct_groups):
  
          if batch_size % n_groups_per_batch != 0:
              raise ValueError(
                  f'batch_size ({batch_size}) must be evenly divisible by '
                  f'n_groups_per_batch ({n_groups_per_batch}).')
  
          if len(group_ids) < batch_size:
              raise ValueError(
                  f'The dataset has only {len(group_ids)} examples but the '
                  f'batch size is {batch_size}. There must be enough examples '
                  f'to form at least one complete batch.')
  
          self.group_ids = group_ids
          (self.unique_groups, self.group_indices,
           unique_counts) = split_into_groups(group_ids)
  
          self.distinct_groups = distinct_groups
          self.n_groups_per_batch = n_groups_per_batch
          self.n_points_per_group = batch_size // n_groups_per_batch
  
          self.dataset_size = len(group_ids)
          self.num_batches = self.dataset_size // batch_size
  
          if uniform_over_groups:  # Sample uniformly over groups
              self.group_prob = None
          else:  # Sample a group proportionately to its size
              unique_counts = unique_counts.numpy()
              self.group_prob = unique_counts / unique_counts.sum()
  
      def __iter__(self):
          for _ in range(self.num_batches):
              # Note that we are selecting group indices rather than groups
              groups_for_batch = np.random.choice(
                  len(self.unique_groups),
                  size=self.n_groups_per_batch,
                  replace=(not self.distinct_groups),
                  p=self.group_prob)
  
              # replace is False if the group is larger than the sample size
              sampled_ids = [
                  np.random.choice(
                      self.group_indices[group],
                      size=self.n_points_per_group,
                      replace=(len(self.group_indices[group])
                               <= self.n_points_per_group),
                      p=None)
                  for group in groups_for_batch]
  
              # Flatten
              sampled_ids = np.concatenate(sampled_ids)
              yield sampled_ids
  
      def __len__(self):
          return self.num_batches
  