
"""simple docstring"""



from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


UpperCAmelCase__          = logging.get_logger(__name__)


if is_torch_available():
 import torch
 from torch import nn
 from torch.utils.data import Dataset

 from transformers import Trainer

 class       lowerCAmelCase__				(						A_	):







      def __init__(							self			:			Optional[Any]    ,      _lowerCamelCase			:			int = 101			):
        _snake_case					     =       length







      def __len__(							self			:			List[Any]			):
        return self.length







      def __getitem__(							self			:			str    ,      _lowerCamelCase			:			Optional[Any]			):
        return i






 class       lowerCAmelCase__				:







      def __call__(							self			:			str    ,      _lowerCamelCase			:			Optional[Any]			):
        return {"input_ids": torch.tensor(_lowerCamelCase			), "labels": torch.tensor(_lowerCamelCase			)}






 class       lowerCAmelCase__				(						nn.Module	):







      def __init__(							self			:			Tuple			):
        super().__init__()
        # Add some (unused) params otherwise DDP will complain.
        _snake_case					     =       nn.Linear(120    ,      80			)







      def      lowercase	(							self			:			int    ,      _lowerCamelCase			:			Any    ,      _lowerCamelCase			:			Union[str, Any]=None			):
        if labels is not None:
          return torch.tensor(0.0    ,      device=input_ids.device			), input_ids
        else:
          return input_ids

class       lowerCAmelCase__				(						A_	):







     @require_torch_neuroncore
     def      lowercase	(							self			:			Union[str, Any]			):
       _snake_case					     =       f'''--nproc_per_node=2
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        '''.split()
       _snake_case					     =       self.get_auto_remove_tmp_dir()
       _snake_case					     =       f'''--output_dir {output_dir}'''.split()
       _snake_case					     =       ['''torchrun'''] + distributed_args + args
       execute_subprocess_async(_lowerCamelCase    ,      env=self.get_env()			)
       # successful return here == success - any errors would have caused an error in the sub-call






class       lowerCAmelCase__				(						A_	):







     @require_torch_multi_gpu
     def      lowercase	(							self			:			List[str]			):
       _snake_case					     =       f'''--nproc_per_node={torch.cuda.device_count()}
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        '''.split()
       _snake_case					     =       self.get_auto_remove_tmp_dir()
       _snake_case					     =       f'''--output_dir {output_dir}'''.split()
       _snake_case					     =       ['''torchrun'''] + distributed_args + args
       execute_subprocess_async(_lowerCamelCase    ,      env=self.get_env()			)
       # successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
 # The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
 #
 # PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

 UpperCAmelCase__          = HfArgumentParser((TrainingArguments,))
 UpperCAmelCase__          = parser.parse_args_into_dataclasses()[0]

 logger.warning(
     F"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
     F"distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}"
 )

 # Essentially, what we want to verify in the distributed case is that we get all samples back,
 # in the right order. (this is crucial for prediction for instance)
 for dataset_length in [101, 40, 7]:
  UpperCAmelCase__          = DummyDataset(dataset_length)
  def 					_UpperCAmelCase      ( __lowerCamelCase   :       EvalPrediction		)     ->					Dict:
    _snake_case					     =       list(range(len(__lowerCamelCase		)		)		)
    _snake_case					     =       p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
    if not success and training_args.local_rank == 0:
      logger.warning(
          '''Predictions and/or labels do not match expected results:\n  - predictions: '''
          f'''{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}'''		)
    return {"success": success}

  UpperCAmelCase__          = Trainer(
      model=DummyModel(),
      args=training_args,
      data_collator=DummyDataCollator(),
      eval_dataset=dataset,
      compute_metrics=compute_metrics,
  )
  UpperCAmelCase__          = trainer.evaluate()
  logger.info(metrics)
  if metrics["eval_success"] is not True:
   logger.error(metrics)
   exit(1)

  UpperCAmelCase__          = trainer.predict(dataset)
  logger.info(p.metrics)
  if p.metrics["test_success"] is not True:
   logger.error(p.metrics)
   exit(1)

  UpperCAmelCase__          = 2

  UpperCAmelCase__          = trainer.evaluate()
  logger.info(metrics)
  if metrics["eval_success"] is not True:
   logger.error(metrics)
   exit(1)

  UpperCAmelCase__          = trainer.predict(dataset)
  logger.info(p.metrics)
  if p.metrics["test_success"] is not True:
   logger.error(p.metrics)
   exit(1)

  UpperCAmelCase__          = None





