code: def load_data_tf_datasets(
      dataset_name, target_num_train_per_class, target_num_val, seed
  ):
      """Load and preprocess data from tensorflow_datasets."""
      logging.info(
          "Loading and preprocessing data from tensorflow datasets..."
      )
      # Load train data.
      ds = tfds.load(
          dataset_name, split=tfds.Split.TRAIN, batch_size=-1
      )
      ds = tfds.as_numpy(ds)
      train_inputs, train_labels = ds["image"], ds["label"]
      # Load test data.
      ds = tfds.load(dataset_name, split=tfds.Split.TEST, batch_size=-1)
      ds = tfds.as_numpy(ds)
      test_inputs, test_labels = ds["image"], ds["label"]
  
      # Remove extra dimensions of size 1.
      train_labels = np.squeeze(train_labels)
      test_labels = np.squeeze(test_labels)
  
      logging.info("Splitting data...")
      data = split_train_val_unlabeled(
          train_inputs,
          train_labels,
          target_num_train_per_class,
          target_num_val,
          seed,
      )
      train_inputs = data[0]
      train_labels = data[1]
      val_inputs = data[2]
      val_labels = data[3]
      unlabeled_inputs = data[4]
      unlabeled_labels = data[5]
  
      logging.info("Converting data to Dataset format...")
      data = Dataset.build_from_splits(
          name=dataset_name,
          inputs_train=train_inputs,
          labels_train=train_labels,
          inputs_val=val_inputs,
          labels_val=val_labels,
          inputs_test=test_inputs,
          labels_test=test_labels,
          inputs_unlabeled=unlabeled_inputs,
          labels_unlabeled=unlabeled_labels,
          feature_preproc_fn=convert_image,
      )
      return data
  
file_path: neural-structured-learning/research/gam/gam/data/loaders.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import json
  import logging
  import os
  import pickle
  import sys
  from .dataset import Dataset
  from .dataset import PlanetoidDataset
  import networkx
  import numpy
  from .preprocessing import convert_image
  from .preprocessing import split_train_val_unlabeled
  from scipy import sparse
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: load_data_tf_datasets(dataset_name, target_num_train_per_class, target_num_val, seed) and returns: data