


import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE			: str				   =					logging.get_logger(__name__)

SCREAMING_SNAKE_CASE			: Dict				   =					{
    "post_extract_proj": "feature_projection.projection",
    "encoder.pos_conv.0": "encoder.pos_conv_embed.conv",
    "self_attn.k_proj": "encoder.layers.*.attention.k_proj",
    "self_attn.v_proj": "encoder.layers.*.attention.v_proj",
    "self_attn.q_proj": "encoder.layers.*.attention.q_proj",
    "self_attn.out_proj": "encoder.layers.*.attention.out_proj",
    "self_attn_layer_norm": "encoder.layers.*.layer_norm",
    "fc1": "encoder.layers.*.feed_forward.intermediate_dense",
    "fc2": "encoder.layers.*.feed_forward.output_dense",
    "final_layer_norm": "encoder.layers.*.final_layer_norm",
    "encoder.layer_norm": "encoder.layer_norm",
    "w2v_model.layer_norm": "feature_projection.layer_norm",
    "quantizer.weight_proj": "quantizer.weight_proj",
    "quantizer.vars": "quantizer.codevectors",
    "project_q": "project_q",
    "final_proj": "project_hid",
    "w2v_encoder.proj": "lm_head",
    "mask_emb": "masked_spec_embed",
}
SCREAMING_SNAKE_CASE			: Optional[int]				   =					[
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]



def 							UpperCamelCase_( lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)       ->		Dict:
    for attribute in key.split('.'					):
        _lowercase			:				Optional[int] 						=						getattr(lowerCamelCase_				,					lowerCamelCase_					)

    if weight_type is not None:
        _lowercase			:				Tuple 						=						getattr(lowerCamelCase_				,					lowerCamelCase_					).shape
    else:
        _lowercase			:				int 						=						hf_pointer.shape

    assert hf_shape == value.shape, (
        F'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
        F''' {value.shape} for {full_name}'''
    )

    if weight_type == "weight":
        _lowercase			:				Union[str, Any] 						=						value
    elif weight_type == "weight_g":
        _lowercase			:				str 						=						value
    elif weight_type == "weight_v":
        _lowercase			:				Optional[int] 						=						value
    elif weight_type == "bias":
        _lowercase			:				str 						=						value
    else:
        _lowercase			:				Union[str, Any] 						=						value

    logger.info(F'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''					)



def 							UpperCamelCase_( lowerCamelCase_				,					lowerCamelCase_					)       ->		Any:
    _lowercase			:				str 						=						[]
    _lowercase			:				List[Any] 						=						fairseq_model.state_dict()

    _lowercase			:				Dict 						=						hf_model.feature_extractor
    _lowercase			:				Tuple 						=						hf_model.adapter

    for name, value in fairseq_dict.items():
        _lowercase			:				List[Any] 						=						False
        if "conv_layers" in name:
            load_conv_layer(
                lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					hf_model.config.feat_extract_norm == 'group'				,					)
            _lowercase			:				Optional[Any] 						=						True
        elif any(x in name for x in ['adaptor', 'w2v_encoder.proj.', 'w2v_proj_ln.']					):
            load_adapter(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
            _lowercase			:				Optional[int] 						=						True
        else:
            for key, mapped_key in MAPPING.items():
                if key in name or key.split('w2v_model.'					)[-1] == name.split('.'					)[0]:
                    _lowercase			:				Tuple 						=						True
                    if "*" in mapped_key:
                        _lowercase			:				Optional[Any] 						=						name.split(lowerCamelCase_					)[0].split('.'					)[-2]
                        _lowercase			:				Dict 						=						mapped_key.replace('*'				,					lowerCamelCase_					)
                    if "weight_g" in name:
                        _lowercase			:				List[Any] 						=						'weight_g'
                    elif "weight_v" in name:
                        _lowercase			:				Tuple 						=						'weight_v'
                    elif "bias" in name:
                        _lowercase			:				Dict 						=						'bias'
                    elif "weight" in name:
                        _lowercase			:				Optional[Any] 						=						'weight'
                    else:
                        _lowercase			:				str 						=						None
                    set_recursively(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
                continue
        if not is_used:
            unused_weights.append(lowerCamelCase_					)

    logger.warning(F'''Unused weights: {unused_weights}'''					)



def 							UpperCamelCase_( lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)       ->		Union[str, Any]:
    _lowercase			:				int 						=						full_name.split('conv_layers.'					)[-1]
    _lowercase			:				str 						=						name.split('.'					)
    _lowercase			:				str 						=						int(items[0]					)
    _lowercase			:				List[str] 						=						int(items[1]					)

    if type_id == 0:
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
                F'''{full_name} has size {value.shape}, but'''
                F''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
            )
            _lowercase			:				int 						=						value
            logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''					)
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
                F'''{full_name} has size {value.shape}, but'''
                F''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
            )
            _lowercase			:				Optional[Any] 						=						value
            logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''					)
    elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
                F'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
                " found."
            )
            _lowercase			:				Union[str, Any] 						=						value
            logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''					)
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
                F'''{full_name} has size {value.shape}, but'''
                F''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
            )
            _lowercase			:				Any 						=						value
            logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''					)
    else:
        unused_weights.append(lowerCamelCase_					)



