




"""simple docstring"""

import argparse
import os
from pathlib import Path
from typing import Dict

import tensorflow as tf
import torch
from tqdm import tqdm

from transformers import PegasusConfig, PegasusForConditionalGeneration, PegasusTokenizer
from transformers.models.pegasus.configuration_pegasus import DEFAULTS, task_specific_params


a							      =     [
    # replace left string with right string to get the relevant state_dict key (identical state dict to bart)
    ['''memory_attention''', '''encoder_attn'''],
    ['''attention''', '''attn'''],
    ['''/''', '''.'''],
    ['''.LayerNorm.gamma''', '''_layer_norm.weight'''],
    ['''.LayerNorm.beta''', '''_layer_norm.bias'''],
    ['''r.layer_''', '''r.layers.'''],
    ['''output_proj''', '''out_proj'''],
    ['''ffn.dense_1.''', '''fc2.'''],
    ['''ffn.dense.''', '''fc1.'''],
    ['''ffn_layer_norm''', '''final_layer_norm'''],
    ['''kernel''', '''weight'''],
    ['''encoder_layer_norm.''', '''encoder.layer_norm.'''],
    ['''decoder_layer_norm.''', '''decoder.layer_norm.'''],
    ['''embeddings.weights''', '''shared.weight'''],
]






def 		_snake_case	(   _snake_case				:   str				)       ->   str:






      '''simple docstring'''


      for pegasus_name, hf_name in PATTERNS:
            _A     =  k.replace(__lowerCamelCase							,    __lowerCamelCase				)
      return k






def 		_snake_case	(   _snake_case				:   dict							,    _snake_case				:   dict				)       ->   int:






      '''simple docstring'''


      _A     =  DEFAULTS.copy()
      cfg_kwargs.update(__lowerCamelCase				)
      _A     =  PegasusConfig(**__lowerCamelCase				)
      _A     =  PegasusForConditionalGeneration(__lowerCamelCase				)
      _A     =  torch_model.model.state_dict()
      _A     =  {}
      for k, v in tf_weights.items():
            _A     =  rename_state_dict_key(__lowerCamelCase				)
            if new_k not in sd:
                  raise ValueError(F'''could not find new key {new_k} in state dict. (converted from {k})'''				)

            if "dense" in k or "proj" in new_k:
                  _A     =  v.T
            _A     =  torch.tensor(__lowerCamelCase							,    dtype=sd[new_k].dtype				)
            assert v.shape == sd[new_k].shape, F'''{new_k}, {k}, {v.shape}, {sd[new_k].shape}'''
      # make sure embedding.padding_idx is respected
      _A     =  torch.zeros_like(mapping['shared.weight'][cfg.pad_token_id + 1]				)
      _A     =  mapping['shared.weight']
      _A     =  mapping['shared.weight']
      _A     =  {k: torch.zeros_like(__lowerCamelCase				) for k, v in sd.items() if k.endswith('bias'				) and k not in mapping}
      mapping.update(**__lowerCamelCase				)
      _A     =  torch_model.model.load_state_dict(__lowerCamelCase							,    strict=__lowerCamelCase				)
      _A     =  [
          k for k in missing if k not in ['encoder.embed_positions.weight', 'decoder.embed_positions.weight']
      ]
      assert unexpected_missing == [], F'''no matches found for the following torch keys {unexpected_missing}'''
      assert extra == [], F'''no matches found for the following tf keys {extra}'''
      return torch_model






def 		_snake_case	(   _snake_case				:   Optional[Any]="./ckpt/aeslc/model.ckpt-32000"				)       ->   Optional[Any]:






      '''simple docstring'''


      _A     =  tf.train.list_variables(__lowerCamelCase				)
      _A     =  {}
      _A     =  ['Adafactor', 'global_step']
      for name, shape in tqdm(__lowerCamelCase							,    desc='converting tf checkpoint to dict'				):
            _A     =  any(pat in name for pat in ignore_name				)
            if skip_key:
                  continue
            _A     =  tf.train.load_variable(__lowerCamelCase							,    __lowerCamelCase				)
            _A     =  array
      return tf_weights






