


'''simple docstring'''



# XXX: we want transformers master here - in the absense of conftest manipulating sys.path:
# hack it in for now:
import sys
from pathlib import Path


A__						:   List[str]				=    Path(__file__).resolve().parents[3] / """src"""
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(42)

A__						:   List[str]				=    {"""base""": """patrickvonplaten/wav2vec2_tiny_random""", """robust""": """patrickvonplaten/wav2vec2_tiny_random_robust"""}

A__						:   Tuple				=    """zero2"""
A__						:   str				=    """zero3"""
A__						:   int				=    [ZEROa, ZEROa]



def      UpperCAmelCase__   (      UpperCAmelCase_       :							Any ,       UpperCAmelCase_       :							Any ,       UpperCAmelCase_       :							str				)	->						int:
     # customize the test name generator function as we want both params to appear in the sub-test
     # name, as by default it shows only the first param
     __lowerCamelCase  :      Dict														=						parameterized.to_safe_name('_'.join(str(UpperCAmelCase_				) for x in param.args				)				)
     return F'{func.__name__}_{param_based_name}'


# Cartesian-product of zero stages with models to test
A__						:   Union[str, Any]				=    list(itertools.product(stages, models.keys()))





@slow
@require_deepspeed
@require_torch_gpu
class   UpperCAmelCase_    (_UpperCAmelCase    ):





    """simple docstring"""






    @parameterized.expand(SCREAMING_SNAKE_CASE_     ,	name_func=SCREAMING_SNAKE_CASE_      )
    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )							->	Dict:
         self.run_and_check(
             stage=SCREAMING_SNAKE_CASE_     ,	model=SCREAMING_SNAKE_CASE_     ,	distributed=SCREAMING_SNAKE_CASE_     ,	fpaa=SCREAMING_SNAKE_CASE_     ,	)







    @require_torch_multi_gpu
    @parameterized.expand(SCREAMING_SNAKE_CASE_     ,	name_func=SCREAMING_SNAKE_CASE_      )
    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )							->	Tuple:
         self.run_and_check(
             stage=SCREAMING_SNAKE_CASE_     ,	model=SCREAMING_SNAKE_CASE_     ,	distributed=SCREAMING_SNAKE_CASE_     ,	fpaa=SCREAMING_SNAKE_CASE_     ,	)







    @parameterized.expand(SCREAMING_SNAKE_CASE_     ,	name_func=SCREAMING_SNAKE_CASE_      )
    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )							->	List[str]:
         self.run_and_check(
             stage=SCREAMING_SNAKE_CASE_     ,	model=SCREAMING_SNAKE_CASE_     ,	distributed=SCREAMING_SNAKE_CASE_     ,	fpaa=SCREAMING_SNAKE_CASE_     ,	)







    @require_torch_multi_gpu
    @parameterized.expand(SCREAMING_SNAKE_CASE_     ,	name_func=SCREAMING_SNAKE_CASE_      )
    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )							->	List[Any]:
         self.run_and_check(
             stage=SCREAMING_SNAKE_CASE_     ,	model=SCREAMING_SNAKE_CASE_     ,	distributed=SCREAMING_SNAKE_CASE_     ,	fpaa=SCREAMING_SNAKE_CASE_     ,	)







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_      )							->	Optional[Any]:
         # XXX: run_asr is premature and doesn't save any results
         # so all we check for now is that the process didn't fail
         pass







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_ = 10     ,	SCREAMING_SNAKE_CASE_ = True     ,	SCREAMING_SNAKE_CASE_ = True     ,	SCREAMING_SNAKE_CASE_ = True     ,	)							->	str:
         __lowerCamelCase  :      List[Any]														=						models[model]

         __lowerCamelCase  :      Union[str, Any]														=						self.run_trainer(
             stage=SCREAMING_SNAKE_CASE_     ,	model_name=SCREAMING_SNAKE_CASE_     ,	eval_steps=SCREAMING_SNAKE_CASE_     ,	num_train_epochs=1     ,	distributed=SCREAMING_SNAKE_CASE_     ,	fpaa=SCREAMING_SNAKE_CASE_     ,	)

         self.do_checks(SCREAMING_SNAKE_CASE_      )

         return output_dir







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_ = 10     ,	SCREAMING_SNAKE_CASE_ = 1     ,	SCREAMING_SNAKE_CASE_ = True     ,	SCREAMING_SNAKE_CASE_ = True     ,	)							->	Optional[Any]:
         __lowerCamelCase  :      Optional[int]														=						self.get_auto_remove_tmp_dir('./xxx'     ,	after=SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      int														=						f'\n            --model_name_or_path {model_name}\n            --dataset_name hf-internal-testing/librispeech_asr_dummy\n            --dataset_config_name clean\n            --train_split_name validation\n            --validation_split_name validation\n            --output_dir {output_dir}\n            --num_train_epochs {str(SCREAMING_SNAKE_CASE_      )}\n            --per_device_train_batch_size 2\n            --per_device_eval_batch_size 2\n            --evaluation_strategy steps\n            --learning_rate 5e-4\n            --warmup_steps 8\n            --orthography timit\n            --preprocessing_num_workers 1\n            --group_by_length\n            --freeze_feature_extractor\n            --report_to none\n            --save_steps 0\n            --eval_steps {eval_steps}\n            --report_to none\n        '.split()

         if fpaa:
              args.extend(['--fp16']      )

         # currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
         # hence the separate config files
         __lowerCamelCase  :      str														=						f'--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json'.split()
         __lowerCamelCase  :      Any														=						[f'{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py']
         __lowerCamelCase  :      Optional[int]														=						self.get_launcher(SCREAMING_SNAKE_CASE_      )

         __lowerCamelCase  :      Union[str, Any]														=						launcher + script + args + ds_args
         # keep for quick debug
         # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
         execute_subprocess_async(SCREAMING_SNAKE_CASE_     ,	env=self.get_env()      )

         return output_dir







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_=False      )							->	Optional[int]:
         # 1. explicitly set --num_nodes=1 just in case these tests end up run on a multi-node setup
         # - it won't be able to handle that
         # 2. for now testing with just 2 gpus max (since some quality tests may give different
         # results with mode gpus because we use very little data)
         __lowerCamelCase  :      List[Any]														=						min(2     ,	get_gpu_count()      ) if distributed else 1
         return f'deepspeed --num_nodes 1 --num_gpus {num_gpus}'.split()










