





import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
lowercase_					     =	logging.get_logger(__name__)




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     Union[str, Any]:
  lowercase__			    =						os.path.abspath(_SCREAMING_SNAKE_CASE	)
  logger.info(F"""Converting TensorFlow checkpoint from {tf_path}"""	)
  # Load weights from TF model
  lowercase__			    =						tf.train.list_variables(_SCREAMING_SNAKE_CASE	)
  lowercase__			    =						[]
  lowercase__			    =						[]
  lowercase__			    =						[]
  for full_name, shape in init_vars:
    # logger.info(f"Loading TF weight {name} with shape {shape}")
    lowercase__			    =						full_name.split('/'	)
    if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
      logger.info(F"""Skipping non-model layer {full_name}"""	)
      continue
    if "optimizer" in full_name:
      logger.info(F"""Skipping optimization layer {full_name}"""	)
      continue
    if name[0] == "model":
      # ignore initial 'model'
      lowercase__			    =						name[1:]
    # figure out how many levels deep the name is
    lowercase__			    =						0
    for _name in name:
      if _name.startswith('layer_with_weights'	):
        depth += 1
      else:
        break
    layer_depth.append(_SCREAMING_SNAKE_CASE	)
    # read data
    lowercase__			    =						tf.train.load_variable(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)
    names.append('/'.join(_SCREAMING_SNAKE_CASE	)	)
    arrays.append(_SCREAMING_SNAKE_CASE	)
  logger.info(F"""Read a total of {len(_SCREAMING_SNAKE_CASE	):,} layers"""	)

  # Sanity check
  if len(set(_SCREAMING_SNAKE_CASE	)	) != 1:
    raise ValueError(F"""Found layer names with different depths (layer depth {list(set(_SCREAMING_SNAKE_CASE	)	)})"""	)
  lowercase__			    =						list(set(_SCREAMING_SNAKE_CASE	)	)[0]
  if layer_depth != 1:
    raise ValueError(
        'The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP'
        ' heads.'	)

  # convert layers
  logger.info('Converting weights...'	)
  for full_name, array in zip(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	):
    lowercase__			    =						full_name.split('/'	)
    lowercase__			    =						model
    lowercase__			    =						[]
    for i, m_name in enumerate(_SCREAMING_SNAKE_CASE	):
      if m_name == ".ATTRIBUTES":
        # variable names end with .ATTRIBUTES/VARIABLE_VALUE
        break
      if m_name.startswith('layer_with_weights'	):
        lowercase__			    =						int(m_name.split('-'	)[-1]	)
        if layer_num <= 2:
          # embedding layers
          # layer_num 0: word_embeddings
          # layer_num 1: position_embeddings
          # layer_num 2: token_type_embeddings
          continue
        elif layer_num == 3:
          # embedding LayerNorm
          trace.extend(['embeddings', 'LayerNorm']	)
          lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'embeddings'	)
          lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'LayerNorm'	)
        elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
          # encoder layers
          trace.extend(['encoder', 'layer', str(layer_num - 4	)]	)
          lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'encoder'	)
          lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'layer'	)
          lowercase__			    =						pointer[layer_num - 4]
        elif layer_num == config.num_hidden_layers + 4:
          # pooler layer
          trace.extend(['pooler', 'dense']	)
          lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'pooler'	)
          lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'dense'	)
      elif m_name == "embeddings":
        trace.append('embeddings'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'embeddings'	)
        if layer_num == 0:
          trace.append('word_embeddings'	)
          lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'word_embeddings'	)
        elif layer_num == 1:
          trace.append('position_embeddings'	)
          lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'position_embeddings'	)
        elif layer_num == 2:
          trace.append('token_type_embeddings'	)
          lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'token_type_embeddings'	)
        else:
          raise ValueError(F"""Unknown embedding layer with name {full_name}"""	)
        trace.append('weight'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'weight'	)
      elif m_name == "_attention_layer":
        # self-attention layer
        trace.extend(['attention', 'self']	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'attention'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'self'	)
      elif m_name == "_attention_layer_norm":
        # output attention norm
        trace.extend(['attention', 'output', 'LayerNorm']	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'attention'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'output'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'LayerNorm'	)
      elif m_name == "_attention_output_dense":
        # output attention dense
        trace.extend(['attention', 'output', 'dense']	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'attention'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'output'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'dense'	)
      elif m_name == "_output_dense":
        # output dense
        trace.extend(['output', 'dense']	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'output'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'dense'	)
      elif m_name == "_output_layer_norm":
        # output dense
        trace.extend(['output', 'LayerNorm']	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'output'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'LayerNorm'	)
      elif m_name == "_key_dense":
        # attention key
        trace.append('key'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'key'	)
      elif m_name == "_query_dense":
        # attention query
        trace.append('query'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'query'	)
      elif m_name == "_value_dense":
        # attention value
        trace.append('value'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'value'	)
      elif m_name == "_intermediate_dense":
        # attention intermediate dense
        trace.extend(['intermediate', 'dense']	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'intermediate'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'dense'	)
      elif m_name == "_output_layer_norm":
        # output layer norm
        trace.append('output'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'output'	)
      # weights & biases
      elif m_name in ["bias", "beta"]:
        trace.append('bias'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'bias'	)
      elif m_name in ["kernel", "gamma"]:
        trace.append('weight'	)
        lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		'weight'	)
      else:
        logger.warning(F"""Ignored {m_name}"""	)
        # for certain layers reshape is necessary
    lowercase__			    =						'.'.join(_SCREAMING_SNAKE_CASE	)
    if re.match(R'(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)'	,		_SCREAMING_SNAKE_CASE	) or re.match(
        R'(\S+)\.attention\.output\.dense\.weight'	,		_SCREAMING_SNAKE_CASE	):
      lowercase__			    =						array.reshape(pointer.data.shape	)
    if "kernel" in full_name:
      lowercase__			    =						array.transpose()
    if pointer.shape == array.shape:
      lowercase__			    =						torch.from_numpy(_SCREAMING_SNAKE_CASE	)
    else:
      raise ValueError(
          F"""Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:"""
          F""" {array.shape}"""	)
    logger.info(F"""Successfully set variable {full_name} to PyTorch layer {trace}"""	)
  return model




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     Dict:
  # Instantiate model
  logger.info(F"""Loading model based on config from {config_path}..."""	)
  lowercase__			    =						BertConfig.from_json_file(_SCREAMING_SNAKE_CASE	)
  lowercase__			    =						BertModel(_SCREAMING_SNAKE_CASE	)

  # Load weights from checkpoint
  logger.info(F"""Loading weights from checkpoint {tf_checkpoint_path}..."""	)
  load_tfa_weights_in_bert(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)

  # Save pytorch-model
  logger.info(F"""Saving PyTorch model to {pytorch_dump_path}..."""	)
  torch.save(model.state_dict()	,		_SCREAMING_SNAKE_CASE	)