def 		_snake_case	(   _snake_case				:   str							,    _snake_case				:   str				)       ->   List[Any]:






      '''simple docstring'''


      _A     =  Path(__lowerCamelCase				).parent.name
      _A     =  task_specific_params[F'''summarization_{dataset}''']['max_position_embeddings']
      _A     =  PegasusTokenizer.from_pretrained('sshleifer/pegasus'							,    model_max_length=__lowerCamelCase				)
      assert tok.model_max_length == desired_max_model_length
      tok.save_pretrained(__lowerCamelCase				)

      # convert model
      _A     =  get_tf_weights_as_numpy(__lowerCamelCase				)
      _A     =  task_specific_params[F'''summarization_{dataset}''']
      if dataset == "large":
            _A     =  task_specific_params
      _A     =  convert_pegasus(__lowerCamelCase							,    __lowerCamelCase				)
      torch_model.save_pretrained(__lowerCamelCase				)
      _A     =  torch_model.state_dict()
      sd.pop('model.decoder.embed_positions.weight'				)
      sd.pop('model.encoder.embed_positions.weight'				)
      torch.save(__lowerCamelCase							,    Path(__lowerCamelCase				) / 'pytorch_model.bin'				)


if __name__ == "__main__":
       a							      =     argparse.ArgumentParser()
       # Required parameters
       parser.add_argument('''tf_ckpt_path''', type=str, help='''passed to tf.train.list_variables''')
       parser.add_argument('''save_dir''', default=None, type=str, help='''Path to the output PyTorch model.''')
       a							      =     parser.parse_args()
       if args.save_dir is None:
              a							      =     Path(args.tf_ckpt_path).parent.name
              a							      =     os.path.join('''pegasus''', dataset)
       convert_pegasus_ckpt_to_pytorch(args.tf_ckpt_path, args.save_dir)





'''simple docstring'''






from __future__ import annotations

import unittest

import numpy as np

from transformers import LayoutLMConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
 import tensorflow as tf

 from transformers.models.layoutlm.modeling_tf_layoutlm import (
     TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
     TFLayoutLMForMaskedLM,
     TFLayoutLMForQuestionAnswering,
     TFLayoutLMForSequenceClassification,
     TFLayoutLMForTokenClassification,
     TFLayoutLMModel,
 )






