




import unittest

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

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       RemBertConfig,
       RemBertForCausalLM,
       RemBertForMaskedLM,
       RemBertForMultipleChoice,
       RemBertForQuestionAnswering,
       RemBertForSequenceClassification,
       RemBertForTokenClassification,
       RemBertModel,
   )
   from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST

class   UpperCAmelCase    :






   def __init__(						self :							Optional[int]    ,       UpperCamelCase_ :							str    ,       UpperCamelCase_ :							int=13    ,       UpperCamelCase_ :							List[Any]=7    ,       UpperCamelCase_ :							str=True    ,       UpperCamelCase_ :							int=True    ,       UpperCamelCase_ :							Any=True    ,       UpperCamelCase_ :							Optional[int]=True    ,       UpperCamelCase_ :							int=99    ,       UpperCamelCase_ :							Tuple=32    ,       UpperCamelCase_ :							Dict=18    ,       UpperCamelCase_ :							Optional[Any]=43    ,       UpperCamelCase_ :							Optional[int]=5    ,       UpperCamelCase_ :							Union[str, Any]=4    ,       UpperCamelCase_ :							List[str]=37    ,       UpperCamelCase_ :							List[str]="gelu"    ,       UpperCamelCase_ :							int=0.1    ,       UpperCamelCase_ :							Any=0.1    ,       UpperCamelCase_ :							int=5_12    ,       UpperCamelCase_ :							Optional[int]=16    ,       UpperCamelCase_ :							int=2    ,       UpperCamelCase_ :							Tuple=0.02    ,       UpperCamelCase_ :							List[str]=3    ,       UpperCamelCase_ :							int=4    ,       UpperCamelCase_ :							Optional[Any]=None    ,       )->      List[str]:
     _lowercase:			Optional[Any]      =     parent
     _lowercase:			Optional[int]      =     batch_size
     _lowercase:			Dict      =     seq_length
     _lowercase:			str      =     is_training
     _lowercase:			Any      =     use_input_mask
     _lowercase:			Optional[Any]      =     use_token_type_ids
     _lowercase:			str      =     use_labels
     _lowercase:			List[str]      =     vocab_size
     _lowercase:			str      =     hidden_size
     _lowercase:			Any      =     input_embedding_size
     _lowercase:			Any      =     output_embedding_size
     _lowercase:			Dict      =     num_hidden_layers
     _lowercase:			Dict      =     num_attention_heads
     _lowercase:			List[Any]      =     intermediate_size
     _lowercase:			Any      =     hidden_act
     _lowercase:			Optional[Any]      =     hidden_dropout_prob
     _lowercase:			Optional[int]      =     attention_probs_dropout_prob
     _lowercase:			Optional[Any]      =     max_position_embeddings
     _lowercase:			Union[str, Any]      =     type_vocab_size
     _lowercase:			List[str]      =     type_sequence_label_size
     _lowercase:			Tuple      =     initializer_range
     _lowercase:			List[str]      =     num_labels
     _lowercase:			Optional[int]      =     num_choices
     _lowercase:			List[str]      =     scope






   def 			_a    (						self :							Dict  )->      Union[str, Any]:
     _lowercase:			List[Any]      =     ids_tensor([self.batch_size, self.seq_length]    ,       self.vocab_size  )

     _lowercase:			Tuple      =     None
     if self.use_input_mask:
       _lowercase:			Optional[Any]      =     random_attention_mask([self.batch_size, self.seq_length]  )

     _lowercase:			Optional[Any]      =     None
     if self.use_token_type_ids:
       _lowercase:			List[str]      =     ids_tensor([self.batch_size, self.seq_length]    ,       self.type_vocab_size  )

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

     _lowercase:			str      =     RemBertConfig(
         vocab_size=self.vocab_size    ,       hidden_size=self.hidden_size    ,       input_embedding_size=self.input_embedding_size    ,       output_embedding_size=self.output_embedding_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=UpperCamelCase_    ,       initializer_range=self.initializer_range    ,       )

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






   def 			_a    (						self :							Union[str, Any]  )->      Any:
     (
         (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       
     ):			Any      =     self.prepare_config_and_inputs()

     _lowercase:			Dict      =     True
     _lowercase:			str      =     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
     _lowercase:			Any      =     ids_tensor([self.batch_size, self.seq_length]    ,       vocab_size=2  )

     return (
         config,
         input_ids,
         token_type_ids,
         input_mask,
         sequence_labels,
         token_labels,
         choice_labels,
         encoder_hidden_states,
         encoder_attention_mask,
     )






   def 			_a    (						self :							Optional[Any]    ,       UpperCamelCase_ :							List[Any]    ,       UpperCamelCase_ :							List[Any]    ,       UpperCamelCase_ :							str    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							List[str]    ,       UpperCamelCase_ :							List[str]    ,       UpperCamelCase_ :							Union[str, Any]  )->      Optional[int]:
     _lowercase:			int      =     RemBertModel(config=UpperCamelCase_  )
     model.to(UpperCamelCase_  )
     model.eval()
     _lowercase:			str      =     model(UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_  )
     _lowercase:			Tuple      =     model(UpperCamelCase_    ,       token_type_ids=UpperCamelCase_  )
     _lowercase:			Optional[int]      =     model(UpperCamelCase_  )
     self.parent.assertEqual(result.last_hidden_state.shape    ,       (self.batch_size, self.seq_length, self.hidden_size)  )






   def 			_a    (						self :							List[str]    ,       UpperCamelCase_ :							Union[str, Any]    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Tuple    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Dict    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Dict    ,       UpperCamelCase_ :							List[Any]    ,       UpperCamelCase_ :							Union[str, Any]    ,       )->      Optional[Any]:
     _lowercase:			List[str]      =     True
     _lowercase:			Optional[int]      =     RemBertModel(UpperCamelCase_  )
     model.to(UpperCamelCase_  )
     model.eval()
     _lowercase:			Dict      =     model(
         UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_    ,       encoder_hidden_states=UpperCamelCase_    ,       encoder_attention_mask=UpperCamelCase_    ,       )
     _lowercase:			Union[str, Any]      =     model(
         UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_    ,       encoder_hidden_states=UpperCamelCase_    ,       )
     _lowercase:			List[Any]      =     model(UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_  )
     self.parent.assertEqual(result.last_hidden_state.shape    ,       (self.batch_size, self.seq_length, self.hidden_size)  )






   def 			_a    (						self :							Any    ,       UpperCamelCase_ :							Union[str, Any]    ,       UpperCamelCase_ :							List[str]    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							Union[str, Any]    ,       UpperCamelCase_ :							List[str]    ,       UpperCamelCase_ :							List[Any]    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Optional[int]    ,       UpperCamelCase_ :							List[str]    ,       )->      Any:
     _lowercase:			Any      =     RemBertForCausalLM(config=UpperCamelCase_  )
     model.to(UpperCamelCase_  )
     model.eval()
     _lowercase:			Dict      =     model(UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_    ,       labels=UpperCamelCase_  )
     self.parent.assertEqual(result.logits.shape    ,       (self.batch_size, self.seq_length, self.vocab_size)  )






   def 			_a    (						self :							List[Any]    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							Dict    ,       UpperCamelCase_ :							Dict    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							Any    ,       UpperCamelCase_ :							Union[str, Any]  )->      Optional[int]:
     _lowercase:			Union[str, Any]      =     RemBertForMaskedLM(config=UpperCamelCase_  )
     model.to(UpperCamelCase_  )
     model.eval()
     _lowercase:			Optional[int]      =     model(UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_    ,       labels=UpperCamelCase_  )
     self.parent.assertEqual(result.logits.shape    ,       (self.batch_size, self.seq_length, self.vocab_size)  )






   def 			_a    (						self :							int    ,       UpperCamelCase_ :							Tuple    ,       UpperCamelCase_ :							str    ,       UpperCamelCase_ :							List[Any]    ,       UpperCamelCase_ :							Any    ,       UpperCamelCase_ :							Optional[int]    ,       UpperCamelCase_ :							Dict    ,       UpperCamelCase_ :							List[Any]    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							Any    ,       )->      Tuple:
     _lowercase:			str      =     True
     _lowercase:			List[str]      =     True
     _lowercase:			List[str]      =     RemBertForCausalLM(config=UpperCamelCase_  )
     model.to(UpperCamelCase_  )
     model.eval()

     # first forward pass
     _lowercase:			str      =     model(
         UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       encoder_hidden_states=UpperCamelCase_    ,       encoder_attention_mask=UpperCamelCase_    ,       use_cache=UpperCamelCase_    ,       )
     _lowercase:			Dict      =     outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     _lowercase:			Optional[int]      =     ids_tensor((self.batch_size, 3)    ,       config.vocab_size  )
     _lowercase:			Optional[int]      =     ids_tensor((self.batch_size, 3)    ,       vocab_size=2  )

     # append to next input_ids and
     _lowercase:			Any      =     torch.cat([input_ids, next_tokens]    ,       dim=-1  )
     _lowercase:			str      =     torch.cat([input_mask, next_mask]    ,       dim=-1  )

     _lowercase:			str      =     model(
         UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       encoder_hidden_states=UpperCamelCase_    ,       encoder_attention_mask=UpperCamelCase_    ,       output_hidden_states=UpperCamelCase_    ,       )['hidden_states'][0]
     _lowercase:			Optional[int]      =     model(
         UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       encoder_hidden_states=UpperCamelCase_    ,       encoder_attention_mask=UpperCamelCase_    ,       past_key_values=UpperCamelCase_    ,       output_hidden_states=UpperCamelCase_    ,       )['hidden_states'][0]

     # select random slice
     _lowercase:			Optional[int]      =     ids_tensor((1,)    ,       output_from_past.shape[-1]  ).item()
     _lowercase:			str      =     output_from_no_past[:, -3:, random_slice_idx].detach()
     _lowercase:			List[Any]      =     output_from_past[:, :, random_slice_idx].detach()

     self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]  )

     # test that outputs are equal for slice
     self.parent.assertTrue(torch.allclose(UpperCamelCase_    ,       UpperCamelCase_    ,       atol=1E-3  )  )






   def 			_a    (						self :							int    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							str    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							List[str]    ,       UpperCamelCase_ :							List[str]    ,       UpperCamelCase_ :							Tuple    ,       UpperCamelCase_ :							Optional[int]  )->      Union[str, Any]:
     _lowercase:			List[str]      =     RemBertForQuestionAnswering(config=UpperCamelCase_  )
     model.to(UpperCamelCase_  )
     model.eval()
     _lowercase:			List[Any]      =     model(
         UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_    ,       start_positions=UpperCamelCase_    ,       end_positions=UpperCamelCase_    ,       )
     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 			_a    (						self :							List[Any]    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							List[str]    ,       UpperCamelCase_ :							List[Any]    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							Dict  )->      Optional[int]:
     _lowercase:			int      =     self.num_labels
     _lowercase:			str      =     RemBertForSequenceClassification(UpperCamelCase_  )
     model.to(UpperCamelCase_  )
     model.eval()
     _lowercase:			int      =     model(UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_    ,       labels=UpperCamelCase_  )
     self.parent.assertEqual(result.logits.shape    ,       (self.batch_size, self.num_labels)  )






   def 			_a    (						self :							Tuple    ,       UpperCamelCase_ :							Any    ,       UpperCamelCase_ :							List[Any]    ,       UpperCamelCase_ :							Optional[int]    ,       UpperCamelCase_ :							Any    ,       UpperCamelCase_ :							str    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							Dict  )->      Optional[Any]:
     _lowercase:			List[str]      =     self.num_labels
     _lowercase:			int      =     RemBertForTokenClassification(config=UpperCamelCase_  )
     model.to(UpperCamelCase_  )
     model.eval()
     _lowercase:			Dict      =     model(UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_    ,       labels=UpperCamelCase_  )
     self.parent.assertEqual(result.logits.shape    ,       (self.batch_size, self.seq_length, self.num_labels)  )






   def 			_a    (						self :							Union[str, Any]    ,       UpperCamelCase_ :							Any    ,       UpperCamelCase_ :							Any    ,       UpperCamelCase_ :							Optional[Any]    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Any    ,       UpperCamelCase_ :							Any    ,       UpperCamelCase_ :							Dict  )->      Dict:
     _lowercase:			int      =     self.num_choices
     _lowercase:			Any      =     RemBertForMultipleChoice(config=UpperCamelCase_  )
     model.to(UpperCamelCase_  )
     model.eval()
     _lowercase:			Any      =     input_ids.unsqueeze(1  ).expand(-1    ,       self.num_choices    ,       -1  ).contiguous()
     _lowercase:			Any      =     token_type_ids.unsqueeze(1  ).expand(-1    ,       self.num_choices    ,       -1  ).contiguous()
     _lowercase:			Union[str, Any]      =     input_mask.unsqueeze(1  ).expand(-1    ,       self.num_choices    ,       -1  ).contiguous()
     _lowercase:			Optional[int]      =     model(
         UpperCamelCase_    ,       attention_mask=UpperCamelCase_    ,       token_type_ids=UpperCamelCase_    ,       labels=UpperCamelCase_    ,       )
     self.parent.assertEqual(result.logits.shape    ,       (self.batch_size, self.num_choices)  )






   def 			_a    (						self :							str  )->      Dict:
     _lowercase:			Tuple      =     self.prepare_config_and_inputs()
     (
         (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       
     ):			Any      =     config_and_inputs
     _lowercase:			Dict      =     {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
     return config, inputs_dict

@require_torch
class   UpperCAmelCase    (      A ,       A ,       unittest.TestCase ):
   __SCREAMING_SNAKE_CASE		:				List[str]      					=       (
       (
           RemBertModel,
           RemBertForMaskedLM,
           RemBertForCausalLM,
           RemBertForMultipleChoice,
           RemBertForQuestionAnswering,
           RemBertForSequenceClassification,
           RemBertForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   __SCREAMING_SNAKE_CASE		:				str      					=       (RemBertForCausalLM,) if is_torch_available() else ()
   __SCREAMING_SNAKE_CASE		:				Union[str, Any]      					=       (
       {
           "feature-extraction": RemBertModel,
           "fill-mask": RemBertForMaskedLM,
           "question-answering": RemBertForQuestionAnswering,
           "text-classification": RemBertForSequenceClassification,
           "text-generation": RemBertForCausalLM,
           "token-classification": RemBertForTokenClassification,
           "zero-shot": RemBertForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )






   def 			_a    (						self :							List[Any]  )->      Dict:
     _lowercase:			str      =     RemBertModelTester(self  )
     _lowercase:			List[str]      =     ConfigTester(self    ,       config_class=UpperCamelCase_    ,       hidden_size=37  )






   def 			_a    (						self :							Dict  )->      str:
     self.config_tester.run_common_tests()






   def 			_a    (						self :							Optional[Any]  )->      List[str]:
     _lowercase:			Optional[int]      =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*UpperCamelCase_  )






   def 			_a    (						self :							str  )->      Dict:
     _lowercase:			Tuple      =     self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
       _lowercase:			Dict      =     type
       self.model_tester.create_and_check_model(*UpperCamelCase_  )






   def 			_a    (						self :							Dict  )->      Dict:
     _lowercase:			List[str]      =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase_  )






   def 			_a    (						self :							Tuple  )->      Tuple:
     _lowercase:			Tuple      =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase_  )






   def 			_a    (						self :							Tuple  )->      Optional[int]:
     _lowercase:			Optional[Any]      =     self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_decoder_model_past_large_inputs(*UpperCamelCase_  )






   def 			_a    (						self :							List[str]  )->      List[str]:
     _lowercase:			Union[str, Any]      =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*UpperCamelCase_  )






   def 			_a    (						self :							Union[str, Any]  )->      Tuple:
     _lowercase:			Optional[Any]      =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase_  )






   def 			_a    (						self :							Union[str, Any]  )->      Dict:
     _lowercase:			Optional[int]      =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*UpperCamelCase_  )






   def 			_a    (						self :							Dict  )->      Dict:
     _lowercase:			Dict      =     self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase_  )






   def 			_a    (						self :							Any  )->      Union[str, Any]:
     # This regression test was failing with PyTorch < 1.3
     (
         (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       (
         _lowercase
     )  ,       
     ):			int      =     self.model_tester.prepare_config_and_inputs_for_decoder()

     _lowercase:			str      =     None

     self.model_tester.create_and_check_model_as_decoder(
         UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       )






   @slow
   def 			_a    (						self :							Tuple  )->      Dict:
     for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _lowercase:			str      =     RemBertModel.from_pretrained(UpperCamelCase_  )
       self.assertIsNotNone(UpperCamelCase_  )

@require_torch
class   UpperCAmelCase    (      unittest.TestCase ):






   @slow
   def 			_a    (						self :							Dict  )->      Optional[int]:
     # Test exact values at the last hidden layer
     _lowercase:			List[Any]      =     RemBertModel.from_pretrained('google/rembert'  )
     _lowercase:			Optional[Any]      =     torch.tensor([[3_12, 5_64_98, 3_13, 21_25, 3_13]]  )
     _lowercase:			Any      =     torch.tensor([[0, 0, 0, 1, 1]]  )
     with torch.no_grad():
       _lowercase:			Optional[int]      =     model(UpperCamelCase_    ,       token_type_ids=UpperCamelCase_    ,       output_hidden_states=UpperCamelCase_  )

     _lowercase:			List[str]      =     11_52

     _lowercase:			Union[str, Any]      =     torch.Size((1, 5, hidden_size)  )
     self.assertEqual(output['last_hidden_state'].shape    ,       UpperCamelCase_  )

     _lowercase:			Optional[Any]      =     torch.tensor(
         [
             [
                 [0.0_754, -0.2_022, 0.1_904],
                 [-0.3_354, -0.3_692, -0.4_791],
                 [-0.2_314, -0.6_729, -0.0_749],
                 [-0.0_396, -0.3_105, -0.4_234],
                 [-0.1_571, -0.0_525, 0.5_353],
             ]
         ]  )

     # Running on the original tf implementation gives slightly different results here.
     # Not clear why this variations is present
     # TODO: Find reason for discrepancy
     # expected_original_implementation = [[
     #     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
     #     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
     #     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
     #     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
     #     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
     # ]]

     self.assertTrue(torch.allclose(output['last_hidden_state'][:, :, :3]    ,       UpperCamelCase_    ,       atol=1E-4  )  )






import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()

class   UpperCAmelCase    (      A ,       unittest.TestCase ):
   __SCREAMING_SNAKE_CASE		:				int      					=       KandinskyPipeline
   __SCREAMING_SNAKE_CASE		:				Union[str, Any]      					=       [
       "prompt",
       "image_embeds",
       "negative_image_embeds",
   ]
   __SCREAMING_SNAKE_CASE		:				Tuple      					=       ["prompt", "negative_prompt", "image_embeds", "negative_image_embeds"]
   __SCREAMING_SNAKE_CASE		:				Tuple      					=       [
       "generator",
       "height",
       "width",
       "latents",
       "guidance_scale",
       "negative_prompt",
       "num_inference_steps",
       "return_dict",
       "guidance_scale",
       "num_images_per_prompt",
       "output_type",
       "return_dict",
   ]
   __SCREAMING_SNAKE_CASE		:				Dict      					=       False






   @property
   def 			_a    (						self :							str  )->      Dict:
     return 32






   @property
   def 			_a    (						self :							Tuple  )->      Any:
     return 32






   @property
   def 			_a    (						self :							Union[str, Any]  )->      List[str]:
     return self.time_input_dim






   @property
   def 			_a    (						self :							Optional[Any]  )->      Tuple:
     return self.time_input_dim * 4






   @property
   def 			_a    (						self :							int  )->      Dict:
     return 1_00






   @property
   def 			_a    (						self :							Tuple  )->      Tuple:
     _lowercase:			List[Any]      =     XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'  )
     return tokenizer






   @property
   def 			_a    (						self :							List[str]  )->      Optional[int]:
     torch.manual_seed(0  )
     _lowercase:			int      =     MCLIPConfig(
         numDims=self.cross_attention_dim    ,       transformerDimensions=self.text_embedder_hidden_size    ,       hidden_size=self.text_embedder_hidden_size    ,       intermediate_size=37    ,       num_attention_heads=4    ,       num_hidden_layers=5    ,       vocab_size=10_05    ,       )

     _lowercase:			Optional[Any]      =     MultilingualCLIP(UpperCamelCase_  )
     _lowercase:			int      =     text_encoder.eval()

     return text_encoder






   @property
   def 			_a    (						self :							Optional[int]  )->      Optional[Any]:
     torch.manual_seed(0  )

     _lowercase:			Optional[Any]      =     {
         'in_channels': 4,
         # Out channels is double in channels because predicts mean and variance
         'out_channels': 8,
         'addition_embed_type': 'text_image',
         'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
         'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
         'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
         'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
         'layers_per_block': 1,
         'encoder_hid_dim': self.text_embedder_hidden_size,
         'encoder_hid_dim_type': 'text_image_proj',
         'cross_attention_dim': self.cross_attention_dim,
         'attention_head_dim': 4,
         'resnet_time_scale_shift': 'scale_shift',
         'class_embed_type': None,
     }

     _lowercase:			str      =     UNetaDConditionModel(**UpperCamelCase_  )
     return model






   @property
   def 			_a    (						self :							Dict  )->      Optional[Any]:
     return {
         "block_out_channels": [32, 64],
         "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
         "in_channels": 3,
         "latent_channels": 4,
         "layers_per_block": 1,
         "norm_num_groups": 8,
         "norm_type": "spatial",
         "num_vq_embeddings": 12,
         "out_channels": 3,
         "up_block_types": [
             "AttnUpDecoderBlock2D",
             "UpDecoderBlock2D",
         ],
         "vq_embed_dim": 4,
     }






   @property
   def 			_a    (						self :							List[str]  )->      int:
     torch.manual_seed(0  )
     _lowercase:			List[str]      =     VQModel(**self.dummy_movq_kwargs  )
     return model






   def 			_a    (						self :							int  )->      List[Any]:
     _lowercase:			List[str]      =     self.dummy_text_encoder
     _lowercase:			List[Any]      =     self.dummy_tokenizer
     _lowercase:			Optional[Any]      =     self.dummy_unet
     _lowercase:			List[str]      =     self.dummy_movq

     _lowercase:			Optional[Any]      =     DDIMScheduler(
         num_train_timesteps=10_00    ,       beta_schedule='linear'    ,       beta_start=0.00_085    ,       beta_end=0.012    ,       clip_sample=UpperCamelCase_    ,       set_alpha_to_one=UpperCamelCase_    ,       steps_offset=1    ,       prediction_type='epsilon'    ,       thresholding=UpperCamelCase_    ,       )

     _lowercase:			List[Any]      =     {
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'unet': unet,
         'scheduler': scheduler,
         'movq': movq,
     }
     return components






   def 			_a    (						self :							str    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Optional[int]=0  )->      Union[str, Any]:
     _lowercase:			str      =     floats_tensor((1, self.cross_attention_dim)    ,       rng=random.Random(UpperCamelCase_  )  ).to(UpperCamelCase_  )
     _lowercase:			Optional[int]      =     floats_tensor((1, self.cross_attention_dim)    ,       rng=random.Random(seed + 1  )  ).to(UpperCamelCase_  )
     if str(UpperCamelCase_  ).startswith('mps'  ):
       _lowercase:			Any      =     torch.manual_seed(UpperCamelCase_  )
     else:
       _lowercase:			Tuple      =     torch.Generator(device=UpperCamelCase_  ).manual_seed(UpperCamelCase_  )
     _lowercase:			List[str]      =     {
         'prompt': 'horse',
         'image_embeds': image_embeds,
         'negative_image_embeds': negative_image_embeds,
         'generator': generator,
         'height': 64,
         'width': 64,
         'guidance_scale': 4.0,
         'num_inference_steps': 2,
         'output_type': 'np',
     }
     return inputs






   def 			_a    (						self :							List[str]  )->      List[str]:
     _lowercase:			str      =     'cpu'

     _lowercase:			Optional[Any]      =     self.get_dummy_components()

     _lowercase:			int      =     self.pipeline_class(**UpperCamelCase_  )
     _lowercase:			List[str]      =     pipe.to(UpperCamelCase_  )

     pipe.set_progress_bar_config(disable=UpperCamelCase_  )

     _lowercase:			Tuple      =     pipe(**self.get_dummy_inputs(UpperCamelCase_  )  )
     _lowercase:			Union[str, Any]      =     output.images

     _lowercase:			Union[str, Any]      =     pipe(
         **self.get_dummy_inputs(UpperCamelCase_  )    ,       return_dict=UpperCamelCase_    ,       )[0]

     _lowercase:			int      =     image[0, -3:, -3:, -1]
     _lowercase:			List[str]      =     image_from_tuple[0, -3:, -3:, -1]

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

     _lowercase:			Union[str, Any]      =     np.array(
         [0.328_663, 1.0, 0.23_216_873, 1.0, 0.92_717_564, 0.4_639_046, 0.96_894_777, 0.31_713_378, 0.6_293_953]  )

     assert (
         np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
     ), f' expected_slice {expected_slice}, but got {image_slice.flatten()}'
     assert (
         np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1E-2
     ), f' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'

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






   def 			_a    (						self :							Optional[int]  )->      Any:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






   def 			_a    (						self :							List[Any]  )->      Optional[int]:
     _lowercase:			Optional[int]      =     load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
         '/kandinsky/kandinsky_text2img_cat_fp16.npy'  )

     _lowercase:			List[Any]      =     KandinskyPriorPipeline.from_pretrained(
         'kandinsky-community/kandinsky-2-1-prior'    ,       torch_dtype=torch.floataa  )
     pipe_prior.to(UpperCamelCase_  )

     _lowercase:			str      =     KandinskyPipeline.from_pretrained('kandinsky-community/kandinsky-2-1'    ,       torch_dtype=torch.floataa  )
     _lowercase:			List[str]      =     pipeline.to(UpperCamelCase_  )
     pipeline.set_progress_bar_config(disable=UpperCamelCase_  )

     _lowercase:			Dict      =     'red cat, 4k photo'

     _lowercase:			List[Any]      =     torch.Generator(device='cuda'  ).manual_seed(0  )
     _lowercase  ,       _lowercase:			Dict      =     pipe_prior(
         UpperCamelCase_    ,       generator=UpperCamelCase_    ,       num_inference_steps=5    ,       negative_prompt=''    ,       ).to_tuple()

     _lowercase:			str      =     torch.Generator(device='cuda'  ).manual_seed(0  )
     _lowercase:			Optional[int]      =     pipeline(
         UpperCamelCase_    ,       image_embeds=UpperCamelCase_    ,       negative_image_embeds=UpperCamelCase_    ,       generator=UpperCamelCase_    ,       num_inference_steps=1_00    ,       output_type='np'    ,       )

     _lowercase:			Optional[int]      =     output.images[0]

     assert image.shape == (5_12, 5_12, 3)

     assert_mean_pixel_difference(UpperCamelCase_    ,       UpperCamelCase_  )
