

'''simple docstring'''





from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin


@dataclass
class  UpperCAmelCase (							UpperCAmelCase__    ):

		'''simple docstring'''






		SCREAMING_SNAKE_CASE_           =					42


class  UpperCAmelCase (							UpperCAmelCase__  , UpperCAmelCase__    ):

		'''simple docstring'''






		@register_to_config
		def __init__(					self   ,			SCREAMING_SNAKE_CASE_ = 16   ,			SCREAMING_SNAKE_CASE_ = 88   ,			SCREAMING_SNAKE_CASE_ = None   ,			SCREAMING_SNAKE_CASE_ = None   ,			SCREAMING_SNAKE_CASE_ = 1   ,			SCREAMING_SNAKE_CASE_ = 0.0   ,			SCREAMING_SNAKE_CASE_ = 32   ,			SCREAMING_SNAKE_CASE_ = None   ,			SCREAMING_SNAKE_CASE_ = False   ,			SCREAMING_SNAKE_CASE_ = None   ,			SCREAMING_SNAKE_CASE_ = "geglu"   ,			SCREAMING_SNAKE_CASE_ = True   ,			SCREAMING_SNAKE_CASE_ = True   ,			)  ->     Tuple:






							'''simple docstring'''



							super().__init__()
							lowerCamelCase_  				=      num_attention_heads
							lowerCamelCase_  				=      attention_head_dim
							lowerCamelCase_  				=      num_attention_heads * attention_head_dim

							lowerCamelCase_  				=      in_channels

							lowerCamelCase_  				=      torch.nn.GroupNorm(num_groups=SCREAMING_SNAKE_CASE_   ,			num_channels=SCREAMING_SNAKE_CASE_   ,			eps=1E-6   ,			affine=SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      nn.Linear(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)

							# 3. Define transformers blocks
							lowerCamelCase_  				=      nn.ModuleList(
							    [
							        BasicTransformerBlock(
							            SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			dropout=SCREAMING_SNAKE_CASE_   ,			cross_attention_dim=SCREAMING_SNAKE_CASE_   ,			activation_fn=SCREAMING_SNAKE_CASE_   ,			attention_bias=SCREAMING_SNAKE_CASE_   ,			double_self_attention=SCREAMING_SNAKE_CASE_   ,			norm_elementwise_affine=SCREAMING_SNAKE_CASE_   ,			)
							        for d in range(SCREAMING_SNAKE_CASE_		)
							    ]		)

							lowerCamelCase_  				=      nn.Linear(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)






		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_=None   ,			SCREAMING_SNAKE_CASE_=None   ,			SCREAMING_SNAKE_CASE_=None   ,			SCREAMING_SNAKE_CASE_=1   ,			SCREAMING_SNAKE_CASE_=None   ,			SCREAMING_SNAKE_CASE_ = True   ,			)  ->     Any:






							'''simple docstring'''



							lowerCamelCase_				,lowerCamelCase_				,lowerCamelCase_				,lowerCamelCase_  				=      hidden_states.shape
							lowerCamelCase_  				=      batch_frames // num_frames

							lowerCamelCase_  				=      hidden_states

							lowerCamelCase_  				=      hidden_states[None, :].reshape(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      hidden_states.permute(0   ,			2   ,			1   ,			3   ,			4		)

							lowerCamelCase_  				=      self.norm(SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      hidden_states.permute(0   ,			3   ,			4   ,			2   ,			1		).reshape(batch_size * height * width   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)

							lowerCamelCase_  				=      self.proj_in(SCREAMING_SNAKE_CASE_		)

							# 2. Blocks
							for block in self.transformer_blocks:
												lowerCamelCase_  				=      block(
												    SCREAMING_SNAKE_CASE_   ,			encoder_hidden_states=SCREAMING_SNAKE_CASE_   ,			timestep=SCREAMING_SNAKE_CASE_   ,			cross_attention_kwargs=SCREAMING_SNAKE_CASE_   ,			class_labels=SCREAMING_SNAKE_CASE_   ,			)

							# 3. Output
							lowerCamelCase_  				=      self.proj_out(SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      (
							    hidden_states[None, None, :]
							    .reshape(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)
							    .permute(0   ,			3   ,			4   ,			1   ,			2		)
							    .contiguous()
							)
							lowerCamelCase_  				=      hidden_states.reshape(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)

							lowerCamelCase_  				=      hidden_states + residual

							if not return_dict:
												return (output,)

							return TransformerTemporalModelOutput(sample=SCREAMING_SNAKE_CASE_		)







import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Audio, Features, Value
from .base import TaskTemplate
@dataclass(frozen=lowerCAmelCase__					)
class 						lowerCamelCase_							(					lowerCAmelCase__					):



					'''simple docstring'''







					__UpperCAmelCase							  =					field(default="automatic-speech-recognition"  ,  metadata={"include_in_asdict_even_if_is_default": True}					)
					__UpperCAmelCase							  =					Features({"audio": Audio()}					)
					__UpperCAmelCase							  =					Features({"transcription": Value("string"					)}					)
					__UpperCAmelCase							  =					"audio"
					__UpperCAmelCase							  =					"transcription"






					def A   (					self			,							snake_case_    ) ->      Any:







								'''simple docstring'''

								if self.audio_column not in features:
											raise ValueError(F'Column {self.audio_column} is not present in features.'    )
								if not isinstance(features[self.audio_column]			,							snake_case_    ):
											raise ValueError(F'Column {self.audio_column} is not an Audio type.'    )
								__lowercase			       = copy.deepcopy(self    )
								__lowercase			       = self.input_schema.copy()
								__lowercase			       = features[self.audio_column]
								__lowercase			       = input_schema
								return task_template






					@property
					def A   (					self    ) ->      Dict[str, str]:







								'''simple docstring'''

								return {self.audio_column: "audio", self.transcription_column: "transcription"}


