code: def prepare_traced_trainer(
      model, task_name, load_best_model_at_end=False
  ):
      is_regression = task_name == "stsb"
      metric = load_metric("glue", task_name)
  
      def compute_metrics(p: EvalPrediction):
          preds = (
              p.predictions[0]
              if isinstance(p.predictions, tuple)
              else p.predictions
          )
          preds = (
              np.squeeze(preds)
              if is_regression
              else np.argmax(preds, axis=1)
          )
          result = metric.compute(
              predictions=preds, references=p.label_ids
          )
          result["default"] = result.get(
              "f1", result.get("accuracy", 0.0)
          )
          return result
  
      tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
      train_dataset, validation_datasets = prepare_datasets(
          task_name, tokenizer, None
      )
      merged_validation_dataset = ConcatDataset(
          [d for d in validation_datasets.values()]
      )
      data_collator = DataCollatorWithPadding(tokenizer)
      training_args = TrainingArguments(
          output_dir="./output/trainer",
          do_train=True,
          do_eval=True,
          evaluation_strategy="steps",
          per_device_train_batch_size=32,
          per_device_eval_batch_size=32,
          num_train_epochs=3,
          dataloader_num_workers=12,
          learning_rate=3e-5,
          save_strategy="steps",
          save_total_limit=1,
          metric_for_best_model="default",
          load_best_model_at_end=load_best_model_at_end,
          disable_tqdm=True,
          optim="adamw_torch",
          seed=1024,
      )
      trainer = nni.trace(Trainer)(
          model=model,
          args=training_args,
          data_collator=data_collator,
          train_dataset=train_dataset,
          eval_dataset=merged_validation_dataset,
          tokenizer=tokenizer,
          compute_metrics=compute_metrics,
      )
      return trainer
  
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 3 for calling a global function with example usage: prepare_traced_trainer(model, task_name, load_best_model_at_end) and returns: trainer, result