import tensorflow as tf

from ...tf_utils import shape_list







class __snake_case       (		tf.keras.layers.Layer	):




   def __init__(					self    :       Union[str, Any] ,			_lowercase    :       str ,			_lowercase    :       int ,			_lowercase    :       List[str] ,			_lowercase    :       Union[str, Any] ,			_lowercase    :       int=1 ,			_lowercase    :       List[str]=False ,			**_lowercase    :       List[str]   ):



     """simple docstring"""






     super().__init__(**_lowercase   )

     SCREAMING_SNAKE_CASE__    							=   vocab_size
     SCREAMING_SNAKE_CASE__    							=   d_embed
     SCREAMING_SNAKE_CASE__    							=   d_proj

     SCREAMING_SNAKE_CASE__    							=   cutoffs + [vocab_size]
     SCREAMING_SNAKE_CASE__    							=   [0] + self.cutoffs
     SCREAMING_SNAKE_CASE__    							=   div_val

     SCREAMING_SNAKE_CASE__    							=   self.cutoffs[0]
     SCREAMING_SNAKE_CASE__    							=   len(self.cutoffs   ) - 1
     SCREAMING_SNAKE_CASE__    							=   self.shortlist_size + self.n_clusters
     SCREAMING_SNAKE_CASE__    							=   keep_order

     SCREAMING_SNAKE_CASE__    							=   []
     SCREAMING_SNAKE_CASE__    							=   []




   def 			__a		(					self    :       str ,			_lowercase    :       List[str]   ):



     """simple docstring"""






     if self.n_clusters > 0:
       SCREAMING_SNAKE_CASE__    							=   self.add_weight(
           shape=(self.n_clusters, self.d_embed) ,			initializer="""zeros""" ,			trainable=_lowercase ,			name="""cluster_weight"""   )
       SCREAMING_SNAKE_CASE__    							=   self.add_weight(
           shape=(self.n_clusters,) ,			initializer="""zeros""" ,			trainable=_lowercase ,			name="""cluster_bias"""   )

     if self.div_val == 1:
       for i in range(len(self.cutoffs   )   ):
         if self.d_proj != self.d_embed:
           SCREAMING_SNAKE_CASE__    							=   self.add_weight(
               shape=(self.d_embed, self.d_proj) ,			initializer="""zeros""" ,			trainable=_lowercase ,			name=f"""out_projs_._{i}""" ,			)
           self.out_projs.append(_lowercase   )
         else:
           self.out_projs.append(_lowercase   )
         SCREAMING_SNAKE_CASE__    							=   self.add_weight(
             shape=(self.vocab_size, self.d_embed) ,			initializer="""zeros""" ,			trainable=_lowercase ,			name=f"""out_layers_._{i}_._weight""" ,			)
         SCREAMING_SNAKE_CASE__    							=   self.add_weight(
             shape=(self.vocab_size,) ,			initializer="""zeros""" ,			trainable=_lowercase ,			name=f"""out_layers_._{i}_._bias""" ,			)
         self.out_layers.append((weight, bias)   )
     else:
       for i in range(len(self.cutoffs   )   ):
         SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   self.cutoff_ends[i], self.cutoff_ends[i + 1]
         SCREAMING_SNAKE_CASE__    							=   self.d_embed // (self.div_val**i)

         SCREAMING_SNAKE_CASE__    							=   self.add_weight(
             shape=(d_emb_i, self.d_proj) ,			initializer="""zeros""" ,			trainable=_lowercase ,			name=f"""out_projs_._{i}"""   )
         self.out_projs.append(_lowercase   )
         SCREAMING_SNAKE_CASE__    							=   self.add_weight(
             shape=(r_idx - l_idx, d_emb_i) ,			initializer="""zeros""" ,			trainable=_lowercase ,			name=f"""out_layers_._{i}_._weight""" ,			)
         SCREAMING_SNAKE_CASE__    							=   self.add_weight(
             shape=(r_idx - l_idx,) ,			initializer="""zeros""" ,			trainable=_lowercase ,			name=f"""out_layers_._{i}_._bias""" ,			)
         self.out_layers.append((weight, bias)   )
     super().build(_lowercase   )




   @staticmethod
   def 			__a		(					_lowercase    :       List[Any] ,			_lowercase    :       Optional[int] ,			_lowercase    :       List[str] ,			_lowercase    :       List[Any]=None   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   x
     if proj is not None:
       SCREAMING_SNAKE_CASE__    							=   tf.einsum("""ibd,ed->ibe""" ,			_lowercase ,			_lowercase   )
     return tf.einsum("""ibd,nd->ibn""" ,			_lowercase ,			_lowercase   ) + b




   @staticmethod
   def 			__a		(					_lowercase    :       Union[str, Any] ,			_lowercase    :       Optional[int]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   shape_list(_lowercase   )
     SCREAMING_SNAKE_CASE__    							=   tf.range(lp_size[0] ,			dtype=target.dtype   )
     SCREAMING_SNAKE_CASE__    							=   tf.stack([r, target] ,			1   )
     return tf.gather_nd(_lowercase ,			_lowercase   )




   def 			__a		(					self    :       int ,			_lowercase    :       int ,			_lowercase    :       Optional[int] ,			_lowercase    :       Optional[int]=True ,			_lowercase    :       str=False   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   0
     if self.n_clusters == 0:
       SCREAMING_SNAKE_CASE__    							=   self._logit(_lowercase ,			self.out_layers[0][0] ,			self.out_layers[0][1] ,			self.out_projs[0]   )
       if target is not None:
         SCREAMING_SNAKE_CASE__    							=   tf.nn.sparse_softmax_cross_entropy_with_logits(labels=_lowercase ,			logits=_lowercase   )
       SCREAMING_SNAKE_CASE__    							=   tf.nn.log_softmax(_lowercase ,			axis=-1   )
     else:
       SCREAMING_SNAKE_CASE__    							=   shape_list(_lowercase   )
       SCREAMING_SNAKE_CASE__    							=   []
       SCREAMING_SNAKE_CASE__    							=   tf.zeros(hidden_sizes[:2]   )
       for i in range(len(self.cutoffs   )   ):
         SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   self.cutoff_ends[i], self.cutoff_ends[i + 1]
         if target is not None:
           SCREAMING_SNAKE_CASE__    							=   (target >= l_idx) & (target < r_idx)
           SCREAMING_SNAKE_CASE__    							=   tf.where(_lowercase   )
           SCREAMING_SNAKE_CASE__    							=   tf.boolean_mask(_lowercase ,			_lowercase   ) - l_idx

         if self.div_val == 1:
           SCREAMING_SNAKE_CASE__    							=   self.out_layers[0][0][l_idx:r_idx]
           SCREAMING_SNAKE_CASE__    							=   self.out_layers[0][1][l_idx:r_idx]
         else:
           SCREAMING_SNAKE_CASE__    							=   self.out_layers[i][0]
           SCREAMING_SNAKE_CASE__    							=   self.out_layers[i][1]

         if i == 0:
           SCREAMING_SNAKE_CASE__    							=   tf.concat([cur_W, self.cluster_weight] ,			0   )
           SCREAMING_SNAKE_CASE__    							=   tf.concat([cur_b, self.cluster_bias] ,			0   )

           SCREAMING_SNAKE_CASE__    							=   self._logit(_lowercase ,			_lowercase ,			_lowercase ,			self.out_projs[0]   )
           SCREAMING_SNAKE_CASE__    							=   tf.nn.log_softmax(_lowercase   )
           out.append(head_logprob[..., : self.cutoffs[0]]   )
           if target is not None:
             SCREAMING_SNAKE_CASE__    							=   tf.boolean_mask(_lowercase ,			_lowercase   )
             SCREAMING_SNAKE_CASE__    							=   self._gather_logprob(_lowercase ,			_lowercase   )
         else:
           SCREAMING_SNAKE_CASE__    							=   self._logit(_lowercase ,			_lowercase ,			_lowercase ,			self.out_projs[i]   )
           SCREAMING_SNAKE_CASE__    							=   tf.nn.log_softmax(_lowercase   )
           SCREAMING_SNAKE_CASE__    							=   self.cutoffs[0] + i - 1  # No probability for the head cluster
           SCREAMING_SNAKE_CASE__    							=   head_logprob[..., cluster_prob_idx, None] + tail_logprob
           out.append(_lowercase   )
           if target is not None:
             SCREAMING_SNAKE_CASE__    							=   tf.boolean_mask(_lowercase ,			_lowercase   )
             SCREAMING_SNAKE_CASE__    							=   tf.boolean_mask(_lowercase ,			_lowercase   )
             SCREAMING_SNAKE_CASE__    							=   self._gather_logprob(_lowercase ,			_lowercase   )
             cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
         if target is not None:
           loss += tf.scatter_nd(_lowercase ,			-cur_logprob ,			shape_list(_lowercase   )   )
       SCREAMING_SNAKE_CASE__    							=   tf.concat(_lowercase ,			axis=-1   )

     if target is not None:
       if return_mean:
         SCREAMING_SNAKE_CASE__    							=   tf.reduce_mean(_lowercase   )
       # Add the training-time loss value to the layer using `self.add_loss()`.
       self.add_loss(_lowercase   )

       # Log the loss as a metric (we could log arbitrary metrics,
       # including different metrics for training and inference.
       self.add_metric(_lowercase ,			name=self.name ,			aggregation="""mean""" if return_mean else """"""   )

     return out

from typing import List, Optional

import numpy as np

from ...processing_utils import ProcessorMixin
from ...utils import to_numpy







class __snake_case       (		lowerCamelCase_	):
   lowerCAmelCase_														=							"EncodecFeatureExtractor"
   lowerCAmelCase_														=							("T5Tokenizer", "T5TokenizerFast")




   def __init__(					self    :       List[Any] ,			_lowercase    :       int ,			_lowercase    :       Any   ):



     """simple docstring"""






     super().__init__(_lowercase ,			_lowercase   )
     SCREAMING_SNAKE_CASE__    							=   self.feature_extractor
     SCREAMING_SNAKE_CASE__    							=   False




   def 			__a		(					self    :       Tuple ,			_lowercase    :       str=None ,			_lowercase    :       int=None ,			_lowercase    :       Any=True   ):



     """simple docstring"""






     return self.tokenizer.get_decoder_prompt_ids(task=_lowercase ,			language=_lowercase ,			no_timestamps=_lowercase   )




   def __call__(					self    :       Tuple ,			*_lowercase    :       Tuple ,			**_lowercase    :       Union[str, Any]   ):



     """simple docstring"""






     if self._in_target_context_manager:
       return self.current_processor(*_lowercase ,			**_lowercase   )

     SCREAMING_SNAKE_CASE__    							=   kwargs.pop("""audio""" ,			_lowercase   )
     SCREAMING_SNAKE_CASE__    							=   kwargs.pop("""sampling_rate""" ,			_lowercase   )
     SCREAMING_SNAKE_CASE__    							=   kwargs.pop("""text""" ,			_lowercase   )
     if len(_lowercase   ) > 0:
       SCREAMING_SNAKE_CASE__    							=   args[0]
       SCREAMING_SNAKE_CASE__    							=   args[1:]

     if audio is None and text is None:
       raise ValueError("""You need to specify either an `audio` or `text` input to process."""   )

     if text is not None:
       SCREAMING_SNAKE_CASE__    							=   self.tokenizer(_lowercase ,			**_lowercase   )

     if audio is not None:
       SCREAMING_SNAKE_CASE__    							=   self.feature_extractor(_lowercase ,			*_lowercase ,			sampling_rate=_lowercase ,			**_lowercase   )

     if audio is None:
       return inputs

     elif text is None:
       return audio_inputs

     else:
       SCREAMING_SNAKE_CASE__    							=   audio_inputs["""input_values"""]
       if "padding_mask" in audio_inputs:
         SCREAMING_SNAKE_CASE__    							=   audio_inputs["""padding_mask"""]
       return inputs




   def 			__a		(					self    :       Dict ,			*_lowercase    :       Any ,			**_lowercase    :       Optional[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   kwargs.pop("""audio""" ,			_lowercase   )
     SCREAMING_SNAKE_CASE__    							=   kwargs.pop("""padding_mask""" ,			_lowercase   )

     if len(_lowercase   ) > 0:
       SCREAMING_SNAKE_CASE__    							=   args[0]
       SCREAMING_SNAKE_CASE__    							=   args[1:]

     if audio_values is not None:
       return self._decode_audio(_lowercase ,			padding_mask=_lowercase   )
     else:
       return self.tokenizer.batch_decode(*_lowercase ,			**_lowercase   )




   def 			__a		(					self    :       List[str] ,			*_lowercase    :       Tuple ,			**_lowercase    :       Optional[Any]   ):



     """simple docstring"""






     return self.tokenizer.decode(*_lowercase ,			**_lowercase   )




   def 			__a		(					self    :       List[str] ,			_lowercase    :       Tuple ,			_lowercase    :       Optional = None   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   to_numpy(_lowercase   )
     SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__							,  SCREAMING_SNAKE_CASE__    							=   audio_values.shape

     if padding_mask is None:
       return list(_lowercase   )

     SCREAMING_SNAKE_CASE__    							=   to_numpy(_lowercase   )

     # match the sequence length of the padding mask to the generated audio arrays by padding with the **non-padding**
     # token (so that the generated audio values are **not** treated as padded tokens)
     SCREAMING_SNAKE_CASE__    							=   seq_len - padding_mask.shape[-1]
     SCREAMING_SNAKE_CASE__    							=   1 - self.feature_extractor.padding_value
     SCREAMING_SNAKE_CASE__    							=   np.pad(_lowercase ,			((0, 0), (0, difference)) ,			"""constant""" ,			constant_values=_lowercase   )

     SCREAMING_SNAKE_CASE__    							=   audio_values.tolist()
     for i in range(_lowercase   ):
       SCREAMING_SNAKE_CASE__    							=   np.asarray(audio_values[i]   )[
           padding_mask[i][None, :] != self.feature_extractor.padding_value
       ]
       SCREAMING_SNAKE_CASE__    							=   sliced_audio.reshape(_lowercase ,			-1   )

     return audio_values
