






'''simple docstring'''
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 numpy
 import tensorflow as tf

 from transformers import (
     TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
     TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
     TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
     BertConfig,
     DPRConfig,
     TFDPRContextEncoder,
     TFDPRQuestionEncoder,
     TFDPRReader,
 )


class        _UpperCAmelCase						:




    def __init__(			self  :       Tuple  ,     a  :       Optional[int]  ,     a  :       Union[str, Any]=1_3  ,     a  :       List[str]=7  ,     a  :       Tuple=True  ,     a  :       Optional[Any]=True  ,     a  :       Any=True  ,     a  :       Any=True  ,     a  :       Optional[Any]=9_9  ,     a  :       str=3_2  ,     a  :       Dict=2  ,     a  :       List[Any]=4  ,     a  :       Tuple=3_7  ,     a  :       Tuple="gelu"  ,     a  :       Tuple=0.1  ,     a  :       str=0.1  ,     a  :       Optional[int]=5_1_2  ,     a  :       Optional[int]=1_6  ,     a  :       Tuple=2  ,     a  :       Optional[int]=0.02  ,     a  :       List[Any]=3  ,     a  :       Union[str, Any]=4  ,     a  :       Optional[Any]=None  ,     a  :       Union[str, Any]=0  ,     ):






       '''simple docstring'''






       lowercase_	:    List[str]           =     parent
       lowercase_	:    Any           =     batch_size
       lowercase_	:    Tuple           =     seq_length
       lowercase_	:    str           =     is_training
       lowercase_	:    str           =     use_input_mask
       lowercase_	:    str           =     use_token_type_ids
       lowercase_	:    List[Any]           =     use_labels
       lowercase_	:    int           =     vocab_size
       lowercase_	:    Optional[int]           =     hidden_size
       lowercase_	:    List[Any]           =     num_hidden_layers
       lowercase_	:    Tuple           =     num_attention_heads
       lowercase_	:    Optional[int]           =     intermediate_size
       lowercase_	:    List[Any]           =     hidden_act
       lowercase_	:    str           =     hidden_dropout_prob
       lowercase_	:    List[Any]           =     attention_probs_dropout_prob
       lowercase_	:    Dict           =     max_position_embeddings
       lowercase_	:    Any           =     type_vocab_size
       lowercase_	:    Optional[int]           =     type_sequence_label_size
       lowercase_	:    List[Any]           =     initializer_range
       lowercase_	:    Dict           =     num_labels
       lowercase_	:    int           =     num_choices
       lowercase_	:    Union[str, Any]           =     scope
       lowercase_	:    Dict           =     projection_dim




    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       lowercase_	:    Optional[int]           =     ids_tensor([self.batch_size, self.seq_length]  ,     self.vocab_size			)

       lowercase_	:    Dict           =     None
       if self.use_input_mask:
          # follow test_modeling_tf_ctrl.py
          lowercase_	:    Union[str, Any]           =     random_attention_mask([self.batch_size, self.seq_length]			)

       lowercase_	:    Optional[int]           =     None
       if self.use_token_type_ids:
          lowercase_	:    Dict           =     ids_tensor([self.batch_size, self.seq_length]  ,     self.type_vocab_size			)

       lowercase_	:    Optional[Any]           =     None
       lowercase_	:    Dict           =     None
       lowercase_	:    Optional[Any]           =     None
       if self.use_labels:
          lowercase_	:    Optional[int]           =     ids_tensor([self.batch_size]  ,     self.type_sequence_label_size			)
          lowercase_	:    Optional[int]           =     ids_tensor([self.batch_size, self.seq_length]  ,     self.num_labels			)
          lowercase_	:    Tuple           =     ids_tensor([self.batch_size]  ,     self.num_choices			)

       lowercase_	:    Union[str, Any]           =     BertConfig(
           vocab_size=self.vocab_size  ,     hidden_size=self.hidden_size  ,     num_hidden_layers=self.num_hidden_layers  ,     num_attention_heads=self.num_attention_heads  ,     intermediate_size=self.intermediate_size  ,     hidden_act=self.hidden_act  ,     hidden_dropout_prob=self.hidden_dropout_prob  ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,     max_position_embeddings=self.max_position_embeddings  ,     type_vocab_size=self.type_vocab_size  ,     is_decoder=a  ,     initializer_range=self.initializer_range  ,     )
       lowercase_	:    Any           =     DPRConfig(projection_dim=self.projection_dim  ,     **config.to_dict()			)

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




    def 				lowerCAmelCase__				(			self  :       List[Any]  ,     a  :       Tuple  ,     a  :       Union[str, Any]  ,     a  :       Tuple  ,     a  :       List[Any]  ,     a  :       Dict  ,     a  :       Union[str, Any]  ,     a  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    Tuple           =     TFDPRContextEncoder(config=a			)
       lowercase_	:    List[Any]           =     model(a  ,     attention_mask=a  ,     token_type_ids=a			)
       lowercase_	:    Optional[int]           =     model(a  ,     token_type_ids=a			)
       lowercase_	:    Optional[int]           =     model(a			)
       self.parent.assertEqual(result.pooler_output.shape  ,     (self.batch_size, self.projection_dim or self.hidden_size)			)




    def 				lowerCAmelCase__				(			self  :       List[str]  ,     a  :       Dict  ,     a  :       Union[str, Any]  ,     a  :       List[str]  ,     a  :       Optional[int]  ,     a  :       str  ,     a  :       List[str]  ,     a  :       Any			):






       '''simple docstring'''






       lowercase_	:    str           =     TFDPRQuestionEncoder(config=a			)
       lowercase_	:    Tuple           =     model(a  ,     attention_mask=a  ,     token_type_ids=a			)
       lowercase_	:    Tuple           =     model(a  ,     token_type_ids=a			)
       lowercase_	:    str           =     model(a			)
       self.parent.assertEqual(result.pooler_output.shape  ,     (self.batch_size, self.projection_dim or self.hidden_size)			)




    def 				lowerCAmelCase__				(			self  :       List[Any]  ,     a  :       Optional[Any]  ,     a  :       List[Any]  ,     a  :       Any  ,     a  :       Optional[int]  ,     a  :       Dict  ,     a  :       Optional[Any]  ,     a  :       int			):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     TFDPRReader(config=a			)
       lowercase_	:    Any           =     model(a  ,     attention_mask=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)			)
       self.parent.assertEqual(result.relevance_logits.shape  ,     (self.batch_size,)			)





    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       lowercase_	:    Union[str, Any]           =     self.prepare_config_and_inputs()
       (
           (
           lowercase_
       )    ,		(
           lowercase_
       )    ,		(
           lowercase_
       )    ,		(
           lowercase_
       )    ,		(
           lowercase_
       )    ,		(
           lowercase_
       )    ,		(
           lowercase_
       )    ,		
       )	:    List[str]           =     config_and_inputs
       lowercase_	:    Dict           =     {"input_ids": input_ids}
       return config, inputs_dict


