code: def prepare_datasets(
      task_name: str, tokenizer: BertTokenizerFast, cache_dir: str
  ):
      task_to_keys = {
          "cola": ("sentence", None),
          "mnli": ("premise", "hypothesis"),
          "mrpc": ("sentence1", "sentence2"),
          "qnli": ("question", "sentence"),
          "qqp": ("question1", "question2"),
          "rte": ("sentence1", "sentence2"),
          "sst2": ("sentence", None),
          "stsb": ("sentence1", "sentence2"),
          "wnli": ("sentence1", "sentence2"),
      }
      sentence1_key, sentence2_key = task_to_keys[task_name]
  
      # used to preprocess the raw data
      def preprocess_function(examples):
          # Tokenize the texts
          args = (
              (examples[sentence1_key],)
              if sentence2_key is None
              else (examples[sentence1_key], examples[sentence2_key])
          )
          result = tokenizer(
              *args, padding=False, max_length=128, truncation=True
          )
  
          if "label" in examples:
              # In all cases, rename the column to labels because the model will expect that.
              result["labels"] = examples["label"]
          return result
  
      raw_datasets = load_dataset(
          "glue", task_name, cache_dir=cache_dir
      )
      for key in list(raw_datasets.keys()):
          if "test" in key:
              raw_datasets.pop(key)
  
      processed_datasets = raw_datasets.map(
          preprocess_function,
          batched=True,
          remove_columns=raw_datasets["train"].column_names,
      )
  
      train_dataset = processed_datasets["train"]
      if task_name == "mnli":
          validation_datasets = {
              "validation_matched": processed_datasets[
                  "validation_matched"
              ],
              "validation_mismatched": processed_datasets[
                  "validation_mismatched"
              ],
          }
      else:
          validation_datasets = {
              "validation": processed_datasets["validation"]
          }
  
      return train_dataset, validation_datasets
  
imports: from __future__ import annotations
  from pathlib import Path
  import numpy
  import torch
  from torch.utils.data import ConcatDataset
  import nni
  from datasets import load_dataset, load_metric
  from transformers import BertTokenizerFast, DataCollatorWithPadding, BertForSequenceClassification, EvalPrediction
  from transformers.trainer import Trainer
  from transformers.training_args import TrainingArguments
  from nni.compression.distillation import DynamicLayerwiseDistiller, Adaptive1dLayerwiseDistiller
  from nni.compression.utils import TransformersEvaluator
  from nni.compression.pruning import MovementPruner
  from nni.compression.speedup import ModelSpeedup
  from nni.compression.utils.external.external_replacer import TransformersAttentionReplacer
  from nni.compression.pruning import TaylorPruner, AGPPruner
  from transformers.models.bert.modeling_bert import BertLayer
  from nni.compression.base.setting import PruningSetting
label_desc: Use image node_id 2 for calling a global function with example usage: prepare_datasets(task_name, tokenizer, cache_dir) and returns: train_dataset,  validation_datasets, result