class    __magic_name__				:






       def __init__(     self    ,				snake_case    ,				snake_case=1_3    ,				snake_case=7    ,				snake_case=True    ,				snake_case=True    ,				snake_case=True    ,				snake_case=True    ,				snake_case=9_9    ,				snake_case=3_2    ,				snake_case=2    ,				snake_case=4    ,				snake_case=3_7    ,				snake_case="gelu"    ,				snake_case=0.1    ,				snake_case=0.1    ,				snake_case=5_1_2    ,				snake_case=1_6    ,				snake_case=2    ,				snake_case=0.02    ,				snake_case=3    ,				snake_case=4    ,				snake_case=None    ,				snake_case=1_0_0_0    ,				)		->    Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  List[str]				       =parent
          _UpperCAmelCase  :  str				       =batch_size
          _UpperCAmelCase  :  List[Any]				       =seq_length
          _UpperCAmelCase  :  List[str]				       =is_training
          _UpperCAmelCase  :  Any				       =use_input_mask
          _UpperCAmelCase  :  List[str]				       =use_token_type_ids
          _UpperCAmelCase  :  Optional[int]				       =use_labels
          _UpperCAmelCase  :  Optional[Any]				       =vocab_size
          _UpperCAmelCase  :  str				       =hidden_size
          _UpperCAmelCase  :  str				       =num_hidden_layers
          _UpperCAmelCase  :  Tuple				       =num_attention_heads
          _UpperCAmelCase  :  List[Any]				       =intermediate_size
          _UpperCAmelCase  :  str				       =hidden_act
          _UpperCAmelCase  :  str				       =hidden_dropout_prob
          _UpperCAmelCase  :  Any				       =attention_probs_dropout_prob
          _UpperCAmelCase  :  int				       =max_position_embeddings
          _UpperCAmelCase  :  List[Any]				       =type_vocab_size
          _UpperCAmelCase  :  List[str]				       =type_sequence_label_size
          _UpperCAmelCase  :  Union[str, Any]				       =initializer_range
          _UpperCAmelCase  :  Tuple				       =num_labels
          _UpperCAmelCase  :  int				       =num_choices
          _UpperCAmelCase  :  List[str]				       =scope
          _UpperCAmelCase  :  Any				       =range_bbox






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




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       =ids_tensor([self.batch_size, self.seq_length]    ,				self.vocab_size)

          # convert bbox to numpy since TF does not support item assignment
          _UpperCAmelCase  :  List[str]				       =ids_tensor([self.batch_size, self.seq_length, 4]    ,				self.range_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]:
                   _UpperCAmelCase  :  Optional[Any]				       =bbox[i, j, 3]
                   _UpperCAmelCase  :  Optional[Any]				       =bbox[i, j, 1]
                   _UpperCAmelCase  :  str				       =t
                if bbox[i, j, 2] < bbox[i, j, 0]:
                   _UpperCAmelCase  :  Optional[int]				       =bbox[i, j, 2]
                   _UpperCAmelCase  :  Tuple				       =bbox[i, j, 0]
                   _UpperCAmelCase  :  Dict				       =t
          _UpperCAmelCase  :  Union[str, Any]				       =tf.convert_to_tensor(snake_case)

          _UpperCAmelCase  :  str				       =None
          if self.use_input_mask:
             _UpperCAmelCase  :  Optional[Any]				       =random_attention_mask([self.batch_size, self.seq_length])

          _UpperCAmelCase  :  Union[str, Any]				       =None
          if self.use_token_type_ids:
             _UpperCAmelCase  :  Optional[Any]				       =ids_tensor([self.batch_size, self.seq_length]    ,				self.type_vocab_size)

          _UpperCAmelCase  :  Optional[int]				       =None
          _UpperCAmelCase  :  Tuple				       =None
          _UpperCAmelCase  :  int				       =None
          if self.use_labels:
             _UpperCAmelCase  :  str				       =ids_tensor([self.batch_size]    ,				self.type_sequence_label_size)
             _UpperCAmelCase  :  List[str]				       =ids_tensor([self.batch_size, self.seq_length]    ,				self.num_labels)
             _UpperCAmelCase  :  Tuple				       =ids_tensor([self.batch_size]    ,				self.num_choices)

          _UpperCAmelCase  :  Optional[int]				       =LayoutLMConfig(
              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    ,				)

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






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  int				       =TFLayoutLMModel(config=snake_case)

          _UpperCAmelCase  :  Dict				       =model(snake_case    ,				snake_case    ,				attention_mask=snake_case    ,				token_type_ids=snake_case)
          _UpperCAmelCase  :  Dict				       =model(snake_case    ,				snake_case    ,				token_type_ids=snake_case)
          _UpperCAmelCase  :  int				       =model(snake_case    ,				snake_case)
          self.parent.assertEqual(result.last_hidden_state.shape    ,				(self.batch_size, self.seq_length, self.hidden_size))
          self.parent.assertEqual(result.pooler_output.shape    ,				(self.batch_size, self.hidden_size))






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case)		->    List[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[Any]				       =TFLayoutLMForMaskedLM(config=snake_case)

          _UpperCAmelCase  :  Optional[Any]				       =model(snake_case    ,				snake_case    ,				attention_mask=snake_case    ,				token_type_ids=snake_case    ,				labels=snake_case)
          self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.vocab_size))






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case)		->    Any:




          '''simple docstring'''
          _UpperCAmelCase  :  Dict				       =self.num_labels
          _UpperCAmelCase  :  Dict				       =TFLayoutLMForSequenceClassification(config=snake_case)

          _UpperCAmelCase  :  Union[str, Any]				       =model(snake_case    ,				snake_case    ,				attention_mask=snake_case    ,				token_type_ids=snake_case)
          self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.num_labels))






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case)		->    Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase  :  str				       =self.num_labels
          _UpperCAmelCase  :  Any				       =TFLayoutLMForTokenClassification(config=snake_case)

          _UpperCAmelCase  :  str				       =model(snake_case    ,				snake_case    ,				attention_mask=snake_case    ,				token_type_ids=snake_case    ,				labels=snake_case)
          self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.num_labels))






       def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case    ,				snake_case)		->    str:




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       =TFLayoutLMForQuestionAnswering(config=snake_case)

          _UpperCAmelCase  :  Union[str, Any]				       =model(snake_case    ,				snake_case    ,				attention_mask=snake_case    ,				token_type_ids=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)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       =self.prepare_config_and_inputs()
          (
              (
              _UpperCAmelCase
          )       ,			(
              _UpperCAmelCase
          )       ,			(
              _UpperCAmelCase
          )       ,			(
              _UpperCAmelCase
          )       ,			(
              _UpperCAmelCase
          )       ,			(
              _UpperCAmelCase
          )       ,			(
              _UpperCAmelCase
          )       ,			(
              _UpperCAmelCase
          )       ,			
          )  :  Optional[int]				       =config_and_inputs
          _UpperCAmelCase  :  Any				       ={
              'input_ids': input_ids,
              'bbox': bbox,
              'token_type_ids': token_type_ids,
              'attention_mask': input_mask,
          }
          return config, inputs_dict






