


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  snake_case_     :
					def __init__(   self     :int    ,__snake_case     :str    ,__snake_case     :int=13    ,__snake_case     :str=7    ,__snake_case     :Any=True    ,__snake_case     :str=True    ,__snake_case     :Dict=True    ,__snake_case     :int=True    ,__snake_case     :Dict=True    ,__snake_case     :Optional[Any]=False    ,__snake_case     :Optional[Any]=False    ,__snake_case     :Optional[int]=False    ,__snake_case     :Dict=2    ,__snake_case     :Any=99    ,__snake_case     :Optional[Any]=0    ,__snake_case     :Tuple=32    ,__snake_case     :Optional[Any]=5    ,__snake_case     :Dict=4    ,__snake_case     :int=0.1    ,__snake_case     :Dict=0.1    ,__snake_case     :int=5_12    ,__snake_case     :int=2    ,__snake_case     :Optional[int]=0.02    ,__snake_case     :Optional[int]=2    ,__snake_case     :Any=4    ,__snake_case     :Tuple="last"    ,__snake_case     :Optional[Any]=True    ,__snake_case     :str=None    ,__snake_case     :List[Any]=0    ,)      ->					Optional[Any]:
												a__    =  parent
												a__    =  batch_size
												a__    =  seq_length
												a__    =  is_training
												a__    =  use_input_lengths
												a__    =  use_token_type_ids
												a__    =  use_labels
												a__    =  gelu_activation
												a__    =  sinusoidal_embeddings
												a__    =  causal
												a__    =  asm
												a__    =  n_langs
												a__    =  vocab_size
												a__    =  n_special
												a__    =  hidden_size
												a__    =  num_hidden_layers
												a__    =  num_attention_heads
												a__    =  hidden_dropout_prob
												a__    =  attention_probs_dropout_prob
												a__    =  max_position_embeddings
												a__    =  type_sequence_label_size
												a__    =  initializer_range
												a__    =  num_labels
												a__    =  num_choices
												a__    =  summary_type
												a__    =  use_proj
												a__    =  scope
												a__    =  bos_token_id
					def 				lowerCamelCase__(   self     :str   )      ->					List[Any]:
												a__    =  ids_tensor([self.batch_size, self.seq_length]    ,self.vocab_size   )
												a__    =  random_attention_mask([self.batch_size, self.seq_length]   )

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

												a__    =  None
												if self.use_token_type_ids:
																			a__    =  ids_tensor([self.batch_size, self.seq_length]    ,self.n_langs   )

												a__    =  None
												a__    =  None
												a__    =  None
												if self.use_labels:
																			a__    =  ids_tensor([self.batch_size]    ,self.type_sequence_label_size   )
																			a__    =  ids_tensor([self.batch_size, self.seq_length]    ,self.num_labels   )
																			a__    =  ids_tensor([self.batch_size]    ,2   ).float()
																			a__    =  ids_tensor([self.batch_size]    ,self.num_choices   )

												a__    =  self.get_config()

												return (
												    config,
												    input_ids,
												    token_type_ids,
												    input_lengths,
												    sequence_labels,
												    token_labels,
												    is_impossible_labels,
												    choice_labels,
												    input_mask,
												)
					def 				lowerCamelCase__(   self     :Optional[Any]   )      ->					Dict:
												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 				lowerCamelCase__(   self     :Optional[Any]    ,__snake_case     :List[Any]    ,__snake_case     :Optional[int]    ,__snake_case     :Union[str, Any]    ,__snake_case     :List[str]    ,__snake_case     :Optional[int]    ,__snake_case     :Union[str, Any]    ,__snake_case     :Optional[int]    ,__snake_case     :List[str]    ,__snake_case     :Optional[Any]    ,)      ->					str:
												a__    =  XLMModel(config=__snake_case   )
												model.to(__snake_case   )
												model.eval()
												a__    =  model(__snake_case    ,lengths=__snake_case    ,langs=__snake_case   )
												a__    =  model(__snake_case    ,langs=__snake_case   )
												a__    =  model(__snake_case   )
												self.parent.assertEqual(result.last_hidden_state.shape    ,(self.batch_size, self.seq_length, self.hidden_size)   )
					def 				lowerCamelCase__(   self     :Dict    ,__snake_case     :Dict    ,__snake_case     :Optional[int]    ,__snake_case     :int    ,__snake_case     :List[str]    ,__snake_case     :Tuple    ,__snake_case     :Union[str, Any]    ,__snake_case     :str    ,__snake_case     :Tuple    ,__snake_case     :str    ,)      ->					str:
												a__    =  XLMWithLMHeadModel(__snake_case   )
												model.to(__snake_case   )
												model.eval()

												a__    =  model(__snake_case    ,token_type_ids=__snake_case    ,labels=__snake_case   )
												self.parent.assertEqual(result.loss.shape    ,()   )
												self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.seq_length, self.vocab_size)   )
					def 				lowerCamelCase__(   self     :Optional[int]    ,__snake_case     :Union[str, Any]    ,__snake_case     :Any    ,__snake_case     :str    ,__snake_case     :Tuple    ,__snake_case     :Union[str, Any]    ,__snake_case     :List[str]    ,__snake_case     :Optional[Any]    ,__snake_case     :Optional[Any]    ,__snake_case     :List[str]    ,)      ->					List[str]:
												a__    =  XLMForQuestionAnsweringSimple(__snake_case   )
												model.to(__snake_case   )
												model.eval()

												a__    =  model(__snake_case   )

												a__    =  model(__snake_case    ,start_positions=__snake_case    ,end_positions=__snake_case   )
												a__    =  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 				lowerCamelCase__(   self     :Dict    ,__snake_case     :Optional[int]    ,__snake_case     :Optional[Any]    ,__snake_case     :List[Any]    ,__snake_case     :Tuple    ,__snake_case     :Dict    ,__snake_case     :Tuple    ,__snake_case     :Any    ,__snake_case     :Optional[int]    ,__snake_case     :Union[str, Any]    ,)      ->					Union[str, Any]:
												a__    =  XLMForQuestionAnswering(__snake_case   )
												model.to(__snake_case   )
												model.eval()

												a__    =  model(__snake_case   )

												a__    =  model(
												    __snake_case    ,start_positions=__snake_case    ,end_positions=__snake_case    ,cls_index=__snake_case    ,is_impossible=__snake_case    ,p_mask=__snake_case    ,)

												a__    =  model(
												    __snake_case    ,start_positions=__snake_case    ,end_positions=__snake_case    ,cls_index=__snake_case    ,is_impossible=__snake_case    ,)

												(a__    )    =  result_with_labels.to_tuple()

												a__    =  model(__snake_case    ,start_positions=__snake_case    ,end_positions=__snake_case   )

												(a__    )    =  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 				lowerCamelCase__(   self     :Optional[Any]    ,__snake_case     :str    ,__snake_case     :int    ,__snake_case     :List[str]    ,__snake_case     :Union[str, Any]    ,__snake_case     :Union[str, Any]    ,__snake_case     :Optional[int]    ,__snake_case     :List[Any]    ,__snake_case     :List[str]    ,__snake_case     :Union[str, Any]    ,)      ->					Union[str, Any]:
												a__    =  XLMForSequenceClassification(__snake_case   )
												model.to(__snake_case   )
												model.eval()

												a__    =  model(__snake_case   )
												a__    =  model(__snake_case    ,labels=__snake_case   )
												self.parent.assertEqual(result.loss.shape    ,()   )
												self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.type_sequence_label_size)   )
					def 				lowerCamelCase__(   self     :List[Any]    ,__snake_case     :Any    ,__snake_case     :Dict    ,__snake_case     :Dict    ,__snake_case     :Any    ,__snake_case     :List[Any]    ,__snake_case     :Any    ,__snake_case     :List[str]    ,__snake_case     :str    ,__snake_case     :Optional[Any]    ,)      ->					Tuple:
												a__    =  self.num_labels
												a__    =  XLMForTokenClassification(__snake_case   )
												model.to(__snake_case   )
												model.eval()

												a__    =  model(__snake_case    ,attention_mask=__snake_case    ,labels=__snake_case   )
												self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.seq_length, self.num_labels)   )
					def 				lowerCamelCase__(   self     :Any    ,__snake_case     :Optional[int]    ,__snake_case     :List[Any]    ,__snake_case     :List[Any]    ,__snake_case     :List[Any]    ,__snake_case     :List[Any]    ,__snake_case     :Union[str, Any]    ,__snake_case     :Optional[int]    ,__snake_case     :Tuple    ,__snake_case     :int    ,)      ->					List[Any]:
												a__    =  self.num_choices
												a__    =  XLMForMultipleChoice(config=__snake_case   )
												model.to(__snake_case   )
												model.eval()
												a__    =  input_ids.unsqueeze(1   ).expand(-1    ,self.num_choices    ,-1   ).contiguous()
												a__    =  token_type_ids.unsqueeze(1   ).expand(-1    ,self.num_choices    ,-1   ).contiguous()
												a__    =  input_mask.unsqueeze(1   ).expand(-1    ,self.num_choices    ,-1   ).contiguous()
												a__    =  model(
												    __snake_case    ,attention_mask=__snake_case    ,token_type_ids=__snake_case    ,labels=__snake_case    ,)
												self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.num_choices)   )



					def 				lowerCamelCase__(   self     :Optional[int]   )      ->					Optional[int]:
												a__    =  self.prepare_config_and_inputs()
												(
												    a__    
												)    =  config_and_inputs
												a__    =  {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'lengths': input_lengths}
												return config, inputs_dict







