



"""simple docstring"""






from __future__ import annotations

import copy
import inspect
import unittest

import numpy as np

from transformers import is_tf_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow
from transformers.utils import cached_property

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


if is_tf_available():
     import tensorflow as tf

     from transformers import (
         TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST,
         TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
         TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
         TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
         TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
         LayoutLMvaConfig,
         TFLayoutLMvaForQuestionAnswering,
         TFLayoutLMvaForSequenceClassification,
         TFLayoutLMvaForTokenClassification,
         TFLayoutLMvaModel,
     )

if is_vision_available():
     from PIL import Image

     from transformers import LayoutLMvaImageProcessor






class 	__snake_case	:




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






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			parent
         _lowerCamelCase						:    Tuple         =			batch_size
         _lowerCamelCase						:    List[Any]         =			num_channels
         _lowerCamelCase						:    Dict         =			image_size
         _lowerCamelCase						:    Tuple         =			patch_size
         _lowerCamelCase						:    str         =			is_training
         _lowerCamelCase						:    Optional[Any]         =			use_input_mask
         _lowerCamelCase						:    int         =			use_token_type_ids
         _lowerCamelCase						:    str         =			use_labels
         _lowerCamelCase						:    Dict         =			vocab_size
         _lowerCamelCase						:    List[Any]         =			hidden_size
         _lowerCamelCase						:    List[str]         =			num_hidden_layers
         _lowerCamelCase						:    Dict         =			num_attention_heads
         _lowerCamelCase						:    Union[str, Any]         =			intermediate_size
         _lowerCamelCase						:    str         =			hidden_act
         _lowerCamelCase						:    Dict         =			hidden_dropout_prob
         _lowerCamelCase						:    Any         =			attention_probs_dropout_prob
         _lowerCamelCase						:    int         =			max_position_embeddings
         _lowerCamelCase						:    Optional[int]         =			type_vocab_size
         _lowerCamelCase						:    Tuple         =			type_sequence_label_size
         _lowerCamelCase						:    int         =			initializer_range
         _lowerCamelCase						:    Optional[int]         =			coordinate_size
         _lowerCamelCase						:    List[Any]         =			shape_size
         _lowerCamelCase						:    Tuple         =			num_labels
         _lowerCamelCase						:    List[Any]         =			num_choices
         _lowerCamelCase						:    Optional[Any]         =			scope
         _lowerCamelCase						:    List[str]         =			range_bbox

         # LayoutLMv3's sequence length equals the number of text tokens + number of patches + 1 (we add 1 for the CLS token)
         _lowerCamelCase						:    List[str]         =			text_seq_length
         _lowerCamelCase						:    str         =			(image_size // patch_size) ** 2 + 1
         _lowerCamelCase						:    Optional[Any]         =			self.text_seq_length + self.image_seq_length




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			ids_tensor([self.batch_size, self.text_seq_length]						, self.vocab_size		)

         _lowerCamelCase						:    str         =			ids_tensor([self.batch_size, self.text_seq_length, 4]						, self.range_bbox		)
         _lowerCamelCase						:    Optional[int]         =			bbox.numpy()
         # Ensure that bbox is legal
         for i in range(bbox.shape[0]		):
                for j in range(bbox.shape[1]		):
                       if bbox[i, j, 3] < bbox[i, j, 1]:
                              _lowerCamelCase						:    Union[str, Any]         =			bbox[i, j, 3]
                              _lowerCamelCase						:    Union[str, Any]         =			bbox[i, j, 1]
                              _lowerCamelCase						:    Any         =			tmp_coordinate
                       if bbox[i, j, 2] < bbox[i, j, 0]:
                              _lowerCamelCase						:    Optional[Any]         =			bbox[i, j, 2]
                              _lowerCamelCase						:    Tuple         =			bbox[i, j, 0]
                              _lowerCamelCase						:    Optional[Any]         =			tmp_coordinate
         _lowerCamelCase						:    Dict         =			tf.constant(__lowerCAmelCase		)

         _lowerCamelCase						:    Tuple         =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

         _lowerCamelCase						:    Any         =			None
         if self.use_input_mask:
                _lowerCamelCase						:    str         =			random_attention_mask([self.batch_size, self.text_seq_length]		)

         _lowerCamelCase						:    List[Any]         =			None
         if self.use_token_type_ids:
                _lowerCamelCase						:    Any         =			ids_tensor([self.batch_size, self.text_seq_length]						, self.type_vocab_size		)

         _lowerCamelCase						:    str         =			None
         _lowerCamelCase						:    List[Any]         =			None
         if self.use_labels:
                _lowerCamelCase						:    Union[str, Any]         =			ids_tensor([self.batch_size]						, self.type_sequence_label_size		)
                _lowerCamelCase						:    str         =			ids_tensor([self.batch_size, self.text_seq_length]						, self.num_labels		)

         _lowerCamelCase						:    List[str]         =			LayoutLMvaConfig(
             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						, initializer_range=self.initializer_range						, coordinate_size=self.coordinate_size						, shape_size=self.shape_size						, input_size=self.image_size						, patch_size=self.patch_size						, )

         return config, input_ids, bbox, pixel_values, token_type_ids, input_mask, sequence_labels, token_labels




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						Dict		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			TFLayoutLMvaModel(config=__lowerCAmelCase		)

         # text + image
         _lowerCamelCase						:    Optional[int]         =			model(__lowerCAmelCase						, pixel_values=__lowerCAmelCase						, training=__lowerCAmelCase		)
         _lowerCamelCase						:    List[str]         =			model(
             __lowerCAmelCase						, bbox=__lowerCAmelCase						, pixel_values=__lowerCAmelCase						, attention_mask=__lowerCAmelCase						, token_type_ids=__lowerCAmelCase						, training=__lowerCAmelCase						, )
         _lowerCamelCase						:    Optional[int]         =			model(__lowerCAmelCase						, bbox=__lowerCAmelCase						, pixel_values=__lowerCAmelCase						, training=__lowerCAmelCase		)

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

         # text only
         _lowerCamelCase						:    Union[str, Any]         =			model(__lowerCAmelCase						, training=__lowerCAmelCase		)
         self.parent.assertEqual(
             result.last_hidden_state.shape						, (self.batch_size, self.text_seq_length, self.hidden_size)		)

         # image only
         _lowerCamelCase						:    Optional[Any]         =			model({'''pixel_values''': pixel_values}						, training=__lowerCAmelCase		)
         self.parent.assertEqual(
             result.last_hidden_state.shape						, (self.batch_size, self.image_seq_length, self.hidden_size)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			self.num_labels
         _lowerCamelCase						:    Optional[int]         =			TFLayoutLMvaForSequenceClassification(config=__lowerCAmelCase		)
         _lowerCamelCase						:    List[Any]         =			model(
             __lowerCAmelCase						, bbox=__lowerCAmelCase						, pixel_values=__lowerCAmelCase						, attention_mask=__lowerCAmelCase						, token_type_ids=__lowerCAmelCase						, labels=__lowerCAmelCase						, training=__lowerCAmelCase						, )
         self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_labels)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    str         =			self.num_labels
         _lowerCamelCase						:    str         =			TFLayoutLMvaForTokenClassification(config=__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			model(
             __lowerCAmelCase						, bbox=__lowerCAmelCase						, pixel_values=__lowerCAmelCase						, attention_mask=__lowerCAmelCase						, token_type_ids=__lowerCAmelCase						, labels=__lowerCAmelCase						, training=__lowerCAmelCase						, )
         self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.text_seq_length, self.num_labels)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			2
         _lowerCamelCase						:    Dict         =			TFLayoutLMvaForQuestionAnswering(config=__lowerCAmelCase		)
         _lowerCamelCase						:    List[Any]         =			model(
             __lowerCAmelCase						, bbox=__lowerCAmelCase						, pixel_values=__lowerCAmelCase						, attention_mask=__lowerCAmelCase						, token_type_ids=__lowerCAmelCase						, start_positions=__lowerCAmelCase						, end_positions=__lowerCAmelCase						, training=__lowerCAmelCase						, )
         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 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			self.prepare_config_and_inputs()
         (_lowerCamelCase)						:    Dict         =			config_and_inputs
         _lowerCamelCase						:    List[Any]         =			{
             """input_ids""": input_ids,
             """bbox""": bbox,
             """pixel_values""": pixel_values,
             """token_type_ids""": token_type_ids,
             """attention_mask""": input_mask,
         }
         return config, inputs_dict






@require_tf
class 	__snake_case	(     __lowercase  ,					__lowercase  ,					unittest.TestCase):
  snake_case__    :       Optional[int]        =						(
      (
          TFLayoutLMvaModel,
          TFLayoutLMvaForQuestionAnswering,
          TFLayoutLMvaForSequenceClassification,
          TFLayoutLMvaForTokenClassification,
      )
      if is_tf_available()
      else ()
  )
  snake_case__    :       Union[str, Any]        =						(
      {'''document-question-answering''': TFLayoutLMvaForQuestionAnswering, '''feature-extraction''': TFLayoutLMvaModel}
      if is_tf_available()
      else {}
  )

  snake_case__    :       Dict        =						False
  snake_case__    :       int        =						False
  snake_case__    :       Dict        =						False




  def 						SCREAMING_SNAKE_CASE    (			self					:						int						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						List[str]		):






         """simple docstring"""






         return True




  def 						SCREAMING_SNAKE_CASE    (			self					:						int						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						int=False		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			copy.deepcopy(__lowerCAmelCase		)

         if model_class in get_values(__lowerCAmelCase		):
                _lowerCamelCase						:    Union[str, Any]         =			{
                    k: tf.tile(tf.expand_dims(__lowerCAmelCase						, 1		)						, (1, self.model_tester.num_choices) + (1,) * (v.ndim - 1)		)
                    if isinstance(__lowerCAmelCase						, tf.Tensor		) and v.ndim > 0
                    else v
                    for k, v in inputs_dict.items()
                }

         if return_labels:
                if model_class in get_values(__lowerCAmelCase		):
                       _lowerCamelCase						:    str         =			tf.ones(self.model_tester.batch_size						, dtype=tf.intaa		)
                elif model_class in get_values(__lowerCAmelCase		):
                       _lowerCamelCase						:    Any         =			tf.zeros(self.model_tester.batch_size						, dtype=tf.intaa		)
                       _lowerCamelCase						:    Dict         =			tf.zeros(self.model_tester.batch_size						, dtype=tf.intaa		)
                elif model_class in get_values(__lowerCAmelCase		):
                       _lowerCamelCase						:    Dict         =			tf.zeros(self.model_tester.batch_size						, dtype=tf.intaa		)
                elif model_class in get_values(__lowerCAmelCase		):
                       _lowerCamelCase						:    int         =			tf.zeros(
                           (self.model_tester.batch_size, self.model_tester.text_seq_length)						, dtype=tf.intaa		)

         return inputs_dict




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    str         =			TFLayoutLMvaModelTester(self		)
         _lowerCamelCase						:    int         =			ConfigTester(self						, config_class=__lowerCAmelCase						, hidden_size=3_7		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         self.config_tester.run_common_tests()




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes:
                _lowerCamelCase						:    str         =			model_class(__lowerCAmelCase		)
                if getattr(__lowerCAmelCase						, '''hf_compute_loss'''						, __lowerCAmelCase		):
                       # The number of elements in the loss should be the same as the number of elements in the label
                       _lowerCamelCase						:    str         =			self._prepare_for_class(inputs_dict.copy()						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)
                       _lowerCamelCase						:    Any         =			prepared_for_class[
                           sorted(prepared_for_class.keys() - inputs_dict.keys()						, reverse=__lowerCAmelCase		)[0]
                       ]
                       _lowerCamelCase						:    List[str]         =			added_label.shape.as_list()[:1]

                       # Test that model correctly compute the loss with kwargs
                       _lowerCamelCase						:    Any         =			self._prepare_for_class(inputs_dict.copy()						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)
                       _lowerCamelCase						:    Tuple         =			prepared_for_class.pop('''input_ids'''		)

                       _lowerCamelCase						:    Union[str, Any]         =			model(__lowerCAmelCase						, **__lowerCAmelCase		)[0]
                       self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]		)

                       # Test that model correctly compute the loss when we mask some positions
                       _lowerCamelCase						:    Union[str, Any]         =			self._prepare_for_class(inputs_dict.copy()						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)
                       _lowerCamelCase						:    str         =			prepared_for_class.pop('''input_ids'''		)
                       if "labels" in prepared_for_class:
                              _lowerCamelCase						:    str         =			prepared_for_class["""labels"""].numpy()
                              if len(labels.shape		) > 1 and labels.shape[1] != 1:
                                     _lowerCamelCase						:    Dict         =			-1_0_0
                                     _lowerCamelCase						:    str         =			tf.convert_to_tensor(__lowerCAmelCase		)
                                     _lowerCamelCase						:    Optional[Any]         =			model(__lowerCAmelCase						, **__lowerCAmelCase		)[0]
                                     self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]		)
                                     self.assertTrue(not np.any(np.isnan(loss.numpy()		)		)		)

                # Test that model correctly compute the loss with a dict
                       _lowerCamelCase						:    Optional[int]         =			self._prepare_for_class(inputs_dict.copy()						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)
                       _lowerCamelCase						:    Tuple         =			model(__lowerCAmelCase		)[0]
                       self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]		)

                       # Test that model correctly compute the loss with a tuple
                       _lowerCamelCase						:    str         =			self._prepare_for_class(inputs_dict.copy()						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)

                       # Get keys that were added with the _prepare_for_class function
                       _lowerCamelCase						:    Tuple         =			prepared_for_class.keys() - inputs_dict.keys()
                       _lowerCamelCase						:    Optional[Any]         =			inspect.signature(model.call		).parameters
                       _lowerCamelCase						:    int         =			list(signature.keys()		)

                       # Create a dictionary holding the location of the tensors in the tuple
                       _lowerCamelCase						:    Union[str, Any]         =			{0: """input_ids"""}
                       for label_key in label_keys:
                              _lowerCamelCase						:    int         =			signature_names.index(__lowerCAmelCase		)
                              _lowerCamelCase						:    Optional[int]         =			label_key
                       _lowerCamelCase						:    Optional[int]         =			sorted(tuple_index_mapping.items()		)
                       # Initialize a list with their default values, update the values and convert to a tuple
                       _lowerCamelCase						:    Any         =			[]

                       for name in signature_names:
                              if name != "kwargs":
                                     list_input.append(signature[name].default		)

                       for index, value in sorted_tuple_index_mapping:
                              _lowerCamelCase						:    List[str]         =			prepared_for_class[value]

                       _lowerCamelCase						:    str         =			tuple(__lowerCAmelCase		)

                       # Send to model
                       _lowerCamelCase						:    List[Any]         =			model(tuple_input[:-1]		)[0]

                       self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         (
             _lowerCamelCase      
         )						:    str         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         (
             _lowerCamelCase      
         )						:    List[Any]         =			self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
                _lowerCamelCase						:    Tuple         =			type
                self.model_tester.create_and_check_model(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




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






         """simple docstring"""






         (
             _lowerCamelCase      
         )						:    int         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(
             __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         (
             _lowerCamelCase      
         )						:    int         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(
             __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         (
             _lowerCamelCase      
         )						:    int         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(
             __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  @slow
  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         for model_name in TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                _lowerCamelCase						:    str         =			TFLayoutLMvaModel.from_pretrained(__lowerCAmelCase		)
                self.assertIsNotNone(__lowerCAmelCase		)






def 	snake_case_  (				):


       '''simple docstring'''





       _lowerCamelCase						:    int         =			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						)
       return image







@require_tf
class 	__snake_case	(     unittest.TestCase):




  @cached_property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         return LayoutLMvaImageProcessor(apply_ocr=__lowerCAmelCase		) if is_vision_available() else None




  @slow
  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			TFLayoutLMvaModel.from_pretrained('''microsoft/layoutlmv3-base'''		)

         _lowerCamelCase						:    str         =			self.default_image_processor
         _lowerCamelCase						:    Union[str, Any]         =			prepare_img()
         _lowerCamelCase						:    List[Any]         =			image_processor(images=__lowerCAmelCase						, return_tensors='''tf'''		).pixel_values

         _lowerCamelCase						:    Tuple         =			tf.constant([[1, 2]]		)
         _lowerCamelCase						:    Tuple         =			tf.expand_dims(tf.constant([[1, 2, 3, 4], [5, 6, 7, 8]]		)						, axis=0		)

         # forward pass
         _lowerCamelCase						:    List[Any]         =			model(input_ids=__lowerCAmelCase						, bbox=__lowerCAmelCase						, pixel_values=__lowerCAmelCase						, training=__lowerCAmelCase		)

         # verify the logits
         _lowerCamelCase						:    List[str]         =			(1, 1_9_9, 7_6_8)
         self.assertEqual(outputs.last_hidden_state.shape						, __lowerCAmelCase		)

         _lowerCamelCase						:    Tuple         =			tf.constant(
             [[-0.05_29, 0.36_18, 0.16_32], [-0.15_87, -0.16_67, -0.04_00], [-0.15_57, -0.16_71, -0.05_05]]		)

         self.assertTrue(np.allclose(outputs.last_hidden_state[0, :3, :3]						, __lowerCAmelCase						, atol=1E-4		)		)







"""simple docstring"""






from dataclasses import dataclass
from typing import Dict, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin






@dataclass
class 	__snake_case	(     _lowercase):
  snake_case__    :       torch.FloatTensor






class 	__snake_case	(     _lowercase  ,					_lowercase):




  @register_to_config
  def __init__(			self					:						int						, __lowerCAmelCase					:						int = 3_2						, __lowerCAmelCase					:						int = 6_4						, __lowerCAmelCase					:						int = 2_0						, __lowerCAmelCase					:						int = 7_6_8						, __lowerCAmelCase					:						Tuple=7_7						, __lowerCAmelCase					:						Tuple=4						, __lowerCAmelCase					:						float = 0.0						, __lowerCAmelCase					:						str = "silu"						, __lowerCAmelCase					:						Optional[str] = None						, __lowerCAmelCase					:						Optional[str] = None						, __lowerCAmelCase					:						Optional[str] = "linear"						, __lowerCAmelCase					:						Optional[str] = "prd"						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[int] = None						, __lowerCAmelCase					:						Optional[int] = None						, ):






         """simple docstring"""






         super().__init__()
         _lowerCamelCase						:    Any         =			num_attention_heads
         _lowerCamelCase						:    Optional[int]         =			attention_head_dim
         _lowerCamelCase						:    Dict         =			num_attention_heads * attention_head_dim
         _lowerCamelCase						:    Optional[Any]         =			additional_embeddings

         _lowerCamelCase						:    List[Any]         =			time_embed_dim or inner_dim
         _lowerCamelCase						:    Optional[Any]         =			embedding_proj_dim or embedding_dim
         _lowerCamelCase						:    Tuple         =			clip_embed_dim or embedding_dim

         _lowerCamelCase						:    int         =			Timesteps(__lowerCAmelCase						, __lowerCAmelCase						, 0		)
         _lowerCamelCase						:    str         =			TimestepEmbedding(__lowerCAmelCase						, __lowerCAmelCase						, out_dim=__lowerCAmelCase						, act_fn=__lowerCAmelCase		)

         _lowerCamelCase						:    Union[str, Any]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)

         if embedding_proj_norm_type is None:
                _lowerCamelCase						:    Dict         =			None
         elif embedding_proj_norm_type == "layer":
                _lowerCamelCase						:    Any         =			nn.LayerNorm(__lowerCAmelCase		)
         else:
                raise ValueError(f'''unsupported embedding_proj_norm_type: {embedding_proj_norm_type}'''		)

         _lowerCamelCase						:    List[str]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)

         if encoder_hid_proj_type is None:
                _lowerCamelCase						:    Any         =			None
         elif encoder_hid_proj_type == "linear":
                _lowerCamelCase						:    List[str]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)
         else:
                raise ValueError(f'''unsupported encoder_hid_proj_type: {encoder_hid_proj_type}'''		)

         _lowerCamelCase						:    Optional[int]         =			nn.Parameter(torch.zeros(1						, num_embeddings + additional_embeddings						, __lowerCAmelCase		)		)

         if added_emb_type == "prd":
                _lowerCamelCase						:    Optional[Any]         =			nn.Parameter(torch.zeros(1						, 1						, __lowerCAmelCase		)		)
         elif added_emb_type is None:
                _lowerCamelCase						:    Union[str, Any]         =			None
         else:
                raise ValueError(
                    f'''`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `\'prd\'` or `None`.'''		)

         _lowerCamelCase						:    Tuple         =			nn.ModuleList(
             [
                 BasicTransformerBlock(
                     __lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase						, dropout=__lowerCAmelCase						, activation_fn='''gelu'''						, attention_bias=__lowerCAmelCase						, )
                 for d in range(__lowerCAmelCase		)
             ]		)

         if norm_in_type == "layer":
                _lowerCamelCase						:    Any         =			nn.LayerNorm(__lowerCAmelCase		)
         elif norm_in_type is None:
                _lowerCamelCase						:    Any         =			None
         else:
                raise ValueError(f'''Unsupported norm_in_type: {norm_in_type}.'''		)

         _lowerCamelCase						:    str         =			nn.LayerNorm(__lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			nn.Linear(__lowerCAmelCase						, __lowerCAmelCase		)

         _lowerCamelCase						:    Optional[int]         =			torch.full(
             [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]						, -1_00_00.0		)
         causal_attention_mask.triu_(1		)
         _lowerCamelCase						:    Dict         =			causal_attention_mask[None, ...]
         self.register_buffer('''causal_attention_mask'''						, __lowerCAmelCase						, persistent=__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			nn.Parameter(torch.zeros(1						, __lowerCAmelCase		)		)
         _lowerCamelCase						:    Tuple         =			nn.Parameter(torch.zeros(1						, __lowerCAmelCase		)		)




  @property
  # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			{}

         def fn_recursive_add_processors(__lowerCAmelCase					:						str						, __lowerCAmelCase					:						torch.nn.Module						, __lowerCAmelCase					:						Dict[str, AttentionProcessor]		):
                if hasattr(__lowerCAmelCase						, '''set_processor'''		):
                       _lowerCamelCase						:    str         =			module.processor

                for sub_name, child in module.named_children():
                       fn_recursive_add_processors(f'''{name}.{sub_name}'''						, __lowerCAmelCase						, __lowerCAmelCase		)

                return processors

         for name, module in self.named_children():
                fn_recursive_add_processors(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)

         return processors




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						Union[AttentionProcessor, Dict[str, AttentionProcessor]]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			len(self.attn_processors.keys()		)

         if isinstance(__lowerCAmelCase						, __lowerCAmelCase		) and len(__lowerCAmelCase		) != count:
                raise ValueError(
                    f'''A dict of processors was passed, but the number of processors {len(__lowerCAmelCase		)} does not match the'''
                    f''' number of attention layers: {count}. Please make sure to pass {count} processor classes.'''		)

         def fn_recursive_attn_processor(__lowerCAmelCase					:						str						, __lowerCAmelCase					:						torch.nn.Module						, __lowerCAmelCase					:						Tuple		):
                if hasattr(__lowerCAmelCase						, '''set_processor'''		):
                       if not isinstance(__lowerCAmelCase						, __lowerCAmelCase		):
                              module.set_processor(__lowerCAmelCase		)
                       else:
                              module.set_processor(processor.pop(f'''{name}.processor'''		)		)

                for sub_name, child in module.named_children():
                       fn_recursive_attn_processor(f'''{name}.{sub_name}'''						, __lowerCAmelCase						, __lowerCAmelCase		)

         for name, module in self.named_children():
                fn_recursive_attn_processor(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         self.set_attn_processor(AttnProcessor()		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Union[torch.Tensor, float, int]						, __lowerCAmelCase					:						torch.FloatTensor						, __lowerCAmelCase					:						Optional[torch.FloatTensor] = None						, __lowerCAmelCase					:						Optional[torch.BoolTensor] = None						, __lowerCAmelCase					:						bool = True						, ):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			hidden_states.shape[0]

         _lowerCamelCase						:    Union[str, Any]         =			timestep
         if not torch.is_tensor(__lowerCAmelCase		):
                _lowerCamelCase						:    Dict         =			torch.tensor([timesteps]						, dtype=torch.long						, device=hidden_states.device		)
         elif torch.is_tensor(__lowerCAmelCase		) and len(timesteps.shape		) == 0:
                _lowerCamelCase						:    str         =			timesteps[None].to(hidden_states.device		)

         # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         _lowerCamelCase						:    Tuple         =			timesteps * torch.ones(__lowerCAmelCase						, dtype=timesteps.dtype						, device=timesteps.device		)

         _lowerCamelCase						:    Tuple         =			self.time_proj(__lowerCAmelCase		)

         # timesteps does not contain any weights and will always return f32 tensors
         # but time_embedding might be fp16, so we need to cast here.
         _lowerCamelCase						:    Union[str, Any]         =			timesteps_projected.to(dtype=self.dtype		)
         _lowerCamelCase						:    Any         =			self.time_embedding(__lowerCAmelCase		)

         if self.embedding_proj_norm is not None:
                _lowerCamelCase						:    str         =			self.embedding_proj_norm(__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			self.embedding_proj(__lowerCAmelCase		)
         if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
                _lowerCamelCase						:    Optional[Any]         =			self.encoder_hidden_states_proj(__lowerCAmelCase		)
         elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
                raise ValueError('''`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set'''		)

         _lowerCamelCase						:    str         =			self.proj_in(__lowerCAmelCase		)

         _lowerCamelCase						:    Tuple         =			self.positional_embedding.to(hidden_states.dtype		)

         _lowerCamelCase						:    Tuple         =			[]
         _lowerCamelCase						:    Dict         =			0

         if encoder_hidden_states is not None:
                additional_embeds.append(__lowerCAmelCase		)
                additional_embeddings_len += encoder_hidden_states.shape[1]

         if len(proj_embeddings.shape		) == 2:
                _lowerCamelCase						:    str         =			proj_embeddings[:, None, :]

         if len(hidden_states.shape		) == 2:
                _lowerCamelCase						:    Dict         =			hidden_states[:, None, :]

         _lowerCamelCase						:    int         =			additional_embeds + [
             proj_embeddings,
             time_embeddings[:, None, :],
             hidden_states,
         ]

         if self.prd_embedding is not None:
                _lowerCamelCase						:    int         =			self.prd_embedding.to(hidden_states.dtype		).expand(__lowerCAmelCase						, -1						, -1		)
                additional_embeds.append(__lowerCAmelCase		)

         _lowerCamelCase						:    str         =			torch.cat(
             __lowerCAmelCase						, dim=1						, )

         # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
         _lowerCamelCase						:    Dict         =			additional_embeddings_len + proj_embeddings.shape[1] + 1
         if positional_embeddings.shape[1] < hidden_states.shape[1]:
                _lowerCamelCase						:    List[Any]         =			F.pad(
                    __lowerCAmelCase						, (
                        0,
                        0,
                        additional_embeddings_len,
                        self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
                    )						, value=0.0						, )

         _lowerCamelCase						:    Any         =			hidden_states + positional_embeddings

         if attention_mask is not None:
                _lowerCamelCase						:    Union[str, Any]         =			(1 - attention_mask.to(hidden_states.dtype		)) * -1_00_00.0
                _lowerCamelCase						:    Tuple         =			F.pad(__lowerCAmelCase						, (0, self.additional_embeddings)						, value=0.0		)
                _lowerCamelCase						:    Any         =			(attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype		)
                _lowerCamelCase						:    int         =			attention_mask.repeat_interleave(self.config.num_attention_heads						, dim=0		)

         if self.norm_in is not None:
                _lowerCamelCase						:    List[str]         =			self.norm_in(__lowerCAmelCase		)

         for block in self.transformer_blocks:
                _lowerCamelCase						:    List[str]         =			block(__lowerCAmelCase						, attention_mask=__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			self.norm_out(__lowerCAmelCase		)

         if self.prd_embedding is not None:
                _lowerCamelCase						:    Any         =			hidden_states[:, -1]
         else:
                _lowerCamelCase						:    Dict         =			hidden_states[:, additional_embeddings_len:]

         _lowerCamelCase						:    str         =			self.proj_to_clip_embeddings(__lowerCAmelCase		)

         if not return_dict:
                return (predicted_image_embedding,)

         return PriorTransformerOutput(predicted_image_embedding=__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			(prior_latents * self.clip_std) + self.clip_mean
         return prior_latents