@require_tf
class    __magic_name__				(						lowerCAmelCase   ,lowerCAmelCase   ,unittest.TestCase							):
       UpperCAmelCase		 =(
           (
               TFLayoutLMModel,
               TFLayoutLMForMaskedLM,
               TFLayoutLMForTokenClassification,
               TFLayoutLMForSequenceClassification,
               TFLayoutLMForQuestionAnswering,
           )
           if is_tf_available()
           else ()
       )
       UpperCAmelCase		 =(
           {
               "feature-extraction": TFLayoutLMModel,
               "fill-mask": TFLayoutLMForMaskedLM,
               "text-classification": TFLayoutLMForSequenceClassification,
               "token-classification": TFLayoutLMForTokenClassification,
               "zero-shot": TFLayoutLMForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       UpperCAmelCase		 =False
       UpperCAmelCase		 =True
       UpperCAmelCase		 =1_0






       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  List[str]				       =TFLayoutLMModelTester(self)
          _UpperCAmelCase  :  int				       =ConfigTester(self    ,				config_class=snake_case    ,				hidden_size=3_7)






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




          '''simple docstring'''
          self.config_tester.run_common_tests()






       def      lowerCAmelCase						(     self)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  Tuple				       =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case)






       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          _UpperCAmelCase  :  List[str]				       =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*snake_case)






       def      lowerCAmelCase						(     self)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[int]				       =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*snake_case)






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




          '''simple docstring'''
          _UpperCAmelCase  :  str				       =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*snake_case)






       def      lowerCAmelCase						(     self)		->    str:




          '''simple docstring'''
          _UpperCAmelCase  :  Dict				       =self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*snake_case)






       @slow
       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          for model_name in TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _UpperCAmelCase  :  str				       =TFLayoutLMModel.from_pretrained(snake_case)
             self.assertIsNotNone(snake_case)






       @unittest.skip('Onnx compliancy broke with TF 2.10')
       def      lowerCAmelCase						(     self)		->    int:




          '''simple docstring'''
          pass






