





'''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


_lowerCamelCase   				=			logging.get_logger(__name__)




class   UpperCamelCase_       (     UpperCamelCase__    ):
							lowerCamelCase_           =					["input_features", "is_longer"]


							def __init__(							self    :str							,				__A    :Tuple=64							,				__A    :List[str]=4_8000							,				__A    :Optional[Any]=480							,				__A    :Optional[int]=10							,				__A    :Dict=1024							,				__A    :Union[str, Any]=0.0							,				__A    :Optional[int]=False							,				__A    :float = 0							,				__A    :float = 1_4000							,				__A    :int = None							,				__A    :str = "fusion"							,				__A    :str = "repeatpad"							,				**__A    :Any							,				)   ->   Optional[int]:



									"""simple docstring"""



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


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      copy.deepcopy(self.__dict__		)
									SCREAMING_SNAKE_CASE__			=      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      _snake_case  (							self    :Any							,				__A    :np.array							,				__A    :Optional[np.array] = None		)   ->   np.ndarray:



									"""simple docstring"""



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


							def      _snake_case  (							self    :Optional[Any]							,				__A    :Any							,				__A    :Optional[int]							,				__A    :str		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      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
											SCREAMING_SNAKE_CASE__			=      [0]
									if len(ranges[2]		) == 0:
											# if the audio is too short, we just use the first chunk
											SCREAMING_SNAKE_CASE__			=      [0]
									# randomly choose index for each part
									SCREAMING_SNAKE_CASE__			=      np.random.choice(ranges[0]		)
									SCREAMING_SNAKE_CASE__			=      np.random.choice(ranges[1]		)
									SCREAMING_SNAKE_CASE__			=      np.random.choice(ranges[2]		)

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

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


							def      _snake_case  (							self    :Tuple							,				__A    :np.array							,				__A    :Tuple							,				__A    :Tuple							,				__A    :Tuple		)   ->   np.array:



									"""simple docstring"""



									if waveform.shape[0] > max_length:
											if truncation == "rand_trunc":
													SCREAMING_SNAKE_CASE__			=      True
													# random crop to max_length (for compatibility) -> this should be handled by self.pad
													SCREAMING_SNAKE_CASE__			=      len(__A		) - max_length
													SCREAMING_SNAKE_CASE__			=      np.random.randint(0							,				overflow + 1		)
													SCREAMING_SNAKE_CASE__			=      waveform[idx : idx + max_length]
													SCREAMING_SNAKE_CASE__			=      self._np_extract_fbank_features(__A							,				self.mel_filters_slaney		)[None, :]
											elif truncation == "fusion":
													SCREAMING_SNAKE_CASE__			=      self._np_extract_fbank_features(__A							,				self.mel_filters		)
													SCREAMING_SNAKE_CASE__			=      max_length // self.hop_length + 1  # the +1 related to how the spectrogram is computed
													SCREAMING_SNAKE_CASE__			=      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.
															SCREAMING_SNAKE_CASE__			=      np.stack([mel, mel, mel, mel]							,				axis=0		)
															SCREAMING_SNAKE_CASE__			=      False
													else:
															SCREAMING_SNAKE_CASE__			=      self._random_mel_fusion(__A							,				__A							,				__A		)
															SCREAMING_SNAKE_CASE__			=      True
											else:
													raise NotImplementedError(f'''data_truncating {truncation} not implemented'''		)

									else:
											SCREAMING_SNAKE_CASE__			=      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":
															SCREAMING_SNAKE_CASE__			=      int(max_length / len(__A		)		)
															SCREAMING_SNAKE_CASE__			=      np.stack(np.tile(__A							,				n_repeat + 1		)		)[:max_length]
													if padding == "repeatpad":
															SCREAMING_SNAKE_CASE__			=      int(max_length / len(__A		)		)
															SCREAMING_SNAKE_CASE__			=      np.stack(np.tile(__A							,				__A		)		)
													SCREAMING_SNAKE_CASE__			=      np.pad(__A							,				(0, max_length - waveform.shape[0])							,				mode="""constant"""							,				constant_values=0		)

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

									return input_mel, longer


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      truncation if truncation is not None else self.truncation
									SCREAMING_SNAKE_CASE__			=      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."""		)

									SCREAMING_SNAKE_CASE__			=      isinstance(__A							,				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}'''		)
									SCREAMING_SNAKE_CASE__			=      is_batched_numpy or (
									    isinstance(__A							,				(list, tuple)		) and (isinstance(raw_speech[0]							,				(np.ndarray, tuple, list)		))
									)

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

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

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

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

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

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

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

									SCREAMING_SNAKE_CASE__			=      {"""input_features""": input_mel, """is_longer""": is_longer}
									SCREAMING_SNAKE_CASE__			=      BatchFeature(__A		)

									if return_tensors is not None:
											SCREAMING_SNAKE_CASE__			=      input_features.convert_to_tensors(__A		)

									return input_features