def 							UpperCamelCase_( lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)       ->		Dict:
    _lowercase			:				Optional[Any] 						=						full_name.split('adaptor.'					)[-1]
    _lowercase			:				List[str] 						=						name.split('.'					)

    if items[1].isdigit():
        _lowercase			:				Any 						=						int(items[1]					)
    else:
        _lowercase			:				Optional[int] 						=						None

    if "adaptor" not in full_name:
        if "proj_ln" in full_name:
            # has to be layer norm
            if "bias" in name:
                assert (
                    value.shape == adapter.proj_layer_norm.bias.data.shape
                ), F'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'''
                _lowercase			:				Dict 						=						value
                logger.info(F'''Adapter proj layer norm bias was initialized from {full_name}.'''					)
            if "weight" in name:
                assert (
                    value.shape == adapter.proj_layer_norm.weight.data.shape
                ), F'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'''
                _lowercase			:				List[Any] 						=						value
        else:
            # has to be projection layer
            if "bias" in name:
                assert (
                    value.shape == adapter.proj.bias.data.shape
                ), F'''{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'''
                _lowercase			:				List[str] 						=						value
                logger.info(F'''Adapter proj layer bias was initialized from {full_name}.'''					)
            if "weight" in name:
                assert (
                    value.shape == adapter.proj.weight.data.shape
                ), F'''{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'''
                _lowercase			:				Any 						=						value
                logger.info(F'''Adapter proj layer weight was initialized from {full_name}.'''					)
    elif isinstance(lowerCamelCase_				,					lowerCamelCase_					):
        if "bias" in name:
            assert (
                value.shape == adapter.layers[layer_id].conv.bias.data.shape
            ), F'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'''
            _lowercase			:				List[str] 						=						value
            logger.info(F'''Adapter layer {layer_id} bias was initialized from {full_name}.'''					)
        elif "weight" in name:
            assert (
                value.shape == adapter.layers[layer_id].conv.weight.data.shape
            ), F'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'''
            _lowercase			:				str 						=						value
            logger.info(F'''Adapter layer {layer_id} bias was initialized from {full_name}.'''					)
    else:
        unused_weights.append(lowerCamelCase_					)



def 							UpperCamelCase_( lowerCamelCase_					)       ->		List[str]:
    _lowercase						,							_lowercase			:				List[str] 						=						emb.weight.shape
    _lowercase			:				Tuple 						=						nn.Linear(lowerCamelCase_				,					lowerCamelCase_				,					bias=lowerCamelCase_					)
    _lowercase			:				List[Any] 						=						emb.weight.data
    return lin_layer



