






import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ASTConfig
from transformers.testing_utils import require_torch, require_torchaudio, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_torchaudio_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch
		from torch import nn

		from transformers import ASTForAudioClassification, ASTModel
		from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import (
		    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		)


if is_torchaudio_available():
		import torchaudio

		from transformers import ASTFeatureExtractor



class 				__UpperCAmelCase		:







			def __init__(     self:  int				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Union[str, Any]=13				,     UpperCAmelCase_:  str=2				,     UpperCAmelCase_:  Optional[Any]=24				,     UpperCAmelCase_:  Optional[Any]=16				,     UpperCAmelCase_:  int=True				,     UpperCAmelCase_:  Optional[int]=True				,     UpperCAmelCase_:  Union[str, Any]=32				,     UpperCAmelCase_:  Any=5				,     UpperCAmelCase_:  str=4				,     UpperCAmelCase_:  int=37				,     UpperCAmelCase_:  Optional[Any]="gelu"				,     UpperCAmelCase_:  Optional[int]=0.1				,     UpperCAmelCase_:  Any=0.1				,     UpperCAmelCase_:  Optional[int]=10				,     UpperCAmelCase_:  Union[str, Any]=0.02				,     UpperCAmelCase_:  List[Any]=None				,     UpperCAmelCase_:  Optional[int]=2				,     UpperCAmelCase_:  Union[str, Any]=2				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  parent
									_SCREAMING_SNAKE_CASE     				=  batch_size
									_SCREAMING_SNAKE_CASE     				=  patch_size
									_SCREAMING_SNAKE_CASE     				=  max_length
									_SCREAMING_SNAKE_CASE     				=  num_mel_bins
									_SCREAMING_SNAKE_CASE     				=  is_training
									_SCREAMING_SNAKE_CASE     				=  use_labels
									_SCREAMING_SNAKE_CASE     				=  hidden_size
									_SCREAMING_SNAKE_CASE     				=  num_hidden_layers
									_SCREAMING_SNAKE_CASE     				=  num_attention_heads
									_SCREAMING_SNAKE_CASE     				=  intermediate_size
									_SCREAMING_SNAKE_CASE     				=  hidden_act
									_SCREAMING_SNAKE_CASE     				=  hidden_dropout_prob
									_SCREAMING_SNAKE_CASE     				=  attention_probs_dropout_prob
									_SCREAMING_SNAKE_CASE     				=  type_sequence_label_size
									_SCREAMING_SNAKE_CASE     				=  initializer_range
									_SCREAMING_SNAKE_CASE     				=  scope
									_SCREAMING_SNAKE_CASE     				=  frequency_stride
									_SCREAMING_SNAKE_CASE     				=  time_stride

									# in AST, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
									_SCREAMING_SNAKE_CASE     				=  (self.num_mel_bins - self.patch_size) // self.frequency_stride + 1
									_SCREAMING_SNAKE_CASE     				=  (self.max_length - self.patch_size) // self.time_stride + 1
									_SCREAMING_SNAKE_CASE     				=  frequency_out_dimension * time_out_dimension
									_SCREAMING_SNAKE_CASE     				=  num_patches + 2







			def 							UpperCamelCase				(     self:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]					)

									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_labels:
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size]				,     self.type_sequence_label_size					)

									_SCREAMING_SNAKE_CASE     				=  self.get_config()

									return config, input_values, labels







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									return ASTConfig(
									    patch_size=self.patch_size				,     max_length=self.max_length				,     num_mel_bins=self.num_mel_bins				,     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				,     is_decoder=UpperCAmelCase_				,     initializer_range=self.initializer_range				,     frequency_stride=self.frequency_stride				,     time_stride=self.time_stride				,     )







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  ASTModel(config=UpperCAmelCase_					)
									model.to(UpperCAmelCase_					)
									model.eval()
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.last_hidden_state.shape				,     (self.batch_size, self.seq_length, self.hidden_size)					)







			def 							UpperCamelCase				(     self:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.prepare_config_and_inputs()
									(
									    (
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					
									)     				=  config_and_inputs
									_SCREAMING_SNAKE_CASE     				=  {"""input_values""": input_values}
									return config, inputs_dict



@require_torch
class 				__UpperCAmelCase		(_UpperCAmelCase	,_UpperCAmelCase	,unittest.TestCase   ):
			__snake_case    :				Tuple									=							(
			    (
			        ASTModel,
			        ASTForAudioClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			__snake_case    :				List[Any]									=							(
			    {"audio-classification": ASTForAudioClassification, "feature-extraction": ASTModel}
			    if is_torch_available()
			    else {}
			)
			__snake_case    :				List[Any]									=							False
			__snake_case    :				Any									=							False
			__snake_case    :				List[Any]									=							False
			__snake_case    :				List[Any]									=							False







			def 							UpperCamelCase				(     self:  Any				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Any				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Optional[Any]					):







									'''simple docstring'''




									if pipeline_test_casse_name == "AudioClassificationPipelineTests":
															return True

									return False







			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  ASTModelTester(self					)
									_SCREAMING_SNAKE_CASE     				=  ConfigTester(self				,     config_class=UpperCAmelCase_				,     has_text_modality=UpperCAmelCase_				,     hidden_size=37					)







			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									self.config_tester.run_common_tests()







			@unittest.skip(reason="""AST does not use inputs_embeds"""					)
			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									pass







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_SCREAMING_SNAKE_CASE     				=  model_class(UpperCAmelCase_					)
															self.assertIsInstance(model.get_input_embeddings()				,     (nn.Module)					)
															_SCREAMING_SNAKE_CASE     				=  model.get_output_embeddings()
															self.assertTrue(x is None or isinstance(UpperCAmelCase_				,     nn.Linear					)					)







			def 							UpperCamelCase				(     self:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_SCREAMING_SNAKE_CASE     				=  model_class(UpperCAmelCase_					)
															_SCREAMING_SNAKE_CASE     				=  inspect.signature(model.forward					)
															# signature.parameters is an OrderedDict => so arg_names order is deterministic
															_SCREAMING_SNAKE_CASE     				=  [*signature.parameters.keys()]

															_SCREAMING_SNAKE_CASE     				=  ["""input_values"""]
															self.assertListEqual(arg_names[:1]				,     UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*UpperCAmelCase_					)







			@slow
			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  ASTModel.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)

def 	__lowerCamelCase      (		)							->  List[str]:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  hf_hub_download(
						    repo_id="""nielsr/audio-spectogram-transformer-checkpoint"""	,filename="""sample_audio.flac"""	,repo_type="""dataset"""		)

						_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  torchaudio.load(snake_case__		)

						return audio, sampling_rate




@require_torch
@require_torchaudio
class 				__UpperCAmelCase		(unittest.TestCase   ):







			@cached_property
			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									return (
									    ASTFeatureExtractor.from_pretrained("""MIT/ast-finetuned-audioset-10-10-0.4593"""					)
									    if is_torchaudio_available()
									    else None
									)







			@slow
			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.default_feature_extractor
									_SCREAMING_SNAKE_CASE     				=  ASTForAudioClassification.from_pretrained("""MIT/ast-finetuned-audioset-10-10-0.4593"""					).to(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  self.default_feature_extractor
									_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  prepare_audio()
									_SCREAMING_SNAKE_CASE     				=  audio.squeeze().numpy()
									_SCREAMING_SNAKE_CASE     				=  feature_extractor(UpperCAmelCase_				,     sampling_rate=UpperCAmelCase_				,     return_tensors="""pt"""					).to(UpperCAmelCase_					)

									# forward pass
									with torch.no_grad():
															_SCREAMING_SNAKE_CASE     				=  model(**UpperCAmelCase_					)

									# verify the logits
									_SCREAMING_SNAKE_CASE     				=  torch.Size((1, 527)					)
									self.assertEqual(outputs.logits.shape				,     UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  torch.tensor([-0.87_60, -7.00_42, -8.66_02]					).to(UpperCAmelCase_					)

									self.assertTrue(torch.allclose(outputs.logits[0, :3]				,     UpperCAmelCase_				,     atol=1E-4					)					)









import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


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

def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Any:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  []
						for i in range(encoder_config.num_hidden_layers		):
												# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
												rename_keys.append(
												    (F'encoder.deit.blocks.{i}.norm1.weight', F'encoder.encoder.layer.{i}.layernorm_before.weight')		)
												rename_keys.append((F'encoder.deit.blocks.{i}.norm1.bias', F'encoder.encoder.layer.{i}.layernorm_before.bias')		)
												rename_keys.append(
												    (F'encoder.deit.blocks.{i}.attn.proj.weight', F'encoder.encoder.layer.{i}.attention.output.dense.weight')		)
												rename_keys.append(
												    (F'encoder.deit.blocks.{i}.attn.proj.bias', F'encoder.encoder.layer.{i}.attention.output.dense.bias')		)
												rename_keys.append(
												    (F'encoder.deit.blocks.{i}.norm2.weight', F'encoder.encoder.layer.{i}.layernorm_after.weight')		)
												rename_keys.append((F'encoder.deit.blocks.{i}.norm2.bias', F'encoder.encoder.layer.{i}.layernorm_after.bias')		)
												rename_keys.append(
												    (F'encoder.deit.blocks.{i}.mlp.fc1.weight', F'encoder.encoder.layer.{i}.intermediate.dense.weight')		)
												rename_keys.append(
												    (F'encoder.deit.blocks.{i}.mlp.fc1.bias', F'encoder.encoder.layer.{i}.intermediate.dense.bias')		)
												rename_keys.append(
												    (F'encoder.deit.blocks.{i}.mlp.fc2.weight', F'encoder.encoder.layer.{i}.output.dense.weight')		)
												rename_keys.append((F'encoder.deit.blocks.{i}.mlp.fc2.bias', F'encoder.encoder.layer.{i}.output.dense.bias')		)

						# cls token, position embeddings and patch embeddings of encoder
						rename_keys.extend(
						    [
						        ("""encoder.deit.cls_token""", """encoder.embeddings.cls_token"""),
						        ("""encoder.deit.pos_embed""", """encoder.embeddings.position_embeddings"""),
						        ("""encoder.deit.patch_embed.proj.weight""", """encoder.embeddings.patch_embeddings.projection.weight"""),
						        ("""encoder.deit.patch_embed.proj.bias""", """encoder.embeddings.patch_embeddings.projection.bias"""),
						        ("""encoder.deit.norm.weight""", """encoder.layernorm.weight"""),
						        ("""encoder.deit.norm.bias""", """encoder.layernorm.bias"""),
						    ]		)

						return rename_keys

def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Dict:






						"""simple docstring"""





						for i in range(encoder_config.num_hidden_layers		):
												# queries, keys and values (only weights, no biases)
												_SCREAMING_SNAKE_CASE     				=  state_dict.pop(F'encoder.deit.blocks.{i}.attn.qkv.weight'		)

												_SCREAMING_SNAKE_CASE     				=  in_proj_weight[
												    : encoder_config.hidden_size, :
												]
												_SCREAMING_SNAKE_CASE     				=  in_proj_weight[
												    encoder_config.hidden_size : encoder_config.hidden_size * 2, :
												]
												_SCREAMING_SNAKE_CASE     				=  in_proj_weight[
												    -encoder_config.hidden_size :, :
												]

def 	__lowerCamelCase      (		snake_case__	,snake_case__	,snake_case__		)							->  Union[str, Any]:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  dct.pop(snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  val

def 	__lowerCamelCase      (		snake_case__		)							->  Union[str, Any]:






						"""simple docstring"""





						if "handwritten" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  """https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg"""  # industry
												# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-12.jpg" # have
												# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-10.jpg" # let
												# url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"  #
												# url = "https://fki.tic.heia-fr.ch/static/img/a01-122.jpg"
						elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  """https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg"""
						_SCREAMING_SNAKE_CASE     				=  Image.open(requests.get(snake_case__	,stream=snake_case__		).raw		).convert("""RGB"""		)
						return im

@torch.no_grad()
def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Tuple:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  ViTConfig(image_size=3_84	,qkv_bias=snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  TrOCRConfig()

						# size of the architecture
						if "base" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  7_68
						elif "large" in checkpoint_url:
												# use ViT-large encoder
												_SCREAMING_SNAKE_CASE     				=  10_24
												_SCREAMING_SNAKE_CASE     				=  40_96
												_SCREAMING_SNAKE_CASE     				=  24
												_SCREAMING_SNAKE_CASE     				=  16
												_SCREAMING_SNAKE_CASE     				=  10_24
						else:
												raise ValueError("""Should either find 'base' or 'large' in checkpoint URL"""		)

						# the large-printed + stage1 checkpoints uses sinusoidal position embeddings, no layernorm afterwards
						if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  False
												_SCREAMING_SNAKE_CASE     				=  """relu"""
												_SCREAMING_SNAKE_CASE     				=  10_24
												_SCREAMING_SNAKE_CASE     				=  True
												_SCREAMING_SNAKE_CASE     				=  False
												_SCREAMING_SNAKE_CASE     				=  False

						# load HuggingFace model
						_SCREAMING_SNAKE_CASE     				=  ViTModel(snake_case__	,add_pooling_layer=snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  TrOCRForCausalLM(snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  VisionEncoderDecoderModel(encoder=snake_case__	,decoder=snake_case__		)
						model.eval()

						# load state_dict of original model, rename some keys
						_SCREAMING_SNAKE_CASE     				=  torch.hub.load_state_dict_from_url(snake_case__	,map_location="""cpu"""	,check_hash=snake_case__		)["""model"""]

						_SCREAMING_SNAKE_CASE     				=  create_rename_keys(snake_case__	,snake_case__		)
						for src, dest in rename_keys:
												rename_key(snake_case__	,snake_case__	,snake_case__		)
						read_in_q_k_v(snake_case__	,snake_case__		)

						# remove parameters we don't need
						del state_dict["encoder.deit.head.weight"]
						del state_dict["encoder.deit.head.bias"]
						del state_dict["decoder.version"]

						# add prefix to decoder keys
						for key, val in state_dict.copy().items():
												_SCREAMING_SNAKE_CASE     				=  state_dict.pop(snake_case__		)
												if key.startswith("""decoder"""		) and "output_projection" not in key:
																		_SCREAMING_SNAKE_CASE     				=  val
												else:
																		_SCREAMING_SNAKE_CASE     				=  val

    # load state dict
						model.load_state_dict(snake_case__		)

						# Check outputs on an image
						_SCREAMING_SNAKE_CASE     				=  ViTImageProcessor(size=encoder_config.image_size		)
						_SCREAMING_SNAKE_CASE     				=  RobertaTokenizer.from_pretrained("""roberta-large"""		)
						_SCREAMING_SNAKE_CASE     				=  TrOCRProcessor(snake_case__	,snake_case__		)

						_SCREAMING_SNAKE_CASE     				=  processor(images=prepare_img(snake_case__		)	,return_tensors="""pt"""		).pixel_values

						# verify logits
						_SCREAMING_SNAKE_CASE     				=  torch.tensor([[model.config.decoder.decoder_start_token_id]]		)
						_SCREAMING_SNAKE_CASE     				=  model(pixel_values=snake_case__	,decoder_input_ids=snake_case__		)
						_SCREAMING_SNAKE_CASE     				=  outputs.logits

						_SCREAMING_SNAKE_CASE     				=  torch.Size([1, 1, 5_02_65]		)
						if "trocr-base-handwritten" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  torch.tensor(
												    [-1.4_502, -4.6_683, -0.5_347, -2.9_291, 9.1_435, -3.0_571, 8.9_764, 1.7_560, 8.7_358, -1.5_311]		)
						elif "trocr-large-handwritten" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  torch.tensor(
												    [-2.6_437, -1.3_129, -2.2_596, -5.3_455, 6.3_539, 1.7_604, 5.4_991, 1.4_702, 5.6_113, 2.0_170]		)
						elif "trocr-base-printed" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  torch.tensor(
												    [-5.6_816, -5.8_388, 1.1_398, -6.9_034, 6.8_505, -2.4_393, 1.2_284, -1.0_232, -1.9_661, -3.9_210]		)
						elif "trocr-large-printed" in checkpoint_url:
												_SCREAMING_SNAKE_CASE     				=  torch.tensor(
												    [-6.0_162, -7.0_959, 4.4_155, -5.1_063, 7.0_468, -3.1_631, 2.6_466, -0.3_081, -0.8_106, -1.7_535]		)

						if "stage1" not in checkpoint_url:
												assert logits.shape == expected_shape, "Shape of logits not as expected"
												assert torch.allclose(logits[0, 0, :10]	,snake_case__	,atol=1e-3		), "First elements of logits not as expected"

						Path(snake_case__		).mkdir(exist_ok=snake_case__		)
						print(F'Saving model to {pytorch_dump_folder_path}'		)
						model.save_pretrained(snake_case__		)
						print(F'Saving processor to {pytorch_dump_folder_path}'		)
						processor.save_pretrained(snake_case__		)


if __name__ == "__main__":
		UpperCamelCase			    =     argparse.ArgumentParser()

		parser.add_argument(
		    '''--checkpoint_url''',
		    default='''https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt''',
		    type=str,
		    help='''URL to the original PyTorch checkpoint (.pth file).''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
		)
		UpperCamelCase			    =     parser.parse_args()
		convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)