"""simple docstring"""



import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
 UpperCAmelCase__          = argparse.ArgumentParser(
     description=(
         'Extraction some layers of the full BertForMaskedLM or RObertaForMaskedLM for Transfer Learned'
         ' Distillation'
     )
 )
 parser.add_argument('--model_type', default='bert', choices=['bert'])
 parser.add_argument('--model_name', default='bert-base-uncased', type=str)
 parser.add_argument('--dump_checkpoint', default='serialization_dir/tf_bert-base-uncased_0247911.pth', type=str)
 parser.add_argument('--vocab_transform', action='store_true')
 UpperCAmelCase__          = parser.parse_args()

 if args.model_type == "bert":
  UpperCAmelCase__          = BertForMaskedLM.from_pretrained(args.model_name)
  UpperCAmelCase__          = 'bert'
 else:
  raise ValueError('args.model_type should be "bert".')

 UpperCAmelCase__          = model.state_dict()
 UpperCAmelCase__          = {}

 for w in ["word_embeddings", "position_embeddings"]:
  UpperCAmelCase__          = state_dict[F"{prefix}.embeddings.{w}.weight"]
 for w in ["weight", "bias"]:
  UpperCAmelCase__          = state_dict[F"{prefix}.embeddings.LayerNorm.{w}"]

 UpperCAmelCase__          = 0
 for teacher_idx in [0, 2, 4, 7, 9, 11]:
  for w in ["weight", "bias"]:
   UpperCAmelCase__          = state_dict[
       F"{prefix}.encoder.layer.{teacher_idx}.attention.self.query.{w}"
   ]
   UpperCAmelCase__          = state_dict[
       F"{prefix}.encoder.layer.{teacher_idx}.attention.self.key.{w}"
   ]
   UpperCAmelCase__          = state_dict[
       F"{prefix}.encoder.layer.{teacher_idx}.attention.self.value.{w}"
   ]

   UpperCAmelCase__          = state_dict[
       F"{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}"
   ]
   UpperCAmelCase__          = state_dict[
       F"{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}"
   ]

   UpperCAmelCase__          = state_dict[
       F"{prefix}.encoder.layer.{teacher_idx}.intermediate.dense.{w}"
   ]
   UpperCAmelCase__          = state_dict[
       F"{prefix}.encoder.layer.{teacher_idx}.output.dense.{w}"
   ]
   UpperCAmelCase__          = state_dict[
       F"{prefix}.encoder.layer.{teacher_idx}.output.LayerNorm.{w}"
   ]
  std_idx += 1

 UpperCAmelCase__          = state_dict['cls.predictions.decoder.weight']
 UpperCAmelCase__          = state_dict['cls.predictions.bias']
 if args.vocab_transform:
  for w in ["weight", "bias"]:
   UpperCAmelCase__          = state_dict[F"cls.predictions.transform.dense.{w}"]
   UpperCAmelCase__          = state_dict[F"cls.predictions.transform.LayerNorm.{w}"]

 print(F"N layers selected for distillation: {std_idx}")
 print(F"Number of params transferred for distillation: {len(compressed_sd.keys())}")

 print(F"Save transferred checkpoint to {args.dump_checkpoint}.")
 torch.save(compressed_sd, args.dump_checkpoint)