def lowerCamelCase__     (		):







   '''simple docstring'''




   _UpperCAmelCase  :  int				       =tf.convert_to_tensor([[1_0_1,1_0_1_9,1_0_1_4,1_0_1_6,1_0_3_7,1_2_8_4_9,4_7_4_7,1_0_0_4,1_4_2_4_6,2_2_7_8,5_4_3_9,4_5_2_4,5_0_0_2,2_9_3_0,2_1_9_3,2_9_3_0,4_3_4_1,3_2_0_8,1_0_0_5,1_0_5_5,2_1_7_1,2_8_4_8,1_1_3_0_0,3_5_3_1,1_0_2],[1_0_1,4_0_7_0,4_0_3_4,7_0_2_0,1_0_2_4,3_0_5_8,1_0_1_5,1_0_1_3,2_8_6_1,1_0_1_3,6_0_7_0,1_9_2_7_4,2_7_7_2,6_2_0_5,2_7_8_1_4,1_6_1_4_7,1_6_1_4_7,4_3_4_3,2_0_4_7,1_0_2_8_3,1_0_9_6_9,1_4_3_8_9,1_0_1_2,2_3_3_8,1_0_2]]				)  # noqa: E231
   _UpperCAmelCase  :  Optional[int]				       =tf.convert_to_tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],]				)  # noqa: E231
   _UpperCAmelCase  :  Dict				       =tf.convert_to_tensor([[[0,0,0,0],[4_2_3,2_3_7,4_4_0,2_5_1],[4_2_7,2_7_2,4_4_1,2_8_7],[4_1_9,1_1_5,4_3_7,1_2_9],[9_6_1,8_8_5,9_9_2,9_1_2],[2_5_6,3_8,3_3_0,5_8],[2_5_6,3_8,3_3_0,5_8],[3_3_6,4_2,3_5_3,5_7],[3_6_0,3_9,4_0_1,5_6],[3_6_0,3_9,4_0_1,5_6],[4_1_1,3_9,4_7_1,5_9],[4_7_9,4_1,5_2_8,5_9],[5_3_3,3_9,6_3_0,6_0],[6_7,1_1_3,1_3_4,1_3_1],[1_4_1,1_1_5,2_0_9,1_3_2],[6_8,1_4_9,1_3_3,1_6_6],[1_4_1,1_4_9,1_8_7,1_6_4],[1_9_5,1_4_8,2_8_7,1_6_5],[1_9_5,1_4_8,2_8_7,1_6_5],[1_9_5,1_4_8,2_8_7,1_6_5],[2_9_5,1_4_8,3_4_9,1_6_5],[4_4_1,1_4_9,4_9_2,1_6_6],[4_9_7,1_4_9,5_4_6,1_6_4],[6_4,2_0_1,1_2_5,2_1_8],[1_0_0_0,1_0_0_0,1_0_0_0,1_0_0_0]],[[0,0,0,0],[6_6_2,1_5_0,7_5_4,1_6_6],[6_6_5,1_9_9,7_4_2,2_1_1],[5_1_9,2_1_3,5_5_4,2_2_8],[5_1_9,2_1_3,5_5_4,2_2_8],[1_3_4,4_3_3,1_8_7,4_5_4],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[3_1_4,4_6_9,3_7_6,4_8_2],[5_0_4,6_8_4,5_8_2,7_0_6],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[6_1_0,7_4_9,6_5_2,7_6_5],[1_3_0,6_5_9,1_6_8,6_7_2],[1_7_6,6_5_7,2_3_7,6_7_2],[2_3_8,6_5_7,3_1_2,6_7_2],[4_4_3,6_5_3,6_2_8,6_7_2],[4_4_3,6_5_3,6_2_8,6_7_2],[7_1_6,3_0_1,8_2_5,3_1_7],[1_0_0_0,1_0_0_0,1_0_0_0,1_0_0_0]]]				)  # noqa: E231
   _UpperCAmelCase  :  List[Any]				       =tf.convert_to_tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]				)  # noqa: E231
   # these are sequence labels (i.e. at the token level)
   _UpperCAmelCase  :  List[Any]				       =tf.convert_to_tensor([[-1_0_0,1_0,1_0,1_0,9,1,-1_0_0,7,7,-1_0_0,7,7,4,2,5,2,8,8,-1_0_0,-1_0_0,5,0,3,2,-1_0_0],[-1_0_0,1_2,1_2,1_2,-1_0_0,1_2,1_0,-1_0_0,-1_0_0,-1_0_0,-1_0_0,1_0,1_2,9,-1_0_0,-1_0_0,-1_0_0,1_0,1_0,1_0,9,1_2,-1_0_0,1_0,-1_0_0]]				)  # noqa: E231
   # fmt: on

   return input_ids, attention_mask, bbox, token_type_ids, labels