@require_tf
class        _UpperCAmelCase						(		snake_case				,					snake_case				,					unittest.TestCase			):
    __lowerCamelCase:   Optional[Any]		       =    (
        (
            TFDPRContextEncoder,
            TFDPRQuestionEncoder,
            TFDPRReader,
        )
        if is_tf_available()
        else ()
    )
    __lowerCamelCase:   List[Any]		       =    {'feature-extraction': TFDPRQuestionEncoder} if is_tf_available() else {}

    __lowerCamelCase:   Optional[Any]		       =    False
    __lowerCamelCase:   List[Any]		       =    False
    __lowerCamelCase:   List[Any]		       =    False
    __lowerCamelCase:   List[Any]		       =    False
    __lowerCamelCase:   Union[str, Any]		       =    False




    def 				lowerCAmelCase__				(			self  :       Tuple			):






       '''simple docstring'''






       lowercase_	:    List[str]           =     TFDPRModelTester(self			)
       lowercase_	:    Optional[Any]           =     ConfigTester(self  ,     config_class=a  ,     hidden_size=3_7			)




    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       self.config_tester.run_common_tests()




    def 				lowerCAmelCase__				(			self  :       Dict			):






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_dpr_context_encoder(*a			)




    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       lowercase_	:    List[str]           =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_dpr_question_encoder(*a			)




    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       lowercase_	:    Tuple           =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_dpr_reader(*a			)





    @slow
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowercase_	:    int           =     TFDPRContextEncoder.from_pretrained(a			)
          self.assertIsNotNone(a			)

       for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowercase_	:    int           =     TFDPRContextEncoder.from_pretrained(a			)
          self.assertIsNotNone(a			)

       for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowercase_	:    Dict           =     TFDPRQuestionEncoder.from_pretrained(a			)
          self.assertIsNotNone(a			)

       for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowercase_	:    Any           =     TFDPRReader.from_pretrained(a			)
          self.assertIsNotNone(a			)


@require_tf
class        _UpperCAmelCase						(		unittest.TestCase			):




    @slow
    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    Any           =     TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"			)

       lowercase_	:    Union[str, Any]           =     tf.constant(
           [[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_0_3, 2_0_2_6, 3_8_9_9, 1_0_1_4_0, 1_0_2_9, 1_0_2]]			)  # [CLS] hello, is my dog cute? [SEP]
       lowercase_	:    List[Any]           =     model(a			)[0]  # embedding shape = (1, 768)
       # compare the actual values for a slice.
       lowercase_	:    Any           =     tf.constant(
           [
               [
                   0.0323_6253,
                   0.1275_3335,
                   0.1681_8509,
                   0.0027_9786,
                   0.389_6933,
                   0.2426_4945,
                   0.217_8971,
                   -0.0233_5227,
                   -0.0848_1959,
                   -0.1432_4117,
               ]
           ]			)
       self.assertTrue(numpy.allclose(output[:, :1_0].numpy()  ,     expected_slice.numpy()  ,     atol=1e-4			)			)