@torch.no_grad()
def 							UpperCamelCase_( lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					)       ->		Union[str, Any]:
    _lowercase			:				Any 						=						WavaVecaConfig.from_pretrained(
        lowerCamelCase_				,					add_adapter=lowerCamelCase_				,					adapter_stride=lowerCamelCase_				,					adapter_kernel_size=lowerCamelCase_				,					use_auth_token=lowerCamelCase_				,					output_hidden_size=lowerCamelCase_				,					)
    _lowercase			:				int 						=						MBartConfig.from_pretrained(lowerCamelCase_					)

    # load model
    _lowercase						,							_lowercase						,							_lowercase			:				Dict 						=						fairseq.checkpoint_utils.load_model_ensemble_and_task(
        [checkpoint_path]				,					arg_overrides={
            'config_yaml': config_yaml_path,
            'data': '/'.join(dict_path.split('/'					)[:-1]					),
            'w2v_path': checkpoint_path,
            'load_pretrained_decoder_from': None,
        }				,					)
    _lowercase			:				Dict 						=						model[0].eval()

    # load feature extractor
    _lowercase			:				Tuple 						=						WavaVecaFeatureExtractor.from_pretrained(lowerCamelCase_				,					use_auth_token=lowerCamelCase_					)

    # set weights for wav2vec2 encoder
    _lowercase			:				str 						=						WavaVecaModel(lowerCamelCase_					)

    recursively_load_weights_wavaveca(model.encoder				,					lowerCamelCase_					)

    # load decoder weights
    _lowercase			:				List[str] 						=						MBartForCausalLM(lowerCamelCase_					)
    _lowercase						,							_lowercase			:				Optional[int] 						=						hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()				,					strict=lowerCamelCase_					)
    logger.warning(F'''The following keys are missing when loading the decoder weights: {missing_keys}'''					)
    logger.warning(F'''The following keys are unexpected when loading the decoder weights: {unexpected_keys}'''					)

    _lowercase			:				Optional[Any] 						=						SpeechEncoderDecoderModel(encoder=lowerCamelCase_				,					decoder=lowerCamelCase_					)
    _lowercase			:				int 						=						False

    _lowercase			:				str 						=						MBartaaTokenizer(lowerCamelCase_					)
    tokenizer.save_pretrained(lowerCamelCase_					)

    _lowercase			:				Tuple 						=						hf_wavavec.config.to_dict()
    _lowercase			:				int 						=						tokenizer.pad_token_id
    _lowercase			:				Optional[Any] 						=						tokenizer.bos_token_id
    _lowercase			:				int 						=						tokenizer.eos_token_id
    _lowercase			:				Tuple 						=						'mbart50'
    _lowercase			:				Tuple 						=						'wav2vec2'

    _lowercase			:				Dict 						=						tokenizer.eos_token_id
    _lowercase			:				List[Any] 						=						25_0004
    _lowercase			:				int 						=						tokenizer.eos_token_id

    _lowercase			:				Any 						=						SpeechEncoderDecoderConfig.from_dict(lowerCamelCase_					)

    hf_wavavec.save_pretrained(lowerCamelCase_					)
    feature_extractor.save_pretrained(lowerCamelCase_					)


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE			: Union[str, Any]				   =					argparse.ArgumentParser()
 parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
 parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint")
 parser.add_argument("--dict_path", default=None, type=str, help="Path to dict of fine-tuned model")
 parser.add_argument("--config_yaml_path", default=None, type=str, help="Path to yaml file of fine-tuned model")
 parser.add_argument(
     "--encoder_config_path",
     default="facebook/wav2vec2-xls-r-1b",
     type=str,
     help="Path to hf encoder wav2vec2 checkpoint config",
 )
 parser.add_argument(
     "--decoder_config_path",
     default="facebook/mbart-large-50-one-to-many-mmt",
     type=str,
     help="Path to hf decoder checkpoint config",
 )
 parser.add_argument("--add_adapter", default=True, type=bool, help="whethere to add model adapter layers")
 parser.add_argument("--adapter_stride", default=2, type=int, help="stride of adapter layers")
 parser.add_argument("--adapter_kernel_size", default=3, type=int, help="kernel size of adapter layers")
 parser.add_argument("--encoder_output_dim", default=1024, type=int, help="encoder output dim")
 parser.add_argument("--start_token_id", default=250004, type=int, help="`decoder_start_token_id` of model config")

 SCREAMING_SNAKE_CASE			: Dict				   =					parser.parse_args()
 convert_wavaveca_checkpoint(
     args.checkpoint_path,
     args.pytorch_dump_folder_path,
     args.dict_path,
     args.config_yaml_path,
     encoder_config_path=args.encoder_config_path,
     decoder_config_path=args.decoder_config_path,
     add_adapter=args.add_adapter,
     adapter_kernel_size=args.adapter_kernel_size,
     adapter_stride=args.adapter_stride,
     decoder_start_token_id=args.start_token_id,
     encoder_output_dim=args.encoder_output_dim,
 )






import unittest