@require_tf
class    __magic_name__				(						unittest.TestCase							):






       @slow
       def      lowerCAmelCase						(     self)		->    Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase  :  Optional[Any]				       =TFLayoutLMModel.from_pretrained('microsoft/layoutlm-base-uncased')

          _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase  :  Tuple				       =prepare_layoutlm_batch_inputs()

          # forward pass
          _UpperCAmelCase  :  Dict				       =model(input_ids=snake_case    ,				bbox=snake_case    ,				attention_mask=snake_case    ,				token_type_ids=snake_case)

          # test the sequence output on [0, :3, :3]
          _UpperCAmelCase  :  List[Any]				       =tf.convert_to_tensor(
              [[0.17_85, -0.19_47, -0.04_25], [-0.32_54, -0.28_07, 0.25_53], [-0.53_91, -0.33_22, 0.33_64]]    ,				)

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

          # test the pooled output on [1, :3]
          _UpperCAmelCase  :  List[Any]				       =tf.convert_to_tensor([-0.65_80, -0.02_14, 0.85_52])

          self.assertTrue(np.allclose(outputs.pooler_output[1, :3]    ,				snake_case    ,				atol=1E-3))






       @slow
       def      lowerCAmelCase						(     self)		->    Dict:




          '''simple docstring'''
          # initialize model with randomly initialized sequence classification head
          _UpperCAmelCase  :  str				       =TFLayoutLMForSequenceClassification.from_pretrained('microsoft/layoutlm-base-uncased'    ,				num_labels=2)

          _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase  :  int				       =prepare_layoutlm_batch_inputs()

          # forward pass
          _UpperCAmelCase  :  Optional[Any]				       =model(
              input_ids=snake_case    ,				bbox=snake_case    ,				attention_mask=snake_case    ,				token_type_ids=snake_case    ,				labels=tf.convert_to_tensor([1, 1])    ,				)

          # test whether we get a loss as a scalar
          _UpperCAmelCase  :  Optional[int]				       =outputs.loss
          _UpperCAmelCase  :  int				       =(2,)
          self.assertEqual(loss.shape    ,				snake_case)

          # test the shape of the logits
          _UpperCAmelCase  :  int				       =outputs.logits
          _UpperCAmelCase  :  Union[str, Any]				       =(2, 2)
          self.assertEqual(logits.shape    ,				snake_case)






       @slow
       def      lowerCAmelCase						(     self)		->    int:




          '''simple docstring'''
          # initialize model with randomly initialized token classification head
          _UpperCAmelCase  :  Optional[int]				       =TFLayoutLMForTokenClassification.from_pretrained('microsoft/layoutlm-base-uncased'    ,				num_labels=1_3)

          _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase  :  Optional[Any]				       =prepare_layoutlm_batch_inputs()

          # forward pass
          _UpperCAmelCase  :  str				       =model(
              input_ids=snake_case    ,				bbox=snake_case    ,				attention_mask=snake_case    ,				token_type_ids=snake_case    ,				labels=snake_case)

          # test the shape of the logits
          _UpperCAmelCase  :  str				       =outputs.logits
          _UpperCAmelCase  :  Optional[Any]				       =tf.convert_to_tensor((2, 2_5, 1_3))
          self.assertEqual(logits.shape    ,				snake_case)






       @slow
       def      lowerCAmelCase						(     self)		->    Optional[int]:




          '''simple docstring'''
          # initialize model with randomly initialized token classification head
          _UpperCAmelCase  :  Optional[Any]				       =TFLayoutLMForQuestionAnswering.from_pretrained('microsoft/layoutlm-base-uncased')

          _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase  :  Optional[Any]				       =prepare_layoutlm_batch_inputs()

          # forward pass
          _UpperCAmelCase  :  Any				       =model(input_ids=snake_case    ,				bbox=snake_case    ,				attention_mask=snake_case    ,				token_type_ids=snake_case)

          # test the shape of the logits
          _UpperCAmelCase  :  Optional[int]				       =tf.convert_to_tensor((2, 2_5))
          self.assertEqual(outputs.start_logits.shape    ,				snake_case)
          self.assertEqual(outputs.end_logits.shape    ,				snake_case)

