



import os
import re
import warnings
from shutil import copyfile
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer


if TYPE_CHECKING:
 from ...tokenization_utils_base import TextInput
from ...utils import logging


a_   :int        =      logging.get_logger(__name__)

a_   :str        =      {"vocab_file": "spiece.model"}

a_   :str        =      {
    "vocab_file": {
        "t5-small": "https://huggingface.co/t5-small/resolve/main/spiece.model",
        "t5-base": "https://huggingface.co/t5-base/resolve/main/spiece.model",
        "t5-large": "https://huggingface.co/t5-large/resolve/main/spiece.model",
        "t5-3b": "https://huggingface.co/t5-3b/resolve/main/spiece.model",
        "t5-11b": "https://huggingface.co/t5-11b/resolve/main/spiece.model",
    }
}


# TODO(PVP) - this should be removed in Transformers v5
a_   :Dict        =      {
    "t5-small": 512,
    "t5-base": 512,
    "t5-large": 512,
    "t5-3b": 512,
    "t5-11b": 512,
}

a_   :Optional[int]        =      "▁"




class 					snake_case__						(       lowerCAmelCase_     ):






     """simple docstring"""


     _SCREAMING_SNAKE_CASE						=	VOCAB_FILES_NAMES
     _SCREAMING_SNAKE_CASE						=	PRETRAINED_VOCAB_FILES_MAP
     _SCREAMING_SNAKE_CASE						=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     _SCREAMING_SNAKE_CASE						=	["""input_ids""", """attention_mask"""]






     def __init__(       self  :  Dict, _snake_case  :  List[Any], _snake_case  :  str="</s>", _snake_case  :  List[Any]="<unk>", _snake_case  :  Union[str, Any]="<pad>", _snake_case  :  List[Any]=1_0_0, _snake_case  :  str=None, _snake_case  :  Optional[Dict[str, Any]] = None, _snake_case  :  Optional[Any]=True, **_snake_case  :  List[str], )   ->None:
            # Add extra_ids to the special token list
            if extra_ids > 0 and additional_special_tokens is None:
                   snake_case__							:	Any   	=	[F'''<extra_id_{i}>''' for i in range(_snake_case					)]
            elif extra_ids > 0 and additional_special_tokens is not None:
                   # Check that we have the right number of extra_id special tokens
                   snake_case__							:	Any   	=	len(set(filter(lambda _snake_case					: bool('extra_id' in str(_snake_case					)					), _snake_case					)					)					)
                   if extra_tokens != extra_ids:
                          raise ValueError(
                              F'''Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are'''
                              ' provided to T5Tokenizer. In this case the additional_special_tokens must include the extra_ids'
                              ' tokens'					)
            if legacy:
                   logger.warning_once(
                       F'''You are using the legacy behaviour of the {self.__class__}. This means that tokens that come after special tokens will not be properly handled. We recommend you to'''
                       ' read the related pull request available at https://github.com/huggingface/transformers/pull/24565'					)

            snake_case__							:	List[str]   	=	legacy
            snake_case__							:	str   	=	{} if sp_model_kwargs is None else sp_model_kwargs

            super().__init__(
                eos_token=_snake_case, unk_token=_snake_case, pad_token=_snake_case, extra_ids=_snake_case, additional_special_tokens=_snake_case, sp_model_kwargs=self.sp_model_kwargs, legacy=_snake_case, **_snake_case, )

            snake_case__							:	Dict   	=	vocab_file
            snake_case__							:	List[str]   	=	extra_ids

            snake_case__							:	Union[str, Any]   	=	spm.SentencePieceProcessor(**self.sp_model_kwargs					)
            self.sp_model.Load(_snake_case					)






     @staticmethod
     def    lowercase_   (       _snake_case  :  Optional[int], _snake_case  :  Optional[int], _snake_case  :  List[Any]					)   ->Tuple:
            if pretrained_model_name_or_path in TaTokenizer.max_model_input_sizes:
                   snake_case__							:	List[str]   	=	TaTokenizer.max_model_input_sizes[pretrained_model_name_or_path]
                   if init_max_model_length is not None and init_max_model_length != max_model_length:
                          return init_max_model_length
                   elif init_max_model_length is None:
                          warnings.warn(
                              'This tokenizer was incorrectly instantiated with a model max length of'
                              F''' {deprecated_max_model_length} which will be corrected in Transformers v5.\nFor now, this'''
                              ' behavior is kept to avoid breaking backwards compatibility when padding/encoding with'
                              ' `truncation is True`.\n- Be aware that you SHOULD NOT rely on'
                              F''' {pretrained_model_name_or_path} automatically truncating your input to'''
                              F''' {deprecated_max_model_length} when padding/encoding.\n- If you want to encode/pad to sequences'''
                              F''' longer than {deprecated_max_model_length} you can either instantiate this tokenizer with'''
                              ' `model_max_length` or pass `max_length` when encoding/padding.\n- To avoid this warning, please'
                              ' instantiate this tokenizer with `model_max_length` set to your preferred value.', _snake_case, )

            return max_model_length






     @property
     def    lowercase_   (       self  :  Optional[Any]					)   ->Tuple:
            return self.sp_model.get_piece_size() + self._extra_ids






     def    lowercase_   (       self  :  Optional[Any]					)   ->List[Any]:
            snake_case__							:	List[str]   	=	{self.convert_ids_to_tokens(_snake_case					): i for i in range(self.vocab_size					)}
            vocab.update(self.added_tokens_encoder					)
            return vocab






     def    lowercase_   (       self  :  List[Any], _snake_case  :  List[int], _snake_case  :  Optional[List[int]] = None, _snake_case  :  bool = False					)   ->List[int]:
            if already_has_special_tokens:
                   return super().get_special_tokens_mask(
                       token_ids_a=_snake_case, token_ids_a=_snake_case, already_has_special_tokens=_snake_case					)

            # normal case: some special tokens
            if token_ids_a is None:
                   return ([0] * len(_snake_case					)) + [1]
            return ([0] * len(_snake_case					)) + [1] + ([0] * len(_snake_case					)) + [1]






     def    lowercase_   (       self  :  Tuple					)   ->Union[str, Any]:
            return list(
                set(filter(lambda _snake_case					: bool(re.search(R'<extra_id_\d+>', _snake_case					)					) is not None, self.additional_special_tokens					)					)					)






     def    lowercase_   (       self  :  Any					)   ->str:
            return [self._convert_token_to_id(_snake_case					) for token in self.get_sentinel_tokens()]






     def    lowercase_   (       self  :  Dict, _snake_case  :  List[int]					)   ->List[int]:
            if len(_snake_case					) > 0 and token_ids[-1] == self.eos_token_id:
                   warnings.warn(
                       F'''This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated'''
                       ' eos tokens being added.'					)
                   return token_ids
            else:
                   return token_ids + [self.eos_token_id]






     def    lowercase_   (       self  :  Optional[Any], _snake_case  :  List[int], _snake_case  :  Optional[List[int]] = None					)   ->List[int]:
            snake_case__							:	List[Any]   	=	[self.eos_token_id]

            if token_ids_a is None:
                   return len(token_ids_a + eos					) * [0]
            return len(token_ids_a + eos + token_ids_a + eos					) * [0]






     def    lowercase_   (       self  :  int, _snake_case  :  List[int], _snake_case  :  Optional[List[int]] = None					)   ->List[int]:
            snake_case__							:	Optional[Any]   	=	self._add_eos_if_not_present(_snake_case					)
            if token_ids_a is None:
                   return token_ids_a
            else:
                   snake_case__							:	Optional[int]   	=	self._add_eos_if_not_present(_snake_case					)
                   return token_ids_a + token_ids_a






     def __getstate__(       self  :  Optional[Any]					)   ->Union[str, Any]:
            snake_case__							:	List[Any]   	=	self.__dict__.copy()
            snake_case__							:	Union[str, Any]   	=	None
            return state






     def __setstate__(       self  :  Optional[int], _snake_case  :  Optional[int]					)   ->Dict:
            snake_case__							:	List[Any]   	=	d

            # for backward compatibility
            if not hasattr(self, 'sp_model_kwargs'					):
                   snake_case__							:	List[Any]   	=	{}

            snake_case__							:	List[str]   	=	spm.SentencePieceProcessor(**self.sp_model_kwargs					)
            self.sp_model.Load(self.vocab_file					)






     def    lowercase_   (       self  :  Tuple, _snake_case  :  "TextInput", **_snake_case  :  Union[str, Any]					)   ->List[str]:
            # Replace the SPIECE_UNDERLINE with a space to make sure SPIECE_UNDERLINE is only used at
            # the beginning of the text
            if not self.legacy:
                   snake_case__							:	Optional[int]   	=	SPIECE_UNDERLINE + text.replace(_snake_case, ' '					)
            return super().tokenize(_snake_case, **_snake_case					)






     def    lowercase_   (       self  :  int, _snake_case  :  Optional[Any], **_snake_case  :  Optional[Any]					)   ->int:
            if not self.legacy:
                   snake_case__							:	str   	=	text.startswith(_snake_case					)
                   if is_first:
                          snake_case__							:	Any   	=	text[1:]

            snake_case__							:	str   	=	self.sp_model.encode(_snake_case, out_type=_snake_case					)

            if not self.legacy and not is_first and not text.startswith(' '					) and tokens[0].startswith(_snake_case					):
                   snake_case__							:	Dict   	=	([tokens[0][1:]] if len(tokens[0]					) > 1 else []) + tokens[1:]
            return tokens






     def    lowercase_   (       self  :  Union[str, Any], _snake_case  :  List[Any]					)   ->int:
            if token.startswith('<extra_id_'					):
                   snake_case__							:	str   	=	re.match(R'<extra_id_(\d+)>', _snake_case					)
                   snake_case__							:	Tuple   	=	int(match.group(1					)					)
                   return self.vocab_size - num - 1
            return self.sp_model.piece_to_id(_snake_case					)






     def    lowercase_   (       self  :  int, _snake_case  :  int					)   ->int:
            if index < self.sp_model.get_piece_size():
                   snake_case__							:	List[Any]   	=	self.sp_model.IdToPiece(_snake_case					)
            else:
                   snake_case__							:	Tuple   	=	F'''<extra_id_{self.vocab_size - 1 - index}>'''
            return token






     def    lowercase_   (       self  :  Optional[Any], _snake_case  :  Union[str, Any]					)   ->Optional[int]:
            snake_case__							:	Optional[Any]   	=	[]
            snake_case__							:	Tuple   	=	''
            snake_case__							:	Dict   	=	False
            for token in tokens:
                   # make sure that special tokens are not decoded using sentencepiece model
                   if token in self.all_special_tokens:
                          if not prev_is_special:
                                 out_string += " "
                          out_string += self.sp_model.decode(_snake_case					) + token
                          snake_case__							:	Any   	=	True
                          snake_case__							:	Dict   	=	[]
                   else:
                          current_sub_tokens.append(_snake_case					)
                          snake_case__							:	Union[str, Any]   	=	False
            out_string += self.sp_model.decode(_snake_case					)
            return out_string.strip()






     def    lowercase_   (       self  :  int, _snake_case  :  str, _snake_case  :  Optional[str] = None					)   ->Tuple[str]:
            if not os.path.isdir(_snake_case					):
                   logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''					)
                   return
            snake_case__							:	List[Any]   	=	os.path.join(
                _snake_case, (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']					)

            if os.path.abspath(self.vocab_file					) != os.path.abspath(_snake_case					) and os.path.isfile(self.vocab_file					):
                   copyfile(self.vocab_file, _snake_case					)
            elif not os.path.isfile(self.vocab_file					):
                   with open(_snake_case, 'wb'					) as fi:
                          snake_case__							:	Optional[Any]   	=	self.sp_model.serialized_model_proto()
                          fi.write(_snake_case					)

            return (out_vocab_file,)











import argparse
import logging
import os
from datetime import datetime

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, RandomSampler, TensorDataset
from tqdm import tqdm

from transformers import GPTaLMHeadModel


a_   :Optional[Any]        =      logging.getLogger(__name__)
def        lowercase_							(A : List[Any]		,					A : List[Any]      ):
       # save results
       if os.path.exists(A      ):
              if os.path.exists(os.path.join(A		,					'config.json'      )      ) and os.path.isfile(
                  os.path.join(A		,					'config.json'      )      ):
                     os.remove(os.path.join(A		,					'config.json'      )      )
              if os.path.exists(os.path.join(A		,					'pytorch_model.bin'      )      ) and os.path.isfile(
                  os.path.join(A		,					'pytorch_model.bin'      )      ):
                     os.remove(os.path.join(A		,					'pytorch_model.bin'      )      )
       else:
              os.makedirs(A      )
       model.save_pretrained(A      )
def        lowercase_							(A : Any		,					A : Optional[Any]=False      ):
       snake_case__							:	str   	=	2
       if unlogit:
              snake_case__							:	Dict   	=	torch.pow(A		,					A      )
       snake_case__							:	Any   	=	p * torch.log(A      )
       snake_case__							:	Tuple   	=	0
       return -plogp.sum(dim=-1      )
def        lowercase_							(A : List[str]      ):
       logger.info('lv, h >\t' + '\t'.join(F'''{x + 1}''' for x in range(len(A      )      )      )      )
       for row in range(len(A      )      ):
              if tensor.dtype != torch.long:
                     logger.info(F'''layer {row + 1}:\t''' + '\t'.join(F'''{x:.5f}''' for x in tensor[row].cpu().data      )      )
              else:
                     logger.info(F'''layer {row + 1}:\t''' + '\t'.join(F'''{x:d}''' for x in tensor[row].cpu().data      )      )
def        lowercase_							(A : Tuple		,					A : Optional[Any]		,					A : str		,					A : int=True		,					A : Optional[int]=True		,					A : Any=None		,					A : int=False      ):
       snake_case__	,				snake_case__							:	Optional[Any]   	=	model.config.num_hidden_layers, model.config.num_attention_heads
       snake_case__							:	int   	=	torch.zeros(A		,					A      ).to(args.device      )
       snake_case__							:	Any   	=	torch.zeros(A		,					A      ).to(args.device      )

       if head_mask is None:
              snake_case__							:	Dict   	=	torch.ones(A		,					A      ).to(args.device      )

       head_mask.requires_grad_(requires_grad=A      )
       # If actually pruned attention multi-head, set head mask to None to avoid shape mismatch
       if actually_pruned:
              snake_case__							:	Optional[int]   	=	None

       snake_case__							:	List[Any]   	=	0.0
       snake_case__							:	str   	=	0.0
       for step, inputs in enumerate(tqdm(A		,					desc='Iteration'		,					disable=args.local_rank not in [-1, 0]      )      ):
              snake_case__							:	Union[str, Any]   	=	tuple(t.to(args.device      ) for t in inputs      )
              ((snake_case__)	,				)							:	Optional[Any]   	=	inputs

              # Do a forward pass (not with torch.no_grad() since we need gradients for importance score - see below)
              snake_case__							:	Union[str, Any]   	=	model(A		,					labels=A		,					head_mask=A      )
              #  (loss), lm_logits, presents, (all hidden_states), (attentions)
              snake_case__	,				snake_case__	,				snake_case__							:	Dict   	=	(
                  outputs[0],
                  outputs[1],
                  outputs[-1],
              )  # Loss and logits are the first, attention the last
              loss.backward()  # Backpropagate to populate the gradients in the head mask
              total_loss += loss.detach().cpu().numpy()
              if compute_entropy:
                     for layer, attn in enumerate(A      ):
                            snake_case__							:	Optional[Any]   	=	entropy(attn.detach()		,					A      )
                            attn_entropy[layer] += masked_entropy.sum(-1      ).sum(0      ).sum(0      ).detach()

              if compute_importance:
                     head_importance += head_mask.grad.abs().detach()
              tot_tokens += torch.ones_like(A      ).float().detach().sum().data

       # Normalize
       attn_entropy /= tot_tokens
       head_importance /= tot_tokens
       # Layerwise importance normalization
       if not args.dont_normalize_importance_by_layer:
              snake_case__							:	Union[str, Any]   	=	2
              snake_case__							:	List[Any]   	=	torch.pow(torch.pow(A		,					A      ).sum(-1      )		,					1 / exponent      )
              head_importance /= norm_by_layer.unsqueeze(-1      ) + 1e-20

       if not args.dont_normalize_global_importance:
              snake_case__							:	Tuple   	=	(head_importance - head_importance.min()) / (head_importance.max() - head_importance.min())

       # Print matrices
       if compute_entropy:
              logger.info('Attention entropies'      )
              print_ad_tensor(A      )
       if compute_importance:
              logger.info('Head importance scores'      )
              print_ad_tensor(A      )
       logger.info('Head ranked by importance scores'      )
       snake_case__							:	Tuple   	=	torch.zeros(head_importance.numel()		,					dtype=torch.long		,					device=args.device      )
       snake_case__							:	Union[str, Any]   	=	torch.arange(
           head_importance.numel()		,					device=args.device      )
       snake_case__							:	str   	=	head_ranks.view_as(A      )
       print_ad_tensor(A      )
       return attn_entropy, head_importance, total_loss
def        lowercase_							(A : Optional[int]		,					A : Dict		,					A : Optional[int]      ):
       snake_case__	,				snake_case__	,				snake_case__							:	Any   	=	compute_heads_importance(A		,					A		,					A		,					compute_entropy=A      )
       snake_case__							:	Tuple   	=	1 / loss  # instead of downsteam score use the LM loss
       logger.info('Pruning: original score: %f, threshold: %f'		,					A		,					original_score * args.masking_threshold      )

       snake_case__							:	Optional[Any]   	=	torch.ones_like(A      )
       snake_case__							:	Union[str, Any]   	=	max(1		,					int(new_head_mask.numel() * args.masking_amount      )      )

       snake_case__							:	Dict   	=	original_score
       while current_score >= original_score * args.masking_threshold:
              snake_case__							:	int   	=	new_head_mask.clone().detach()  # save current head mask
              # heads from least important to most - keep only not-masked heads
              snake_case__							:	List[Any]   	=	float('Inf'      )
              snake_case__							:	Union[str, Any]   	=	head_importance.view(-1      ).sort()[1]

              if len(A      ) <= num_to_mask:
                     print('BREAK BY num_to_mask'      )
                     break

              # mask heads
              snake_case__							:	int   	=	current_heads_to_mask[:num_to_mask]
              logger.info('Heads to mask: %s'		,					str(current_heads_to_mask.tolist()      )      )
              snake_case__							:	int   	=	new_head_mask.view(-1      )
              snake_case__							:	int   	=	0.0
              snake_case__							:	Union[str, Any]   	=	new_head_mask.view_as(A      )
              snake_case__							:	List[str]   	=	new_head_mask.clone().detach()
              print_ad_tensor(A      )

              # Compute metric and head importance again
              snake_case__	,				snake_case__	,				snake_case__							:	Any   	=	compute_heads_importance(
                  A		,					A		,					A		,					compute_entropy=A		,					head_mask=A      )
              snake_case__							:	Dict   	=	1 / loss
              logger.info(
                  'Masking: current score: %f, remaining heads %d (%.1f percents)'		,					A		,					new_head_mask.sum()		,					new_head_mask.sum() / new_head_mask.numel() * 1_0_0		,					)

       logger.info('Final head mask'      )
       print_ad_tensor(A      )
       np.save(os.path.join(args.output_dir		,					'head_mask.npy'      )		,					head_mask.detach().cpu().numpy()      )

       return head_mask
def        lowercase_							(A : List[str]		,					A : Tuple		,					A : Optional[Any]		,					A : int      ):
       snake_case__							:	Any   	=	datetime.now()
       snake_case__	,				snake_case__	,				snake_case__							:	str   	=	compute_heads_importance(
           A		,					A		,					A		,					compute_entropy=A		,					compute_importance=A		,					head_mask=A      )
       snake_case__							:	Tuple   	=	1 / loss
       snake_case__							:	Dict   	=	datetime.now() - before_time

       snake_case__							:	Union[str, Any]   	=	sum(p.numel() for p in model.parameters()      )
       snake_case__							:	Optional[Any]   	=	{
           layer: (1 - head_mask[layer].long()).nonzero().squeeze().tolist() for layer in range(len(A      )      )
       }

       for k, v in heads_to_prune.items():
              if isinstance(A		,					A      ):
                     snake_case__							:	Any   	=	[
                         v,
                     ]

       assert sum(len(A      ) for h in heads_to_prune.values()      ) == (1 - head_mask.long()).sum().item()
       model.prune_heads(A      )
       snake_case__							:	Dict   	=	sum(p.numel() for p in model.parameters()      )

       snake_case__							:	Tuple   	=	datetime.now()
       snake_case__	,				snake_case__	,				snake_case__							:	Dict   	=	compute_heads_importance(
           A		,					A		,					A		,					compute_entropy=A		,					compute_importance=A		,					head_mask=A		,					actually_pruned=A		,					)

       snake_case__							:	Any   	=	1 / loss
       snake_case__							:	int   	=	datetime.now() - before_time

       logger.info(
           'Pruning: original num of params: %.2e, after pruning %.2e (%.1f percents)'		,					A		,					A		,					pruned_num_params / original_num_params * 1_0_0		,					)
       logger.info('Pruning: score with masking: %f score with pruning: %f'		,					A		,					A      )
       logger.info('Pruning: speed ratio (original timing / new timing): %f percents'		,					original_time / new_time * 1_0_0      )
       save_model(A		,					args.output_dir      )
def        lowercase_							():
       snake_case__							:	str   	=	argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           '--data_dir'		,					default=A		,					type=A		,					required=A		,					help='The input data dir. Should contain the .tsv files (or other data files) for the task.'		,					)
       parser.add_argument(
           '--model_name_or_path'		,					default=A		,					type=A		,					required=A		,					help='Path to pretrained model or model identifier from huggingface.co/models'		,					)
       parser.add_argument(
           '--output_dir'		,					default=A		,					type=A		,					required=A		,					help='The output directory where the model predictions and checkpoints will be written.'		,					)

       # Other parameters
       parser.add_argument(
           '--config_name'		,					default=''		,					type=A		,					help='Pretrained config name or path if not the same as model_name_or_path'		,					)
       parser.add_argument(
           '--tokenizer_name'		,					default=''		,					type=A		,					help='Pretrained tokenizer name or path if not the same as model_name_or_path'		,					)
       parser.add_argument(
           '--cache_dir'		,					default=A		,					type=A		,					help='Where do you want to store the pre-trained models downloaded from s3'		,					)
       parser.add_argument(
           '--data_subset'		,					type=A		,					default=-1		,					help='If > 0: limit the data to a subset of data_subset instances.'      )
       parser.add_argument(
           '--overwrite_output_dir'		,					action='store_true'		,					help='Whether to overwrite data in output directory'      )
       parser.add_argument(
           '--overwrite_cache'		,					action='store_true'		,					help='Overwrite the cached training and evaluation sets'      )

       parser.add_argument(
           '--dont_normalize_importance_by_layer'		,					action='store_true'		,					help='Don\'t normalize importance score by layers'      )
       parser.add_argument(
           '--dont_normalize_global_importance'		,					action='store_true'		,					help='Don\'t normalize all importance scores between 0 and 1'		,					)

       parser.add_argument(
           '--try_masking'		,					action='store_true'		,					help='Whether to try to mask head until a threshold of accuracy.'      )
       parser.add_argument(
           '--masking_threshold'		,					default=0.9		,					type=A		,					help='masking threshold in term of metrics (stop masking when metric < threshold * original metric value).'		,					)
       parser.add_argument(
           '--masking_amount'		,					default=0.1		,					type=A		,					help='Amount to heads to masking at each masking step.'      )
       parser.add_argument('--metric_name'		,					default='acc'		,					type=A		,					help='Metric to use for head masking.'      )

       parser.add_argument(
           '--max_seq_length'		,					default=1_2_8		,					type=A		,					help=(
               'The maximum total input sequence length after WordPiece tokenization. \n'
               'Sequences longer than this will be truncated, sequences shorter padded.'
           )		,					)
       parser.add_argument('--batch_size'		,					default=1		,					type=A		,					help='Batch size.'      )

       parser.add_argument('--seed'		,					type=A		,					default=4_2      )
       parser.add_argument('--local_rank'		,					type=A		,					default=-1		,					help='local_rank for distributed training on gpus'      )
       parser.add_argument('--no_cuda'		,					action='store_true'		,					help='Whether not to use CUDA when available'      )
       parser.add_argument('--server_ip'		,					type=A		,					default=''		,					help='Can be used for distant debugging.'      )
       parser.add_argument('--server_port'		,					type=A		,					default=''		,					help='Can be used for distant debugging.'      )
       snake_case__							:	Optional[int]   	=	parser.parse_args()

       if args.server_ip and args.server_port:
              # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
              import ptvsd

              print('Waiting for debugger attach'      )
              ptvsd.enable_attach(address=(args.server_ip, args.server_port)		,					redirect_output=A      )
              ptvsd.wait_for_attach()

       # Setup devices and distributed training
       if args.local_rank == -1 or args.no_cuda:
              snake_case__							:	List[Any]   	=	torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu'      )
              snake_case__							:	Optional[Any]   	=	0 if args.no_cuda else torch.cuda.device_count()
       else:
              torch.cuda.set_device(args.local_rank      )
              snake_case__							:	int   	=	torch.device('cuda'		,					args.local_rank      )
              snake_case__							:	List[str]   	=	1
              torch.distributed.init_process_group(backend='nccl'      )  # Initializes the distributed backend

       # Setup logging
       logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN      )
       logger.info('device: {} n_gpu: {}, distributed: {}'.format(args.device		,					args.n_gpu		,					bool(args.local_rank != -1      )      )      )

       snake_case__							:	Any   	=	GPTaLMHeadModel.from_pretrained(args.model_name_or_path      )

       # Distributed and parallel training
       model.to(args.device      )
       if args.local_rank != -1:
              snake_case__							:	List[str]   	=	nn.parallel.DistributedDataParallel(
                  A		,					device_ids=[args.local_rank]		,					output_device=args.local_rank		,					find_unused_parameters=A      )
       elif args.n_gpu > 1:
              snake_case__							:	Optional[int]   	=	nn.DataParallel(A      )

       # Print/save training arguments
       os.makedirs(args.output_dir		,					exist_ok=A      )
       torch.save(A		,					os.path.join(args.output_dir		,					'run_args.bin'      )      )
       logger.info('Training/evaluation parameters %s'		,					A      )

       # Prepare dataset
       snake_case__							:	Optional[Any]   	=	np.concatenate(
           [
               np.loadtxt(args.data_dir		,					dtype=np.intaa      ),
           ]      )
       snake_case__							:	List[str]   	=	(torch.from_numpy(A      ),)
       snake_case__							:	int   	=	TensorDataset(*A      )
       snake_case__							:	Union[str, Any]   	=	RandomSampler(A      )
       snake_case__							:	Any   	=	DataLoader(A		,					sampler=A		,					batch_size=args.batch_size      )

       # Compute head entropy and importance score
       compute_heads_importance(A		,					A		,					A      )

       # Try head masking (set heads to zero until the score goes under a threshole)
       # and head pruning (remove masked heads and see the effect on the network)
       if args.try_masking and args.masking_threshold > 0.0 and args.masking_threshold < 1.0:
              snake_case__							:	Dict   	=	mask_heads(A		,					A		,					A      )
              prune_heads(A		,					A		,					A		,					A      )


if __name__ == "__main__":
 main()






