"""simple docstring"""



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      _A	:




    """simple docstring"""




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

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

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

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

         __snake_case				=							self.get_config()

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

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

         __snake_case				=							model(A_    )

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

         __snake_case				=							model(A_    )

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

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

         ((__snake_case)		,	)				=							result_with_labels.to_tuple()

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

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

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

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







    def  lowercase					(		self							:     List[str]    )    ->       Any:
         __snake_case				=							self.prepare_config_and_inputs()
         (
             (
             __snake_case
         )		,	(
             __snake_case
         )		,	(
             __snake_case
         )		,	(
             __snake_case
         )		,	(
             __snake_case
         )		,	(
             __snake_case
         )		,	(
             __snake_case
         )		,	(
             __snake_case
         )		,	(
             __snake_case
         )		,	
         )				=							config_and_inputs
         __snake_case				=							{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''lengths''': input_lengths}
         return config, inputs_dict






@require_torch
class      _A	(						_UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, unittest.TestCase					):




    """simple docstring"""




    UpperCamelCase_	:					Union[str, Any]							=					(
        (
            XLMModel,
            XLMWithLMHeadModel,
            XLMForQuestionAnswering,
            XLMForSequenceClassification,
            XLMForQuestionAnsweringSimple,
            XLMForTokenClassification,
            XLMForMultipleChoice,
        )
        if is_torch_available()
        else ()
    )
    UpperCamelCase_	:					List[str]							=					(
        (XLMWithLMHeadModel,) if is_torch_available() else ()
    )  # TODO (PVP): Check other models whether language generation is also applicable
    UpperCamelCase_	:					List[str]							=					(
        {
            '''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  lowercase					(		self							:     int	,      A_							:     Optional[int]	,      A_							:     Union[str, Any]	,      A_							:     str	,      A_							:     Tuple	,      A_							:     List[str]    )    ->       Optional[int]:
         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  lowercase					(		self							:     Optional[int]	,      A_							:     str	,      A_							:     Tuple	,      A_							:     List[Any]=False    )    ->       Optional[int]:
         __snake_case				=							super()._prepare_for_class(A_	,      A_	,      return_labels=A_    )

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

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

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

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

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







    @slow
    def  lowercase					(		self							:     int    )    ->       List[str]:
         for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              __snake_case				=							XLMModel.from_pretrained(A_    )
              self.assertIsNotNone(A_    )






@require_torch
class      _A	(						unittest.TestCase					):




    """simple docstring"""




    @slow
    def  lowercase					(		self							:     Any    )    ->       int:
         __snake_case				=							XLMWithLMHeadModel.from_pretrained('''xlm-mlm-en-2048'''    )
         model.to(A_    )
         __snake_case				=							torch.tensor([[14, 447]]	,      dtype=torch.long	,      device=A_    )  # the president
         __snake_case				=							[
             14,
             447,
             14,
             447,
             14,
             447,
             14,
             447,
             14,
             447,
             14,
             447,
             14,
             447,
             14,
             447,
             14,
             447,
             14,
             447,
         ]  # 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
         __snake_case				=							model.generate(A_	,      do_sample=A_    )
         self.assertListEqual(output_ids[0].cpu().numpy().tolist()	,      A_    )
"""simple docstring"""



import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()



@slow
@require_torch_gpu
class      _A	(						unittest.TestCase					):




    """simple docstring"""




    def  lowercase					(		self							:     Tuple    )    ->       str:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()
    def  lowercase					(		self							:     int    )    ->       Tuple:
         __snake_case				=							StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''    )
         __snake_case				=							sd_pipe.to(A_    )
         sd_pipe.set_progress_bar_config(disable=A_    )

         sd_pipe.set_scheduler('''sample_euler'''    )

         __snake_case				=							'''A painting of a squirrel eating a burger'''
         __snake_case				=							torch.manual_seed(0    )
         __snake_case				=							sd_pipe([prompt]	,      generator=A_	,      guidance_scale=9.0	,      num_inference_steps=20	,      output_type='''np'''    )

         __snake_case				=							output.images

         __snake_case				=							image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         __snake_case				=							np.array([0.04_47, 0.04_92, 0.04_68, 0.04_08, 0.03_83, 0.04_08, 0.03_54, 0.03_80, 0.03_39]    )

         assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2
    def  lowercase					(		self							:     Optional[Any]    )    ->       Tuple:
         __snake_case				=							StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''    )
         __snake_case				=							sd_pipe.to(A_    )
         sd_pipe.set_progress_bar_config(disable=A_    )

         sd_pipe.set_scheduler('''sample_euler'''    )

         __snake_case				=							'''A painting of a squirrel eating a burger'''
         __snake_case				=							torch.manual_seed(0    )
         __snake_case				=							sd_pipe([prompt]	,      generator=A_	,      guidance_scale=9.0	,      num_inference_steps=20	,      output_type='''np'''    )

         __snake_case				=							output.images

         __snake_case				=							image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         __snake_case				=							np.array([0.12_37, 0.13_20, 0.14_38, 0.13_59, 0.13_90, 0.11_32, 0.12_77, 0.11_75, 0.11_12]    )

         assert np.abs(image_slice.flatten() - expected_slice    ).max() < 5E-1







    def  lowercase					(		self							:     List[str]    )    ->       Optional[Any]:
         __snake_case				=							StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''    )
         __snake_case				=							sd_pipe.to(A_    )
         sd_pipe.set_progress_bar_config(disable=A_    )

         sd_pipe.set_scheduler('''sample_dpmpp_2m'''    )

         __snake_case				=							'''A painting of a squirrel eating a burger'''
         __snake_case				=							torch.manual_seed(0    )
         __snake_case				=							sd_pipe(
             [prompt]	,      generator=A_	,      guidance_scale=7.5	,      num_inference_steps=15	,      output_type='''np'''	,      use_karras_sigmas=A_	,      )

         __snake_case				=							output.images

         __snake_case				=							image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         __snake_case				=							np.array(
             [0.11_38_16_89, 0.12_11_29_21, 0.1_38_94_57, 0.12_54_96_06, 0.1_24_49_64, 0.10_83_15_17, 0.11_56_28_66, 0.10_86_78_16, 0.10_49_90_48]    )

         assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2