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





       '''simple docstring'''



       def __init__(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase=2     ,       __lowerCAmelCase=3     ,       __lowerCAmelCase=4     ,       __lowerCAmelCase=2     ,       __lowerCAmelCase=7     ,       __lowerCAmelCase=True     ,       __lowerCAmelCase=True     ,       __lowerCAmelCase=True     ,       __lowerCAmelCase=True     ,       __lowerCAmelCase=99     ,       __lowerCAmelCase=36     ,       __lowerCAmelCase=2     ,       __lowerCAmelCase=4     ,       __lowerCAmelCase=37     ,       __lowerCAmelCase="gelu"     ,       __lowerCAmelCase=0.1     ,       __lowerCAmelCase=0.1     ,       __lowerCAmelCase=512     ,       __lowerCAmelCase=16     ,       __lowerCAmelCase=2     ,       __lowerCAmelCase=0.0_2     ,       __lowerCAmelCase=6     ,       __lowerCAmelCase=6     ,       __lowerCAmelCase=3     ,       __lowerCAmelCase=4     ,       __lowerCAmelCase=None     ,       __lowerCAmelCase=1000     ,       )  ->  List[Any]:
        lowercase__    :   Union[str, Any]						= parent
        lowercase__    :   Optional[Any]						= batch_size
        lowercase__    :   str						= num_channels
        lowercase__    :   Optional[Any]						= image_size
        lowercase__    :   Any						= patch_size
        lowercase__    :   int						= is_training
        lowercase__    :   List[Any]						= use_input_mask
        lowercase__    :   Any						= use_token_type_ids
        lowercase__    :   int						= use_labels
        lowercase__    :   Optional[Any]						= vocab_size
        lowercase__    :   Union[str, Any]						= hidden_size
        lowercase__    :   Optional[int]						= num_hidden_layers
        lowercase__    :   Dict						= num_attention_heads
        lowercase__    :   str						= intermediate_size
        lowercase__    :   Optional[Any]						= hidden_act
        lowercase__    :   str						= hidden_dropout_prob
        lowercase__    :   Any						= attention_probs_dropout_prob
        lowercase__    :   Dict						= max_position_embeddings
        lowercase__    :   Dict						= type_vocab_size
        lowercase__    :   Dict						= type_sequence_label_size
        lowercase__    :   Optional[Any]						= initializer_range
        lowercase__    :   Tuple						= coordinate_size
        lowercase__    :   Optional[int]						= shape_size
        lowercase__    :   List[str]						= num_labels
        lowercase__    :   Any						= num_choices
        lowercase__    :   List[str]						= scope
        lowercase__    :   Tuple						= range_bbox

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

       def _lowerCAmelCase(						self			)  ->  Dict:
        lowercase__    :   List[str]						= ids_tensor([self.batch_size, self.text_seq_length]     ,       self.vocab_size			)

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

        lowercase__    :   Any						= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

        lowercase__    :   List[Any]						= None
        if self.use_input_mask:
         lowercase__    :   Optional[int]						= random_attention_mask([self.batch_size, self.text_seq_length]			)

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

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

        lowercase__    :   Any						= 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     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  Any:
        lowercase__    :   Union[str, Any]						= TFLayoutLMvaModel(config=__lowerCAmelCase			)

        # text + image
        lowercase__    :   Tuple						= model(__lowerCAmelCase     ,       pixel_values=__lowerCAmelCase     ,       training=__lowerCAmelCase			)
        lowercase__    :   Any						= model(
            __lowerCAmelCase     ,       bbox=__lowerCAmelCase     ,       pixel_values=__lowerCAmelCase     ,       attention_mask=__lowerCAmelCase     ,       token_type_ids=__lowerCAmelCase     ,       training=__lowerCAmelCase     ,       )
        lowercase__    :   str						= 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
        lowercase__    :   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
        lowercase__    :   List[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 _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  int:
        lowercase__    :   Optional[Any]						= self.num_labels
        lowercase__    :   Dict						= TFLayoutLMvaForSequenceClassification(config=__lowerCAmelCase			)
        lowercase__    :   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 _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  Union[str, Any]:
        lowercase__    :   Tuple						= self.num_labels
        lowercase__    :   List[Any]						= TFLayoutLMvaForTokenClassification(config=__lowerCAmelCase			)
        lowercase__    :   str						= 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 _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  str:
        lowercase__    :   Dict						= 2
        lowercase__    :   Optional[int]						= TFLayoutLMvaForQuestionAnswering(config=__lowerCAmelCase			)
        lowercase__    :   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 _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   str						= self.prepare_config_and_inputs()
        ((lowercase__)			,							(lowercase__)			,							(lowercase__)			,							(lowercase__)			,							(lowercase__)			,							(lowercase__)			,							(lowercase__)			,							(lowercase__))    :   List[str]						= config_and_inputs
        lowercase__    :   Optional[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       UpperCAmelCase  (						a__   ,						a__   ,						unittest.TestCase ):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE 		=   (
           (
               TFLayoutLMvaModel,
               TFLayoutLMvaForQuestionAnswering,
               TFLayoutLMvaForSequenceClassification,
               TFLayoutLMvaForTokenClassification,
           )
           if is_tf_available()
           else ()
       )
       SCREAMING_SNAKE_CASE 		=   (
           {"document-question-answering": TFLayoutLMvaForQuestionAnswering, "feature-extraction": TFLayoutLMvaModel}
           if is_tf_available()
           else {}
       )

       SCREAMING_SNAKE_CASE 		=   False
       SCREAMING_SNAKE_CASE 		=   False
       SCREAMING_SNAKE_CASE 		=   False

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  Tuple:
        return True

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase=False			)  ->  dict:
        lowercase__    :   List[str]						= copy.deepcopy(__lowerCAmelCase			)

        if model_class in get_values(__lowerCAmelCase			):
         lowercase__    :   Optional[int]						= {
             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			):
          lowercase__    :   Any						= tf.ones(self.model_tester.batch_size     ,       dtype=tf.intaa			)
         elif model_class in get_values(__lowerCAmelCase			):
          lowercase__    :   str						= tf.zeros(self.model_tester.batch_size     ,       dtype=tf.intaa			)
          lowercase__    :   Optional[Any]						= tf.zeros(self.model_tester.batch_size     ,       dtype=tf.intaa			)
         elif model_class in get_values(__lowerCAmelCase			):
          lowercase__    :   Optional[Any]						= tf.zeros(self.model_tester.batch_size     ,       dtype=tf.intaa			)
         elif model_class in get_values(__lowerCAmelCase			):
          lowercase__    :   Union[str, Any]						= tf.zeros(
              (self.model_tester.batch_size, self.model_tester.text_seq_length)     ,       dtype=tf.intaa			)

        return inputs_dict

       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        lowercase__    :   List[str]						= TFLayoutLMvaModelTester(self			)
        lowercase__    :   Union[str, Any]						= ConfigTester(self     ,       config_class=__lowerCAmelCase     ,       hidden_size=37			)

       def _lowerCAmelCase(						self			)  ->  Any:
        self.config_tester.run_common_tests()

       def _lowerCAmelCase(						self			)  ->  List[str]:
        lowercase__			,							lowercase__    :   List[Any]						= self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
         lowercase__    :   Dict						= 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
          lowercase__    :   Union[str, Any]						= self._prepare_for_class(inputs_dict.copy()     ,       __lowerCAmelCase     ,       return_labels=__lowerCAmelCase			)
          lowercase__    :   Dict						= prepared_for_class[
              sorted(prepared_for_class.keys() - inputs_dict.keys()     ,       reverse=__lowerCAmelCase			)[0]
          ]
          lowercase__    :   List[Any]						= added_label.shape.as_list()[:1]

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

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

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

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

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

          for index, value in sorted_tuple_index_mapping:
           lowercase__    :   List[Any]						= prepared_for_class[value]

          lowercase__    :   Optional[int]						= tuple(__lowerCAmelCase			)

          # Send to model
          lowercase__    :   Optional[Any]						= model(tuple_input[:-1]			)[0]

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

       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        (
            (
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							
        )    :   Any						= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(__lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)

       def _lowerCAmelCase(						self			)  ->  Tuple:
        (
            (
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							
        )    :   List[str]						= self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
         lowercase__    :   Tuple						= type
         self.model_tester.create_and_check_model(__lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)

       def _lowerCAmelCase(						self			)  ->  int:
        (
            (
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							
        )    :   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 _lowerCAmelCase(						self			)  ->  Tuple:
        (
            (
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							
        )    :   Optional[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 _lowerCAmelCase(						self			)  ->  Tuple:
        (
            (
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							
        )    :   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 _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        for model_name in TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowercase__    :   int						= TFLayoutLMvaModel.from_pretrained(__lowerCAmelCase			)
         self.assertIsNotNone(__lowerCAmelCase			)


def       __UpperCamelCase	(       ):
 lowercase__    :   List[str]						= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
 return image





@require_tf
class       UpperCAmelCase  (						unittest.TestCase ):





       '''simple docstring'''



       @cached_property
       def _lowerCAmelCase(						self			)  ->  str:
        return LayoutLMvaImageProcessor(apply_ocr=__lowerCAmelCase			) if is_vision_available() else None


       @slow
       def _lowerCAmelCase(						self			)  ->  Dict:
        lowercase__    :   Any						= TFLayoutLMvaModel.from_pretrained('''microsoft/layoutlmv3-base'''			)

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

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

        # forward pass
        lowercase__    :   Optional[int]						= model(input_ids=__lowerCAmelCase     ,       bbox=__lowerCAmelCase     ,       pixel_values=__lowerCAmelCase     ,       training=__lowerCAmelCase			)

        # verify the logits
        lowercase__    :   Union[str, Any]						= (1, 199, 768)
        self.assertEqual(outputs.last_hidden_state.shape     ,       __lowerCAmelCase			)

        lowercase__    :   List[Any]						= tf.constant(
            [[-0.0_5_2_9, 0.3_6_1_8, 0.1_6_3_2], [-0.1_5_8_7, -0.1_6_6_7, -0.0_4_0_0], [-0.1_5_5_7, -0.1_6_7_1, -0.0_5_0_5]]			)

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


'''simple docstring'''

import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase="attention"					):
 lowercase__    :   Dict						= params[F"""{prefix}/layers_{i}/{layer_name}/key/kernel"""]
 lowercase__    :   Dict						= params[F"""{prefix}/layers_{i}/{layer_name}/out/kernel"""]
 lowercase__    :   int						= params[F"""{prefix}/layers_{i}/{layer_name}/query/kernel"""]
 lowercase__    :   Any						= params[F"""{prefix}/layers_{i}/{layer_name}/value/kernel"""]
 return k, o, q, v
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase=False					):
 if split_mlp_wi:
  lowercase__    :   Tuple						= params[F"""{prefix}/layers_{i}/mlp/wi_0/kernel"""]
  lowercase__    :   Tuple						= params[F"""{prefix}/layers_{i}/mlp/wi_1/kernel"""]
  lowercase__    :   Dict						= (wi_a, wi_a)
 else:
  lowercase__    :   List[Any]						= params[F"""{prefix}/layers_{i}/mlp/wi/kernel"""]

 lowercase__    :   int						= params[F"""{prefix}/layers_{i}/mlp/wo/kernel"""]
 return wi, wo
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase					):
 return params[F"""{prefix}/layers_{i}/{layer_name}/scale"""]
def       __UpperCamelCase	(       UpperCAmelCase		,					*, UpperCAmelCase		,					UpperCAmelCase					):
 lowercase__    :   Union[str, Any]						= traverse_util.flatten_dict(variables['''target''']					)
 lowercase__    :   Any						= {'''/'''.join(UpperCAmelCase					): v for k, v in old.items()}

 # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
 lowercase__    :   List[str]						= '''encoder/layers_0/mlp/wi_0/kernel''' in old
 print('''Split MLP:'''		,					UpperCAmelCase					)

 lowercase__    :   Union[str, Any]						= collections.OrderedDict()

 # Shared embeddings.
 lowercase__    :   str						= old['''token_embedder/embedding''']

 # Encoder.
 for i in range(UpperCAmelCase					):
  # Block i, layer 0 (Self Attention).
  lowercase__    :   List[Any]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''encoder'''		,					'''pre_attention_layer_norm'''					)
  lowercase__			,							lowercase__			,							lowercase__			,							lowercase__    :   Tuple						= tax_attention_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''encoder'''		,					'''attention'''					)
  lowercase__    :   Dict						= layer_norm
  lowercase__    :   Any						= k.T
  lowercase__    :   Optional[int]						= o.T
  lowercase__    :   Optional[int]						= q.T
  lowercase__    :   Any						= v.T

  # Block i, layer 1 (MLP).
  lowercase__    :   List[str]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''encoder'''		,					'''pre_mlp_layer_norm'''					)
  lowercase__			,							lowercase__    :   List[Any]						= tax_mlp_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''encoder'''		,					UpperCAmelCase					)
  lowercase__    :   int						= layer_norm
  if split_mlp_wi:
   lowercase__    :   Union[str, Any]						= wi[0].T
   lowercase__    :   List[str]						= wi[1].T
  else:
   lowercase__    :   int						= wi.T
  lowercase__    :   List[str]						= wo.T

 lowercase__    :   Dict						= old[
     '''encoder/relpos_bias/rel_embedding'''
 ].T
 lowercase__    :   str						= old['''encoder/encoder_norm/scale''']

 if not is_encoder_only:
  # Decoder.
  for i in range(UpperCAmelCase					):
   # Block i, layer 0 (Self Attention).
   lowercase__    :   Optional[Any]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''pre_self_attention_layer_norm'''					)
   lowercase__			,							lowercase__			,							lowercase__			,							lowercase__    :   Optional[Any]						= tax_attention_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''self_attention'''					)
   lowercase__    :   List[str]						= layer_norm
   lowercase__    :   List[Any]						= k.T
   lowercase__    :   List[Any]						= o.T
   lowercase__    :   List[str]						= q.T
   lowercase__    :   Optional[Any]						= v.T

   # Block i, layer 1 (Cross Attention).
   lowercase__    :   List[Any]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''pre_cross_attention_layer_norm'''					)
   lowercase__			,							lowercase__			,							lowercase__			,							lowercase__    :   Any						= tax_attention_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''encoder_decoder_attention'''					)
   lowercase__    :   Any						= layer_norm
   lowercase__    :   Optional[int]						= k.T
   lowercase__    :   Optional[Any]						= o.T
   lowercase__    :   Tuple						= q.T
   lowercase__    :   List[Any]						= v.T

   # Block i, layer 2 (MLP).
   lowercase__    :   List[str]						= tax_layer_norm_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					'''pre_mlp_layer_norm'''					)
   lowercase__			,							lowercase__    :   List[Any]						= tax_mlp_lookup(UpperCAmelCase		,					UpperCAmelCase		,					'''decoder'''		,					UpperCAmelCase					)
   lowercase__    :   Optional[int]						= layer_norm
   if split_mlp_wi:
    lowercase__    :   List[str]						= wi[0].T
    lowercase__    :   str						= wi[1].T
   else:
    lowercase__    :   str						= wi.T
   lowercase__    :   Optional[Any]						= wo.T

  lowercase__    :   Any						= old['''decoder/decoder_norm/scale''']
  lowercase__    :   Optional[Any]						= old[
      '''decoder/relpos_bias/rel_embedding'''
  ].T

  # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
  if "decoder/logits_dense/kernel" in old:
   lowercase__    :   Dict						= old['''decoder/logits_dense/kernel'''].T

 return new
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase					):
 lowercase__    :   int						= collections.OrderedDict([(k, torch.from_numpy(v.copy()					)) for (k, v) in converted_params.items()]					)

 # Add what is missing.
 if "encoder.embed_tokens.weight" not in state_dict:
  lowercase__    :   Optional[int]						= state_dict['''shared.weight''']

 if not is_encoder_only:
  if "decoder.embed_tokens.weight" not in state_dict:
   lowercase__    :   List[Any]						= state_dict['''shared.weight''']

  if "lm_head.weight" not in state_dict:  # For old 1.0 models.
   print('''Using shared word embeddings as lm_head.'''					)
   lowercase__    :   Union[str, Any]						= state_dict['''shared.weight''']

 return state_dict
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase					):
 lowercase__    :   List[str]						= checkpoints.load_tax_checkpoint(UpperCAmelCase					)
 lowercase__    :   str						= convert_tax_to_pytorch(UpperCAmelCase		,					num_layers=config.num_layers		,					is_encoder_only=UpperCAmelCase					)
 lowercase__    :   List[Any]						= make_state_dict(UpperCAmelCase		,					UpperCAmelCase					)
 model.load_state_dict(UpperCAmelCase		,					strict=UpperCAmelCase					)
