



# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import warnings
from typing import List
from unittest.mock import Mock

import torch
from torch.utils.data import DataLoader, IterableDataset, TensorDataset

from accelerate.accelerator import Accelerator
from accelerate.utils.dataclasses import DistributedType



class 			__lowerCAmelCase  (						lowerCAmelCase__ ):




 def __init__(    self						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     data




 def __iter__(    self					):



    '''simple docstring'''
    for element in self.data:
       yield element





def 		a__      (		_UpperCamelCase      :			Any=True				):
   __lowerCamelCase    				=     Accelerator(even_batches=_UpperCamelCase				)
   assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
   return accelerator





def 		a__      (		_UpperCamelCase      :			Accelerator			,_UpperCamelCase      :			int			,_UpperCamelCase      :			int			,_UpperCamelCase      :			bool = False				):
   if iterable:
      __lowerCamelCase    				=     DummyIterableDataset(torch.as_tensor(range(_UpperCamelCase				)				)				)
   else:
      __lowerCamelCase    				=     TensorDataset(torch.as_tensor(range(_UpperCamelCase				)				)				)

   __lowerCamelCase    				=     DataLoader(_UpperCamelCase			,batch_size=_UpperCamelCase				)
   __lowerCamelCase    				=     accelerator.prepare(_UpperCamelCase				)

   return dl





def 		a__      (		_UpperCamelCase      :			Accelerator			,_UpperCamelCase      :			int			,_UpperCamelCase      :			int			,_UpperCamelCase      :			List[int]			,_UpperCamelCase      :			List[int]			,):
   __lowerCamelCase    				=     create_dataloader(accelerator=_UpperCamelCase			,dataset_size=_UpperCamelCase			,batch_size=_UpperCamelCase				)

   __lowerCamelCase    				=     [len(batch[0]				) for batch in dl]

   if accelerator.process_index == 0:
      assert batch_sizes == process_0_expected_batch_sizes
   elif accelerator.process_index == 1:
      assert batch_sizes == process_1_expected_batch_sizes





def 		a__      (		):
   __lowerCamelCase    				=     create_accelerator()

   # without padding, we would expect a different number of batches
   verify_dataloader_batch_sizes(
       _UpperCamelCase			,dataset_size=3			,batch_size=1			,process_0_expected_batch_sizes=[1, 1]			,process_1_expected_batch_sizes=[1, 1]			,)

   # without padding, we would expect the same number of batches, but different sizes
   verify_dataloader_batch_sizes(
       _UpperCamelCase			,dataset_size=7			,batch_size=2			,process_0_expected_batch_sizes=[2, 2]			,process_1_expected_batch_sizes=[2, 2]			,)





def 		a__      (		):
   __lowerCamelCase    				=     create_accelerator(even_batches=_UpperCamelCase				)

   verify_dataloader_batch_sizes(
       _UpperCamelCase			,dataset_size=3			,batch_size=1			,process_0_expected_batch_sizes=[1, 1]			,process_1_expected_batch_sizes=[1]			,)

   verify_dataloader_batch_sizes(
       _UpperCamelCase			,dataset_size=7			,batch_size=2			,process_0_expected_batch_sizes=[2, 2]			,process_1_expected_batch_sizes=[2, 1]			,)





def 		a__      (		):
   __lowerCamelCase    				=     create_accelerator(even_batches=_UpperCamelCase				)

   __lowerCamelCase    				=     torch.nn.Linear(1			,1				)
   __lowerCamelCase    				=     accelerator.prepare(_UpperCamelCase				)

   __lowerCamelCase    				=     create_dataloader(_UpperCamelCase			,dataset_size=3			,batch_size=1				)

   __lowerCamelCase    				=     []
   with accelerator.join_uneven_inputs([ddp_model]				):
      for batch_idx, batch in enumerate(_UpperCamelCase				):
         __lowerCamelCase    				=     ddp_model(batch[0].float()				)
         __lowerCamelCase    				=     output.sum()
         loss.backward()
         batch_idxs.append(_UpperCamelCase				)

   accelerator.wait_for_everyone()

   if accelerator.process_index == 0:
      assert batch_idxs == [0, 1]
   elif accelerator.process_index == 1:
      assert batch_idxs == [0]





def 		a__      (		_UpperCamelCase      :			Tuple				):
   with warnings.catch_warnings(record=_UpperCamelCase				) as w:
      with accelerator.join_uneven_inputs([Mock()]				):
         pass

      assert issubclass(w[-1].category			,_UpperCamelCase				)
      assert "only supported for multi-GPU" in str(w[-1].message				)





def 		a__      (		):
   __lowerCamelCase    				=     True
   __lowerCamelCase    				=     False
   __lowerCamelCase    				=     create_accelerator(even_batches=_UpperCamelCase				)
   __lowerCamelCase    				=     torch.nn.Linear(1			,1				)
   __lowerCamelCase    				=     accelerator.prepare(_UpperCamelCase				)
   __lowerCamelCase    				=     create_dataloader(_UpperCamelCase			,dataset_size=3			,batch_size=1				)
   __lowerCamelCase    				=     create_dataloader(_UpperCamelCase			,dataset_size=3			,batch_size=1				)

   with accelerator.join_uneven_inputs([ddp_model]			,even_batches=_UpperCamelCase				):
      __lowerCamelCase    				=     train_dl.batch_sampler.even_batches
      __lowerCamelCase    				=     valid_dl.batch_sampler.even_batches

   assert train_dl_overridden_value == overridden_even_batches
   assert valid_dl_overridden_value == overridden_even_batches
   assert train_dl.batch_sampler.even_batches == default_even_batches
   assert valid_dl.batch_sampler.even_batches == default_even_batches





def 		a__      (		):
   __lowerCamelCase    				=     True
   __lowerCamelCase    				=     False
   __lowerCamelCase    				=     create_accelerator(even_batches=_UpperCamelCase				)
   __lowerCamelCase    				=     torch.nn.Linear(1			,1				)
   __lowerCamelCase    				=     accelerator.prepare(_UpperCamelCase				)
   create_dataloader(_UpperCamelCase			,dataset_size=3			,batch_size=1			,iterable=_UpperCamelCase				)
   __lowerCamelCase    				=     create_dataloader(_UpperCamelCase			,dataset_size=3			,batch_size=1				)

   with warnings.catch_warnings():
      warnings.filterwarnings('''ignore'''				)
      try:
         with accelerator.join_uneven_inputs([ddp_model]			,even_batches=_UpperCamelCase				):
            __lowerCamelCase    				=     batch_dl.batch_sampler.even_batches
      except AttributeError:
         # ensure attribute error is not raised when processing iterable dl
         raise AssertionError

   assert batch_dl_overridden_value == overridden_even_batches
   assert batch_dl.batch_sampler.even_batches == default_even_batches





def 		a__      (		):
   __lowerCamelCase    				=     create_accelerator()
   __lowerCamelCase    				=     torch.nn.Linear(1			,1				)
   __lowerCamelCase    				=     accelerator.prepare(_UpperCamelCase				)
   create_dataloader(_UpperCamelCase			,dataset_size=3			,batch_size=1			,iterable=_UpperCamelCase				)

   with warnings.catch_warnings(record=_UpperCamelCase				) as w:
      with accelerator.join_uneven_inputs([ddp_model]			,even_batches=_UpperCamelCase				):
         pass

      assert issubclass(w[-1].category			,_UpperCamelCase				)
      assert "only supported for map-style datasets" in str(w[-1].message				)





def 		a__      (		):
   __lowerCamelCase    				=     create_accelerator()

   accelerator.print('''Test that even_batches variable ensures uniform batches across processes'''				)
   test_default_ensures_even_batch_sizes()

   accelerator.print('''Run tests with even_batches disabled'''				)
   test_can_disable_even_batches()

   accelerator.print('''Test joining uneven inputs'''				)
   test_can_join_uneven_inputs()

   accelerator.print('''Test overriding even_batches when joining uneven inputs'''				)
   test_join_can_override_even_batches()

   accelerator.print('''Test overriding even_batches for mixed dataloader types'''				)
   test_join_can_override_for_mixed_type_dataloaders()

   accelerator.print('''Test overriding even_batches raises a warning for iterable dataloaders'''				)
   test_join_raises_warning_for_iterable_when_overriding_even_batches()

   accelerator.print('''Test join with non DDP distributed raises warning'''				)
   __lowerCamelCase    				=     accelerator.state.distributed_type
   __lowerCamelCase    				=     DistributedType.FSDP
   test_join_raises_warning_for_non_ddp_distributed(_UpperCamelCase				)
   __lowerCamelCase    				=     original_state


if __name__ == "__main__":
  main()











import argparse
from pathlib import Path

import fairseq
import torch
from fairseq.models.xmod import XMODModel as FairseqXmodModel
from packaging import version

from transformers import XmodConfig, XmodForMaskedLM, XmodForSequenceClassification
from transformers.utils import logging


if version.parse(fairseq.__version__) < version.parse("""0.12.2"""):
  raise Exception("""requires fairseq >= 0.12.2""")
if version.parse(fairseq.__version__) > version.parse("""2"""):
  raise Exception("""requires fairseq < v2""")

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

a_								=       """Hello, World!"""
a_								=       """en_XX"""





def 		a__      (		_UpperCamelCase      :			str			,_UpperCamelCase      :			str			,_UpperCamelCase      :			bool				):
   __lowerCamelCase    				=     Path('''data_bin'''				)
   __lowerCamelCase    				=     FairseqXmodModel.from_pretrained(
       model_name_or_path=str(Path(_UpperCamelCase				).parent				)			,checkpoint_file=Path(_UpperCamelCase				).name			,_name='''xmod_base'''			,arch='''xmod_base'''			,task='''multilingual_masked_lm'''			,data_name_or_path=str(_UpperCamelCase				)			,bpe='''sentencepiece'''			,sentencepiece_model=str(Path(_UpperCamelCase				).parent / '''sentencepiece.bpe.model'''				)			,src_dict=str(data_dir / '''dict.txt'''				)			,)
   xmod.eval()  # disable dropout
   print(_UpperCamelCase				)

   __lowerCamelCase    				=     xmod.model.encoder.sentence_encoder
   __lowerCamelCase    				=     XmodConfig(
       vocab_size=xmod_sent_encoder.embed_tokens.num_embeddings			,hidden_size=xmod.cfg.model.encoder_embed_dim			,num_hidden_layers=xmod.cfg.model.encoder_layers			,num_attention_heads=xmod.cfg.model.encoder_attention_heads			,intermediate_size=xmod.cfg.model.encoder_ffn_embed_dim			,max_position_embeddings=5_14			,type_vocab_size=1			,layer_norm_eps=1e-5			,pre_norm=xmod.cfg.model.encoder_normalize_before			,adapter_reduction_factor=getattr(xmod.cfg.model			,'''bottleneck'''			,2				)			,adapter_layer_norm=xmod.cfg.model.adapter_layer_norm			,adapter_reuse_layer_norm=xmod.cfg.model.adapter_reuse_layer_norm			,ln_before_adapter=xmod.cfg.model.ln_before_adapter			,languages=xmod.cfg.model.languages			,)
   if classification_head:
      __lowerCamelCase    				=     xmod.model.classification_heads['''mnli'''].out_proj.weight.shape[0]

   print('''Our X-MOD config:'''			,_UpperCamelCase				)

   __lowerCamelCase    				=     XmodForSequenceClassification(_UpperCamelCase				) if classification_head else XmodForMaskedLM(_UpperCamelCase				)
   model.eval()

   # Now let's copy all the weights.
   # Embeddings
   __lowerCamelCase    				=     xmod_sent_encoder.embed_tokens.weight
   __lowerCamelCase    				=     xmod_sent_encoder.embed_positions.weight
   __lowerCamelCase    				=     torch.zeros_like(
       model.roberta.embeddings.token_type_embeddings.weight				)  # just zero them out b/c xmod doesn't use them.

   __lowerCamelCase    				=     xmod_sent_encoder.layernorm_embedding.weight
   __lowerCamelCase    				=     xmod_sent_encoder.layernorm_embedding.bias

   for i in range(config.num_hidden_layers				):
      # Encoder: start of layer
      __lowerCamelCase    				=     model.roberta.encoder.layer[i]
      __lowerCamelCase    				=     xmod_sent_encoder.layers[i]

      # self attention
      __lowerCamelCase    				=     layer.attention.self
      if not (
          xmod_layer.self_attn.k_proj.weight.data.shape
          == xmod_layer.self_attn.q_proj.weight.data.shape
          == xmod_layer.self_attn.v_proj.weight.data.shape
          == torch.Size((config.hidden_size, config.hidden_size)				)
      ):
         raise AssertionError('''Dimensions of self-attention weights do not match.'''				)

      __lowerCamelCase    				=     xmod_layer.self_attn.q_proj.weight
      __lowerCamelCase    				=     xmod_layer.self_attn.q_proj.bias
      __lowerCamelCase    				=     xmod_layer.self_attn.k_proj.weight
      __lowerCamelCase    				=     xmod_layer.self_attn.k_proj.bias
      __lowerCamelCase    				=     xmod_layer.self_attn.v_proj.weight
      __lowerCamelCase    				=     xmod_layer.self_attn.v_proj.bias

      # self-attention output
      __lowerCamelCase    				=     layer.attention.output
      if self_output.dense.weight.shape != xmod_layer.self_attn.out_proj.weight.shape:
         raise AssertionError('''Dimensions of self-attention output weights do not match.'''				)
      __lowerCamelCase    				=     xmod_layer.self_attn.out_proj.weight
      __lowerCamelCase    				=     xmod_layer.self_attn.out_proj.bias
      __lowerCamelCase    				=     xmod_layer.self_attn_layer_norm.weight
      __lowerCamelCase    				=     xmod_layer.self_attn_layer_norm.bias

      # intermediate
      __lowerCamelCase    				=     layer.intermediate
      if intermediate.dense.weight.shape != xmod_layer.fca.weight.shape:
         raise AssertionError('''Dimensions of intermediate weights do not match.'''				)
      __lowerCamelCase    				=     xmod_layer.fca.weight
      __lowerCamelCase    				=     xmod_layer.fca.bias

      # output
      __lowerCamelCase    				=     layer.output
      if bert_output.dense.weight.shape != xmod_layer.fca.weight.shape:
         raise AssertionError('''Dimensions of feed-forward weights do not match.'''				)
      __lowerCamelCase    				=     xmod_layer.fca.weight
      __lowerCamelCase    				=     xmod_layer.fca.bias
      __lowerCamelCase    				=     xmod_layer.final_layer_norm.weight
      __lowerCamelCase    				=     xmod_layer.final_layer_norm.bias
      if bert_output.adapter_layer_norm is not None:
         __lowerCamelCase    				=     xmod_layer.adapter_layer_norm.weight
         __lowerCamelCase    				=     xmod_layer.adapter_layer_norm.bias

      if sorted(bert_output.adapter_modules.keys()				) != sorted(xmod_layer.adapter_modules.keys()				):
         raise AssertionError('''Lists of language adapters do not match.'''				)
      for lang_code, adapter in xmod_layer.adapter_modules.items():
         __lowerCamelCase    				=     bert_output.adapter_modules[lang_code]
         __lowerCamelCase    				=     xmod_layer.adapter_modules[lang_code]
         __lowerCamelCase    				=     from_adapter.fca.weight
         __lowerCamelCase    				=     from_adapter.fca.bias
         __lowerCamelCase    				=     from_adapter.fca.weight
         __lowerCamelCase    				=     from_adapter.fca.bias

      # end of layer

   if xmod_sent_encoder.layer_norm is not None:
      __lowerCamelCase    				=     xmod_sent_encoder.layer_norm.weight
      __lowerCamelCase    				=     xmod_sent_encoder.layer_norm.bias

   if classification_head:
      __lowerCamelCase    				=     xmod.model.classification_heads['''mnli'''].dense.weight
      __lowerCamelCase    				=     xmod.model.classification_heads['''mnli'''].dense.bias
      __lowerCamelCase    				=     xmod.model.classification_heads['''mnli'''].out_proj.weight
      __lowerCamelCase    				=     xmod.model.classification_heads['''mnli'''].out_proj.bias
   else:
      # LM Head
      __lowerCamelCase    				=     xmod.model.encoder.lm_head.dense.weight
      __lowerCamelCase    				=     xmod.model.encoder.lm_head.dense.bias
      __lowerCamelCase    				=     xmod.model.encoder.lm_head.layer_norm.weight
      __lowerCamelCase    				=     xmod.model.encoder.lm_head.layer_norm.bias
      __lowerCamelCase    				=     xmod.model.encoder.lm_head.weight
      __lowerCamelCase    				=     xmod.model.encoder.lm_head.bias

   # Let's check that we get the same results.
   __lowerCamelCase    				=     xmod.encode(_UpperCamelCase				).unsqueeze(0				)  # batch of size 1
   model.roberta.set_default_language(_UpperCamelCase				)

   __lowerCamelCase    				=     model(_UpperCamelCase				)[0]
   if classification_head:
      __lowerCamelCase    				=     xmod.model.classification_heads['''mnli'''](xmod.extract_features(_UpperCamelCase				)				)
   else:
      __lowerCamelCase    				=     xmod.model(_UpperCamelCase			,lang_id=[SAMPLE_LANGUAGE]				)[0]
   print(our_output.shape			,their_output.shape				)
   __lowerCamelCase    				=     torch.max(torch.abs(our_output - their_output				)				).item()
   print(F"""max_absolute_diff = {max_absolute_diff}"""				)  # ~ 1e-7
   __lowerCamelCase    				=     torch.allclose(_UpperCamelCase			,_UpperCamelCase			,atol=1e-3				)
   print('''Do both models output the same tensors?'''			,'''🔥''' if success else '''💩'''				)
   if not success:
      raise Exception('''Something went wRoNg'''				)

   Path(_UpperCamelCase				).mkdir(parents=_UpperCamelCase			,exist_ok=_UpperCamelCase				)
   print(F"""Saving model to {pytorch_dump_folder_path}"""				)
   model.save_pretrained(_UpperCamelCase				)


if __name__ == "__main__":
  a_								=       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--xmod_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."""
  )
  a_								=       parser.parse_args()
  convert_xmod_checkpoint_to_pytorch(
      args.xmod_checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
  )






