code: def pruning_ffn():
      model: BertForSequenceClassification = torch.load(
          "./output/pruning/attn_pruned_model.pth"
      )
      teacher_model: BertForSequenceClassification = (
          build_finetuning_model(
              "mnli", f"./output/bert_finetuned/{task_name}.bin"
          )
      )
      # create ffn config list, here simply use a linear function related to the number of retained heads to determine the sparse ratio
      config_list = []
      for name, module in model.named_modules():
          if isinstance(module, BertLayer):
              retained_head_num = (
                  module.attention.self.num_attention_heads
              )
              ori_head_num = (
                  len(module.attention.pruned_heads) + retained_head_num
              )
              ffn_sparse_ratio = (
                  1 - retained_head_num / ori_head_num / 2
              )
              config_list.append(
                  {
                      "op_names": [f"{name}.intermediate.dense"],
                      "sparse_ratio": ffn_sparse_ratio,
                  }
              )
  
      trainer = prepare_traced_trainer(model, task_name)
      teacher_model.eval().to(trainer.args.device)
      # create a distiller for restoring the accuracy
      distiller = dynamic_distiller(model, teacher_model, trainer)
      # fusion compress: TaylorPruner + DynamicLayerwiseDistiller
      taylor_pruner = TaylorPruner.from_compressor(
          distiller, config_list, 1000
      )
      # fusion compress: AGPPruner(TaylorPruner) + DynamicLayerwiseDistiller
      agp_pruner = AGPPruner(taylor_pruner, 1000, 36)
      agp_pruner.compress(None, 3)
      agp_pruner.unwrap_model()
      distiller.unwrap_teacher_model()
  
      masks = agp_pruner.get_masks()
      Path("./output/pruning/").mkdir(parents=True, exist_ok=True)
      torch.save(masks, "./output/pruning/ffn_masks.pth")
      torch.save(model, "./output/pruning/ffn_masked_model.pth")
  
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 11 for calling a global function with example usage: pruning_ffn() without return types