





'''simple docstring'''



import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()


class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase					,    __UpperCamelCase					,    unittest.TestCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						Optional[int]      =   AutoencoderKL
     UpperCAmelCase__       :						Union[str, Any]      =   '''sample'''
     UpperCAmelCase__       :						Optional[int]      =   1E-2

     @property
     def 	UpperCamelCase(					self	):
       _snake_case     		=				4
       _snake_case     		=				3
       _snake_case     		=				(32, 32)

       _snake_case     		=				floats_tensor((batch_size, num_channels) + sizes	).to(lowerCamelCase	)

       return {"sample": image}

     @property
     def 	UpperCamelCase(					self	):
       return (3, 32, 32)

     @property
     def 	UpperCamelCase(					self	):
       return (3, 32, 32)

     def 	UpperCamelCase(					self	):
       _snake_case     		=				{
           "block_out_channels": [32, 64],
           "in_channels": 3,
           "out_channels": 3,
           "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
           "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"],
           "latent_channels": 4,
       }
       _snake_case     		=				self.dummy_input
       return init_dict, inputs_dict

     def 	UpperCamelCase(					self	):
       pass

     def 	UpperCamelCase(					self	):
       pass

     @unittest.skipIf(torch_device == "mps"						,   "Gradient checkpointing skipped on MPS"	)
     def 	UpperCamelCase(					self	):
       # enable deterministic behavior for gradient checkpointing
       _snake_case		,		_snake_case     		=				self.prepare_init_args_and_inputs_for_common()
       _snake_case     		=				self.model_class(**lowerCamelCase	)
       model.to(lowerCamelCase	)

       assert not model.is_gradient_checkpointing and model.training

       _snake_case     		=				model(**lowerCamelCase	).sample
       # run the backwards pass on the model. For backwards pass, for simplicity purpose,
       # we won't calculate the loss and rather backprop on out.sum()
       model.zero_grad()

       _snake_case     		=				torch.randn_like(lowerCamelCase	)
       _snake_case     		=				(out - labels).mean()
       loss.backward()

       # re-instantiate the model now enabling gradient checkpointing
       _snake_case     		=				self.model_class(**lowerCamelCase	)
       # clone model
       model_a.load_state_dict(model.state_dict()	)
       model_a.to(lowerCamelCase	)
       model_a.enable_gradient_checkpointing()

       assert model_a.is_gradient_checkpointing and model_a.training

       _snake_case     		=				model_a(**lowerCamelCase	).sample
       # run the backwards pass on the model. For backwards pass, for simplicity purpose,
       # we won't calculate the loss and rather backprop on out.sum()
       model_a.zero_grad()
       _snake_case     		=				(out_a - labels).mean()
       loss_a.backward()

       # compare the output and parameters gradients
       self.assertTrue((loss - loss_a).abs() < 1e-5	)
       _snake_case     		=				dict(model.named_parameters()	)
       _snake_case     		=				dict(model_a.named_parameters()	)
       for name, param in named_params.items():
         self.assertTrue(torch_all_close(param.grad.data						,   named_params_a[name].grad.data						,   atol=5e-5	)	)

     def 	UpperCamelCase(					self	):
       _snake_case		,		_snake_case     		=				AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy"						,   output_loading_info=lowerCamelCase	)
       self.assertIsNotNone(lowerCamelCase	)
       self.assertEqual(len(loading_info["missing_keys"]	)						,   0	)

       model.to(lowerCamelCase	)
       _snake_case     		=				model(**self.dummy_input	)

       assert image is not None, "Make sure output is not None"






     def 	UpperCamelCase(					self	):
       _snake_case     		=				AutoencoderKL.from_pretrained("fusing/autoencoder-kl-dummy"	)
       _snake_case     		=				model.to(lowerCamelCase	)
       model.eval()

       if torch_device == "mps":
         _snake_case     		=				torch.manual_seed(0	)
       else:
         _snake_case     		=				torch.Generator(device=lowerCamelCase	).manual_seed(0	)

       _snake_case     		=				torch.randn(
           1						,   model.config.in_channels						,   model.config.sample_size						,   model.config.sample_size						,   generator=torch.manual_seed(0	)						,   )
       _snake_case     		=				image.to(lowerCamelCase	)
       with torch.no_grad():
         _snake_case     		=				model(lowerCamelCase						,   sample_posterior=lowerCamelCase						,   generator=lowerCamelCase	).sample

       _snake_case     		=				output[0, -1, -3:, -3:].flatten().cpu()

       # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
       # the expected output slices are not the same for CPU and GPU.
       if torch_device == "mps":
         _snake_case     		=				torch.tensor(
             [
                 -4.0078e-01,
                 -3.8323e-04,
                 -1.2681e-01,
                 -1.1462e-01,
                 2.0095e-01,
                 1.0893e-01,
                 -8.8247e-02,
                 -3.0361e-01,
                 -9.8644e-03,
             ]	)
       elif torch_device == "cpu":
         _snake_case     		=				torch.tensor(
             [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]	)
       else:
         _snake_case     		=				torch.tensor(
             [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]	)

       self.assertTrue(torch_all_close(lowerCamelCase						,   lowerCamelCase						,   rtol=1e-2	)	)







