




"""simple docstring"""



import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


SCREAMING_SNAKE_CASE    =	logging.get_logger(__name__)




class      UpperCAmelCase_  (						A_							):
    lowercase__						 =				['''input_features''', '''is_longer''']




    def __init__(   self						:	str     ,			snake_case_						:	Any=64     ,			snake_case_						:	Tuple=48_000     ,			snake_case_						:	Optional[Any]=480     ,			snake_case_						:	Any=10     ,			snake_case_						:	Tuple=1_024     ,			snake_case_						:	Union[str, Any]=0.0     ,			snake_case_						:	List[Any]=False     ,			snake_case_						:	float = 0     ,			snake_case_						:	float = 14_000     ,			snake_case_						:	int = None     ,			snake_case_						:	str = "fusion"     ,			snake_case_						:	str = "repeatpad"     ,			**snake_case_						:	List[str]     ,			)				->							Optional[int]:




     '''simple docstring'''

     super().__init__(
         feature_size=snake_case_     ,			sampling_rate=snake_case_     ,			padding_value=snake_case_     ,			return_attention_mask=snake_case_     ,			**snake_case_     ,			)
     A__	=    top_db
     A__	=    truncation
     A__	=    padding
     A__	=    fft_window_size
     A__	=    (fft_window_size >> 1) + 1
     A__	=    hop_length
     A__	=    max_length_s
     A__	=    max_length_s * sampling_rate
     A__	=    sampling_rate
     A__	=    frequency_min
     A__	=    frequency_max
     A__	=    mel_filter_bank(
         num_frequency_bins=self.nb_frequency_bins     ,			num_mel_filters=snake_case_     ,			min_frequency=snake_case_     ,			max_frequency=snake_case_     ,			sampling_rate=snake_case_     ,			norm=snake_case_     ,			mel_scale="htk"     ,			)
     A__	=    mel_filter_bank(
         num_frequency_bins=self.nb_frequency_bins     ,			num_mel_filters=snake_case_     ,			min_frequency=snake_case_     ,			max_frequency=snake_case_     ,			sampling_rate=snake_case_     ,			norm="slaney"     ,			mel_scale="slaney"     ,			)




    def 							__magic_name__							(   self						:	Tuple				)				->							Dict[str, Any]:




     '''simple docstring'''

     A__	=    copy.deepcopy(self.__dict__				)
     A__	=    self.__class__.__name__
     if "mel_filters" in output:
      del output["mel_filters"]
     if "mel_filters_slaney" in output:
      del output["mel_filters_slaney"]
     return output




    def 							__magic_name__							(   self						:	List[Any]     ,			snake_case_						:	np.array     ,			snake_case_						:	Optional[np.array] = None				)				->							np.ndarray:




     '''simple docstring'''

     A__	=    spectrogram(
         snake_case_     ,			window_function(self.fft_window_size     ,			"hann"				)     ,			frame_length=self.fft_window_size     ,			hop_length=self.hop_length     ,			power=2.0     ,			mel_filters=snake_case_     ,			log_mel="dB"     ,			)
     return log_mel_spectrogram.T




    def 							__magic_name__							(   self						:	Tuple     ,			snake_case_						:	Optional[Any]     ,			snake_case_						:	int     ,			snake_case_						:	List[str]				)				->							Any:




     '''simple docstring'''

     A__	=    np.array_split(list(range(0     ,			total_frames - chunk_frames + 1				)				)     ,			3				)
     if len(ranges[1]				) == 0:
      # if the audio is too short, we just use the first chunk
      A__	=    [0]
     if len(ranges[2]				) == 0:
      # if the audio is too short, we just use the first chunk
      A__	=    [0]
     # randomly choose index for each part
     A__	=    np.random.choice(ranges[0]				)
     A__	=    np.random.choice(ranges[1]				)
     A__	=    np.random.choice(ranges[2]				)

     A__	=    mel[idx_front : idx_front + chunk_frames, :]
     A__	=    mel[idx_middle : idx_middle + chunk_frames, :]
     A__	=    mel[idx_back : idx_back + chunk_frames, :]

     A__	=    torch.tensor(mel[None, None, :]				)
     A__	=    torch.nn.functional.interpolate(
         snake_case_     ,			size=[chunk_frames, 64]     ,			mode="bilinear"     ,			align_corners=snake_case_				)
     A__	=    mel_shrink[0][0].numpy()
     A__	=    np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back]     ,			axis=0				)
     return mel_fusion




    def 							__magic_name__							(   self						:	Optional[Any]     ,			snake_case_						:	np.array     ,			snake_case_						:	List[str]     ,			snake_case_						:	str     ,			snake_case_						:	Dict				)				->							np.array:




     '''simple docstring'''

     if waveform.shape[0] > max_length:
      if truncation == "rand_trunc":
       A__	=    True
       # random crop to max_length (for compatibility) -> this should be handled by self.pad
       A__	=    len(snake_case_				) - max_length
       A__	=    np.random.randint(0     ,			overflow + 1				)
       A__	=    waveform[idx : idx + max_length]
       A__	=    self._np_extract_fbank_features(snake_case_     ,			self.mel_filters_slaney				)[None, :]
      elif truncation == "fusion":
       A__	=    self._np_extract_fbank_features(snake_case_     ,			self.mel_filters				)
       A__	=    max_length // self.hop_length + 1  # the +1 related to how the spectrogram is computed
       A__	=    mel.shape[0]
       if chunk_frames == total_frames:
        # there is a corner case where the audio length is larger than max_length but smaller than max_length+hop_length.
        # In this case, we just use the whole audio.
        A__	=    np.stack([mel, mel, mel, mel]     ,			axis=0				)
        A__	=    False
       else:
        A__	=    self._random_mel_fusion(snake_case_     ,			snake_case_     ,			snake_case_				)
        A__	=    True
      else:
       raise NotImplementedError(F"""data_truncating {truncation} not implemented"""				)

     else:
      A__	=    False
      # only use repeat as a new possible value for padding. you repeat the audio before applying the usual max_length padding
      if waveform.shape[0] < max_length:
       if padding == "repeat":
        A__	=    int(max_length / len(snake_case_				)				)
        A__	=    np.stack(np.tile(snake_case_     ,			n_repeat + 1				)				)[:max_length]
       if padding == "repeatpad":
        A__	=    int(max_length / len(snake_case_				)				)
        A__	=    np.stack(np.tile(snake_case_     ,			snake_case_				)				)
       A__	=    np.pad(snake_case_     ,			(0, max_length - waveform.shape[0])     ,			mode="constant"     ,			constant_values=0				)

      if truncation == "fusion":
       A__	=    self._np_extract_fbank_features(snake_case_     ,			self.mel_filters				)
       A__	=    np.stack([input_mel, input_mel, input_mel, input_mel]     ,			axis=0				)
      else:
       A__	=    self._np_extract_fbank_features(snake_case_     ,			self.mel_filters_slaney				)[None, :]

     return input_mel, longer




    def __call__(   self						:	str     ,			snake_case_						:	Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]     ,			snake_case_						:	str = None     ,			snake_case_						:	Optional[str] = None     ,			snake_case_						:	Optional[int] = None     ,			snake_case_						:	Optional[int] = None     ,			snake_case_						:	Optional[Union[str, TensorType]] = None     ,			**snake_case_						:	Union[str, Any]     ,			)				->							BatchFeature:




     '''simple docstring'''

     A__	=    truncation if truncation is not None else self.truncation
     A__	=    padding if padding else self.padding

     if sampling_rate is not None:
      if sampling_rate != self.sampling_rate:
       raise ValueError(
           F"""The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"""
           F""" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"""
           F""" was sampled with {self.sampling_rate} and not {sampling_rate}."""				)
     else:
      logger.warning(
          "It is strongly recommended to pass the `sampling_rate` argument to this function. "
          "Failing to do so can result in silent errors that might be hard to debug."				)

     A__	=    isinstance(snake_case_     ,			np.ndarray				) and len(raw_speech.shape				) > 1
     if is_batched_numpy and len(raw_speech.shape				) > 2:
      raise ValueError(F"""Only mono-channel audio is supported for input to {self}"""				)
     A__	=    is_batched_numpy or (
         isinstance(snake_case_     ,			(list, tuple)				) and (isinstance(raw_speech[0]     ,			(np.ndarray, tuple, list)				))
     )

     if is_batched:
      A__	=    [np.asarray(snake_case_     ,			dtype=np.floataa				) for speech in raw_speech]
     elif not is_batched and not isinstance(snake_case_     ,			np.ndarray				):
      A__	=    np.asarray(snake_case_     ,			dtype=np.floataa				)
     elif isinstance(snake_case_     ,			np.ndarray				) and raw_speech.dtype is np.dtype(np.floataa				):
      A__	=    raw_speech.astype(np.floataa				)

     # always return batch
     if not is_batched:
      A__	=    [np.asarray(snake_case_				)]

     # convert to mel spectrogram, truncate and pad if needed.
     A__	=    [
         self._get_input_mel(snake_case_     ,			max_length if max_length else self.nb_max_samples     ,			snake_case_     ,			snake_case_				)
         for waveform in raw_speech
     ]

     A__	=    []
     A__	=    []
     for mel, longer in padded_inputs:
      input_mel.append(snake_case_				)
      is_longer.append(snake_case_				)

     if truncation == "fusion" and sum(snake_case_				) == 0:
      # if no audio is longer than 10s, then randomly select one audio to be longer
      A__	=    np.random.randint(0     ,			len(snake_case_				)				)
      A__	=    True

     if isinstance(input_mel[0]     ,			snake_case_				):
      A__	=    [np.asarray(snake_case_     ,			dtype=np.floataa				) for feature in input_mel]

     # is_longer is a list of bool
     A__	=    [[longer] for longer in is_longer]

     A__	=    {"input_features": input_mel, "is_longer": is_longer}
     A__	=    BatchFeature(snake_case_				)

     if return_tensors is not None:
      A__	=    input_features.convert_to_tensors(snake_case_				)

     return input_features











