


"""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     :    Tuple    , UpperCamelCase_     :    Any    , UpperCamelCase_     :    List[Any]=2    , UpperCamelCase_     :    Dict=3    , UpperCamelCase_     :    List[str]=4    , UpperCamelCase_     :    Any=2    , UpperCamelCase_     :    List[str]=7    , UpperCamelCase_     :    Dict=True    , UpperCamelCase_     :    int=True    , UpperCamelCase_     :    int=True    , UpperCamelCase_     :    Union[str, Any]=True    , UpperCamelCase_     :    List[str]=9_9    , UpperCamelCase_     :    Optional[Any]=3_6    , UpperCamelCase_     :    Dict=2    , UpperCamelCase_     :    Optional[Any]=4    , UpperCamelCase_     :    Tuple=3_7    , UpperCamelCase_     :    str="gelu"    , UpperCamelCase_     :    Union[str, Any]=0.1    , UpperCamelCase_     :    Optional[int]=0.1    , UpperCamelCase_     :    Dict=5_1_2    , UpperCamelCase_     :    Dict=1_6    , UpperCamelCase_     :    List[Any]=2    , UpperCamelCase_     :    str=0.02    , UpperCamelCase_     :    Tuple=6    , UpperCamelCase_     :    List[Any]=6    , UpperCamelCase_     :    Optional[int]=3    , UpperCamelCase_     :    Any=4    , UpperCamelCase_     :    Union[str, Any]=None    , UpperCamelCase_     :    str=1_0_0_0    , ):
           lowerCAmelCase    :		Tuple				=      parent
           lowerCAmelCase    :		Dict				=      batch_size
           lowerCAmelCase    :		Optional[int]				=      num_channels
           lowerCAmelCase    :		List[Any]				=      image_size
           lowerCAmelCase    :		Optional[Any]				=      patch_size
           lowerCAmelCase    :		Union[str, Any]				=      is_training
           lowerCAmelCase    :		Optional[int]				=      use_input_mask
           lowerCAmelCase    :		List[Any]				=      use_token_type_ids
           lowerCAmelCase    :		Dict				=      use_labels
           lowerCAmelCase    :		str				=      vocab_size
           lowerCAmelCase    :		Any				=      hidden_size
           lowerCAmelCase    :		List[str]				=      num_hidden_layers
           lowerCAmelCase    :		str				=      num_attention_heads
           lowerCAmelCase    :		Optional[int]				=      intermediate_size
           lowerCAmelCase    :		Tuple				=      hidden_act
           lowerCAmelCase    :		Optional[Any]				=      hidden_dropout_prob
           lowerCAmelCase    :		Optional[int]				=      attention_probs_dropout_prob
           lowerCAmelCase    :		Union[str, Any]				=      max_position_embeddings
           lowerCAmelCase    :		int				=      type_vocab_size
           lowerCAmelCase    :		List[Any]				=      type_sequence_label_size
           lowerCAmelCase    :		List[Any]				=      initializer_range
           lowerCAmelCase    :		List[str]				=      coordinate_size
           lowerCAmelCase    :		Any				=      shape_size
           lowerCAmelCase    :		Tuple				=      num_labels
           lowerCAmelCase    :		Union[str, Any]				=      num_choices
           lowerCAmelCase    :		List[str]				=      scope
           lowerCAmelCase    :		Any				=      range_bbox

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






     def    lowerCamelCase__   (		self     :    List[str] ):
           lowerCAmelCase    :		List[str]				=      ids_tensor([self.batch_size, self.text_seq_length]    , self.vocab_size )

           lowerCAmelCase    :		Union[str, Any]				=      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    :		List[str]				=      bbox[i, j, 3]
                             lowerCAmelCase    :		str				=      bbox[i, j, 1]
                             lowerCAmelCase    :		Dict				=      tmp_coordinate
                       if bbox[i, j, 2] < bbox[i, j, 0]:
                             lowerCAmelCase    :		int				=      bbox[i, j, 2]
                             lowerCAmelCase    :		Union[str, Any]				=      bbox[i, j, 0]
                             lowerCAmelCase    :		List[str]				=      tmp_coordinate
           lowerCAmelCase    :		Dict				=      tf.constant(UpperCamelCase_ )

           lowerCAmelCase    :		str				=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

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

           lowerCAmelCase    :		List[str]				=      None
           lowerCAmelCase    :		Optional[Any]				=      None
           if self.use_labels:
                 lowerCAmelCase    :		Tuple				=      ids_tensor([self.batch_size]    , self.type_sequence_label_size )
                 lowerCAmelCase    :		Any				=      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    lowerCamelCase__   (		self     :    Dict    , UpperCamelCase_     :    Tuple    , UpperCamelCase_     :    Optional[int]    , UpperCamelCase_     :    Any    , UpperCamelCase_     :    List[Any]    , UpperCamelCase_     :    List[str]    , UpperCamelCase_     :    List[Any] ):
           lowerCAmelCase    :		List[Any]				=      TFLayoutLMvaModel(config=UpperCamelCase_ )

           # text + image
           lowerCAmelCase    :		List[Any]				=      model(UpperCamelCase_    , pixel_values=UpperCamelCase_    , training=UpperCamelCase_ )
           lowerCAmelCase    :		str				=      model(
               UpperCamelCase_    , bbox=UpperCamelCase_    , pixel_values=UpperCamelCase_    , attention_mask=UpperCamelCase_    , token_type_ids=UpperCamelCase_    , training=UpperCamelCase_    , )
           lowerCAmelCase    :		Dict				=      model(UpperCamelCase_    , bbox=UpperCamelCase_    , pixel_values=UpperCamelCase_    , training=UpperCamelCase_ )

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

           # text only
           lowerCAmelCase    :		Optional[int]				=      model(UpperCamelCase_    , training=UpperCamelCase_ )
           self.parent.assertEqual(
               result.last_hidden_state.shape    , (self.batch_size, self.text_seq_length, self.hidden_size) )

           # image only
           lowerCAmelCase    :		int				=      model({'''pixel_values''': pixel_values}    , training=UpperCamelCase_ )
           self.parent.assertEqual(
               result.last_hidden_state.shape    , (self.batch_size, self.image_seq_length, self.hidden_size) )






     def    lowerCamelCase__   (		self     :    Any    , UpperCamelCase_     :    Dict    , UpperCamelCase_     :    List[str]    , UpperCamelCase_     :    int    , UpperCamelCase_     :    List[Any]    , UpperCamelCase_     :    Union[str, Any]    , UpperCamelCase_     :    str    , UpperCamelCase_     :    Optional[int] ):
           lowerCAmelCase    :		int				=      self.num_labels
           lowerCAmelCase    :		str				=      TFLayoutLMvaForSequenceClassification(config=UpperCamelCase_ )
           lowerCAmelCase    :		Dict				=      model(
               UpperCamelCase_    , bbox=UpperCamelCase_    , pixel_values=UpperCamelCase_    , attention_mask=UpperCamelCase_    , token_type_ids=UpperCamelCase_    , labels=UpperCamelCase_    , training=UpperCamelCase_    , )
           self.parent.assertEqual(result.logits.shape    , (self.batch_size, self.num_labels) )






     def    lowerCamelCase__   (		self     :    Any    , UpperCamelCase_     :    Union[str, Any]    , UpperCamelCase_     :    List[str]    , UpperCamelCase_     :    Optional[Any]    , UpperCamelCase_     :    Dict    , UpperCamelCase_     :    int    , UpperCamelCase_     :    int    , UpperCamelCase_     :    Union[str, Any] ):
           lowerCAmelCase    :		Union[str, Any]				=      self.num_labels
           lowerCAmelCase    :		Tuple				=      TFLayoutLMvaForTokenClassification(config=UpperCamelCase_ )
           lowerCAmelCase    :		List[Any]				=      model(
               UpperCamelCase_    , bbox=UpperCamelCase_    , pixel_values=UpperCamelCase_    , attention_mask=UpperCamelCase_    , token_type_ids=UpperCamelCase_    , labels=UpperCamelCase_    , training=UpperCamelCase_    , )
           self.parent.assertEqual(result.logits.shape    , (self.batch_size, self.text_seq_length, self.num_labels) )






     def    lowerCamelCase__   (		self     :    Any    , UpperCamelCase_     :    int    , UpperCamelCase_     :    str    , UpperCamelCase_     :    int    , UpperCamelCase_     :    Union[str, Any]    , UpperCamelCase_     :    Optional[int]    , UpperCamelCase_     :    Tuple    , UpperCamelCase_     :    Optional[Any] ):
           lowerCAmelCase    :		Optional[int]				=      2
           lowerCAmelCase    :		int				=      TFLayoutLMvaForQuestionAnswering(config=UpperCamelCase_ )
           lowerCAmelCase    :		Tuple				=      model(
               UpperCamelCase_    , bbox=UpperCamelCase_    , pixel_values=UpperCamelCase_    , attention_mask=UpperCamelCase_    , token_type_ids=UpperCamelCase_    , start_positions=UpperCamelCase_    , end_positions=UpperCamelCase_    , training=UpperCamelCase_    , )
           self.parent.assertEqual(result.start_logits.shape    , (self.batch_size, self.seq_length) )
           self.parent.assertEqual(result.end_logits.shape    , (self.batch_size, self.seq_length) )






     def    lowerCamelCase__   (		self     :    Any ):
           lowerCAmelCase    :		Union[str, Any]				=      self.prepare_config_and_inputs()
           ((lowerCAmelCase),					(lowerCAmelCase),					(lowerCAmelCase),					(lowerCAmelCase),					(lowerCAmelCase),					(lowerCAmelCase),					(lowerCAmelCase),					(lowerCAmelCase))    :		Optional[Any]				=      config_and_inputs
           lowerCAmelCase    :		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 						snake_case_(				a__							,    a__							,    unittest.TestCase	):
     __UpperCamelCase     =    (
         (
             TFLayoutLMvaModel,
             TFLayoutLMvaForQuestionAnswering,
             TFLayoutLMvaForSequenceClassification,
             TFLayoutLMvaForTokenClassification,
         )
         if is_tf_available()
         else ()
     )
     __UpperCamelCase     =    (
         {'''document-question-answering''': TFLayoutLMvaForQuestionAnswering, '''feature-extraction''': TFLayoutLMvaModel}
         if is_tf_available()
         else {}
     )

     __UpperCamelCase     =    False
     __UpperCamelCase     =    False
     __UpperCamelCase     =    False






     def    lowerCamelCase__   (		self     :    int    , UpperCamelCase_     :    int    , UpperCamelCase_     :    Tuple    , UpperCamelCase_     :    List[Any]    , UpperCamelCase_     :    List[Any]    , UpperCamelCase_     :    List[str] ):
           return True






     def    lowerCamelCase__   (		self     :    List[Any]    , UpperCamelCase_     :    Dict    , UpperCamelCase_     :    int    , UpperCamelCase_     :    Optional[Any]=False ):
           lowerCAmelCase    :		List[Any]				=      copy.deepcopy(UpperCamelCase_ )

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

           if return_labels:
                 if model_class in get_values(UpperCamelCase_ ):
                       lowerCAmelCase    :		List[Any]				=      tf.ones(self.model_tester.batch_size    , dtype=tf.intaa )
                 elif model_class in get_values(UpperCamelCase_ ):
                       lowerCAmelCase    :		Optional[int]				=      tf.zeros(self.model_tester.batch_size    , dtype=tf.intaa )
                       lowerCAmelCase    :		Union[str, Any]				=      tf.zeros(self.model_tester.batch_size    , dtype=tf.intaa )
                 elif model_class in get_values(UpperCamelCase_ ):
                       lowerCAmelCase    :		List[Any]				=      tf.zeros(self.model_tester.batch_size    , dtype=tf.intaa )
                 elif model_class in get_values(UpperCamelCase_ ):
                       lowerCAmelCase    :		Optional[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] ):
           lowerCAmelCase    :		Any				=      TFLayoutLMvaModelTester(self )
           lowerCAmelCase    :		Tuple				=      ConfigTester(self    , config_class=UpperCamelCase_    , hidden_size=3_7 )






     def    lowerCamelCase__   (		self     :    Any ):
           self.config_tester.run_common_tests()






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

                       # Test that model correctly compute the loss with kwargs
                       lowerCAmelCase    :		Dict				=      self._prepare_for_class(inputs_dict.copy()    , UpperCamelCase_    , return_labels=UpperCamelCase_ )
                       lowerCAmelCase    :		List[str]				=      prepared_for_class.pop('''input_ids''' )

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

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

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

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

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

                       lowerCAmelCase    :		Optional[Any]				=      tuple(UpperCamelCase_ )

                       # Send to model
                       lowerCAmelCase    :		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] ):
           (
               (
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					
           )    :		Union[str, Any]				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_ )






     def    lowerCamelCase__   (		self     :    List[str] ):
           (
               (
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					
           )    :		int				=      self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                 lowerCAmelCase    :		List[Any]				=      type
                 self.model_tester.create_and_check_model(UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_ )






     def    lowerCamelCase__   (		self     :    Optional[int] ):
           (
               (
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					
           )    :		Any				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(
               UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_ )






     def    lowerCamelCase__   (		self     :    List[str] ):
           (
               (
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					
           )    :		int				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(
               UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_ )






     def    lowerCamelCase__   (		self     :    int ):
           (
               (
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					(
               lowerCAmelCase
           ),					
           )    :		str				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(
               UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_ )






     @slow
     def    lowerCamelCase__   (		self     :    str ):
           for model_name in TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCAmelCase    :		Tuple				=      TFLayoutLMvaModel.from_pretrained(UpperCamelCase_ )
                 self.assertIsNotNone(UpperCamelCase_ )



def    _snake_case  (  ):
      lowerCAmelCase    :		Optional[int]				=      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''	)
      return image






@require_tf
class 						snake_case_(				unittest.TestCase	):






     @cached_property
     def    lowerCamelCase__   (		self     :    Tuple ):
           return LayoutLMvaImageProcessor(apply_ocr=UpperCamelCase_ ) if is_vision_available() else None






     @slow
     def    lowerCamelCase__   (		self     :    List[Any] ):
           lowerCAmelCase    :		Any				=      TFLayoutLMvaModel.from_pretrained('''microsoft/layoutlmv3-base''' )

           lowerCAmelCase    :		Optional[Any]				=      self.default_image_processor
           lowerCAmelCase    :		Any				=      prepare_img()
           lowerCAmelCase    :		Optional[Any]				=      image_processor(images=UpperCamelCase_    , return_tensors='''tf''' ).pixel_values

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

           # forward pass
           lowerCAmelCase    :		str				=      model(input_ids=UpperCamelCase_    , bbox=UpperCamelCase_    , pixel_values=UpperCamelCase_    , training=UpperCamelCase_ )

           # verify the logits
           lowerCAmelCase    :		Any				=      (1, 1_9_9, 7_6_8)
           self.assertEqual(outputs.last_hidden_state.shape    , UpperCamelCase_ )

           lowerCAmelCase    :		Tuple				=      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]    , UpperCamelCase_    , atol=1E-4 ) )






'''simple docstring'''







import argparse
import gc
import json
import os
import shutil
import warnings

import torch

from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer


try:
   from transformers import LlamaTokenizerFast
except ImportError as e:
   warnings.warn(e)
   warnings.warn(
       '''The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion'''
   )
   A__:    Dict					    =       None
A__:    Tuple					    =       {
    '''7B''': 1_1008,
    '''13B''': 1_3824,
    '''30B''': 1_7920,
    '''65B''': 2_2016,
    '''70B''': 2_8672,
}
A__:    Any					    =       {
    '''7B''': 1,
    '''7Bf''': 1,
    '''13B''': 2,
    '''13Bf''': 2,
    '''30B''': 4,
    '''65B''': 8,
    '''70B''': 8,
    '''70Bf''': 8,
}


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[int]    ,_UpperCAmelCase	:	Optional[int]=1    ,_UpperCAmelCase	:	List[str]=256  )  ->			Dict:
       return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3  )  ) + multiple_of - 1) // multiple_of)


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	List[Any]  )  ->			List[str]:
       with open(_UpperCAmelCase    ,"""r"""  ) as f:
              return json.load(_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	str    ,_UpperCAmelCase	:	Optional[Any]  )  ->			Tuple:
       with open(_UpperCAmelCase    ,"""w"""  ) as f:
              json.dump(_UpperCAmelCase    ,_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	str    ,_UpperCAmelCase	:	Optional[Any]    ,_UpperCAmelCase	:	int    ,_UpperCAmelCase	:	List[Any]=True  )  ->			Union[str, Any]:
       os.makedirs(_UpperCAmelCase    ,exist_ok=_UpperCAmelCase  )
       _a    : Union[str, Any]     					=os.path.join(_UpperCAmelCase    ,"""tmp"""  )
       os.makedirs(_UpperCAmelCase    ,exist_ok=_UpperCAmelCase  )

       _a    : int     					=read_json(os.path.join(_UpperCAmelCase    ,"""params.json"""  )  )
       _a    : int     					=NUM_SHARDS[model_size]
       _a    : Dict     					=params["""n_layers"""]
       _a    : Union[str, Any]     					=params["""n_heads"""]
       _a    : List[str]     					=n_heads // num_shards
       _a    : int     					=params["""dim"""]
       _a    : Union[str, Any]     					=dim // n_heads
       _a    : int     					=1_0_0_0_0.0
       _a    : str     					=1.0 / (base ** (torch.arange(0    ,_UpperCAmelCase    ,2  ).float() / dims_per_head))

       if "n_kv_heads" in params:
              _a    : str     					=params["""n_kv_heads"""]  # for GQA / MQA
              _a    : Optional[Any]     					=n_heads_per_shard // num_key_value_heads
              _a    : Optional[int]     					=dim // num_key_value_heads
       else:  # compatibility with other checkpoints
              _a    : str     					=n_heads
              _a    : Any     					=n_heads_per_shard
              _a    : str     					=dim

       # permute for sliced rotary
       def permute(_UpperCAmelCase	:	Tuple    ,_UpperCAmelCase	:	Optional[int]=n_heads    ,_UpperCAmelCase	:	Optional[int]=dim    ,_UpperCAmelCase	:	List[str]=dim  ):
              return w.view(_UpperCAmelCase    ,dima // n_heads // 2    ,2    ,_UpperCAmelCase  ).transpose(1    ,2  ).reshape(_UpperCAmelCase    ,_UpperCAmelCase  )

       print(F"Fetching all parameters from the checkpoint at {input_base_path}."  )
       # Load weights
       if model_size == "7B":
              # Not sharded
              # (The sharded implementation would also work, but this is simpler.)
              _a    : Any     					=torch.load(os.path.join(_UpperCAmelCase    ,"""consolidated.00.pth"""  )    ,map_location="""cpu"""  )
       else:
              # Sharded
              _a    : List[Any]     					=[
                  torch.load(os.path.join(_UpperCAmelCase    ,F"consolidated.{i:02d}.pth"  )    ,map_location="""cpu"""  )
                  for i in range(_UpperCAmelCase  )
              ]
       _a    : Any     					=0
       _a    : Optional[int]     					={"""weight_map""": {}}
       for layer_i in range(_UpperCAmelCase  ):
              _a    : List[str]     					=F"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
              if model_size == "7B":
                     # Unsharded
                     _a    : List[str]     					={
                         F"model.layers.{layer_i}.self_attn.q_proj.weight": permute(
                             loaded[F"layers.{layer_i}.attention.wq.weight"]  ),
                         F"model.layers.{layer_i}.self_attn.k_proj.weight": permute(
                             loaded[F"layers.{layer_i}.attention.wk.weight"]  ),
                         F"model.layers.{layer_i}.self_attn.v_proj.weight": loaded[F"layers.{layer_i}.attention.wv.weight"],
                         F"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[F"layers.{layer_i}.attention.wo.weight"],
                         F"model.layers.{layer_i}.mlp.gate_proj.weight": loaded[F"layers.{layer_i}.feed_forward.w1.weight"],
                         F"model.layers.{layer_i}.mlp.down_proj.weight": loaded[F"layers.{layer_i}.feed_forward.w2.weight"],
                         F"model.layers.{layer_i}.mlp.up_proj.weight": loaded[F"layers.{layer_i}.feed_forward.w3.weight"],
                         F"model.layers.{layer_i}.input_layernorm.weight": loaded[F"layers.{layer_i}.attention_norm.weight"],
                         F"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[F"layers.{layer_i}.ffn_norm.weight"],
                     }
              else:
                     # Sharded
                     # Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share
                     # the same storage object, saving attention_norm and ffn_norm will save other weights too, which is
                     # redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned.

                     _a    : Tuple     					={
                         F"model.layers.{layer_i}.input_layernorm.weight": loaded[0][
                             F"layers.{layer_i}.attention_norm.weight"
                         ].clone(),
                         F"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[0][
                             F"layers.{layer_i}.ffn_norm.weight"
                         ].clone(),
                     }
                     _a    : str     					=permute(
                         torch.cat(
                             [
                                 loaded[i][F"layers.{layer_i}.attention.wq.weight"].view(_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase  )
                                 for i in range(_UpperCAmelCase  )
                             ]    ,dim=0    ,).reshape(_UpperCAmelCase    ,_UpperCAmelCase  )  )
                     _a    : Tuple     					=permute(
                         torch.cat(
                             [
                                 loaded[i][F"layers.{layer_i}.attention.wk.weight"].view(
                                     _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase  )
                                 for i in range(_UpperCAmelCase  )
                             ]    ,dim=0    ,).reshape(_UpperCAmelCase    ,_UpperCAmelCase  )    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,)
                     _a    : Any     					=torch.cat(
                         [
                             loaded[i][F"layers.{layer_i}.attention.wv.weight"].view(
                                 _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase  )
                             for i in range(_UpperCAmelCase  )
                         ]    ,dim=0    ,).reshape(_UpperCAmelCase    ,_UpperCAmelCase  )

                     _a    : List[str]     					=torch.cat(
                         [loaded[i][F"layers.{layer_i}.attention.wo.weight"] for i in range(_UpperCAmelCase  )]    ,dim=1  )
                     _a    : Union[str, Any]     					=torch.cat(
                         [loaded[i][F"layers.{layer_i}.feed_forward.w1.weight"] for i in range(_UpperCAmelCase  )]    ,dim=0  )
                     _a    : Tuple     					=torch.cat(
                         [loaded[i][F"layers.{layer_i}.feed_forward.w2.weight"] for i in range(_UpperCAmelCase  )]    ,dim=1  )
                     _a    : Union[str, Any]     					=torch.cat(
                         [loaded[i][F"layers.{layer_i}.feed_forward.w3.weight"] for i in range(_UpperCAmelCase  )]    ,dim=0  )

              _a    : str     					=inv_freq
              for k, v in state_dict.items():
                     _a    : Any     					=filename
                     param_count += v.numel()
              torch.save(_UpperCAmelCase    ,os.path.join(_UpperCAmelCase    ,_UpperCAmelCase  )  )

       _a    : Union[str, Any]     					=F"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"
       if model_size == "7B":
              # Unsharded
              _a    : List[str]     					={
                  """model.embed_tokens.weight""": loaded["""tok_embeddings.weight"""],
                  """model.norm.weight""": loaded["""norm.weight"""],
                  """lm_head.weight""": loaded["""output.weight"""],
              }
       else:
              _a    : int     					={
                  """model.norm.weight""": loaded[0]["""norm.weight"""],
                  """model.embed_tokens.weight""": torch.cat(
                      [loaded[i]["""tok_embeddings.weight"""] for i in range(_UpperCAmelCase  )]    ,dim=1  ),
                  """lm_head.weight""": torch.cat([loaded[i]["""output.weight"""] for i in range(_UpperCAmelCase  )]    ,dim=0  ),
              }

       for k, v in state_dict.items():
              _a    : Dict     					=filename
              param_count += v.numel()
       torch.save(_UpperCAmelCase    ,os.path.join(_UpperCAmelCase    ,_UpperCAmelCase  )  )

       # Write configs
       _a    : Tuple     					={"""total_size""": param_count * 2}
       write_json(_UpperCAmelCase    ,os.path.join(_UpperCAmelCase    ,"""pytorch_model.bin.index.json"""  )  )
       _a    : Optional[Any]     					=params["""ffn_dim_multiplier"""] if """ffn_dim_multiplier""" in params else 1
       _a    : int     					=params["""multiple_of"""] if """multiple_of""" in params else 256
       _a    : List[Any]     					=LlamaConfig(
           hidden_size=_UpperCAmelCase    ,intermediate_size=compute_intermediate_size(_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase  )    ,num_attention_heads=params["""n_heads"""]    ,num_hidden_layers=params["""n_layers"""]    ,rms_norm_eps=params["""norm_eps"""]    ,num_key_value_heads=_UpperCAmelCase    ,)
       config.save_pretrained(_UpperCAmelCase  )

       # Make space so we can load the model properly now.
       del state_dict
       del loaded
       gc.collect()

       print("""Loading the checkpoint in a Llama model."""  )
       _a    : Any     					=LlamaForCausalLM.from_pretrained(_UpperCAmelCase    ,torch_dtype=torch.floataa    ,low_cpu_mem_usage=_UpperCAmelCase  )
       # Avoid saving this as part of the config.
       del model.config._name_or_path

       print("""Saving in the Transformers format."""  )
       model.save_pretrained(_UpperCAmelCase    ,safe_serialization=_UpperCAmelCase  )
       shutil.rmtree(_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	int    ,_UpperCAmelCase	:	int  )  ->			Optional[Any]:
       # Initialize the tokenizer based on the `spm` model
       _a    : List[str]     					=LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
       print(F"Saving a {tokenizer_class.__name__} to {tokenizer_path}."  )
       _a    : List[Any]     					=tokenizer_class(_UpperCAmelCase  )
       tokenizer.save_pretrained(_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  )  ->			Union[str, Any]:
       _a    : List[str]     					=argparse.ArgumentParser()
       parser.add_argument(
           """--input_dir"""    ,help="""Location of LLaMA weights, which contains tokenizer.model and model folders"""    ,)
       parser.add_argument(
           """--model_size"""    ,choices=["""7B""", """7Bf""", """13B""", """13Bf""", """30B""", """65B""", """70B""", """70Bf""", """tokenizer_only"""]    ,)
       parser.add_argument(
           """--output_dir"""    ,help="""Location to write HF model and tokenizer"""    ,)
       parser.add_argument("""--safe_serialization"""    ,type=_UpperCAmelCase    ,help="""Whether or not to save using `safetensors`."""  )
       _a    : Optional[Any]     					=parser.parse_args()
       if args.model_size != "tokenizer_only":
              write_model(
                  model_path=args.output_dir    ,input_base_path=os.path.join(args.input_dir    ,args.model_size  )    ,model_size=args.model_size    ,safe_serialization=args.safe_serialization    ,)
       _a    : List[Any]     					=os.path.join(args.input_dir    ,"""tokenizer.model"""  )
       write_tokenizer(args.output_dir    ,_UpperCAmelCase  )


if __name__ == "__main__":
   main()