import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCamelCase   				=			logging.get_logger(__name__)

_lowerCamelCase   				=			{
    'google/pix2struct-textcaps-base': (
        'https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json'
    ),
}




class   UpperCamelCase_       (     UpperCamelCase__    ):
							lowerCamelCase_           =					"pix2struct_text_model"
							lowerCamelCase_           =					["past_key_values"]
							lowerCamelCase_           =					{
							    "hidden_size": "hidden_size",
							    "num_attention_heads": "num_heads",
							    "num_hidden_layers": "num_layers",
							}


							def __init__(							self    :Union[str, Any]							,				__A    :Any=5_0244							,				__A    :Optional[Any]=768							,				__A    :Tuple=64							,				__A    :List[str]=2048							,				__A    :int=12							,				__A    :str=12							,				__A    :Any=32							,				__A    :Tuple=128							,				__A    :int=0.1							,				__A    :str=1E-6							,				__A    :Optional[Any]=1.0							,				__A    :Union[str, Any]="gelu_new"							,				__A    :Any=0							,				__A    :List[str]=False							,				__A    :Optional[Any]=0							,				__A    :int=1							,				__A    :Optional[int]=False							,				__A    :Optional[Any]=True							,				**__A    :List[Any]							,				)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      vocab_size
									SCREAMING_SNAKE_CASE__			=      hidden_size
									SCREAMING_SNAKE_CASE__			=      d_kv
									SCREAMING_SNAKE_CASE__			=      d_ff
									SCREAMING_SNAKE_CASE__			=      num_layers
									SCREAMING_SNAKE_CASE__			=      num_heads
									SCREAMING_SNAKE_CASE__			=      relative_attention_num_buckets
									SCREAMING_SNAKE_CASE__			=      relative_attention_max_distance
									SCREAMING_SNAKE_CASE__			=      dropout_rate
									SCREAMING_SNAKE_CASE__			=      layer_norm_epsilon
									SCREAMING_SNAKE_CASE__			=      initializer_factor
									SCREAMING_SNAKE_CASE__			=      use_cache

									SCREAMING_SNAKE_CASE__			=      eos_token_id
									SCREAMING_SNAKE_CASE__			=      decoder_start_token_id

									# for backwards compatibility
									SCREAMING_SNAKE_CASE__			=      dense_act_fn

									super().__init__(
									    pad_token_id=__A							,				eos_token_id=__A							,				decoder_start_token_id=__A							,				tie_word_embeddings=__A							,				is_decoder=__A							,				**__A							,				)


							@classmethod
							def      _snake_case  (							cls    :Optional[int]							,				__A    :Union[str, os.PathLike]							,				**__A    :Optional[int]		)   ->   "PretrainedConfig":



									"""simple docstring"""



									cls._set_token_in_kwargs(__A		)

									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      cls.get_config_dict(__A							,				**__A		)

									# get the text config dict if we are loading from Pix2StructConfig
									if config_dict.get("""model_type"""		) == "pix2struct":
											SCREAMING_SNAKE_CASE__			=      config_dict["""text_config"""]

									if "model_type" in config_dict and hasattr(cls							,				"""model_type"""		) and config_dict["model_type"] != cls.model_type:
											logger.warning(
											    f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
											    f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''		)

									return cls.from_dict(__A							,				**__A		)