from transformers import XLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     XLMForMultipleChoice,
     XLMForQuestionAnswering,
     XLMForQuestionAnsweringSimple,
     XLMForSequenceClassification,
     XLMForTokenClassification,
     XLMModel,
     XLMWithLMHeadModel,
 )
 from transformers.models.xlm.modeling_xlm import XLM_PRETRAINED_MODEL_ARCHIVE_LIST






class _lowerCamelCase:
       def __init__(   self, lowerCamelCase, lowerCamelCase=13, lowerCamelCase=7, lowerCamelCase=True, lowerCamelCase=True, lowerCamelCase=True, lowerCamelCase=True, lowerCamelCase=True, lowerCamelCase=False, lowerCamelCase=False, lowerCamelCase=False, lowerCamelCase=2, lowerCamelCase=99, lowerCamelCase=0, lowerCamelCase=32, lowerCamelCase=5, lowerCamelCase=4, lowerCamelCase=0.1, lowerCamelCase=0.1, lowerCamelCase=5_12, lowerCamelCase=2, lowerCamelCase=0.0_2, lowerCamelCase=2, lowerCamelCase=4, lowerCamelCase="last", lowerCamelCase=True, lowerCamelCase=None, lowerCamelCase=0, )			->  str:




           """simple docstring"""

           _lowercase			:				Union[str, Any] 						=						parent
           _lowercase			:				Optional[Any] 						=						batch_size
           _lowercase			:				List[str] 						=						seq_length
           _lowercase			:				int 						=						is_training
           _lowercase			:				List[str] 						=						use_input_lengths
           _lowercase			:				int 						=						use_token_type_ids
           _lowercase			:				Any 						=						use_labels
           _lowercase			:				Union[str, Any] 						=						gelu_activation
           _lowercase			:				List[str] 						=						sinusoidal_embeddings
           _lowercase			:				str 						=						causal
           _lowercase			:				Optional[int] 						=						asm
           _lowercase			:				Union[str, Any] 						=						n_langs
           _lowercase			:				List[Any] 						=						vocab_size
           _lowercase			:				Any 						=						n_special
           _lowercase			:				Any 						=						hidden_size
           _lowercase			:				str 						=						num_hidden_layers
           _lowercase			:				Union[str, Any] 						=						num_attention_heads
           _lowercase			:				Tuple 						=						hidden_dropout_prob
           _lowercase			:				Optional[int] 						=						attention_probs_dropout_prob
           _lowercase			:				Union[str, Any] 						=						max_position_embeddings
           _lowercase			:				List[str] 						=						type_sequence_label_size
           _lowercase			:				Any 						=						initializer_range
           _lowercase			:				int 						=						num_labels
           _lowercase			:				Optional[int] 						=						num_choices
           _lowercase			:				Optional[Any] 						=						summary_type
           _lowercase			:				Optional[Any] 						=						use_proj
           _lowercase			:				int 						=						scope
           _lowercase			:				List[Any] 						=						bos_token_id
       def      UpperCamelCase	(   self)			->  Optional[int]:




           """simple docstring"""

           _lowercase			:				Tuple 						=						ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
           _lowercase			:				List[str] 						=						random_attention_mask([self.batch_size, self.seq_length])

           _lowercase			:				int 						=						None
           if self.use_input_lengths:
               _lowercase			:				Dict 						=						(
                   ids_tensor([self.batch_size], vocab_size=2) + self.seq_length - 2
               )  # small variation of seq_length

           _lowercase			:				Tuple 						=						None
           if self.use_token_type_ids:
               _lowercase			:				Optional[Any] 						=						ids_tensor([self.batch_size, self.seq_length], self.n_langs)

           _lowercase			:				Tuple 						=						None
           _lowercase			:				int 						=						None
           _lowercase			:				int 						=						None
           if self.use_labels:
               _lowercase			:				str 						=						ids_tensor([self.batch_size], self.type_sequence_label_size)
               _lowercase			:				str 						=						ids_tensor([self.batch_size, self.seq_length], self.num_labels)
               _lowercase			:				Dict 						=						ids_tensor([self.batch_size], 2).float()
               _lowercase			:				Tuple 						=						ids_tensor([self.batch_size], self.num_choices)

           _lowercase			:				Dict 						=						self.get_config()

           return (
               config,
               input_ids,
               token_type_ids,
               input_lengths,
               sequence_labels,
               token_labels,
               is_impossible_labels,
               choice_labels,
               input_mask,
           )
       def      UpperCamelCase	(   self)			->  Union[str, Any]:




           """simple docstring"""

           return XLMConfig(
               vocab_size=self.vocab_size, n_special=self.n_special, emb_dim=self.hidden_size, n_layers=self.num_hidden_layers, n_heads=self.num_attention_heads, dropout=self.hidden_dropout_prob, attention_dropout=self.attention_probs_dropout_prob, gelu_activation=self.gelu_activation, sinusoidal_embeddings=self.sinusoidal_embeddings, asm=self.asm, causal=self.causal, n_langs=self.n_langs, max_position_embeddings=self.max_position_embeddings, initializer_range=self.initializer_range, summary_type=self.summary_type, use_proj=self.use_proj, num_labels=self.num_labels, bos_token_id=self.bos_token_id, )
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, )			->  Tuple:




           """simple docstring"""

           _lowercase			:				List[Any] 						=						XLMModel(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase			:				str 						=						model(lowerCamelCase, lengths=lowerCamelCase, langs=lowerCamelCase)
           _lowercase			:				int 						=						model(lowerCamelCase, langs=lowerCamelCase)
           _lowercase			:				Any 						=						model(lowerCamelCase)
           self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, )			->  List[Any]:




           """simple docstring"""

           _lowercase			:				Dict 						=						XLMWithLMHeadModel(lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()

           _lowercase			:				int 						=						model(lowerCamelCase, token_type_ids=lowerCamelCase, labels=lowerCamelCase)
           self.parent.assertEqual(result.loss.shape, ())
           self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, )			->  str:




           """simple docstring"""

           _lowercase			:				Tuple 						=						XLMForQuestionAnsweringSimple(lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()

           _lowercase			:				Dict 						=						model(lowerCamelCase)

           _lowercase			:				List[str] 						=						model(lowerCamelCase, start_positions=lowerCamelCase, end_positions=lowerCamelCase)
           _lowercase			:				Any 						=						outputs
           self.parent.assertEqual(result.start_logits.shape, (self.batch_size, self.seq_length))
           self.parent.assertEqual(result.end_logits.shape, (self.batch_size, self.seq_length))
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, )			->  Union[str, Any]:




           """simple docstring"""

           _lowercase			:				Tuple 						=						XLMForQuestionAnswering(lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()

           _lowercase			:				Optional[Any] 						=						model(lowerCamelCase)

           _lowercase			:				List[Any] 						=						model(
               lowerCamelCase, start_positions=lowerCamelCase, end_positions=lowerCamelCase, cls_index=lowerCamelCase, is_impossible=lowerCamelCase, p_mask=lowerCamelCase, )

           _lowercase			:				List[str] 						=						model(
               lowerCamelCase, start_positions=lowerCamelCase, end_positions=lowerCamelCase, cls_index=lowerCamelCase, is_impossible=lowerCamelCase, )

           ((_lowercase)						,							)			:				Optional[Any] 						=						result_with_labels.to_tuple()

           _lowercase			:				List[str] 						=						model(lowerCamelCase, start_positions=lowerCamelCase, end_positions=lowerCamelCase)

           ((_lowercase)						,							)			:				Any 						=						result_with_labels.to_tuple()

           self.parent.assertEqual(result_with_labels.loss.shape, ())
           self.parent.assertEqual(result.start_top_log_probs.shape, (self.batch_size, model.config.start_n_top))
           self.parent.assertEqual(result.start_top_index.shape, (self.batch_size, model.config.start_n_top))
           self.parent.assertEqual(
               result.end_top_log_probs.shape, (self.batch_size, model.config.start_n_top * model.config.end_n_top))
           self.parent.assertEqual(
               result.end_top_index.shape, (self.batch_size, model.config.start_n_top * model.config.end_n_top))
           self.parent.assertEqual(result.cls_logits.shape, (self.batch_size,))
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, )			->  int:




           """simple docstring"""

           _lowercase			:				Optional[Any] 						=						XLMForSequenceClassification(lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()

           _lowercase			:				Optional[int] 						=						model(lowerCamelCase)
           _lowercase			:				Optional[int] 						=						model(lowerCamelCase, labels=lowerCamelCase)
           self.parent.assertEqual(result.loss.shape, ())
           self.parent.assertEqual(result.logits.shape, (self.batch_size, self.type_sequence_label_size))
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, )			->  List[str]:




           """simple docstring"""

           _lowercase			:				Any 						=						self.num_labels
           _lowercase			:				str 						=						XLMForTokenClassification(lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()

           _lowercase			:				int 						=						model(lowerCamelCase, attention_mask=lowerCamelCase, labels=lowerCamelCase)
           self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, )			->  Dict:




           """simple docstring"""

           _lowercase			:				Optional[Any] 						=						self.num_choices
           _lowercase			:				Optional[int] 						=						XLMForMultipleChoice(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase			:				Optional[Any] 						=						input_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
           _lowercase			:				int 						=						token_type_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
           _lowercase			:				Optional[Any] 						=						input_mask.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
           _lowercase			:				List[str] 						=						model(
               lowerCamelCase, attention_mask=lowerCamelCase, token_type_ids=lowerCamelCase, labels=lowerCamelCase, )
           self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_choices))




       def      UpperCamelCase	(   self)			->  Tuple:




           """simple docstring"""

           _lowercase			:				Dict 						=						self.prepare_config_and_inputs()
           (
               (
               _lowercase
           )						,							(
               _lowercase
           )						,							(
               _lowercase
           )						,							(
               _lowercase
           )						,							(
               _lowercase
           )						,							(
               _lowercase
           )						,							(
               _lowercase
           )						,							(
               _lowercase
           )						,							(
               _lowercase
           )						,							
           )			:				Optional[Any] 						=						config_and_inputs
           _lowercase			:				List[str] 						=						{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'lengths': input_lengths}
           return config, inputs_dict