if __name__ == "__main__":
      lowercase_					     =	argparse.ArgumentParser()
      parser.add_argument(
          """--tf_checkpoint_path""", type=str, required=True, help="""Path to the TensorFlow 2.x checkpoint path."""
      )
      parser.add_argument(
          """--bert_config_file""",
          type=str,
          required=True,
          help="""The config json file corresponding to the BERT model. This specifies the model architecture.""",
      )
      parser.add_argument(
          """--pytorch_dump_path""",
          type=str,
          required=True,
          help="""Path to the output PyTorch model (must include filename).""",
      )
      lowercase_					     =	parser.parse_args()
      convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)









import unittest

import numpy as np

from transformers.testing_utils import is_flaky, require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import DonutImageProcessor
class        SCREAMING_SNAKE_CASE     (unittest.TestCase   ):


   def __init__(					self    :       Any    ,  a    :       str    ,  a    :       List[Any]=7    ,  a    :       int=3    ,  a    :       int=18    ,  a    :       Optional[Any]=30    ,  a    :       Optional[int]=400    ,  a    :       int=True    ,  a    :       Tuple=None    ,  a    :       Optional[Any]=True    ,  a    :       str=False    ,  a    :       str=True    ,  a    :       int=True    ,  a    :       Tuple=[0.5, 0.5, 0.5]    ,  a    :       Any=[0.5, 0.5, 0.5]    ,  )->     Optional[int]:





     """simple docstring"""







     lowercase__			    =						parent
     lowercase__			    =						batch_size
     lowercase__			    =						num_channels
     lowercase__			    =						image_size
     lowercase__			    =						min_resolution
     lowercase__			    =						max_resolution
     lowercase__			    =						do_resize
     lowercase__			    =						size if size is not None else {'height': 18, 'width': 20}
     lowercase__			    =						do_thumbnail
     lowercase__			    =						do_align_axis
     lowercase__			    =						do_pad
     lowercase__			    =						do_normalize
     lowercase__			    =						image_mean
     lowercase__			    =						image_std


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[Any] )->     Tuple:





     """simple docstring"""







     return {
         "do_resize": self.do_resize,
         "size": self.size,
         "do_thumbnail": self.do_thumbnail,
         "do_align_long_axis": self.do_align_axis,
         "do_pad": self.do_pad,
         "do_normalize": self.do_normalize,
         "image_mean": self.image_mean,
         "image_std": self.image_std,
     }