class   UpperCamelCase_       (     UpperCamelCase__    ):
							lowerCamelCase_           =					"pix2struct_vision_model"


							def __init__(							self    :Optional[int]							,				__A    :int=768							,				__A    :Optional[Any]=768							,				__A    :Union[str, Any]=2048							,				__A    :int=64							,				__A    :Union[str, Any]=12							,				__A    :str=12							,				__A    :Any="gelu_new"							,				__A    :List[Any]=1E-6							,				__A    :Dict=0.0							,				__A    :int=0.0							,				__A    :int=1E-10							,				__A    :Dict=1.0							,				__A    :int=4096							,				__A    :int=32							,				__A    :int=128							,				**__A    :Tuple							,				)   ->   str:



									"""simple docstring"""



									super().__init__(**__A		)

									SCREAMING_SNAKE_CASE__			=      hidden_size
									SCREAMING_SNAKE_CASE__			=      patch_embed_hidden_size
									SCREAMING_SNAKE_CASE__			=      d_ff
									SCREAMING_SNAKE_CASE__			=      dropout_rate
									SCREAMING_SNAKE_CASE__			=      num_hidden_layers
									SCREAMING_SNAKE_CASE__			=      num_attention_heads
									SCREAMING_SNAKE_CASE__			=      initializer_range
									SCREAMING_SNAKE_CASE__			=      initializer_factor
									SCREAMING_SNAKE_CASE__			=      attention_dropout
									SCREAMING_SNAKE_CASE__			=      layer_norm_eps
									SCREAMING_SNAKE_CASE__			=      dense_act_fn
									SCREAMING_SNAKE_CASE__			=      seq_len
									SCREAMING_SNAKE_CASE__			=      relative_attention_num_buckets
									SCREAMING_SNAKE_CASE__			=      relative_attention_max_distance
									SCREAMING_SNAKE_CASE__			=      d_kv


							@classmethod
							def      _snake_case  (							cls    :str							,				__A    :Union[str, os.PathLike]							,				**__A    :str		)   ->   "PretrainedConfig":



									"""simple docstring"""



									cls._set_token_in_kwargs(__A		)

									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      cls.get_config_dict(__A							,				**__A		)

									# get the vision config dict if we are loading from Pix2StructConfig
									if config_dict.get("""model_type"""		) == "pix2struct":
											SCREAMING_SNAKE_CASE__			=      config_dict["""vision_config"""]

									if "model_type" in config_dict and hasattr(cls							,				"""model_type"""		) and config_dict["model_type"] != cls.model_type:
											logger.warning(
											    f'''You are using a model of type {config_dict['model_type']} to instantiate a model of type '''
											    f'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''		)

									return cls.from_dict(__A							,				**__A		)




class   UpperCamelCase_       (     UpperCamelCase__    ):
							lowerCamelCase_           =					"pix2struct"
							lowerCamelCase_           =					True


							def __init__(							self    :str							,				__A    :Optional[Any]=None							,				__A    :List[str]=None							,				__A    :Optional[Any]=1.0							,				__A    :Optional[Any]=0.0_2							,				__A    :Any=False							,				__A    :Tuple=False							,				__A    :Any=True							,				**__A    :Dict							,				)   ->   Union[str, Any]:



									"""simple docstring"""



									super().__init__(tie_word_embeddings=__A							,				is_encoder_decoder=__A							,				**__A		)

									if text_config is None:
											SCREAMING_SNAKE_CASE__			=      {}
											logger.info("""text_config is None. Initializing the Pix2StructTextConfig with default values."""		)

									if vision_config is None:
											SCREAMING_SNAKE_CASE__			=      {}
											logger.info("""vision_config is None. Initializing the Pix2StructVisionConfig with default values."""		)

									SCREAMING_SNAKE_CASE__			=      PixaStructTextConfig(**__A		)
									SCREAMING_SNAKE_CASE__			=      PixaStructVisionConfig(**__A		)

									SCREAMING_SNAKE_CASE__			=      self.text_config.decoder_start_token_id
									SCREAMING_SNAKE_CASE__			=      self.text_config.pad_token_id
									SCREAMING_SNAKE_CASE__			=      self.text_config.eos_token_id

									SCREAMING_SNAKE_CASE__			=      initializer_factor
									SCREAMING_SNAKE_CASE__			=      initializer_range

									SCREAMING_SNAKE_CASE__			=      self.initializer_range
									SCREAMING_SNAKE_CASE__			=      self.initializer_range

									SCREAMING_SNAKE_CASE__			=      is_vqa


							@classmethod
							def      _snake_case  (							cls    :Union[str, Any]							,				__A    :PixaStructTextConfig							,				__A    :PixaStructVisionConfig							,				**__A    :Optional[int]		)   ->   Optional[Any]:



									"""simple docstring"""



									return cls(text_config=text_config.to_dict()							,				vision_config=vision_config.to_dict()							,				**__A		)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      copy.deepcopy(self.__dict__		)
									SCREAMING_SNAKE_CASE__			=      self.text_config.to_dict()
									SCREAMING_SNAKE_CASE__			=      self.vision_config.to_dict()
									SCREAMING_SNAKE_CASE__			=      self.__class__.model_type
									return output