@require_torch
class _lowerCamelCase(		_a,					_a,					_a,					unittest.TestCase  ):
       lowercase_					:      Any       					=     (
           (
               XLMModel,
               XLMWithLMHeadModel,
               XLMForQuestionAnswering,
               XLMForSequenceClassification,
               XLMForQuestionAnsweringSimple,
               XLMForTokenClassification,
               XLMForMultipleChoice,
           )
           if is_torch_available()
           else ()
       )
       lowercase_					:      Optional[int]       					=     (
           (XLMWithLMHeadModel,) if is_torch_available() else ()
       )  # TODO (PVP): Check other models whether language generation is also applicable
       lowercase_					:      Union[str, Any]       					=     (
           {
               """feature-extraction""": XLMModel,
               """fill-mask""": XLMWithLMHeadModel,
               """question-answering""": XLMForQuestionAnsweringSimple,
               """text-classification""": XLMForSequenceClassification,
               """text-generation""": XLMWithLMHeadModel,
               """token-classification""": XLMForTokenClassification,
               """zero-shot""": XLMForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase)			->  Optional[int]:




           """simple docstring"""

           if (
               pipeline_test_casse_name == "QAPipelineTests"
               and tokenizer_name is not None
               and not tokenizer_name.endswith('Fast')
           ):
               # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
               # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
               # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
               return True

           return False
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase=False)			->  Optional[int]:




           """simple docstring"""

           _lowercase			:				Any 						=						super()._prepare_for_class(lowerCamelCase, lowerCamelCase, return_labels=lowerCamelCase)

           if return_labels:
               if model_class.__name__ == "XLMForQuestionAnswering":
                   _lowercase			:				Any 						=						torch.zeros(
                       self.model_tester.batch_size, dtype=torch.long, device=lowerCamelCase)
                   _lowercase			:				Dict 						=						torch.zeros(
                       self.model_tester.batch_size, dtype=torch.long, device=lowerCamelCase)

           return inputs_dict
       def      UpperCamelCase	(   self)			->  int:




           """simple docstring"""

           _lowercase			:				Union[str, Any] 						=						XLMModelTester(self)
           _lowercase			:				List[str] 						=						ConfigTester(self, config_class=lowerCamelCase, emb_dim=37)
       def      UpperCamelCase	(   self)			->  Tuple:




           """simple docstring"""

           self.config_tester.run_common_tests()
       def      UpperCamelCase	(   self)			->  List[Any]:




           """simple docstring"""

           _lowercase			:				Optional[Any] 						=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_xlm_model(*lowerCamelCase)
       def      UpperCamelCase	(   self)			->  List[str]:




           """simple docstring"""

           _lowercase			:				str 						=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_xlm_lm_head(*lowerCamelCase)
       def      UpperCamelCase	(   self)			->  Optional[Any]:




           """simple docstring"""

           _lowercase			:				Optional[Any] 						=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_xlm_simple_qa(*lowerCamelCase)
       def      UpperCamelCase	(   self)			->  int:




           """simple docstring"""

           _lowercase			:				str 						=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_xlm_qa(*lowerCamelCase)
       def      UpperCamelCase	(   self)			->  List[str]:




           """simple docstring"""

           _lowercase			:				Dict 						=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_xlm_sequence_classif(*lowerCamelCase)
       def      UpperCamelCase	(   self)			->  Dict:




           """simple docstring"""

           _lowercase			:				Union[str, Any] 						=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_xlm_token_classif(*lowerCamelCase)
       def      UpperCamelCase	(   self)			->  Any:




           """simple docstring"""

           _lowercase			:				str 						=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_xlm_for_multiple_choice(*lowerCamelCase)
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase=False, lowerCamelCase=1)			->  int:




           """simple docstring"""

           self.assertIsInstance(lowerCamelCase, lowerCamelCase)
           self.assertListEqual(
               [isinstance(lowerCamelCase, lowerCamelCase) for iter_attentions in attentions], [True] * len(lowerCamelCase))
           self.assertEqual(len(lowerCamelCase), (max_length - min_length) * num_beam_groups)

           for idx, iter_attentions in enumerate(lowerCamelCase):
               # adds PAD dummy token
               _lowercase			:				Dict 						=						min_length + idx + 1
               _lowercase			:				int 						=						min_length + idx + 1

               _lowercase			:				Dict 						=						(
                   batch_size * num_beam_groups,
                   config.num_attention_heads,
                   tgt_len,
                   src_len,
               )
               # check attn size
               self.assertListEqual(
                   [layer_attention.shape for layer_attention in iter_attentions], [expected_shape] * len(lowerCamelCase))
       def      UpperCamelCase	(   self, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase, lowerCamelCase=False, lowerCamelCase=1)			->  List[Any]:




           """simple docstring"""

           self.assertIsInstance(lowerCamelCase, lowerCamelCase)
           self.assertListEqual(
               [isinstance(lowerCamelCase, lowerCamelCase) for iter_hidden_states in hidden_states], [True] * len(lowerCamelCase), )
           self.assertEqual(len(lowerCamelCase), (max_length - min_length) * num_beam_groups)

           for idx, iter_hidden_states in enumerate(lowerCamelCase):
               # adds PAD dummy token
               _lowercase			:				int 						=						min_length + idx + 1
               _lowercase			:				int 						=						(batch_size * num_beam_groups, seq_len, config.hidden_size)
               # check hidden size
               self.assertListEqual(
                   [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states], [expected_shape] * len(lowerCamelCase), )
           pass




       @slow
       def      UpperCamelCase	(   self)			->  int:




           """simple docstring"""

           for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _lowercase			:				Dict 						=						XLMModel.from_pretrained(lowerCamelCase)
               self.assertIsNotNone(lowerCamelCase)






@require_torch
class _lowerCamelCase(		unittest.TestCase  ):
       @slow
       def      UpperCamelCase	(   self)			->  Union[str, Any]:




           """simple docstring"""

           _lowercase			:				Tuple 						=						XLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048')
           model.to(lowerCamelCase)
           _lowercase			:				Union[str, Any] 						=						torch.tensor([[14, 4_47]], dtype=torch.long, device=lowerCamelCase)  # the president
           _lowercase			:				Any 						=						[
               14,
               4_47,
               14,
               4_47,
               14,
               4_47,
               14,
               4_47,
               14,
               4_47,
               14,
               4_47,
               14,
               4_47,
               14,
               4_47,
               14,
               4_47,
               14,
               4_47,
           ]  # the president the president the president the president the president the president the president the president the president the president
           # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
           _lowercase			:				str 						=						model.generate(lowerCamelCase, do_sample=lowerCamelCase)
           self.assertListEqual(output_ids[0].cpu().numpy().tolist(), lowerCamelCase)