@require_torch
class  snake_case_     (__A			,		__A			,		__A			,		unittest.TestCase   ):
					UpperCAmelCase__      :					List[str]					   =    (
					    (
					        XLMModel,
					        XLMWithLMHeadModel,
					        XLMForQuestionAnswering,
					        XLMForSequenceClassification,
					        XLMForQuestionAnsweringSimple,
					        XLMForTokenClassification,
					        XLMForMultipleChoice,
					    )
					    if is_torch_available()
					    else ()
					)
					UpperCAmelCase__      :					Any					   =    (
					    (XLMWithLMHeadModel,) if is_torch_available() else ()
					)  # TODO (PVP): Check other models whether language generation is also applicable
					UpperCAmelCase__      :					Optional[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 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :Any    ,__snake_case     :int    ,__snake_case     :List[Any]    ,__snake_case     :List[Any]    ,__snake_case     :int   )      ->					Tuple:





												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 				lowerCamelCase__(   self     :str    ,__snake_case     :Union[str, Any]    ,__snake_case     :Union[str, Any]    ,__snake_case     :List[str]=False   )      ->					List[Any]:
												a__    =  super()._prepare_for_class(__snake_case    ,__snake_case    ,return_labels=__snake_case   )

												if return_labels:
																			if model_class.__name__ == "XLMForQuestionAnswering":
																										a__    =  torch.zeros(
																										    self.model_tester.batch_size    ,dtype=torch.long    ,device=__snake_case   )
																										a__    =  torch.zeros(
																										    self.model_tester.batch_size    ,dtype=torch.long    ,device=__snake_case   )

												return inputs_dict
					def 				lowerCamelCase__(   self     :Dict   )      ->					Union[str, Any]:
												a__    =  XLMModelTester(self   )
												a__    =  ConfigTester(self    ,config_class=__snake_case    ,emb_dim=37   )
					def 				lowerCamelCase__(   self     :List[Any]   )      ->					List[str]:
												self.config_tester.run_common_tests()
					def 				lowerCamelCase__(   self     :Union[str, Any]   )      ->					Optional[Any]:
												a__    =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_xlm_model(*__snake_case   )
					def 				lowerCamelCase__(   self     :int   )      ->					Dict:
												a__    =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_xlm_lm_head(*__snake_case   )
					def 				lowerCamelCase__(   self     :List[Any]   )      ->					Dict:
												a__    =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_xlm_simple_qa(*__snake_case   )
					def 				lowerCamelCase__(   self     :Any   )      ->					str:
												a__    =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_xlm_qa(*__snake_case   )
					def 				lowerCamelCase__(   self     :Dict   )      ->					List[str]:
												a__    =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_xlm_sequence_classif(*__snake_case   )
					def 				lowerCamelCase__(   self     :Optional[int]   )      ->					int:
												a__    =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_xlm_token_classif(*__snake_case   )
					def 				lowerCamelCase__(   self     :Optional[int]   )      ->					List[Any]:
												a__    =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_xlm_for_multiple_choice(*__snake_case   )
					def 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :Optional[int]    ,__snake_case     :List[Any]    ,__snake_case     :List[Any]    ,__snake_case     :Any    ,__snake_case     :Optional[Any]    ,__snake_case     :Union[str, Any]=False    ,__snake_case     :List[str]=1   )      ->					Union[str, Any]:
												self.assertIsInstance(__snake_case    ,__snake_case   )
												self.assertListEqual(
												    [isinstance(__snake_case    ,__snake_case   ) for iter_attentions in attentions]    ,[True] * len(__snake_case   )   )
												self.assertEqual(len(__snake_case   )    ,(max_length - min_length) * num_beam_groups   )

												for idx, iter_attentions in enumerate(__snake_case   ):
																			# adds PAD dummy token
																			a__    =  min_length + idx + 1
																			a__    =  min_length + idx + 1

																			a__    =  (
																			    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(__snake_case   )   )
					def 				lowerCamelCase__(   self     :Tuple    ,__snake_case     :Optional[int]    ,__snake_case     :Optional[int]    ,__snake_case     :Union[str, Any]    ,__snake_case     :str    ,__snake_case     :Union[str, Any]    ,__snake_case     :Dict=False    ,__snake_case     :int=1   )      ->					Any:
												self.assertIsInstance(__snake_case    ,__snake_case   )
												self.assertListEqual(
												    [isinstance(__snake_case    ,__snake_case   ) for iter_hidden_states in hidden_states]    ,[True] * len(__snake_case   )    ,)
												self.assertEqual(len(__snake_case   )    ,(max_length - min_length) * num_beam_groups   )

												for idx, iter_hidden_states in enumerate(__snake_case   ):
																			# adds PAD dummy token
																			a__    =  min_length + idx + 1
																			a__    =  (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(__snake_case   )    ,)
												pass



					@slow
					def 				lowerCamelCase__(   self     :Optional[Any]   )      ->					Optional[Any]:





												for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			a__    =  XLMModel.from_pretrained(__snake_case   )
																			self.assertIsNotNone(__snake_case   )







