

from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

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


if is_tf_available():
 import tensorflow as tf

 from transformers.models.gptj.modeling_tf_gptj import (
     TFGPTJForCausalLM,
     TFGPTJForQuestionAnswering,
     TFGPTJForSequenceClassification,
     TFGPTJModel,
     shape_list,
 )



class 					lowercase		:



 '''simple docstring'''


 def __init__(self , __a				)				->      Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       parent
  UpperCAmelCase__           =       13
  UpperCAmelCase__           =       7
  UpperCAmelCase__           =       True
  UpperCAmelCase__           =       True
  UpperCAmelCase__           =       True
  UpperCAmelCase__           =       True
  UpperCAmelCase__           =       True
  UpperCAmelCase__           =       99
  UpperCAmelCase__           =       32
  UpperCAmelCase__           =       4
  UpperCAmelCase__           =       2
  UpperCAmelCase__           =       4
  UpperCAmelCase__           =       37
  UpperCAmelCase__           =       'gelu'
  UpperCAmelCase__           =       0.1
  UpperCAmelCase__           =       0.1
  UpperCAmelCase__           =       512
  UpperCAmelCase__           =       16
  UpperCAmelCase__           =       2
  UpperCAmelCase__           =       0.02
  UpperCAmelCase__           =       3
  UpperCAmelCase__           =       4
  UpperCAmelCase__           =       None
  UpperCAmelCase__           =       self.vocab_size - 1
  UpperCAmelCase__           =       self.vocab_size - 1
  UpperCAmelCase__           =       self.vocab_size - 1




 def 						UpperCamelCase__							(self				)				->      List[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       ids_tensor([self.batch_size, self.seq_length] , self.vocab_size				)

  UpperCAmelCase__           =       None
  if self.use_input_mask:
   UpperCAmelCase__           =       random_attention_mask([self.batch_size, self.seq_length]				)

  UpperCAmelCase__           =       None
  if self.use_token_type_ids:
   UpperCAmelCase__           =       ids_tensor([self.batch_size, self.seq_length] , self.type_vocab_size				)

  UpperCAmelCase__           =       None
  if self.use_mc_token_ids:
   UpperCAmelCase__           =       ids_tensor([self.batch_size, self.num_choices] , self.seq_length				)

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

  UpperCAmelCase__           =       GPTJConfig(
      vocab_size=self.vocab_size , n_embd=self.hidden_size , n_layer=self.num_hidden_layers , n_head=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 , n_positions=self.max_position_embeddings , type_vocab_size=self.type_vocab_size , initializer_range=self.initializer_range , bos_token_id=self.bos_token_id , eos_token_id=self.eos_token_id , pad_token_id=self.pad_token_id , rotary_dim=self.rotary_dim , return_dict=lowerCamelCase__ , )

  UpperCAmelCase__           =       ids_tensor([self.num_hidden_layers, self.num_attention_heads] , 2				)

  return (
      config,
      input_ids,
      input_mask,
      head_mask,
      token_type_ids,
      mc_token_ids,
      sequence_labels,
      token_labels,
      choice_labels,
  )




 def 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , *__a				)				->      Dict:






  """simple docstring"""





  UpperCAmelCase__           =       TFGPTJModel(config=lowerCamelCase__				)
  UpperCAmelCase__           =       {
      'input_ids': input_ids,
      'attention_mask': input_mask,
      'token_type_ids': token_type_ids,
  }
  UpperCAmelCase__           =       model(lowerCamelCase__				)

  UpperCAmelCase__           =       [input_ids, None, input_mask]  # None is the input for 'past'
  UpperCAmelCase__           =       model(lowerCamelCase__				)

  UpperCAmelCase__           =       model(lowerCamelCase__				)

  self.parent.assertEqual(result.last_hidden_state.shape , (self.batch_size, self.seq_length, self.hidden_size)				)




 def 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , *__a				)				->      str:






  """simple docstring"""





  UpperCAmelCase__           =       TFGPTJModel(config=lowerCamelCase__				)

  # first forward pass
  UpperCAmelCase__           =       model(lowerCamelCase__ , token_type_ids=lowerCamelCase__ , use_cache=lowerCamelCase__				)
  UpperCAmelCase__           =       model(lowerCamelCase__ , token_type_ids=lowerCamelCase__				)
  UpperCAmelCase__           =       model(lowerCamelCase__ , token_type_ids=lowerCamelCase__ , use_cache=lowerCamelCase__				)

  self.parent.assertTrue(len(lowerCamelCase__				) == len(lowerCamelCase__				)				)
  self.parent.assertTrue(len(lowerCamelCase__				) == len(lowerCamelCase__				) + 1				)

  UpperCAmelCase__       ,					UpperCAmelCase__           =       outputs.to_tuple()

  # create hypothetical next token and extent to next_input_ids
  UpperCAmelCase__           =       ids_tensor((self.batch_size, 1) , config.vocab_size				)
  UpperCAmelCase__           =       ids_tensor([self.batch_size, 1] , self.type_vocab_size				)

  # append to next input_ids and token_type_ids
  UpperCAmelCase__           =       tf.concat([input_ids, next_tokens] , axis=-1				)
  UpperCAmelCase__           =       tf.concat([token_type_ids, next_token_types] , axis=-1				)

  UpperCAmelCase__           =       model(lowerCamelCase__ , token_type_ids=lowerCamelCase__				)['last_hidden_state']
  UpperCAmelCase__           =       model(lowerCamelCase__ , token_type_ids=lowerCamelCase__ , past_key_values=lowerCamelCase__				)[
      'last_hidden_state'
  ]

  # select random slice
  UpperCAmelCase__           =       int(ids_tensor((1,) , shape_list(lowerCamelCase__				)[-1]				)				)
  UpperCAmelCase__           =       output_from_no_past[:, -1, random_slice_idx]
  UpperCAmelCase__           =       output_from_past[:, 0, random_slice_idx]

  # test that outputs are equal for slice
  tf.debugging.assert_near(lowerCamelCase__ , lowerCamelCase__ , rtol=1E-6				)




 def 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , *__a				)				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__           =       TFGPTJModel(config=lowerCamelCase__				)

  # create attention mask
  UpperCAmelCase__           =       self.seq_length // 2
  UpperCAmelCase__           =       tf.ones((self.batch_size, half_seq_length) , dtype=tf.intaa				)
  UpperCAmelCase__           =       tf.zeros((self.batch_size, self.seq_length - half_seq_length) , dtype=tf.intaa				)
  UpperCAmelCase__           =       tf.concat([attn_mask_begin, attn_mask_end] , axis=1				)

  # first forward pass
  UpperCAmelCase__       ,					UpperCAmelCase__           =       model(lowerCamelCase__ , attention_mask=lowerCamelCase__				).to_tuple()

  # create hypothetical next token and extent to next_input_ids
  UpperCAmelCase__           =       ids_tensor((self.batch_size, 1) , config.vocab_size				)

  # change a random masked slice from input_ids
  UpperCAmelCase__           =       ids_tensor((1,) , lowerCamelCase__				).numpy() + 1
  UpperCAmelCase__           =       ids_tensor((self.batch_size, self.seq_length) , config.vocab_size				)
  UpperCAmelCase__           =       tf.range(self.seq_length				) == (self.seq_length - random_seq_idx_to_change)
  UpperCAmelCase__           =       tf.transpose(
      tf.broadcast_to(tf.expand_dims(lowerCamelCase__ , -1				) , (self.seq_length, self.batch_size)				)				)
  UpperCAmelCase__           =       tf.where(lowerCamelCase__ , lowerCamelCase__ , lowerCamelCase__				)

  # append to next input_ids and attn_mask
  UpperCAmelCase__           =       tf.concat([input_ids, next_tokens] , axis=-1				)
  UpperCAmelCase__           =       tf.concat([attn_mask, tf.ones((shape_list(lowerCamelCase__				)[0], 1) , dtype=tf.intaa				)] , axis=1				)

  # get two different outputs
  UpperCAmelCase__           =       model(lowerCamelCase__ , attention_mask=lowerCamelCase__				)['last_hidden_state']
  UpperCAmelCase__           =       model(lowerCamelCase__ , past_key_values=lowerCamelCase__ , attention_mask=lowerCamelCase__				)[
      'last_hidden_state'
  ]

  # select random slice
  UpperCAmelCase__           =       int(ids_tensor((1,) , shape_list(lowerCamelCase__				)[-1]				)				)
  UpperCAmelCase__           =       output_from_no_past[:, -1, random_slice_idx]
  UpperCAmelCase__           =       output_from_past[:, 0, random_slice_idx]

  # test that outputs are equal for slice
  tf.debugging.assert_near(lowerCamelCase__ , lowerCamelCase__ , rtol=1E-1_2				)




 def 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , *__a				)				->      Tuple:






  """simple docstring"""





  UpperCAmelCase__           =       TFGPTJModel(config=lowerCamelCase__				)

  UpperCAmelCase__           =       input_ids[:1, :]
  UpperCAmelCase__           =       input_mask[:1, :]
  UpperCAmelCase__           =       token_type_ids[:1, :]
  UpperCAmelCase__           =       1

  # first forward pass
  UpperCAmelCase__           =       model(lowerCamelCase__ , attention_mask=lowerCamelCase__ , token_type_ids=lowerCamelCase__ , use_cache=lowerCamelCase__				)

  UpperCAmelCase__       ,					UpperCAmelCase__           =       outputs.to_tuple()

  # create hypothetical next token and extent to next_input_ids
  UpperCAmelCase__           =       ids_tensor((self.batch_size, 3) , config.vocab_size				)
  UpperCAmelCase__           =       ids_tensor((self.batch_size, 3) , 2				)
  UpperCAmelCase__           =       ids_tensor((self.batch_size, 3) , self.type_vocab_size				)

  # append to next input_ids and token_type_ids
  UpperCAmelCase__           =       tf.concat([input_ids, next_tokens] , axis=-1				)
  UpperCAmelCase__           =       tf.concat([input_mask, next_attn_mask] , axis=-1				)
  UpperCAmelCase__           =       tf.concat([token_type_ids, next_token_types] , axis=-1				)

  UpperCAmelCase__           =       model(
      lowerCamelCase__ , token_type_ids=lowerCamelCase__ , attention_mask=lowerCamelCase__				)['last_hidden_state']
  UpperCAmelCase__           =       model(
      lowerCamelCase__ , token_type_ids=lowerCamelCase__ , attention_mask=lowerCamelCase__ , past_key_values=lowerCamelCase__ , )['last_hidden_state']
  self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]				)

  # select random slice
  UpperCAmelCase__           =       int(ids_tensor((1,) , shape_list(lowerCamelCase__				)[-1]				)				)
  UpperCAmelCase__           =       output_from_no_past[:, -3:, random_slice_idx]
  UpperCAmelCase__           =       output_from_past[:, :, random_slice_idx]

  # test that outputs are equal for slice
  tf.debugging.assert_near(lowerCamelCase__ , lowerCamelCase__ , rtol=1E-3				)




 def 						UpperCamelCase__							(self , __a , __a , __a , __a , __a , *__a				)				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__           =       TFGPTJForCausalLM(config=lowerCamelCase__				)
  UpperCAmelCase__           =       {
      'input_ids': input_ids,
      'attention_mask': input_mask,
      'token_type_ids': token_type_ids,
  }
  UpperCAmelCase__           =       model(lowerCamelCase__				)
  self.parent.assertEqual(result.logits.shape , (self.batch_size, self.seq_length, self.vocab_size)				)




 def 						UpperCamelCase__							(self				)				->      List[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       self.prepare_config_and_inputs()

  (
      (
      UpperCAmelCase__						
  )       ,					(
      UpperCAmelCase__						
  )       ,					(
      UpperCAmelCase__						
  )       ,					(
      UpperCAmelCase__						
  )       ,					(
      UpperCAmelCase__						
  )       ,					(
      UpperCAmelCase__						
  )       ,					(
      UpperCAmelCase__						
  )       ,					(
      UpperCAmelCase__						
  )       ,					(
      UpperCAmelCase__						
  )       ,											
  )           =       config_and_inputs

  UpperCAmelCase__           =       {
      'input_ids': input_ids,
      'token_type_ids': token_type_ids,
      'attention_mask': input_mask,
  }
  return config, inputs_dict







@require_tf
class 					lowercase		( __UpperCAmelCase				,       __UpperCAmelCase				,       __UpperCAmelCase				,       unittest.TestCase      ):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE      							=       (
     (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
     if is_tf_available()
     else ()
 )

 __SCREAMING_SNAKE_CASE      							=       (TFGPTJForCausalLM,) if is_tf_available() else ()
 __SCREAMING_SNAKE_CASE      							=       (
     {
         """feature-extraction""": TFGPTJModel,
         """question-answering""": TFGPTJForQuestionAnswering,
         """text-classification""": TFGPTJForSequenceClassification,
         """text-generation""": TFGPTJForCausalLM,
         """zero-shot""": TFGPTJForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 __SCREAMING_SNAKE_CASE      							=       False
 __SCREAMING_SNAKE_CASE      							=       False
 __SCREAMING_SNAKE_CASE      							=       False
 __SCREAMING_SNAKE_CASE      							=       False




 def 						UpperCamelCase__							(self , __a , __a , __a , __a , __a				)				->      List[Any]:






  """simple docstring"""





  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				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       TFGPTJModelTester(self				)
  UpperCAmelCase__           =       ConfigTester(self , config_class=lowerCamelCase__ , n_embd=37				)




 def 						UpperCamelCase__							(self				)				->      Optional[Any]:






  """simple docstring"""





  self.config_tester.run_common_tests()




 def 						UpperCamelCase__							(self				)				->      Dict:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_gptj_model(*lowerCamelCase__				)




 def 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_gptj_model_past(*lowerCamelCase__				)




 def 						UpperCamelCase__							(self				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_gptj_model_attention_mask_past(*lowerCamelCase__				)




 def 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_gptj_model_past_large_inputs(*lowerCamelCase__				)




 def 						UpperCamelCase__							(self				)				->      Any:






  """simple docstring"""





  UpperCAmelCase__           =       self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_gptj_lm_head_model(*lowerCamelCase__				)




 @slow
 @unittest.skipIf(
     not is_tf_available() or len(tf.config.list_physical_devices('GPU'				)				) > 0 , 'skip testing on GPU for now to avoid GPU OOM.' , )
 def 						UpperCamelCase__							(self				)				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__           =       TFGPTJModel.from_pretrained('EleutherAI/gpt-j-6B' , from_pt=lowerCamelCase__				)
  self.assertIsNotNone(lowerCamelCase__				)




 @unittest.skip(reason='Currently, model embeddings are going to undergo a major refactor.'				)
 def 						UpperCamelCase__							(self				)				->      Union[str, Any]:






  """simple docstring"""





  super().test_resize_token_embeddings()







@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class 					lowercase		( unittest.TestCase      ):



 '''simple docstring'''


 def 						UpperCamelCase__							(self				)				->      Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B' , from_pt=lowerCamelCase__				)
  UpperCAmelCase__           =       tf.convert_to_tensor([[464, 3290]] , dtype=tf.intaa				)  # The dog
  # fmt: off
  # The dog is a man's best friend. It is a loyal companion, and it is a friend
  UpperCAmelCase__           =       [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 15185, 11, 290, 340, 318, 257, 1545]
  # fmt: on
  UpperCAmelCase__           =       model.generate(lowerCamelCase__ , do_sample=lowerCamelCase__				)
  self.assertListEqual(output_ids[0].numpy().tolist() , lowerCamelCase__				)




 def 						UpperCamelCase__							(self				)				->      Union[str, Any]:






  """simple docstring"""





  UpperCAmelCase__           =       AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B' , revision='float16'				)
  UpperCAmelCase__           =       TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B' , revision='float16' , from_pt=lowerCamelCase__				)

  UpperCAmelCase__           =       tokenizer('Today is a nice day and' , return_tensors='tf'				)
  # forces the generation to happen on CPU, to avoid GPU-related quirks
  with tf.device(':/CPU:0'				):
   UpperCAmelCase__           =       model.generate(**lowerCamelCase__ , do_sample=lowerCamelCase__ , seed=[42, 0]				)
  UpperCAmelCase__           =       tokenizer.decode(output_ids[0] , skip_special_tokens=lowerCamelCase__				)

  UpperCAmelCase__           =       'Today is a nice day and I’m going to go for a walk. I’'
  self.assertEqual(lowerCamelCase__ , lowerCamelCase__				)




 def 						UpperCamelCase__							(self				)				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__           =       TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B' , revision='float16' , from_pt=lowerCamelCase__				)
  UpperCAmelCase__           =       AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B' , revision='float16'				)

  UpperCAmelCase__           =       'left'

  # Define PAD Token = EOS Token = 50256
  UpperCAmelCase__           =       tokenizer.eos_token
  UpperCAmelCase__           =       model.config.eos_token_id

  # use different length sentences to test batching
  UpperCAmelCase__           =       [
      'Hello, my dog is a little',
      'Today, I',
  ]
  UpperCAmelCase__           =       [
      'Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia',
      'Today, I’m going to be talking about a topic that’',
  ]
  return model, tokenizer, sentences, expected_output_sentences




 def 						UpperCamelCase__							(self				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__           =       self._get_beam_search_test_objects()

  UpperCAmelCase__           =       tokenizer(lowerCamelCase__ , return_tensors='tf' , padding=lowerCamelCase__				)
  UpperCAmelCase__           =       model.generate(**lowerCamelCase__ , do_sample=lowerCamelCase__ , num_beams=2				)
  UpperCAmelCase__           =       tokenizer.batch_decode(lowerCamelCase__ , skip_special_tokens=lowerCamelCase__				)
  self.assertListEqual(lowerCamelCase__ , lowerCamelCase__				)




 def 						UpperCamelCase__							(self				)				->      Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__           =       self._get_beam_search_test_objects()

  UpperCAmelCase__           =       tokenizer(lowerCamelCase__ , return_tensors='tf' , padding=lowerCamelCase__				)
  UpperCAmelCase__           =       tokenizer(sentences[0] , return_tensors='tf'				)
  UpperCAmelCase__           =       model.generate(**lowerCamelCase__ , do_sample=lowerCamelCase__ , num_beams=2				)
  UpperCAmelCase__           =       (
      shape_list(inputs_non_padded['input_ids']				)[-1]
      - tf.reduce_sum(tf.cast(inputs['attention_mask'][-1] , tf.intaa				)				).numpy()
  )
  UpperCAmelCase__           =       tokenizer(sentences[1] , return_tensors='tf'				)
  UpperCAmelCase__           =       model.generate(
      **lowerCamelCase__ , do_sample=lowerCamelCase__ , num_beams=2 , max_length=model.config.max_length - num_paddings				)
  UpperCAmelCase__           =       tokenizer.decode(output_non_padded[0] , skip_special_tokens=lowerCamelCase__				)
  UpperCAmelCase__           =       tokenizer.decode(output_padded[0] , skip_special_tokens=lowerCamelCase__				)
  self.assertListEqual(lowerCamelCase__ , [non_padded_sentence, padded_sentence]				)




 def 						UpperCamelCase__							(self				)				->      List[str]:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__       ,					UpperCAmelCase__           =       self._get_beam_search_test_objects()

  UpperCAmelCase__           =       tokenizer(lowerCamelCase__ , return_tensors='tf' , padding=lowerCamelCase__				)
  UpperCAmelCase__           =       tf.function(model.generate , jit_compile=lowerCamelCase__				)
  UpperCAmelCase__           =       xla_generate(**lowerCamelCase__ , do_sample=lowerCamelCase__ , num_beams=2				)
  UpperCAmelCase__           =       tokenizer.batch_decode(lowerCamelCase__ , skip_special_tokens=lowerCamelCase__				)
  self.assertListEqual(lowerCamelCase__ , lowerCamelCase__				)






import random
import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
 import torch

if is_vision_available():
 import PIL

 from transformers import FlavaImageProcessor
 from transformers.image_utils import PILImageResampling
 from transformers.models.flava.image_processing_flava import (
     FLAVA_CODEBOOK_MEAN,
     FLAVA_CODEBOOK_STD,
     FLAVA_IMAGE_MEAN,
     FLAVA_IMAGE_STD,
 )
else:
 _UpperCamelCase     =    _UpperCamelCase     =    _UpperCamelCase     =    _UpperCamelCase     =    None



class 					lowercase		( unittest.TestCase      ):



 '''simple docstring'''


 def __init__(self , __a , __a=7 , __a=3 , __a=30 , __a=400 , __a=True , __a=None , __a=True , __a=None , __a=None , __a=True , __a=1 / 255 , __a=True , __a=FLAVA_IMAGE_MEAN , __a=FLAVA_IMAGE_STD , __a=14 , __a=75 , __a=None , __a=16 , __a=0.3 , __a=None , __a=True , __a=None , __a=None , __a=True , __a=None , __a=True , __a=True , __a=FLAVA_CODEBOOK_MEAN , __a=FLAVA_CODEBOOK_STD , )				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       size if size is not None else {'height': 224, 'width': 224}
  UpperCAmelCase__           =       crop_size if crop_size is not None else {'height': 224, 'width': 224}
  UpperCAmelCase__           =       codebook_size if codebook_size is not None else {'height': 112, 'width': 112}
  UpperCAmelCase__           =       codebook_crop_size if codebook_crop_size is not None else {'height': 112, 'width': 112}

  UpperCAmelCase__           =       parent
  UpperCAmelCase__           =       batch_size
  UpperCAmelCase__           =       num_channels
  UpperCAmelCase__           =       do_resize
  UpperCAmelCase__           =       do_rescale
  UpperCAmelCase__           =       rescale_factor
  UpperCAmelCase__           =       min_resolution
  UpperCAmelCase__           =       max_resolution
  UpperCAmelCase__           =       size
  UpperCAmelCase__           =       resample if resample is not None else PILImageResampling.BICUBIC
  UpperCAmelCase__           =       do_normalize
  UpperCAmelCase__           =       image_mean
  UpperCAmelCase__           =       image_std
  UpperCAmelCase__           =       do_center_crop
  UpperCAmelCase__           =       crop_size

  UpperCAmelCase__           =       input_size_patches
  UpperCAmelCase__           =       total_mask_patches
  UpperCAmelCase__           =       mask_group_max_patches
  UpperCAmelCase__           =       mask_group_min_patches
  UpperCAmelCase__           =       mask_group_min_aspect_ratio
  UpperCAmelCase__           =       mask_group_max_aspect_ratio

  UpperCAmelCase__           =       codebook_do_resize
  UpperCAmelCase__           =       codebook_size
  UpperCAmelCase__           =       codebook_resample if codebook_resample is not None else PILImageResampling.LANCZOS
  UpperCAmelCase__           =       codebook_do_center_crop
  UpperCAmelCase__           =       codebook_crop_size
  UpperCAmelCase__           =       codebook_do_map_pixels
  UpperCAmelCase__           =       codebook_do_normalize
  UpperCAmelCase__           =       codebook_image_mean
  UpperCAmelCase__           =       codebook_image_std




 def 						UpperCamelCase__							(self				)				->      Any:






  """simple docstring"""





  return {
      "image_mean": self.image_mean,
      "image_std": self.image_std,
      "do_normalize": self.do_normalize,
      "do_resize": self.do_resize,
      "size": self.size,
      "resample": self.resample,
      "do_rescale": self.do_rescale,
      "rescale_factor": self.rescale_factor,
      "do_center_crop": self.do_center_crop,
      "crop_size": self.crop_size,
      "input_size_patches": self.input_size_patches,
      "total_mask_patches": self.total_mask_patches,
      "mask_group_max_patches": self.mask_group_max_patches,
      "mask_group_min_patches": self.mask_group_min_patches,
      "mask_group_min_aspect_ratio": self.mask_group_min_aspect_ratio,
      "mask_group_max_aspect_ratio": self.mask_group_min_aspect_ratio,
      "codebook_do_resize": self.codebook_do_resize,
      "codebook_size": self.codebook_size,
      "codebook_resample": self.codebook_resample,
      "codebook_do_center_crop": self.codebook_do_center_crop,
      "codebook_crop_size": self.codebook_crop_size,
      "codebook_do_map_pixels": self.codebook_do_map_pixels,
      "codebook_do_normalize": self.codebook_do_normalize,
      "codebook_image_mean": self.codebook_image_mean,
      "codebook_image_std": self.codebook_image_std,
  }




 def 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  return (self.size["height"], self.size["width"])




 def 						UpperCamelCase__							(self				)				->      Optional[Any]:






  """simple docstring"""





  return (
      (self.input_size_patches, self.input_size_patches)
      if not isinstance(self.input_size_patches , __a				)
      else self.input_size_patches
  )




 def 						UpperCamelCase__							(self				)				->      Optional[Any]:






  """simple docstring"""





  return (self.codebook_size["height"], self.codebook_size["width"])







@require_torch
@require_vision
class 					lowercase		( _UpperCamelCase				,       unittest.TestCase      ):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE      							=       FlavaImageProcessor if is_vision_available() else None
 __SCREAMING_SNAKE_CASE      							=       None




 def 						UpperCamelCase__							(self				)				->      Tuple:






  """simple docstring"""





  UpperCAmelCase__           =       FlavaImageProcessingTester(self				)




 @property
 def 						UpperCamelCase__							(self				)				->      Dict:






  """simple docstring"""





  return self.image_processor_tester.prepare_image_processor_dict()




 def 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  UpperCAmelCase__           =       self.image_processing_class(**self.image_processor_dict				)
  self.assertTrue(hasattr(__a , 'image_mean'				)				)
  self.assertTrue(hasattr(__a , 'image_std'				)				)
  self.assertTrue(hasattr(__a , 'do_normalize'				)				)
  self.assertTrue(hasattr(__a , 'do_resize'				)				)
  self.assertTrue(hasattr(__a , 'resample'				)				)
  self.assertTrue(hasattr(__a , 'crop_size'				)				)
  self.assertTrue(hasattr(__a , 'do_center_crop'				)				)
  self.assertTrue(hasattr(__a , 'do_rescale'				)				)
  self.assertTrue(hasattr(__a , 'rescale_factor'				)				)
  self.assertTrue(hasattr(__a , 'masking_generator'				)				)
  self.assertTrue(hasattr(__a , 'codebook_do_resize'				)				)
  self.assertTrue(hasattr(__a , 'codebook_size'				)				)
  self.assertTrue(hasattr(__a , 'codebook_resample'				)				)
  self.assertTrue(hasattr(__a , 'codebook_do_center_crop'				)				)
  self.assertTrue(hasattr(__a , 'codebook_crop_size'				)				)
  self.assertTrue(hasattr(__a , 'codebook_do_map_pixels'				)				)
  self.assertTrue(hasattr(__a , 'codebook_do_normalize'				)				)
  self.assertTrue(hasattr(__a , 'codebook_image_mean'				)				)
  self.assertTrue(hasattr(__a , 'codebook_image_std'				)				)




 def 						UpperCamelCase__							(self				)				->      Tuple:






  """simple docstring"""





  UpperCAmelCase__           =       self.image_processing_class.from_dict(self.image_processor_dict				)
  self.assertEqual(image_processor.size , {'height': 224, 'width': 224}				)
  self.assertEqual(image_processor.crop_size , {'height': 224, 'width': 224}				)
  self.assertEqual(image_processor.codebook_size , {'height': 112, 'width': 112}				)
  self.assertEqual(image_processor.codebook_crop_size , {'height': 112, 'width': 112}				)

  UpperCAmelCase__           =       self.image_processing_class.from_dict(
      self.image_processor_dict , size=42 , crop_size=84 , codebook_size=33 , codebook_crop_size=66				)
  self.assertEqual(image_processor.size , {'height': 42, 'width': 42}				)
  self.assertEqual(image_processor.crop_size , {'height': 84, 'width': 84}				)
  self.assertEqual(image_processor.codebook_size , {'height': 33, 'width': 33}				)
  self.assertEqual(image_processor.codebook_crop_size , {'height': 66, 'width': 66}				)




 def 						UpperCamelCase__							(self				)				->      Any:






  """simple docstring"""





  pass




 def 						UpperCamelCase__							(self				)				->      List[Any]:






  """simple docstring"""





  UpperCAmelCase__           =       self.image_processing_class(**self.image_processor_dict				)
  # create random PIL images
  UpperCAmelCase__           =       prepare_image_inputs(self.image_processor_tester , equal_resolution=__a				)
  for image in image_inputs:
   self.assertIsInstance(__a , PIL.Image.Image				)

  # Test not batched input
  UpperCAmelCase__           =       image_processing(image_inputs[0] , return_tensors='pt'				)

  # Test no bool masked pos
  self.assertFalse('bool_masked_pos' in encoded_images				)

  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_image_size()

  self.assertEqual(
      encoded_images.pixel_values.shape , (1, self.image_processor_tester.num_channels, expected_height, expected_width) , )

  # Test batched
  UpperCAmelCase__           =       image_processing(__a , return_tensors='pt'				)
  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_image_size()

  # Test no bool masked pos
  self.assertFalse('bool_masked_pos' in encoded_images				)

  self.assertEqual(
      encoded_images.pixel_values.shape , (
          self.image_processor_tester.batch_size,
          self.image_processor_tester.num_channels,
          expected_height,
          expected_width,
      ) , )




 def 						UpperCamelCase__							(self , __a , __a				)				->      Any:






  """simple docstring"""





  UpperCAmelCase__           =       self.image_processing_class(**self.image_processor_dict				)
  # create random tensors
  UpperCAmelCase__           =       prepare_image_inputs(self.image_processor_tester , equal_resolution=__a , **__a				)
  for image in image_inputs:
   self.assertIsInstance(__a , __a				)

  # Test not batched input
  UpperCAmelCase__           =       image_processing(image_inputs[0] , return_tensors='pt'				)

  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_image_size()
  self.assertEqual(
      encoded_images.pixel_values.shape , (1, self.image_processor_tester.num_channels, expected_height, expected_width) , )

  UpperCAmelCase__           =       image_processing(__a , return_image_mask=__a , return_tensors='pt'				)

  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_image_size()
  self.assertEqual(
      encoded_images.pixel_values.shape , (
          self.image_processor_tester.batch_size,
          self.image_processor_tester.num_channels,
          expected_height,
          expected_width,
      ) , )

  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_mask_size()
  self.assertEqual(
      encoded_images.bool_masked_pos.shape , (
          self.image_processor_tester.batch_size,
          expected_height,
          expected_width,
      ) , )

  # Test batched
  UpperCAmelCase__           =       image_processing(__a , return_tensors='pt'				).pixel_values

  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_image_size()
  self.assertEqual(
      encoded_images.shape , (
          self.image_processor_tester.batch_size,
          self.image_processor_tester.num_channels,
          expected_height,
          expected_width,
      ) , )

  # Test masking
  UpperCAmelCase__           =       image_processing(__a , return_image_mask=__a , return_tensors='pt'				)

  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_image_size()
  self.assertEqual(
      encoded_images.pixel_values.shape , (
          self.image_processor_tester.batch_size,
          self.image_processor_tester.num_channels,
          expected_height,
          expected_width,
      ) , )

  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_mask_size()
  self.assertEqual(
      encoded_images.bool_masked_pos.shape , (
          self.image_processor_tester.batch_size,
          expected_height,
          expected_width,
      ) , )




 def 						UpperCamelCase__							(self				)				->      Dict:






  """simple docstring"""





  self._test_call_framework(np.ndarray , prepare_kwargs={'numpify': True}				)




 def 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  self._test_call_framework(torch.Tensor , prepare_kwargs={'torchify': True}				)




 def 						UpperCamelCase__							(self				)				->      Dict:






  """simple docstring"""





  random.seed(1234				)
  UpperCAmelCase__           =       self.image_processing_class(**self.image_processor_dict				)
  UpperCAmelCase__           =       prepare_image_inputs(self.image_processor_tester , equal_resolution=__a , torchify=__a				)

  # Test not batched input
  UpperCAmelCase__           =       image_processing(image_inputs[0] , return_image_mask=__a , return_tensors='pt'				)
  self.assertEqual(encoded_images.bool_masked_pos.sum().item() , 75				)




 def 						UpperCamelCase__							(self				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       self.image_processing_class(**self.image_processor_dict				)
  # create random PIL images
  UpperCAmelCase__           =       prepare_image_inputs(self.image_processor_tester , equal_resolution=__a				)
  for image in image_inputs:
   self.assertIsInstance(__a , PIL.Image.Image				)

  # Test not batched input
  UpperCAmelCase__           =       image_processing(image_inputs[0] , return_codebook_pixels=__a , return_tensors='pt'				)
  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_codebook_image_size()
  self.assertEqual(
      encoded_images.codebook_pixel_values.shape , (1, self.image_processor_tester.num_channels, expected_height, expected_width) , )

  # Test batched
  UpperCAmelCase__           =       image_processing(__a , return_codebook_pixels=__a , return_tensors='pt'				)
  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.image_processor_tester.get_expected_codebook_image_size()
  self.assertEqual(
      encoded_images.codebook_pixel_values.shape , (
          self.image_processor_tester.batch_size,
          self.image_processor_tester.num_channels,
          expected_height,
          expected_width,
      ) , )