"""simple docstring"""



import argparse

import torch
from datasets import load_dataset
from donut import DonutModel

from transformers import (
    DonutImageProcessor,
    DonutProcessor,
    DonutSwinConfig,
    DonutSwinModel,
    MBartConfig,
    MBartForCausalLM,
    VisionEncoderDecoderModel,
    XLMRobertaTokenizerFast,
)







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		List[str]:
 A__	=    model.config

 A__	=    DonutSwinConfig(
     image_size=original_config.input_size      ,			patch_size=4      ,			depths=original_config.encoder_layer      ,			num_heads=[4, 8, 16, 32]      ,			window_size=original_config.window_size      ,			embed_dim=1_28      ,			)
 A__	=    MBartConfig(
     is_decoder=lowercase_      ,			is_encoder_decoder=lowercase_      ,			add_cross_attention=lowercase_      ,			decoder_layers=original_config.decoder_layer      ,			max_position_embeddings=original_config.max_position_embeddings      ,			vocab_size=len(
         model.decoder.tokenizer						)      ,			scale_embedding=lowercase_      ,			add_final_layer_norm=lowercase_      ,			)

 return encoder_config, decoder_config







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		int:
 if "encoder.model" in name:
  A__	=    name.replace("encoder.model"      ,			"encoder"						)
 if "decoder.model" in name:
  A__	=    name.replace("decoder.model"      ,			"decoder"						)
 if "patch_embed.proj" in name:
  A__	=    name.replace("patch_embed.proj"      ,			"embeddings.patch_embeddings.projection"						)
 if "patch_embed.norm" in name:
  A__	=    name.replace("patch_embed.norm"      ,			"embeddings.norm"						)
 if name.startswith("encoder"						):
  if "layers" in name:
   A__	=    "encoder." + name
  if "attn.proj" in name:
   A__	=    name.replace("attn.proj"      ,			"attention.output.dense"						)
  if "attn" in name and "mask" not in name:
   A__	=    name.replace("attn"      ,			"attention.self"						)
  if "norm1" in name:
   A__	=    name.replace("norm1"      ,			"layernorm_before"						)
  if "norm2" in name:
   A__	=    name.replace("norm2"      ,			"layernorm_after"						)
  if "mlp.fc1" in name:
   A__	=    name.replace("mlp.fc1"      ,			"intermediate.dense"						)
  if "mlp.fc2" in name:
   A__	=    name.replace("mlp.fc2"      ,			"output.dense"						)

  if name == "encoder.norm.weight":
   A__	=    "encoder.layernorm.weight"
  if name == "encoder.norm.bias":
   A__	=    "encoder.layernorm.bias"

 return name







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		Any:
 for key in orig_state_dict.copy().keys():
  A__	=    orig_state_dict.pop(lowercase_						)

  if "qkv" in key:
   A__	=    key.split("."						)
   A__	=    int(key_split[3]						)
   A__	=    int(key_split[5]						)
   A__	=    model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

   if "weight" in key:
    A__	=    val[:dim, :]
    A__	=    val[dim : dim * 2, :]
    A__	=    val[-dim:, :]
   else:
    A__	=    val[:dim]
    A__	=    val[dim : dim * 2]
    A__	=    val[-dim:]
  elif "attn_mask" in key or key in ["encoder.model.norm.weight", "encoder.model.norm.bias"]:
   # HuggingFace implementation doesn't use attn_mask buffer
   # and model doesn't use final LayerNorms for the encoder
   pass
  else:
   A__	=    val

 return orig_state_dict







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_=None      ,			lowercase_=False						)						->		Dict:
 # load original model
 A__	=    DonutModel.from_pretrained(lowercase_						).eval()

 # load HuggingFace model
 A__,				A__	=    get_configs(lowercase_						)
 A__	=    DonutSwinModel(lowercase_						)
 A__	=    MBartForCausalLM(lowercase_						)
 A__	=    VisionEncoderDecoderModel(encoder=lowercase_      ,			decoder=lowercase_						)
 model.eval()

 A__	=    original_model.state_dict()
 A__	=    convert_state_dict(lowercase_      ,			lowercase_						)
 model.load_state_dict(lowercase_						)

 # verify results on scanned document
 A__	=    load_dataset("hf-internal-testing/example-documents"						)
 A__	=    dataset["test"][0]["image"].convert("RGB"						)

 A__	=    XLMRobertaTokenizerFast.from_pretrained(lowercase_      ,			from_slow=lowercase_						)
 A__	=    DonutImageProcessor(
     do_align_long_axis=original_model.config.align_long_axis      ,			size=original_model.config.input_size[::-1]						)
 A__	=    DonutProcessor(lowercase_      ,			lowercase_						)
 A__	=    processor(lowercase_      ,			return_tensors="pt"						).pixel_values

 if model_name == "naver-clova-ix/donut-base-finetuned-docvqa":
  A__	=    "<s_docvqa><s_question>{user_input}</s_question><s_answer>"
  A__	=    "When is the coffee break?"
  A__	=    task_prompt.replace("{user_input}"      ,			lowercase_						)
 elif model_name == "naver-clova-ix/donut-base-finetuned-rvlcdip":
  A__	=    "<s_rvlcdip>"
 elif model_name in [
     "naver-clova-ix/donut-base-finetuned-cord-v1",
     "naver-clova-ix/donut-base-finetuned-cord-v1-2560",
 ]:
  A__	=    "<s_cord>"
 elif model_name == "naver-clova-ix/donut-base-finetuned-cord-v2":
  A__	=    "s_cord-v2>"
 elif model_name == "naver-clova-ix/donut-base-finetuned-zhtrainticket":
  A__	=    "<s_zhtrainticket>"
 elif model_name in ["naver-clova-ix/donut-proto", "naver-clova-ix/donut-base"]:
  # use a random prompt
  A__	=    "hello world"
 else:
  raise ValueError("Model name not supported"						)
 A__	=    original_model.decoder.tokenizer(lowercase_      ,			add_special_tokens=lowercase_      ,			return_tensors="pt"						)[
     "input_ids"
 ]

 A__	=    original_model.encoder.model.patch_embed(lowercase_						)
 A__,				A__	=    model.encoder.embeddings(lowercase_						)
 assert torch.allclose(lowercase_      ,			lowercase_      ,			atol=1E-3						)

 # verify encoder hidden states
 A__	=    original_model.encoder(lowercase_						)
 A__	=    model.encoder(lowercase_						).last_hidden_state
 assert torch.allclose(lowercase_      ,			lowercase_      ,			atol=1E-2						)

 # verify decoder hidden states
 A__	=    original_model(lowercase_      ,			lowercase_      ,			lowercase_						).logits
 A__	=    model(lowercase_      ,			decoder_input_ids=lowercase_						).logits
 assert torch.allclose(lowercase_      ,			lowercase_      ,			atol=1E-3						)
 print("Looks ok!"						)

 if pytorch_dump_folder_path is not None:
  print(f"""Saving model and processor to {pytorch_dump_folder_path}"""						)
  model.save_pretrained(lowercase_						)
  processor.save_pretrained(lowercase_						)

 if push_to_hub:
  model.push_to_hub("nielsr/" + model_name.split("/"						)[-1]      ,			commit_message="Update model"						)
  processor.push_to_hub("nielsr/" + model_name.split("/"						)[-1]      ,			commit_message="Update model"						)


if __name__ == "__main__":
      SCREAMING_SNAKE_CASE    =	argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          "--model_name",
          default="naver-clova-ix/donut-base-finetuned-docvqa",
          required=False,
          type=str,
          help="Name of the original model you'd like to convert.",
      )
      parser.add_argument(
          "--pytorch_dump_folder_path",
          default=None,
          required=False,
          type=str,
          help="Path to the output PyTorch model directory.",
      )
      parser.add_argument(
          "--push_to_hub",
          action="store_true",
          help="Whether or not to push the converted model and processor to the 🤗 hub.",
      )

      SCREAMING_SNAKE_CASE    =	parser.parse_args()
      convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