@require_torch
class  snake_case_     (unittest.TestCase   ):
					@slow
					def 				lowerCamelCase__(   self     :Tuple   )      ->					Tuple:
												a__    =  XLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'   )
												model.to(__snake_case   )
												a__    =  torch.tensor([[14, 4_47]]    ,dtype=torch.long    ,device=__snake_case   )  # the president
												a__    =  [
												    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
												a__    =  model.generate(__snake_case    ,do_sample=__snake_case   )
												self.assertListEqual(output_ids[0].cpu().numpy().tolist()    ,__snake_case   )









import math
from typing import List, Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils import SchedulerMixin, SchedulerOutput







class  snake_case_     (lowerCamelCase_			,		lowerCamelCase_   ):
					UpperCAmelCase__      :					Optional[Any]					   =    1
					@register_to_config
					def __init__(   self     :Optional[int]    ,__snake_case     :int = 10_00    ,__snake_case     :Optional[Union[np.ndarray, List[float]]] = None   )      ->					int:
												# set `betas`, `alphas`, `timesteps`
												self.set_timesteps(__snake_case   )

												# standard deviation of the initial noise distribution
												a__    =  1.0

												# For now we only support F-PNDM, i.e. the runge-kutta method
												# For more information on the algorithm please take a look at the paper: https://arxiv.org/pdf/2202.09778.pdf
												# mainly at formula (9), (12), (13) and the Algorithm 2.
												a__    =  4

												# running values
												a__    =  []
					def 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :int    ,__snake_case     :Union[str, torch.device] = None   )      ->					Union[str, Any]:
												a__    =  num_inference_steps
												a__    =  torch.linspace(1    ,0    ,num_inference_steps + 1   )[:-1]
												a__    =  torch.cat([steps, torch.tensor([0.0]   )]   )

												if self.config.trained_betas is not None:
																			a__    =  torch.tensor(self.config.trained_betas    ,dtype=torch.floataa   )
												else:
																			a__    =  torch.sin(steps * math.pi / 2   ) ** 2

												a__    =  (1.0 - self.betas**2) ** 0.5

												a__    =  (torch.atana(self.betas    ,self.alphas   ) / math.pi * 2)[:-1]
												a__    =  timesteps.to(__snake_case   )

												a__    =  []
					def 				lowerCamelCase__(   self     :Any    ,__snake_case     :torch.FloatTensor    ,__snake_case     :int    ,__snake_case     :torch.FloatTensor    ,__snake_case     :bool = True    ,)      ->					Union[SchedulerOutput, Tuple]:
												if self.num_inference_steps is None:
																			raise ValueError(
																			    'Number of inference steps is \'None\', you need to run \'set_timesteps\' after creating the scheduler'   )

												a__    =  (self.timesteps == timestep).nonzero().item()
												a__    =  timestep_index + 1

												a__    =  sample * self.betas[timestep_index] + model_output * self.alphas[timestep_index]
												self.ets.append(__snake_case   )

												if len(self.ets   ) == 1:
																			a__    =  self.ets[-1]
												elif len(self.ets   ) == 2:
																			a__    =  (3 * self.ets[-1] - self.ets[-2]) / 2
												elif len(self.ets   ) == 3:
																			a__    =  (23 * self.ets[-1] - 16 * self.ets[-2] + 5 * self.ets[-3]) / 12
												else:
																			a__    =  (1 / 24) * (55 * self.ets[-1] - 59 * self.ets[-2] + 37 * self.ets[-3] - 9 * self.ets[-4])

												a__    =  self._get_prev_sample(__snake_case    ,__snake_case    ,__snake_case    ,__snake_case   )

												if not return_dict:
																			return (prev_sample,)

												return SchedulerOutput(prev_sample=__snake_case   )
					def 				lowerCamelCase__(   self     :Union[str, Any]    ,__snake_case     :torch.FloatTensor    ,*__snake_case     :int    ,**__snake_case     :Optional[int]   )      ->					torch.FloatTensor:
												return sample
					def 				lowerCamelCase__(   self     :Optional[Any]    ,__snake_case     :List[Any]    ,__snake_case     :Optional[int]    ,__snake_case     :Dict    ,__snake_case     :Any   )      ->					Optional[Any]:
												a__    =  self.alphas[timestep_index]
												a__    =  self.betas[timestep_index]

												a__    =  self.alphas[prev_timestep_index]
												a__    =  self.betas[prev_timestep_index]

												a__    =  (sample - sigma * ets) / max(__snake_case    ,1E-8   )
												a__    =  next_alpha * pred + ets * next_sigma

												return prev_sample



					def __len__(   self     :Any   )      ->					Union[str, Any]:
												return self.config.num_train_timesteps