def       __UpperCamelCase	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase = False					):
 lowercase__    :   Dict						= TaConfig.from_json_file(UpperCAmelCase					)
 print(F"""Building PyTorch model from configuration: {config}"""					)
 # Non-v1.1 checkpoints could also use T5Model, but this works for all.
 # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
 if is_encoder_only:
  lowercase__    :   List[str]						= TaEncoderModel(UpperCAmelCase					)
 else:
  lowercase__    :   Tuple						= TaForConditionalGeneration(UpperCAmelCase					)

 # Load weights from tf checkpoint
 load_tax_weights_in_ta(UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase					)

 # Save pytorch-model
 print(F"""Save PyTorch model to {pytorch_dump_path}"""					)
 model.save_pretrained(UpperCAmelCase					)

 # Verify that we can load the checkpoint.
 model.from_pretrained(UpperCAmelCase					)
 print('''Done'''					)


if __name__ == "__main__":
  __a:			Optional[Any]	       =       argparse.ArgumentParser(description="""Converts a native T5X checkpoint into a PyTorch checkpoint.""")
  # Required parameters
  parser.add_argument(
      """--t5x_checkpoint_path""", default=None, type=str, required=True, help="""Path to the T5X checkpoint."""
  )
  parser.add_argument(
      """--config_file""",
      default=None,
      type=str,
      required=True,
      help="""The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.""",
  )
  parser.add_argument(
      """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
  )
  parser.add_argument(
      """--is_encoder_only""", action="""store_true""", help="""Check if the model is encoder-decoder model""", default=False
  )
  __a:			Tuple	       =       parser.parse_args()
  convert_tax_checkpoint_to_pytorch(
      args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
  )