@require_torch
@require_vision
class        SCREAMING_SNAKE_CASE     (UpperCAmelCase    ,			unittest.TestCase   ):
   _UpperCamelCase					:	Optional[Any]				      = DonutImageProcessor if is_vision_available() else None


   def  SCREAMING_SNAKE_CASE_							(					self    :       int )->     List[Any]:





     """simple docstring"""







     lowercase__			    =						DonutImageProcessingTester(self )


   @property
   def  SCREAMING_SNAKE_CASE_							(					self    :       Union[str, Any] )->     Optional[int]:





     """simple docstring"""







     return self.image_processor_tester.prepare_image_processor_dict()


   def  SCREAMING_SNAKE_CASE_							(					self    :       Any )->     int:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class(**self.image_processor_dict )
     self.assertTrue(hasattr(a    ,  'do_resize' ) )
     self.assertTrue(hasattr(a    ,  'size' ) )
     self.assertTrue(hasattr(a    ,  'do_thumbnail' ) )
     self.assertTrue(hasattr(a    ,  'do_align_long_axis' ) )
     self.assertTrue(hasattr(a    ,  'do_pad' ) )
     self.assertTrue(hasattr(a    ,  'do_normalize' ) )
     self.assertTrue(hasattr(a    ,  'image_mean' ) )
     self.assertTrue(hasattr(a    ,  'image_std' ) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[Any] )->     Dict:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class.from_dict(self.image_processor_dict )
     self.assertEqual(image_processor.size    ,  {'height': 18, 'width': 20} )

     lowercase__			    =						self.image_processing_class.from_dict(self.image_processor_dict    ,  size=42 )
     self.assertEqual(image_processor.size    ,  {'height': 42, 'width': 42} )

     # Previous config had dimensions in (width, height) order
     lowercase__			    =						self.image_processing_class.from_dict(self.image_processor_dict    ,  size=(42, 84) )
     self.assertEqual(image_processor.size    ,  {'height': 84, 'width': 42} )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[Any] )->     Dict:





     """simple docstring"""







     pass


   @is_flaky()
   def  SCREAMING_SNAKE_CASE_							(					self    :       str )->     Optional[int]:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class(**self.image_processor_dict )
     # create random PIL images
     lowercase__			    =						prepare_image_inputs(self.image_processor_tester    ,  equal_resolution=a )
     for image in image_inputs:
       self.assertIsInstance(a    ,  Image.Image )

     # Test not batched input
     lowercase__			    =						image_processing(image_inputs[0]    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             1,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )

     # Test batched
     lowercase__			    =						image_processing(a    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )


   @is_flaky()
   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[Any] )->     Tuple:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class(**self.image_processor_dict )
     # create random numpy tensors
     lowercase__			    =						prepare_image_inputs(self.image_processor_tester    ,  equal_resolution=a    ,  numpify=a )
     for image in image_inputs:
       self.assertIsInstance(a    ,  np.ndarray )

     # Test not batched input
     lowercase__			    =						image_processing(image_inputs[0]    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             1,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )

     # Test batched
     lowercase__			    =						image_processing(a    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )


   @is_flaky()
   def  SCREAMING_SNAKE_CASE_							(					self    :       List[str] )->     Dict:





     """simple docstring"""







     lowercase__			    =						self.image_processing_class(**self.image_processor_dict )
     # create random PyTorch tensors
     lowercase__			    =						prepare_image_inputs(self.image_processor_tester    ,  equal_resolution=a    ,  torchify=a )
     for image in image_inputs:
       self.assertIsInstance(a    ,  torch.Tensor )

     # Test not batched input
     lowercase__			    =						image_processing(image_inputs[0]    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             1,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )

     # Test batched
     lowercase__			    =						image_processing(a    ,  return_tensors='pt' ).pixel_values
     self.assertEqual(
         encoded_images.shape    ,  (
             self.image_processor_tester.batch_size,
             self.image_processor_tester.num_channels,
             self.image_processor_tester.size['height'],
             self.image_processor_tester.size['width'],
         )    ,  )