'''simple docstring'''
import unittest

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

from diffusers import KandinskyVaaPriorPipeline, 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        _UpperCAmelCase						(		snake_case				,					unittest.TestCase			):
    __lowerCamelCase:   Dict		       =    KandinskyVaaPriorPipeline
    __lowerCamelCase:   Optional[int]		       =    ['prompt']
    __lowerCamelCase:   Any		       =    ['prompt', 'negative_prompt']
    __lowerCamelCase:   List[Any]		       =    [
        'num_images_per_prompt',
        'generator',
        'num_inference_steps',
        'latents',
        'negative_prompt',
        'guidance_scale',
        'output_type',
        'return_dict',
    ]
    __lowerCamelCase:   List[Any]		       =    False




    @property
    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       return 3_2




    @property
    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       return 3_2




    @property
    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       return self.time_input_dim




    @property
    def 				lowerCAmelCase__				(			self  :       str			):






       '''simple docstring'''






       return self.time_input_dim * 4




    @property
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       return 1_0_0




    @property
    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    Any           =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)
       return tokenizer




    @property
    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       torch.manual_seed(0			)
       lowercase_	:    Tuple           =     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 				lowerCAmelCase__				(			self  :       Tuple			):






       '''simple docstring'''






       torch.manual_seed(0			)

       lowercase_	:    List[str]           =     {
           "num_attention_heads": 2,
           "attention_head_dim": 1_2,
           "embedding_dim": self.text_embedder_hidden_size,
           "num_layers": 1,
       }

       lowercase_	:    Union[str, Any]           =     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
       lowercase_	:    List[Any]           =     nn.Parameter(torch.ones(model.clip_std.shape			)			)
       return model




    @property
    def 				lowerCAmelCase__				(			self  :       Tuple			):






       '''simple docstring'''






       torch.manual_seed(0			)
       lowercase_	:    Dict           =     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  ,     )

       lowercase_	:    Optional[Any]           =     CLIPVisionModelWithProjection(a			)
       return model




    @property
    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       lowercase_	:    List[str]           =     CLIPImageProcessor(
           crop_size=2_2_4  ,     do_center_crop=a  ,     do_normalize=a  ,     do_resize=a  ,     image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]  ,     image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]  ,     resample=3  ,     size=2_2_4  ,     )

       return image_processor




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       lowercase_	:    Any           =     self.dummy_prior
       lowercase_	:    Optional[Any]           =     self.dummy_image_encoder
       lowercase_	:    List[Any]           =     self.dummy_text_encoder
       lowercase_	:    Any           =     self.dummy_tokenizer
       lowercase_	:    Optional[Any]           =     self.dummy_image_processor

       lowercase_	:    List[str]           =     UnCLIPScheduler(
           variance_type="fixed_small_log"  ,     prediction_type="sample"  ,     num_train_timesteps=1_0_0_0  ,     clip_sample=a  ,     clip_sample_range=10.0  ,     )

       lowercase_	:    List[Any]           =     {
           "prior": prior,
           "image_encoder": image_encoder,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "scheduler": scheduler,
           "image_processor": image_processor,
       }

       return components




    def 				lowerCAmelCase__				(			self  :       Any  ,     a  :       Dict  ,     a  :       Dict=0			):






       '''simple docstring'''






       if str(a			).startswith("mps"			):
          lowercase_	:    int           =     torch.manual_seed(a			)
       else:
          lowercase_	:    Optional[Any]           =     torch.Generator(device=a			).manual_seed(a			)
       lowercase_	:    Any           =     {
           "prompt": "horse",
           "generator": generator,
           "guidance_scale": 4.0,
           "num_inference_steps": 2,
           "output_type": "np",
       }
       return inputs




    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       lowercase_	:    str           =     "cpu"

       lowercase_	:    Any           =     self.get_dummy_components()

       lowercase_	:    int           =     self.pipeline_class(**a			)
       lowercase_	:    Any           =     pipe.to(a			)

       pipe.set_progress_bar_config(disable=a			)

       lowercase_	:    Any           =     pipe(**self.get_dummy_inputs(a			)			)
       lowercase_	:    List[Any]           =     output.image_embeds

       lowercase_	:    str           =     pipe(
           **self.get_dummy_inputs(a			)  ,     return_dict=a  ,     )[0]

       lowercase_	:    Any           =     image[0, -1_0:]
       lowercase_	:    Dict           =     image_from_tuple[0, -1_0:]

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

       lowercase_	:    int           =     np.array(
           [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]			)

       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 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       lowercase_	:    int           =     torch_device == "cpu"
       lowercase_	:    Tuple           =     True
       lowercase_	:    str           =     False

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





    @skip_mps
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       lowercase_	:    Any           =     torch_device == "cpu"
       lowercase_	:    int           =     False

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