@slow
class       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase	):
       return F'''gaussian_noise_s={seed}_shape={"_".join([str(lowerCamelCase	) for s in shape]	)}.npy'''

     def 	UpperCamelCase(					self	):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()

     def 	UpperCamelCase(					self						,   lowerCamelCase=0						,   lowerCamelCase=(4, 3, 512, 512)						,   lowerCamelCase=False	):
       _snake_case     		=				torch.floataa if fpaa else torch.floataa
       _snake_case     		=				torch.from_numpy(load_hf_numpy(self.get_file_format(lowerCamelCase						,   lowerCamelCase	)	)	).to(lowerCamelCase	).to(lowerCamelCase	)
       return image

     def 	UpperCamelCase(					self						,   lowerCamelCase="CompVis/stable-diffusion-v1-4"						,   lowerCamelCase=False	):
       _snake_case     		=				"fp16" if fpaa else None
       _snake_case     		=				torch.floataa if fpaa else torch.floataa

       _snake_case     		=				AutoencoderKL.from_pretrained(
           lowerCamelCase						,   subfolder="vae"						,   torch_dtype=lowerCamelCase						,   revision=lowerCamelCase						,   )
       model.to(lowerCamelCase	).eval()

       return model

     def 	UpperCamelCase(					self						,   lowerCamelCase=0	):
       if torch_device == "mps":
         return torch.manual_seed(lowerCamelCase	)
       return torch.Generator(device=lowerCamelCase	).manual_seed(lowerCamelCase	)

     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
             [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
             # fmt: on
         ]	)
     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				self.get_sd_vae_model()
       _snake_case     		=				self.get_sd_image(lowerCamelCase	)
       _snake_case     		=				self.get_generator(lowerCamelCase	)

       with torch.no_grad():
         _snake_case     		=				model(lowerCamelCase						,   generator=lowerCamelCase						,   sample_posterior=lowerCamelCase	).sample

       assert sample.shape == image.shape

       _snake_case     		=				sample[-1, -2:, -2:, :2].flatten().float().cpu()
       _snake_case     		=				torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice	)

       assert torch_all_close(lowerCamelCase						,   lowerCamelCase						,   atol=3e-3	)

     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
             [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
             # fmt: on
         ]	)
     @require_torch_gpu
     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				self.get_sd_vae_model(fpaa=lowerCamelCase	)
       _snake_case     		=				self.get_sd_image(lowerCamelCase						,   fpaa=lowerCamelCase	)
       _snake_case     		=				self.get_generator(lowerCamelCase	)

       with torch.no_grad():
         _snake_case     		=				model(lowerCamelCase						,   generator=lowerCamelCase						,   sample_posterior=lowerCamelCase	).sample

       assert sample.shape == image.shape

       _snake_case     		=				sample[-1, -2:, :2, -2:].flatten().float().cpu()
       _snake_case     		=				torch.tensor(lowerCamelCase	)

       assert torch_all_close(lowerCamelCase						,   lowerCamelCase						,   atol=1e-2	)

     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
             [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
             # fmt: on
         ]	)
     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				self.get_sd_vae_model()
       _snake_case     		=				self.get_sd_image(lowerCamelCase	)

       with torch.no_grad():
         _snake_case     		=				model(lowerCamelCase	).sample

       assert sample.shape == image.shape

       _snake_case     		=				sample[-1, -2:, -2:, :2].flatten().float().cpu()
       _snake_case     		=				torch.tensor(expected_slice_mps if torch_device == "mps" else expected_slice	)

       assert torch_all_close(lowerCamelCase						,   lowerCamelCase						,   atol=3e-3	)

     @parameterized.expand(
         [
             # fmt: off
             [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
             [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
             # fmt: on
         ]	)
     @require_torch_gpu
     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				self.get_sd_vae_model()
       _snake_case     		=				self.get_sd_image(lowerCamelCase						,   shape=(3, 4, 64, 64)	)

       with torch.no_grad():
         _snake_case     		=				model.decode(lowerCamelCase	).sample

       assert list(sample.shape	) == [3, 3, 512, 512]

       _snake_case     		=				sample[-1, -2:, :2, -2:].flatten().cpu()
       _snake_case     		=				torch.tensor(lowerCamelCase	)

       assert torch_all_close(lowerCamelCase						,   lowerCamelCase						,   atol=1e-3	)

     @parameterized.expand(
         [
             # fmt: off
             [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
             [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
             # fmt: on
         ]	)
     @require_torch_gpu
     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				self.get_sd_vae_model(fpaa=lowerCamelCase	)
       _snake_case     		=				self.get_sd_image(lowerCamelCase						,   shape=(3, 4, 64, 64)						,   fpaa=lowerCamelCase	)

       with torch.no_grad():
         _snake_case     		=				model.decode(lowerCamelCase	).sample

       assert list(sample.shape	) == [3, 3, 512, 512]

       _snake_case     		=				sample[-1, -2:, :2, -2:].flatten().float().cpu()
       _snake_case     		=				torch.tensor(lowerCamelCase	)

       assert torch_all_close(lowerCamelCase						,   lowerCamelCase						,   atol=5e-3	)

     @parameterized.expand([(13,), (16,), (27,)]	)
     @require_torch_gpu
     @unittest.skipIf(not is_xformers_available()						,   reason="xformers is not required when using PyTorch 2.0."	)
     def 	UpperCamelCase(					self						,   lowerCamelCase	):
       _snake_case     		=				self.get_sd_vae_model(fpaa=lowerCamelCase	)
       _snake_case     		=				self.get_sd_image(lowerCamelCase						,   shape=(3, 4, 64, 64)						,   fpaa=lowerCamelCase	)

       with torch.no_grad():
         _snake_case     		=				model.decode(lowerCamelCase	).sample

       model.enable_xformers_memory_efficient_attention()
       with torch.no_grad():
         _snake_case     		=				model.decode(lowerCamelCase	).sample

       assert list(sample.shape	) == [3, 3, 512, 512]

       assert torch_all_close(lowerCamelCase						,   lowerCamelCase						,   atol=1e-1	)

     @parameterized.expand([(13,), (16,), (37,)]	)
     @require_torch_gpu
     @unittest.skipIf(not is_xformers_available()						,   reason="xformers is not required when using PyTorch 2.0."	)
     def 	UpperCamelCase(					self						,   lowerCamelCase	):
       _snake_case     		=				self.get_sd_vae_model()
       _snake_case     		=				self.get_sd_image(lowerCamelCase						,   shape=(3, 4, 64, 64)	)

       with torch.no_grad():
         _snake_case     		=				model.decode(lowerCamelCase	).sample

       model.enable_xformers_memory_efficient_attention()
       with torch.no_grad():
         _snake_case     		=				model.decode(lowerCamelCase	).sample

       assert list(sample.shape	) == [3, 3, 512, 512]

       assert torch_all_close(lowerCamelCase						,   lowerCamelCase						,   atol=1e-2	)






     @parameterized.expand(
         [
             # fmt: off
             [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
             [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
             # fmt: on
         ]	)
     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase	):
       _snake_case     		=				self.get_sd_vae_model()
       _snake_case     		=				self.get_sd_image(lowerCamelCase	)
       _snake_case     		=				self.get_generator(lowerCamelCase	)

       with torch.no_grad():
         _snake_case     		=				model.encode(lowerCamelCase	).latent_dist
         _snake_case     		=				dist.sample(generator=lowerCamelCase	)

       assert list(sample.shape	) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

       _snake_case     		=				sample[0, -1, -3:, -3:].flatten().cpu()
       _snake_case     		=				torch.tensor(lowerCamelCase	)

       _snake_case     		=				3e-3 if torch_device != "mps" else 1e-2
       assert torch_all_close(lowerCamelCase						,   lowerCamelCase						,   atol=lowerCamelCase	)









'''simple docstring'''



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()
__magic_name__   :						Union[str, Any]    =			logging.get_logger(__name__)

__magic_name__   :						Tuple    =			"""Hello, World!"""
__magic_name__   :						Union[str, Any]    =			"""en_XX"""




def 				snake_case_      (		SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__					,	SCREAMING_SNAKE_CASE__       ):







  '''simple docstring'''



  _snake_case     		=				Path("data_bin"       )
  _snake_case     		=				FairseqXmodModel.from_pretrained(
      model_name_or_path=str(Path(SCREAMING_SNAKE_CASE__       ).parent       )					,	checkpoint_file=Path(SCREAMING_SNAKE_CASE__       ).name					,	_name="xmod_base"					,	arch="xmod_base"					,	task="multilingual_masked_lm"					,	data_name_or_path=str(SCREAMING_SNAKE_CASE__       )					,	bpe="sentencepiece"					,	sentencepiece_model=str(Path(SCREAMING_SNAKE_CASE__       ).parent / "sentencepiece.bpe.model"       )					,	src_dict=str(data_dir / "dict.txt"       )					,	)
  xmod.eval()  # disable dropout
  print(SCREAMING_SNAKE_CASE__       )

  _snake_case     		=				xmod.model.encoder.sentence_encoder
  _snake_case     		=				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:
    _snake_case     		=				xmod.model.classification_heads["mnli"].out_proj.weight.shape[0]

  print("Our X-MOD config:"					,	SCREAMING_SNAKE_CASE__       )

  _snake_case     		=				XmodForSequenceClassification(SCREAMING_SNAKE_CASE__       ) if classification_head else XmodForMaskedLM(SCREAMING_SNAKE_CASE__       )
  model.eval()

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

  _snake_case     		=				xmod_sent_encoder.layernorm_embedding.weight
  _snake_case     		=				xmod_sent_encoder.layernorm_embedding.bias

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

    # self attention
    _snake_case     		=				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."       )

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

    # self-attention output
    _snake_case     		=				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."       )
    _snake_case     		=				xmod_layer.self_attn.out_proj.weight
    _snake_case     		=				xmod_layer.self_attn.out_proj.bias
    _snake_case     		=				xmod_layer.self_attn_layer_norm.weight
    _snake_case     		=				xmod_layer.self_attn_layer_norm.bias

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

    # output
    _snake_case     		=				layer.output
    if bert_output.dense.weight.shape != xmod_layer.fca.weight.shape:
      raise AssertionError("Dimensions of feed-forward weights do not match."       )
    _snake_case     		=				xmod_layer.fca.weight
    _snake_case     		=				xmod_layer.fca.bias
    _snake_case     		=				xmod_layer.final_layer_norm.weight
    _snake_case     		=				xmod_layer.final_layer_norm.bias
    if bert_output.adapter_layer_norm is not None:
      _snake_case     		=				xmod_layer.adapter_layer_norm.weight
      _snake_case     		=				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():
      _snake_case     		=				bert_output.adapter_modules[lang_code]
      _snake_case     		=				xmod_layer.adapter_modules[lang_code]
      _snake_case     		=				from_adapter.fca.weight
      _snake_case     		=				from_adapter.fca.bias
      _snake_case     		=				from_adapter.fca.weight
      _snake_case     		=				from_adapter.fca.bias

    # end of layer

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

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

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

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

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


if __name__ == "__main__":
 __magic_name__   :						str    =			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."""
 )
 __magic_name__   :						List[Any]    =			parser.parse_args()
 convert_xmod_checkpoint_to_pytorch(
     args.xmod_checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
 )


