
from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import (
       TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFXLMForMultipleChoice,
       TFXLMForQuestionAnsweringSimple,
       TFXLMForSequenceClassification,
       TFXLMForTokenClassification,
       TFXLMModel,
       TFXLMWithLMHeadModel,
       XLMConfig,
   )







class      lowerCamelCase		:





     def __init__(self      :    List[str]				,    _A      :    Any				,    )						->						Any:
         snake_case      =					parent
         snake_case      =					1_3
         snake_case      =					7
         snake_case      =					True
         snake_case      =					True
         snake_case      =					True
         snake_case      =					True
         snake_case      =					True
         snake_case      =					False
         snake_case      =					False
         snake_case      =					False
         snake_case      =					2
         snake_case      =					9_9
         snake_case      =					0
         snake_case      =					3_2
         snake_case      =					2
         snake_case      =					4
         snake_case      =					0.1
         snake_case      =					0.1
         snake_case      =					5_1_2
         snake_case      =					1_6
         snake_case      =					2
         snake_case      =					0.02
         snake_case      =					3
         snake_case      =					4
         snake_case      =					"last"
         snake_case      =					True
         snake_case      =					None
         snake_case      =					0





     def 		UpperCAmelCase(self      :    str   )						->						Optional[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]				,    dtype=tf.floataa   )

         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				,    dtype=tf.floataa   )
             snake_case      =					ids_tensor([self.batch_size]				,    self.num_choices   )

         snake_case      =					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				,    bos_token_id=self.bos_token_id				,    )

         return (
             config,
             input_ids,
             token_type_ids,
             input_lengths,
             sequence_labels,
             token_labels,
             is_impossible_labels,
             choice_labels,
             input_mask,
         )





     def 		UpperCAmelCase(self      :    int				,    _A      :    List[str]				,    _A      :    Tuple				,    _A      :    Any				,    _A      :    List[Any]				,    _A      :    List[Any]				,    _A      :    Dict				,    _A      :    Dict				,    _A      :    int				,    _A      :    Dict				,    )						->						List[Any]:
         snake_case      =					TFXLMModel(config=_A   )
         snake_case      =					{"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
         snake_case      =					model(_A   )

         snake_case      =					[input_ids, input_mask]
         snake_case      =					model(_A   )
         self.parent.assertEqual(result.last_hidden_state.shape				,    (self.batch_size, self.seq_length, self.hidden_size)   )





     def 		UpperCAmelCase(self      :    Dict				,    _A      :    Optional[Any]				,    _A      :    Union[str, Any]				,    _A      :    List[Any]				,    _A      :    str				,    _A      :    Optional[Any]				,    _A      :    str				,    _A      :    List[str]				,    _A      :    Optional[Any]				,    _A      :    Optional[int]				,    )						->						List[str]:
         snake_case      =					TFXLMWithLMHeadModel(_A   )

         snake_case      =					{"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
         snake_case      =					model(_A   )

         snake_case      =					outputs

         self.parent.assertEqual(result.logits.shape				,    (self.batch_size, self.seq_length, self.vocab_size)   )





     def 		UpperCAmelCase(self      :    List[str]				,    _A      :    Optional[Any]				,    _A      :    int				,    _A      :    str				,    _A      :    int				,    _A      :    Optional[int]				,    _A      :    str				,    _A      :    int				,    _A      :    Optional[int]				,    _A      :    int				,    )						->						Optional[int]:
         snake_case      =					TFXLMForQuestionAnsweringSimple(_A   )

         snake_case      =					{"input_ids": input_ids, "lengths": input_lengths}

         snake_case      =					model(_A   )

         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      :    Tuple				,    _A      :    Dict				,    _A      :    Tuple				,    _A      :    Optional[Any]				,    _A      :    Optional[int]				,    _A      :    Union[str, Any]				,    _A      :    Any				,    _A      :    Any				,    _A      :    Optional[Any]				,    _A      :    Dict				,    )						->						Union[str, Any]:
         snake_case      =					TFXLMForSequenceClassification(_A   )

         snake_case      =					{"input_ids": input_ids, "lengths": input_lengths}

         snake_case      =					model(_A   )

         self.parent.assertEqual(result.logits.shape				,    (self.batch_size, self.type_sequence_label_size)   )





     def 		UpperCAmelCase(self      :    List[Any]				,    _A      :    List[Any]				,    _A      :    List[Any]				,    _A      :    Union[str, Any]				,    _A      :    Optional[Any]				,    _A      :    List[Any]				,    _A      :    Any				,    _A      :    str				,    _A      :    Optional[Any]				,    _A      :    str				,    )						->						int:
         snake_case      =					self.num_labels
         snake_case      =					TFXLMForTokenClassification(config=_A   )
         snake_case      =					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
         snake_case      =					model(_A   )
         self.parent.assertEqual(result.logits.shape				,    (self.batch_size, self.seq_length, self.num_labels)   )





     def 		UpperCAmelCase(self      :    Tuple				,    _A      :    Tuple				,    _A      :    Optional[int]				,    _A      :    str				,    _A      :    List[str]				,    _A      :    Optional[int]				,    _A      :    str				,    _A      :    Union[str, Any]				,    _A      :    Union[str, Any]				,    _A      :    List[Any]				,    )						->						Union[str, Any]:
         snake_case      =					self.num_choices
         snake_case      =					TFXLMForMultipleChoice(config=_A   )
         snake_case      =					tf.tile(tf.expand_dims(_A				,    1   )				,    (1, self.num_choices, 1)   )
         snake_case      =					tf.tile(tf.expand_dims(_A				,    1   )				,    (1, self.num_choices, 1)   )
         snake_case      =					tf.tile(tf.expand_dims(_A				,    1   )				,    (1, self.num_choices, 1)   )
         snake_case      =					{
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
             "token_type_ids": multiple_choice_token_type_ids,
         }
         snake_case      =					model(_A   )
         self.parent.assertEqual(result.logits.shape				,    (self.batch_size, self.num_choices)   )





     def 		UpperCAmelCase(self      :    List[str]   )						->						Optional[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,
             "langs": token_type_ids,
             "lengths": input_lengths,
         }
         return config, inputs_dict







@require_tf
class      lowerCamelCase		(  A_							, A_							, unittest.TestCase  ):
     UpperCAmelCase__				:							Union[str, Any]           =		(
         (
             TFXLMModel,
             TFXLMWithLMHeadModel,
             TFXLMForSequenceClassification,
             TFXLMForQuestionAnsweringSimple,
             TFXLMForTokenClassification,
             TFXLMForMultipleChoice,
         )
         if is_tf_available()
         else ()
     )
     UpperCAmelCase__				:							List[Any]           =		(
         (TFXLMWithLMHeadModel,) if is_tf_available() else ()
     )  # TODO (PVP): Check other models whether language generation is also applicable
     UpperCAmelCase__				:							Union[str, Any]           =		(
         {
             "feature-extraction": TFXLMModel,
             "fill-mask": TFXLMWithLMHeadModel,
             "question-answering": TFXLMForQuestionAnsweringSimple,
             "text-classification": TFXLMForSequenceClassification,
             "text-generation": TFXLMWithLMHeadModel,
             "token-classification": TFXLMForTokenClassification,
             "zero-shot": TFXLMForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     UpperCAmelCase__				:							int           =		False
     UpperCAmelCase__				:							Optional[int]           =		False





     def 		UpperCAmelCase(self      :    List[str]				,    _A      :    Optional[Any]				,    _A      :    Optional[int]				,    _A      :    int				,    _A      :    Dict				,    _A      :    Dict   )						->						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 		UpperCAmelCase(self      :    List[Any]   )						->						Dict:
         snake_case      =					TFXLMModelTester(self   )
         snake_case      =					ConfigTester(self				,    config_class=_A				,    emb_dim=3_7   )





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





     def 		UpperCAmelCase(self      :    int   )						->						Optional[int]:
         snake_case      =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_model(*_A   )





     def 		UpperCAmelCase(self      :    Tuple   )						->						Optional[Any]:
         snake_case      =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_lm_head(*_A   )





     def 		UpperCAmelCase(self      :    Optional[Any]   )						->						Optional[Any]:
         snake_case      =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_qa(*_A   )





     def 		UpperCAmelCase(self      :    Tuple   )						->						Union[str, Any]:
         snake_case      =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_sequence_classif(*_A   )





     def 		UpperCAmelCase(self      :    List[Any]   )						->						Dict:
         snake_case      =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_for_token_classification(*_A   )





     def 		UpperCAmelCase(self      :    int   )						->						Union[str, Any]:
         snake_case      =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_for_multiple_choice(*_A   )





     @slow
     def 		UpperCAmelCase(self      :    Tuple   )						->						str:
         for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             snake_case      =					TFXLMModel.from_pretrained(_A   )
             self.assertIsNotNone(_A   )







@require_tf
class      lowerCamelCase		(  unittest.TestCase  ):





     @slow
     def 		UpperCAmelCase(self      :    Union[str, Any]   )						->						List[Any]:
         snake_case      =					TFXLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048"   )
         snake_case      =					tf.convert_to_tensor([[1_4, 4_4_7]]				,    dtype=tf.intaa   )  # the president
         snake_case      =					[
             1_4,
             4_4_7,
             1_4,
             4_4_7,
             1_4,
             4_4_7,
             1_4,
             4_4_7,
             1_4,
             4_4_7,
             1_4,
             4_4_7,
             1_4,
             4_4_7,
             1_4,
             4_4_7,
             1_4,
             4_4_7,
             1_4,
             4_4_7,
         ]  # 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].numpy().tolist()				,    _A   )



import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class      lowerCamelCase		(  A_							, unittest.TestCase  ):
     UpperCAmelCase__				:							Tuple           =		KandinskyPriorPipeline
     UpperCAmelCase__				:							List[str]           =		["prompt"]
     UpperCAmelCase__				:							List[str]           =		["prompt", "negative_prompt"]
     UpperCAmelCase__				:							List[str]           =		[
         "num_images_per_prompt",
         "generator",
         "num_inference_steps",
         "latents",
         "negative_prompt",
         "guidance_scale",
         "output_type",
         "return_dict",
     ]
     UpperCAmelCase__				:							str           =		False





     @property
     def 		UpperCAmelCase(self      :    List[str]   )						->						Union[str, Any]:
         return 3_2





     @property
     def 		UpperCAmelCase(self      :    Any   )						->						Any:
         return 3_2





     @property
     def 		UpperCAmelCase(self      :    int   )						->						Dict:
         return self.time_input_dim





     @property
     def 		UpperCAmelCase(self      :    str   )						->						Tuple:
         return self.time_input_dim * 4





     @property
     def 		UpperCAmelCase(self      :    Dict   )						->						Dict:
         return 1_0_0





     @property
     def 		UpperCAmelCase(self      :    Any   )						->						int:
         snake_case      =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )
         return tokenizer





     @property
     def 		UpperCAmelCase(self      :    Optional[Any]   )						->						Optional[Any]:
         torch.manual_seed(0   )
         snake_case      =					CLIPTextConfig(
             bos_token_id=0				,    eos_token_id=2				,    hidden_size=self.text_embedder_hidden_size				,    projection_dim=self.text_embedder_hidden_size				,    intermediate_size=3_7				,    layer_norm_eps=1E-05				,    num_attention_heads=4				,    num_hidden_layers=5				,    pad_token_id=1				,    vocab_size=1_0_0_0				,    )
         return CLIPTextModelWithProjection(_A   )





     @property
     def 		UpperCAmelCase(self      :    List[Any]   )						->						Tuple:
         torch.manual_seed(0   )

         snake_case      =					{
             "num_attention_heads": 2,
             "attention_head_dim": 1_2,
             "embedding_dim": self.text_embedder_hidden_size,
             "num_layers": 1,
         }

         snake_case      =					PriorTransformer(**_A   )
         # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
         snake_case      =					nn.Parameter(torch.ones(model.clip_std.shape   )   )
         return model





     @property
     def 		UpperCAmelCase(self      :    str   )						->						Union[str, Any]:
         torch.manual_seed(0   )
         snake_case      =					CLIPVisionConfig(
             hidden_size=self.text_embedder_hidden_size				,    image_size=2_2_4				,    projection_dim=self.text_embedder_hidden_size				,    intermediate_size=3_7				,    num_attention_heads=4				,    num_channels=3				,    num_hidden_layers=5				,    patch_size=1_4				,    )

         snake_case      =					CLIPVisionModelWithProjection(_A   )
         return model





     @property
     def 		UpperCAmelCase(self      :    List[Any]   )						->						int:
         snake_case      =					CLIPImageProcessor(
             crop_size=2_2_4				,    do_center_crop=_A				,    do_normalize=_A				,    do_resize=_A				,    image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]				,    image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]				,    resample=3				,    size=2_2_4				,    )

         return image_processor





     def 		UpperCAmelCase(self      :    int   )						->						Union[str, Any]:
         snake_case      =					self.dummy_prior
         snake_case      =					self.dummy_image_encoder
         snake_case      =					self.dummy_text_encoder
         snake_case      =					self.dummy_tokenizer
         snake_case      =					self.dummy_image_processor

         snake_case      =					UnCLIPScheduler(
             variance_type="fixed_small_log"				,    prediction_type="sample"				,    num_train_timesteps=1_0_0_0				,    clip_sample=_A				,    clip_sample_range=10.0				,    )

         snake_case      =					{
             "prior": prior,
             "image_encoder": image_encoder,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "scheduler": scheduler,
             "image_processor": image_processor,
         }

         return components





     def 		UpperCAmelCase(self      :    Optional[int]				,    _A      :    Tuple				,    _A      :    str=0   )						->						Union[str, Any]:
         if str(_A   ).startswith("mps"   ):
             snake_case      =					torch.manual_seed(_A   )
         else:
             snake_case      =					torch.Generator(device=_A   ).manual_seed(_A   )
         snake_case      =					{
             "prompt": "horse",
             "generator": generator,
             "guidance_scale": 4.0,
             "num_inference_steps": 2,
             "output_type": "np",
         }
         return inputs





     def 		UpperCAmelCase(self      :    List[Any]   )						->						int:
         snake_case      =					"cpu"

         snake_case      =					self.get_dummy_components()

         snake_case      =					self.pipeline_class(**_A   )
         snake_case      =					pipe.to(_A   )

         pipe.set_progress_bar_config(disable=_A   )

         snake_case      =					pipe(**self.get_dummy_inputs(_A   )   )
         snake_case      =					output.image_embeds

         snake_case      =					pipe(
             **self.get_dummy_inputs(_A   )				,    return_dict=_A				,    )[0]

         snake_case      =					image[0, -1_0:]
         snake_case      =					image_from_tuple[0, -1_0:]

         assert image.shape == (1, 3_2)

         snake_case      =					np.array(
             [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]   )

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





     @skip_mps
     def 		UpperCAmelCase(self      :    List[Any]   )						->						str:
         snake_case      =					torch_device == "cpu"
         snake_case      =					True
         snake_case      =					False

         self._test_inference_batch_single_identical(
             test_max_difference=_A				,    relax_max_difference=_A				,    test_mean_pixel_difference=_A				,    )





     @skip_mps
     def 		UpperCAmelCase(self      :    Optional[int]   )						->						List[str]:
         snake_case      =					torch_device == "cpu"
         snake_case      =					False

         self._test_attention_slicing_forward_pass(
             test_max_difference=_A				,    test_mean_pixel_difference=_A				,    )

