
'''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  A				:





     def __init__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=2						,			SCREAMING_SNAKE_CASE=3						,			SCREAMING_SNAKE_CASE=4						,			SCREAMING_SNAKE_CASE=2						,			SCREAMING_SNAKE_CASE=7						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=99						,			SCREAMING_SNAKE_CASE=36						,			SCREAMING_SNAKE_CASE=2						,			SCREAMING_SNAKE_CASE=4						,			SCREAMING_SNAKE_CASE=37						,			SCREAMING_SNAKE_CASE="gelu"						,			SCREAMING_SNAKE_CASE=0.1						,			SCREAMING_SNAKE_CASE=0.1						,			SCREAMING_SNAKE_CASE=512						,			SCREAMING_SNAKE_CASE=16						,			SCREAMING_SNAKE_CASE=2						,			SCREAMING_SNAKE_CASE=0.02						,			SCREAMING_SNAKE_CASE=6						,			SCREAMING_SNAKE_CASE=6						,			SCREAMING_SNAKE_CASE=3						,			SCREAMING_SNAKE_CASE=4						,			SCREAMING_SNAKE_CASE=None						,			SCREAMING_SNAKE_CASE=1000						,			) ->				List[Any]:






         """simple docstring"""






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

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





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






         """simple docstring"""






         A       :   Optional[Any]					=						ids_tensor([self.batch_size, self.text_seq_length]						,			self.vocab_size	)

         A       :   List[str]					=						ids_tensor([self.batch_size, self.text_seq_length, 4]						,			self.range_bbox	)
         A       :   List[Any]					=						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]:
                     A       :   List[str]					=						bbox[i, j, 3]
                     A       :   Union[str, Any]					=						bbox[i, j, 1]
                     A       :   str					=						tmp_coordinate
                 if bbox[i, j, 2] < bbox[i, j, 0]:
                     A       :   Union[str, Any]					=						bbox[i, j, 2]
                     A       :   Union[str, Any]					=						bbox[i, j, 0]
                     A       :   int					=						tmp_coordinate
         A       :   Any					=						tf.constant(SCREAMING_SNAKE_CASE__	)

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

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

         A       :   Tuple					=						None
         if self.use_token_type_ids:
             A       :   str					=						ids_tensor([self.batch_size, self.text_seq_length]						,			self.type_vocab_size	)

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

         A       :   Dict					=						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 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[int]:






         """simple docstring"""






         A       :   Union[str, Any]					=						TFLayoutLMvaModel(config=SCREAMING_SNAKE_CASE__	)

         # text + image
         A       :   Optional[Any]					=						model(SCREAMING_SNAKE_CASE__						,			pixel_values=SCREAMING_SNAKE_CASE__						,			training=SCREAMING_SNAKE_CASE__	)
         A       :   Optional[int]					=						model(
             SCREAMING_SNAKE_CASE__						,			bbox=SCREAMING_SNAKE_CASE__						,			pixel_values=SCREAMING_SNAKE_CASE__						,			attention_mask=SCREAMING_SNAKE_CASE__						,			token_type_ids=SCREAMING_SNAKE_CASE__						,			training=SCREAMING_SNAKE_CASE__						,			)
         A       :   Optional[Any]					=						model(SCREAMING_SNAKE_CASE__						,			bbox=SCREAMING_SNAKE_CASE__						,			pixel_values=SCREAMING_SNAKE_CASE__						,			training=SCREAMING_SNAKE_CASE__	)

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

         # text only
         A       :   Optional[Any]					=						model(SCREAMING_SNAKE_CASE__						,			training=SCREAMING_SNAKE_CASE__	)
         self.parent.assertEqual(
             result.last_hidden_state.shape						,			(self.batch_size, self.text_seq_length, self.hidden_size)	)

         # image only
         A       :   List[str]					=						model({'''pixel_values''': pixel_values}						,			training=SCREAMING_SNAKE_CASE__	)
         self.parent.assertEqual(
             result.last_hidden_state.shape						,			(self.batch_size, self.image_seq_length, self.hidden_size)	)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				List[str]:






         """simple docstring"""






         A       :   int					=						self.num_labels
         A       :   Union[str, Any]					=						TFLayoutLMvaForSequenceClassification(config=SCREAMING_SNAKE_CASE__	)
         A       :   str					=						model(
             SCREAMING_SNAKE_CASE__						,			bbox=SCREAMING_SNAKE_CASE__						,			pixel_values=SCREAMING_SNAKE_CASE__						,			attention_mask=SCREAMING_SNAKE_CASE__						,			token_type_ids=SCREAMING_SNAKE_CASE__						,			labels=SCREAMING_SNAKE_CASE__						,			training=SCREAMING_SNAKE_CASE__						,			)
         self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_labels)	)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				List[str]:






         """simple docstring"""






         A       :   List[Any]					=						self.num_labels
         A       :   List[Any]					=						TFLayoutLMvaForTokenClassification(config=SCREAMING_SNAKE_CASE__	)
         A       :   Dict					=						model(
             SCREAMING_SNAKE_CASE__						,			bbox=SCREAMING_SNAKE_CASE__						,			pixel_values=SCREAMING_SNAKE_CASE__						,			attention_mask=SCREAMING_SNAKE_CASE__						,			token_type_ids=SCREAMING_SNAKE_CASE__						,			labels=SCREAMING_SNAKE_CASE__						,			training=SCREAMING_SNAKE_CASE__						,			)
         self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.text_seq_length, self.num_labels)	)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				int:






         """simple docstring"""






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






         """simple docstring"""






         A       :   Dict					=						self.prepare_config_and_inputs()
         (A)       :   List[str]					=						config_and_inputs
         A       :   List[str]					=						{
             '''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  A				(							__snake_case    ,			__snake_case    ,			unittest.TestCase						):
     __magic_name__						=		(
         (
             TFLayoutLMvaModel,
             TFLayoutLMvaForQuestionAnswering,
             TFLayoutLMvaForSequenceClassification,
             TFLayoutLMvaForTokenClassification,
         )
         if is_tf_available()
         else ()
     )
     __magic_name__						=		(
         {'''document-question-answering''': TFLayoutLMvaForQuestionAnswering, '''feature-extraction''': TFLayoutLMvaModel}
         if is_tf_available()
         else {}
     )

     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		False





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[Any]:






         """simple docstring"""






         return True





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=False	) ->				dict:






         """simple docstring"""






         A       :   Tuple					=						copy.deepcopy(SCREAMING_SNAKE_CASE__	)

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

         if return_labels:
             if model_class in get_values(SCREAMING_SNAKE_CASE__	):
                 A       :   Union[str, Any]					=						tf.ones(self.model_tester.batch_size						,			dtype=tf.intaa	)
             elif model_class in get_values(SCREAMING_SNAKE_CASE__	):
                 A       :   Optional[int]					=						tf.zeros(self.model_tester.batch_size						,			dtype=tf.intaa	)
                 A       :   int					=						tf.zeros(self.model_tester.batch_size						,			dtype=tf.intaa	)
             elif model_class in get_values(SCREAMING_SNAKE_CASE__	):
                 A       :   Optional[int]					=						tf.zeros(self.model_tester.batch_size						,			dtype=tf.intaa	)
             elif model_class in get_values(SCREAMING_SNAKE_CASE__	):
                 A       :   Dict					=						tf.zeros(
                     (self.model_tester.batch_size, self.model_tester.text_seq_length)						,			dtype=tf.intaa	)

         return inputs_dict





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






         """simple docstring"""






         A       :   Union[str, Any]					=						TFLayoutLMvaModelTester(self	)
         A       :   Union[str, Any]					=						ConfigTester(self						,			config_class=SCREAMING_SNAKE_CASE__						,			hidden_size=37	)





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






         """simple docstring"""






         self.config_tester.run_common_tests()





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






         """simple docstring"""






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

                 # Test that model correctly compute the loss with kwargs
                 A       :   Any					=						self._prepare_for_class(inputs_dict.copy()						,			SCREAMING_SNAKE_CASE__						,			return_labels=SCREAMING_SNAKE_CASE__	)
                 A       :   Optional[int]					=						prepared_for_class.pop('''input_ids'''	)

                 A       :   List[Any]					=						model(SCREAMING_SNAKE_CASE__						,			**SCREAMING_SNAKE_CASE__	)[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
                 A       :   Tuple					=						self._prepare_for_class(inputs_dict.copy()						,			SCREAMING_SNAKE_CASE__						,			return_labels=SCREAMING_SNAKE_CASE__	)
                 A       :   Tuple					=						prepared_for_class.pop('''input_ids'''	)
                 if "labels" in prepared_for_class:
                     A       :   Optional[Any]					=						prepared_for_class['''labels'''].numpy()
                     if len(labels.shape	) > 1 and labels.shape[1] != 1:
                         A       :   Optional[Any]					=						-100
                         A       :   int					=						tf.convert_to_tensor(SCREAMING_SNAKE_CASE__	)
                         A       :   Any					=						model(SCREAMING_SNAKE_CASE__						,			**SCREAMING_SNAKE_CASE__	)[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
                 A       :   str					=						self._prepare_for_class(inputs_dict.copy()						,			SCREAMING_SNAKE_CASE__						,			return_labels=SCREAMING_SNAKE_CASE__	)
                 A       :   Dict					=						model(SCREAMING_SNAKE_CASE__	)[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
                 A       :   str					=						self._prepare_for_class(inputs_dict.copy()						,			SCREAMING_SNAKE_CASE__						,			return_labels=SCREAMING_SNAKE_CASE__	)

                 # Get keys that were added with the _prepare_for_class function
                 A       :   Union[str, Any]					=						prepared_for_class.keys() - inputs_dict.keys()
                 A       :   Any					=						inspect.signature(model.call	).parameters
                 A       :   Any					=						list(signature.keys()	)

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

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

                 for index, value in sorted_tuple_index_mapping:
                     A       :   int					=						prepared_for_class[value]

                 A       :   str					=						tuple(SCREAMING_SNAKE_CASE__	)

                 # Send to model
                 A       :   Dict					=						model(tuple_input[:-1]	)[0]

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





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






         """simple docstring"""






         (
             A
         )       :   str					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__	)





     def 				__lowerCAmelCase							(							self	) ->				str:






         """simple docstring"""






         (
             A
         )       :   Optional[Any]					=						self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
             A       :   List[str]					=						type
             self.model_tester.create_and_check_model(SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__	)





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         (
             A
         )       :   Optional[int]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(
             SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__	)





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         (
             A
         )       :   Dict					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(
             SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__	)





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






         """simple docstring"""






         (
             A
         )       :   str					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(
             SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__						,			SCREAMING_SNAKE_CASE__	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				int:






         """simple docstring"""






         for model_name in TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             A       :   List[Any]					=						TFLayoutLMvaModel.from_pretrained(SCREAMING_SNAKE_CASE__	)
             self.assertIsNotNone(SCREAMING_SNAKE_CASE__	)







def 		lowerCAmelCase_				(   ):




    '''simple docstring'''



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







@require_tf
class  A				(							unittest.TestCase						):





     @cached_property
     def 				__lowerCAmelCase							(							self	) ->				List[Any]:






         """simple docstring"""






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





     @slow
     def 				__lowerCAmelCase							(							self	) ->				List[Any]:






         """simple docstring"""






         A       :   Union[str, Any]					=						TFLayoutLMvaModel.from_pretrained('''microsoft/layoutlmv3-base'''	)

         A       :   List[str]					=						self.default_image_processor
         A       :   Dict					=						prepare_img()
         A       :   List[str]					=						image_processor(images=SCREAMING_SNAKE_CASE__						,			return_tensors='''tf'''	).pixel_values

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

         # forward pass
         A       :   Any					=						model(input_ids=SCREAMING_SNAKE_CASE__						,			bbox=SCREAMING_SNAKE_CASE__						,			pixel_values=SCREAMING_SNAKE_CASE__						,			training=SCREAMING_SNAKE_CASE__	)

         # verify the logits
         A       :   int					=						(1, 199, 768)
         self.assertEqual(outputs.last_hidden_state.shape						,			SCREAMING_SNAKE_CASE__	)

         A       :   List[Any]					=						tf.constant(
             [[-0.0_529, 0.3_618, 0.1_632], [-0.1_587, -0.1_667, -0.0_400], [-0.1_557, -0.1_671, -0.0_505]]	)

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











"""simple docstring"""
import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()



def      __lowerCamelCase   (							a_     :		Union[str, Any] ,   a_     :		Tuple ,   a_     :		str=None							)			->      Union[str, Any]:
      # set parameter of one layer
      assert torch_layer.weight.shape == weight.shape, f'''{torch_layer} layer.weight does not match'''
      __SCREAMING_SNAKE_CASE   :Dict						=  nn.Parameter(a_							)
      if bias is not None:
            assert torch_layer.bias.shape == bias.shape, f'''{torch_layer} layer.bias does not match'''
            __SCREAMING_SNAKE_CASE   :Optional[int]						=  nn.Parameter(a_							)



def      __lowerCamelCase   (							a_     :		Dict ,   a_     :		str ,   a_     :		Optional[int]							)			->      Any:
      # set torch weights for 1-to-1 comparison
      __SCREAMING_SNAKE_CASE   :List[Any]						=  np.asarray(weights[0]							)
      __SCREAMING_SNAKE_CASE   :Optional[int]						=  np.asarray(weights[1]							)
      __SCREAMING_SNAKE_CASE   :Optional[int]						=  np.asarray(weights[2]							)

      set_param(
          torch_layer.self_attention.query_key ,   torch.tensor(a_							).transpose(1 ,   2							).contiguous().view(-1 ,   a_							) ,   )
      set_param(
          torch_layer.self_attention.value ,   torch.tensor(a_							).transpose(1 ,   2							).contiguous().view(-1 ,   a_							) ,   )
      set_param(
          torch_layer.output.dense ,   torch.tensor(a_							).view(-1 ,   a_							).contiguous().transpose(0 ,   1							) ,   )



def      __lowerCamelCase   (							a_     :		List[Any] ,   a_     :		Dict ,   a_     :		List[str]							)			->      Union[str, Any]:
      # set torch weights for 1-to-1 comparison
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  np.asarray(weights[0]							)
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  np.asarray(weights[1]							)
      __SCREAMING_SNAKE_CASE   :Any						=  np.asarray(weights[2]							)
      __SCREAMING_SNAKE_CASE   :Dict						=  np.asarray(weights[3]							)

      set_param(
          torch_layer.self_attention.query ,   torch.tensor(a_							).transpose(1 ,   2							).contiguous().view(-1 ,   a_							) ,   )
      set_param(
          torch_layer.self_attention.key ,   torch.tensor(a_							).transpose(1 ,   2							).contiguous().view(-1 ,   a_							) ,   )
      set_param(
          torch_layer.self_attention.value ,   torch.tensor(a_							).transpose(1 ,   2							).contiguous().view(-1 ,   a_							) ,   )
      set_param(
          torch_layer.output.dense ,   torch.tensor(a_							).view(-1 ,   a_							).contiguous().transpose(0 ,   1							) ,   )



def      __lowerCamelCase   (							a_     :		Any ,   a_     :		List[str] ,   a_     :		Optional[int]							)			->      Union[str, Any]:
      # layernorm 1
      __SCREAMING_SNAKE_CASE   :Any						=  weights[0][0][0]
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  np.asarray(layer_norm_a[0]							)
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  np.asarray(layer_norm_a[1]							)
      set_param(
          torch_block.attention.layer_norm ,   torch.tensor(a_							) ,   torch.tensor(a_							) ,   )

      # lsh weights + output
      __SCREAMING_SNAKE_CASE   :List[Any]						=  weights[0][1]
      if len(a_							) < 4:
            set_layer_weights_in_torch_lsh(a_ ,   torch_block.attention ,   a_							)
      else:
            set_layer_weights_in_torch_local(a_ ,   torch_block.attention ,   a_							)

      # intermediate weighs
      __SCREAMING_SNAKE_CASE   :List[Any]						=  weights[2][0][1][2]

      # Chunked Feed Forward
      if len(a_							) == 4:
            __SCREAMING_SNAKE_CASE   :List[str]						=  intermediate_weights[2]

      # layernorm 2
      __SCREAMING_SNAKE_CASE   :Tuple						=  np.asarray(intermediate_weights[0][0]							)
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  np.asarray(intermediate_weights[0][1]							)
      set_param(
          torch_block.feed_forward.layer_norm ,   torch.tensor(a_							) ,   torch.tensor(a_							) ,   )

      # intermediate dense
      __SCREAMING_SNAKE_CASE   :int						=  np.asarray(intermediate_weights[1][0]							)
      __SCREAMING_SNAKE_CASE   :int						=  np.asarray(intermediate_weights[1][1]							)
      set_param(
          torch_block.feed_forward.dense.dense ,   torch.tensor(a_							).transpose(0 ,   1							).contiguous() ,   torch.tensor(a_							) ,   )

      # intermediate out
      __SCREAMING_SNAKE_CASE   :str						=  np.asarray(intermediate_weights[4][0]							)
      __SCREAMING_SNAKE_CASE   :str						=  np.asarray(intermediate_weights[4][1]							)
      set_param(
          torch_block.feed_forward.output.dense ,   torch.tensor(a_							).transpose(0 ,   1							).contiguous() ,   torch.tensor(a_							) ,   )



def      __lowerCamelCase   (							a_     :		List[str] ,   a_     :		str ,   a_     :		List[Any]							)			->      Optional[Any]:
      # reformer model
      __SCREAMING_SNAKE_CASE   :Dict						=  torch_model.reformer

      # word embeds
      __SCREAMING_SNAKE_CASE   :List[Any]						=  np.asarray(weights[1]							)
      set_param(
          torch_model_reformer.embeddings.word_embeddings ,   torch.tensor(a_							) ,   )

      if isinstance(weights[3] ,   a_							):
            __SCREAMING_SNAKE_CASE   :List[Any]						=  torch_model_reformer.embeddings.position_embeddings
            for emb_idx in range(len(position_embeddings.weights							)							):
                  __SCREAMING_SNAKE_CASE   :List[str]						=  np.asarray(weights[3][emb_idx][0]							)
                  assert (
                      position_embeddings.weights[emb_idx].shape == emb_weights.shape
                  ), f'''{position_embeddings[emb_idx]} emb does not match'''
                  __SCREAMING_SNAKE_CASE   :str						=  nn.Parameter(torch.tensor(a_							)							)

      __SCREAMING_SNAKE_CASE   :Optional[int]						=  weights[5]
      assert len(torch_model_reformer.encoder.layers							) * 4 == len(
          a_							), "HF and trax model do not have the same number of layers"
      for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers							):
            __SCREAMING_SNAKE_CASE   :Dict						=  trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
            set_block_weights_in_torch(a_ ,   a_ ,   a_							)

      # output layer norm
      __SCREAMING_SNAKE_CASE   :Optional[int]						=  np.asarray(weights[7][0]							)
      __SCREAMING_SNAKE_CASE   :List[Any]						=  np.asarray(weights[7][1]							)
      set_param(
          torch_model_reformer.encoder.layer_norm ,   torch.tensor(a_							) ,   torch.tensor(a_							) ,   )

      # output embeddings
      __SCREAMING_SNAKE_CASE   :Optional[int]						=  np.asarray(weights[9][0]							)
      __SCREAMING_SNAKE_CASE   :str						=  np.asarray(weights[9][1]							)
      set_param(
          torch_model.lm_head.decoder ,   torch.tensor(a_							).transpose(0 ,   1							).contiguous() ,   torch.tensor(a_							) ,   )



def      __lowerCamelCase   (							a_     :		Any ,   a_     :		Dict ,   a_     :		Dict							)			->      Tuple:
      # Initialise PyTorch model
      __SCREAMING_SNAKE_CASE   :List[str]						=  ReformerConfig.from_json_file(a_							)
      print(f'''Building PyTorch model from configuration: {config}'''							)
      __SCREAMING_SNAKE_CASE   :List[Any]						=  ReformerModelWithLMHead(a_							)

      with open(a_ ,   '''rb'''							) as f:
            __SCREAMING_SNAKE_CASE   :Any						=  pickle.load(a_							)['''weights''']

      set_model_weights_in_torch(a_ ,   a_ ,   config.hidden_size							)

      # Save pytorch-model
      print(f'''Save PyTorch model to {pytorch_dump_path}'''							)
      torch.save(model.state_dict() ,   a_							)


if __name__ == "__main__":
  lowerCamelCase_								=		argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--trax_model_pkl_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
  )
  parser.add_argument(
      "--config_file",
      default=None,
      type=str,
      required=True,
      help=(
          "The config json file corresponding to the pre-trained Reformer model. \n"
          "This specifies the model architecture."
      ),
  )
  parser.add_argument(
      "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
  )
  lowerCamelCase_								=		parser.parse_args()
  convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)