"""simple docstring"""


import argparse
import pathlib

import fairseq
import torch
from fairseq.models.roberta import RobertaModel as FairseqRobertaModel
from fairseq.modules import TransformerSentenceEncoderLayer
from packaging import version

from transformers import XLMRobertaConfig, XLMRobertaXLForMaskedLM, XLMRobertaXLForSequenceClassification
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.models.roberta.modeling_roberta import RobertaAttention
from transformers.utils import logging


if version.parse(fairseq.__version__) < version.parse('1.0.0a'):
   raise Exception('requires fairseq >= 1.0.0a')

logging.set_verbosity_info()
__snake_case    					=					logging.get_logger(__name__)

__snake_case    					=					'Hello world! cécé herlolip'


def    _lowerCamelCase   (     lowerCamelCase__	:							str       ,		lowerCamelCase__	:							str       ,		lowerCamelCase__	:							bool  ):
 lowercase__       :       int				=						FairseqRobertaModel.from_pretrained(lowerCamelCase__  )
 roberta.eval()  # disable dropout
 lowercase__       :       Tuple				=						roberta.model.encoder.sentence_encoder
 lowercase__       :       Tuple				=						XLMRobertaConfig(
     vocab_size=roberta_sent_encoder.embed_tokens.num_embeddings       ,		hidden_size=roberta.cfg.model.encoder_embed_dim       ,		num_hidden_layers=roberta.cfg.model.encoder_layers       ,		num_attention_heads=roberta.cfg.model.encoder_attention_heads       ,		intermediate_size=roberta.cfg.model.encoder_ffn_embed_dim       ,		max_position_embeddings=5_14       ,		type_vocab_size=1       ,		layer_norm_eps=1E-5       ,		)
 if classification_head:
  lowercase__       :       Any				=						roberta.model.classification_heads["""mnli"""].out_proj.weight.shape[0]

 print("""Our RoBERTa config:"""       ,		lowerCamelCase__  )

 lowercase__       :       List[Any]				=						XLMRobertaXLForSequenceClassification(lowerCamelCase__  ) if classification_head else XLMRobertaXLForMaskedLM(lowerCamelCase__  )
 model.eval()

 # Now let's copy all the weights.
 # Embeddings
 lowercase__       :       int				=						roberta_sent_encoder.embed_tokens.weight
 lowercase__       :       Union[str, Any]				=						roberta_sent_encoder.embed_positions.weight
 lowercase__       :       int				=						torch.zeros_like(
     model.roberta.embeddings.token_type_embeddings.weight  )  # just zero them out b/c RoBERTa doesn't use them.

 lowercase__       :       int				=						roberta_sent_encoder.layer_norm.weight
 lowercase__       :       List[Any]				=						roberta_sent_encoder.layer_norm.bias

 for i in range(config.num_hidden_layers  ):
  # Encoder: start of layer
  lowercase__       :       BertLayer				=						model.roberta.encoder.layer[i]
  lowercase__       :       TransformerSentenceEncoderLayer				=						roberta_sent_encoder.layers[i]

  lowercase__       :       RobertaAttention				=						layer.attention
  lowercase__       :       str				=						roberta_layer.self_attn_layer_norm.weight
  lowercase__       :       Union[str, Any]				=						roberta_layer.self_attn_layer_norm.bias

  # self attention
  lowercase__       :       BertSelfAttention				=						layer.attention.self
  assert (
      roberta_layer.self_attn.k_proj.weight.data.shape
      == roberta_layer.self_attn.q_proj.weight.data.shape
      == roberta_layer.self_attn.v_proj.weight.data.shape
      == torch.Size((config.hidden_size, config.hidden_size)  )
  )

  lowercase__       :       Optional[Any]				=						roberta_layer.self_attn.q_proj.weight
  lowercase__       :       str				=						roberta_layer.self_attn.q_proj.bias
  lowercase__       :       Optional[int]				=						roberta_layer.self_attn.k_proj.weight
  lowercase__       :       Optional[int]				=						roberta_layer.self_attn.k_proj.bias
  lowercase__       :       int				=						roberta_layer.self_attn.v_proj.weight
  lowercase__       :       Union[str, Any]				=						roberta_layer.self_attn.v_proj.bias

  # self-attention output
  lowercase__       :       BertSelfOutput				=						layer.attention.output
  assert self_output.dense.weight.shape == roberta_layer.self_attn.out_proj.weight.shape
  lowercase__       :       Any				=						roberta_layer.self_attn.out_proj.weight
  lowercase__       :       Optional[int]				=						roberta_layer.self_attn.out_proj.bias

  # this one is final layer norm
  lowercase__       :       Optional[Any]				=						roberta_layer.final_layer_norm.weight
  lowercase__       :       Any				=						roberta_layer.final_layer_norm.bias

  # intermediate
  lowercase__       :       BertIntermediate				=						layer.intermediate
  assert intermediate.dense.weight.shape == roberta_layer.fca.weight.shape
  lowercase__       :       Dict				=						roberta_layer.fca.weight
  lowercase__       :       Any				=						roberta_layer.fca.bias

  # output
  lowercase__       :       BertOutput				=						layer.output
  assert bert_output.dense.weight.shape == roberta_layer.fca.weight.shape
  lowercase__       :       Union[str, Any]				=						roberta_layer.fca.weight
  lowercase__       :       Optional[Any]				=						roberta_layer.fca.bias
  # end of layer

 if classification_head:
  lowercase__       :       Optional[Any]				=						roberta.model.classification_heads["""mnli"""].dense.weight
  lowercase__       :       str				=						roberta.model.classification_heads["""mnli"""].dense.bias
  lowercase__       :       str				=						roberta.model.classification_heads["""mnli"""].out_proj.weight
  lowercase__       :       List[str]				=						roberta.model.classification_heads["""mnli"""].out_proj.bias
 else:
  # LM Head
  lowercase__       :       Tuple				=						roberta.model.encoder.lm_head.dense.weight
  lowercase__       :       int				=						roberta.model.encoder.lm_head.dense.bias
  lowercase__       :       Any				=						roberta.model.encoder.lm_head.layer_norm.weight
  lowercase__       :       Union[str, Any]				=						roberta.model.encoder.lm_head.layer_norm.bias
  lowercase__       :       Dict				=						roberta.model.encoder.lm_head.weight
  lowercase__       :       List[Any]				=						roberta.model.encoder.lm_head.bias

 # Let's check that we get the same results.
 lowercase__       :       torch.Tensor				=						roberta.encode(lowerCamelCase__  ).unsqueeze(0  )  # batch of size 1

 lowercase__       :       Any				=						model(lowerCamelCase__  )[0]
 if classification_head:
  lowercase__       :       Optional[Any]				=						roberta.model.classification_heads["""mnli"""](roberta.extract_features(lowerCamelCase__  )  )
 else:
  lowercase__       :       Tuple				=						roberta.model(lowerCamelCase__  )[0]
 print(our_output.shape       ,		their_output.shape  )
 lowercase__       :       Tuple				=						torch.max(torch.abs(our_output - their_output  )  ).item()
 print(f'''max_absolute_diff = {max_absolute_diff}'''  )  # ~ 1e-7
 lowercase__       :       int				=						torch.allclose(lowerCamelCase__       ,		lowerCamelCase__       ,		atol=1E-3  )
 print("""Do both models output the same tensors?"""       ,		"""🔥""" if success else """💩"""  )
 if not success:
  raise Exception("""Something went wRoNg"""  )

 pathlib.Path(lowerCamelCase__  ).mkdir(parents=lowerCamelCase__       ,		exist_ok=lowerCamelCase__  )
 print(f'''Saving model to {pytorch_dump_folder_path}'''  )
 model.save_pretrained(lowerCamelCase__  )


if __name__ == "__main__":
   __snake_case    					=					argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--roberta_checkpoint_path', default=None, type=str, required=True, help='Path the official PyTorch dump.'
   )
   parser.add_argument(
       '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
   )
   parser.add_argument(
       '--classification_head', action='store_true', help='Whether to convert a final classification head.'
   )
   __snake_case    					=					parser.parse_args()
   convert_xlm_roberta_xl_checkpoint_to_pytorch(
       args.roberta_checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
   )