code: def load_roiterberg_ANI(mode="atomization"):
      """
      Load the ANI dataset.
  
      Parameters
      ----------
      mode: str
        Accepted modes are "relative", "atomization", or "absolute". These settings are used
        to adjust the dynamic range of the model, with absolute having the greatest and relative
        having the lowest. Note that for atomization we approximate the single atom energy
        using a different level of theory
  
  
      Returns
      -------
      tuples
        Elements returned are 3-tuple (a,b,c) where and b are the train and test datasets, respectively,
        and c is an array of indices denoting the group of each
  
      """
      if "ROITBERG_ANI" not in os.environ:
          raise ValueError(
              "Please set environment variable ROITBERG_ANI to where the ani_dgb_s0x.h5 files are."
          )
  
      base_dir = os.environ["ROITBERG_ANI"]
  
      # Number of conformations in each file increases exponentially.
      # Start with a smaller dataset before continuing. Use all of them
      # for production
      hdf5files = [
          "ani_gdb_s01.h5",
          "ani_gdb_s02.h5",
          # 'ani_gdb_s03.h5',
          # 'ani_gdb_s04.h5',
          # 'ani_gdb_s05.h5',
          # 'ani_gdb_s06.h5',
          # 'ani_gdb_s07.h5',
          # 'ani_gdb_s08.h5'
      ]
  
      hdf5files = [os.path.join(base_dir, f) for f in hdf5files]
  
      groups = []
  
      def shard_generator():
          shard_size = 4096 * 64
  
          row_idx = 0
          group_idx = 0
  
          X_cache = []
          y_cache = []
          w_cache = []
          ids_cache = []
  
          for hdf5file in hdf5files:
              adl = pya.anidataloader(hdf5file)
              for data in adl:
                  # Extract the data
                  P = data["path"]
                  R = data["coordinates"]
                  E = data["energies"]
                  S = data["species"]
                  smi = data["smiles"]
  
                  if len(S) > 23:
                      print("skipping:", smi, "due to atom count.")
                      continue
  
                  # Print the data
                  print("Processing: ", P)
                  print("  Smiles:      ", "".join(smi))
                  print("  Symbols:     ", S)
                  print("  Coordinates: ", R.shape)
                  print("  Energies:    ", E.shape)
  
                  Z_padded = np.zeros((23,), dtype=np.float32)
                  nonpadded = convert_species_to_atomic_nums(S)
                  Z_padded[: nonpadded.shape[0]] = nonpadded
  
                  if mode == "relative":
                      offset = np.amin(E)
                  elif mode == "atomization":
                      # self-interaction energies taken from
                      # https://github.com/isayev/ANI1_dataset README
                      atomizationEnergies = {
                          0: 0,
                          1: -0.500607632585,
                          6: -37.8302333826,
                          7: -54.5680045287,
                          8: -75.0362229210,
                      }
  
                      offset = 0
  
                      for z in nonpadded:
                          offset -= atomizationEnergies[z]
                  elif mode == "absolute":
                      offset = 0
                  else:
                      raise Exception("Unsupported mode: ", mode)
  
                  for k in range(len(E)):
                      R_padded = np.zeros((23, 3), dtype=np.float32)
                      R_padded[: R[k].shape[0], : R[k].shape[1]] = R[k]
  
                      X = np.concatenate(
                          [np.expand_dims(Z_padded, 1), R_padded],
                          axis=1,
                      )
  
                      y = E[k] - offset
  
                      if len(X_cache) == shard_size:
                          yield np.array(X_cache), np.array(
                              y_cache
                          ), np.array(w_cache), np.array(ids_cache)
  
                          X_cache = []
                          y_cache = []
                          w_cache = []
                          ids_cache = []
  
                      else:
                          X_cache.append(X)
                          y_cache.append(np.array(y).reshape((1,)))
                          w_cache.append(np.array(1).reshape((1,)))
                          ids_cache.append(row_idx)
                          row_idx += 1
                          groups.append(group_idx)
  
                  group_idx += 1
  
          # flush once more at the end
          if len(X_cache) > 0:
              yield np.array(X_cache), np.array(y_cache), np.array(
                  w_cache
              ), np.array(ids_cache)
  
      tasks = ["ani"]
      dataset = dc.data.DiskDataset.create_dataset(
          shard_generator(), tasks=tasks, data_dir=all_dir
      )
  
      print("Number of groups", np.amax(groups))
      splitter = dc.splits.RandomGroupSplitter(groups)
  
      train_dataset, test_dataset = splitter.train_test_split(
          dataset, train_dir=fold_dir, test_dir=test_dir, frac_train=0.8
      )
  
      return train_dataset, test_dataset, groups
  
file_path: deepchem/examples/roitberg/roitberg.py
imports: import numpy
  import os
  import tensorflow
  import deepchem
  import pyanitools
  import app
label_desc: Use image node_id 2 for calling a global function with example usage: load_roiterberg_ANI(mode) and returns: train_dataset,  test_dataset,  groups