code: def prune_heads(args, model, eval_dataloader, head_mask):
      """This method shows how to prune head (remove heads weights) based on
      the head importance scores as described in Michel et al. (http://arxiv.org/abs/1905.10650)
      """
      # Try pruning and test time speedup
      # Pruning is like masking but we actually remove the masked weights
      before_time = datetime.now()
      _, _, loss = compute_heads_importance(
          args,
          model,
          eval_dataloader,
          compute_entropy=False,
          compute_importance=False,
          head_mask=head_mask,
      )
      score_masking = 1 / loss
      original_time = datetime.now() - before_time
  
      original_num_params = sum(p.numel() for p in model.parameters())
      heads_to_prune = dict(
          (
              layer,
              (1 - head_mask[layer].long())
              .nonzero()
              .squeeze()
              .tolist(),
          )
          for layer in range(len(head_mask))
      )
  
      for k, v in heads_to_prune.items():
          if isinstance(v, int):
              heads_to_prune[k] = [
                  v,
              ]
  
      assert (
          sum(len(h) for h in heads_to_prune.values())
          == (1 - head_mask.long()).sum().item()
      )
      model.prune_heads(heads_to_prune)
      pruned_num_params = sum(p.numel() for p in model.parameters())
  
      before_time = datetime.now()
      _, _, loss = compute_heads_importance(
          args,
          model,
          eval_dataloader,
          compute_entropy=False,
          compute_importance=False,
          head_mask=None,
          actually_pruned=True,
      )
  
      score_pruning = 1 / loss
      new_time = datetime.now() - before_time
  
      logger.info(
          "Pruning: original num of params: %.2e, after pruning %.2e (%.1f percents)",
          original_num_params,
          pruned_num_params,
          pruned_num_params / original_num_params * 100,
      )
      logger.info(
          "Pruning: score with masking: %f score with pruning: %f",
          score_masking,
          score_pruning,
      )
      logger.info(
          "Pruning: speed ratio (original timing / new timing): %f percents",
          original_time / new_time * 100,
      )
      save_model(model, args.output_dir)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/examples/research_projects/bertology/run_prune_gpt.py
imports: import argparse
  import logging
  import os
  from datetime import datetime
  import numpy
  import torch
  from torch import nn
  from torch.utils.data import DataLoader, RandomSampler, TensorDataset
  from tqdm import tqdm
  from transformers import GPT2LMHeadModel
label_desc: Use image node_id 6 for calling a global function with example usage: prune_heads(args, model, eval_dataloader, head_mask) without return types