code: def generate_dataset(args):
      if args.type == "text":
          dataset = load_dataset(args.dataset)
          train_dataset = sample_dataset(
              dataset["train"],
              label_column=args.label_column,
              num_samples=args.shots,
          )
          val_dataset = (
              dataset["validation"]
              if "validaiton" in dataset
              else dataset["test"]
          )
          test_dataset = dataset["test"]
  
          train_dataset.set_format("pandas")
          val_dataset.set_format("pandas")
          test_dataset.set_format("pandas")
  
          train_df = train_dataset[:].drop(columns=["label_text"])
          val_df = val_dataset[:].drop(columns=["label_text"])
          test_df = test_dataset[:].drop(columns=["label_text"])
  
          num_classes = train_dataset[args.label_column].max() + 1
      else:
          dataset = build_dataset(
              args.dataset, args.data_path, args.shots
          )
          train_df = generate_image_df(args, dataset.train_x)
          val_df = generate_image_df(args, dataset.val)
          test_df = generate_image_df(args, dataset.test)
          num_classes = len(dataset.classnames)
  
      return dataset, train_df, val_df, test_df, num_classes
  
imports: import torch
  import pandas
  from tqdm import tqdm
  import torch.nn.functional
  import torch.nn
  from autogluon.multimodal import MultiModalPredictor
  from datasets import list_datasets, load_dataset
  from setfit import sample_dataset
  from imagedatasets import build_dataset
  from torch.utils.data import Dataset
label_desc: Use image node_id 3 for calling a global function with example usage: generate_dataset(args) and returns: dataset,  train_df,  val_df,  test_df,  num_classes