
'''simple docstring'''
import torch
from torch import nn






class    UpperCAmelCase_				(	nn.Module					):





 def __init__(    self ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_=1 ,       lowercase_=False):
     super().__init__()

     snake_case_ :    Any     			=						n_token
     snake_case_ :    str     			=						d_embed
     snake_case_ :    Optional[int]     			=						d_proj

     snake_case_ :    Tuple     			=						cutoffs + [n_token]
     snake_case_ :    int     			=						[0] + self.cutoffs
     snake_case_ :    Optional[int]     			=						div_val

     snake_case_ :    str     			=						self.cutoffs[0]
     snake_case_ :    Optional[Any]     			=						len(self.cutoffs) - 1
     snake_case_ :    Union[str, Any]     			=						self.shortlist_size + self.n_clusters

     if self.n_clusters > 0:
         snake_case_ :    Union[str, Any]     			=						nn.Parameter(torch.zeros(self.n_clusters ,       self.d_embed))
         snake_case_ :    Union[str, Any]     			=						nn.Parameter(torch.zeros(self.n_clusters))

     snake_case_ :    Tuple     			=						nn.ModuleList()
     snake_case_ :    List[Any]     			=						nn.ParameterList()

     if div_val == 1:
         for i in range(len(self.cutoffs)):
             if d_proj != d_embed:
                 self.out_projs.append(nn.Parameter(torch.FloatTensor(UpperCAmelCase__ ,       UpperCAmelCase__)))
             else:
                 self.out_projs.append(UpperCAmelCase__)

         self.out_layers.append(nn.Linear(UpperCAmelCase__ ,       UpperCAmelCase__))
     else:
         for i in range(len(self.cutoffs)):
             snake_case_ :    str     			=						self.cutoff_ends[i], self.cutoff_ends[i + 1]
             snake_case_ :    Optional[Any]     			=						d_embed // (div_val**i)

             self.out_projs.append(nn.Parameter(torch.FloatTensor(UpperCAmelCase__ ,       UpperCAmelCase__)))

             self.out_layers.append(nn.Linear(UpperCAmelCase__ ,       r_idx - l_idx))

     snake_case_ :    Optional[int]     			=						keep_order





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_ ,       lowercase_ ,       lowercase_):
     if proj is None:
         snake_case_ :    Dict     			=						nn.functional.linear(UpperCAmelCase__ ,       UpperCAmelCase__ ,       bias=UpperCAmelCase__)
     else:
         # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
         snake_case_ :    Optional[Any]     			=						nn.functional.linear(UpperCAmelCase__ ,       proj.t().contiguous())
         snake_case_ :    Any     			=						nn.functional.linear(UpperCAmelCase__ ,       UpperCAmelCase__ ,       bias=UpperCAmelCase__)
         # else:
         #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
         #     if bias is not None:
         #         logit = logit + bias

     return logit





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_=None ,       lowercase_=False):

     if labels is not None:
         # Shift so that tokens < n predict n
         snake_case_ :    Tuple     			=						hidden[..., :-1, :].contiguous()
         snake_case_ :    Union[str, Any]     			=						labels[..., 1:].contiguous()
         snake_case_ :    Any     			=						hidden.view(-1 ,       hidden.size(-1))
         snake_case_ :    Any     			=						labels.view(-1)
         if hidden.size(0) != labels.size(0):
             raise RuntimeError("Input and labels should have the same size in the batch dimension.")
     else:
         snake_case_ :    int     			=						hidden.view(-1 ,       hidden.size(-1))

     if self.n_clusters == 0:
         snake_case_ :    str     			=						self._compute_logit(UpperCAmelCase__ ,       self.out_layers[0].weight ,       self.out_layers[0].bias ,       self.out_projs[0])
         if labels is not None:
             snake_case_ :    Tuple     			=						labels != -1_00
             snake_case_ :    Optional[Any]     			=						torch.zeros_like(UpperCAmelCase__ ,       dtype=hidden.dtype ,       device=hidden.device)
             snake_case_ :    Any     			=						(
                 -nn.functional.log_softmax(UpperCAmelCase__ ,       dim=-1)[mask].gather(1 ,       labels[mask].unsqueeze(1)).squeeze(1)
             )
         else:
             snake_case_ :    Any     			=						nn.functional.log_softmax(UpperCAmelCase__ ,       dim=-1)
     else:
         # construct weights and biases
         snake_case_ :    Any     			=						[], []
         for i in range(len(self.cutoffs)):
             if self.div_val == 1:
                 snake_case_ :    Optional[int]     			=						self.cutoff_ends[i], self.cutoff_ends[i + 1]
                 snake_case_ :    List[Any]     			=						self.out_layers[0].weight[l_idx:r_idx]
                 snake_case_ :    str     			=						self.out_layers[0].bias[l_idx:r_idx]
             else:
                 snake_case_ :    Union[str, Any]     			=						self.out_layers[i].weight
                 snake_case_ :    List[str]     			=						self.out_layers[i].bias

             if i == 0:
                 snake_case_ :    List[str]     			=						torch.cat([weight_i, self.cluster_weight] ,       dim=0)
                 snake_case_ :    int     			=						torch.cat([bias_i, self.cluster_bias] ,       dim=0)

             weights.append(UpperCAmelCase__)
             biases.append(UpperCAmelCase__)

         snake_case_ :    Dict     			=						weights[0], biases[0], self.out_projs[0]

         snake_case_ :    Dict     			=						self._compute_logit(UpperCAmelCase__ ,       UpperCAmelCase__ ,       UpperCAmelCase__ ,       UpperCAmelCase__)
         snake_case_ :    Optional[int]     			=						nn.functional.log_softmax(UpperCAmelCase__ ,       dim=1)

         if labels is None:
             snake_case_ :    List[str]     			=						hidden.new_empty((head_logit.size(0), self.n_token))
         else:
             snake_case_ :    Optional[Any]     			=						torch.zeros_like(UpperCAmelCase__ ,       dtype=hidden.dtype ,       device=hidden.device)

         snake_case_ :    Optional[int]     			=						0
         snake_case_ :    Dict     			=						[0] + self.cutoffs
         for i in range(len(UpperCAmelCase__) - 1):
             snake_case_ :    Optional[int]     			=						cutoff_values[i], cutoff_values[i + 1]

             if labels is not None:
                 snake_case_ :    Optional[int]     			=						(labels >= l_idx) & (labels < r_idx)
                 snake_case_ :    Optional[Any]     			=						mask_i.nonzero().squeeze()

                 if indices_i.numel() == 0:
                     continue

                 snake_case_ :    Any     			=						labels.index_select(0 ,       UpperCAmelCase__) - l_idx
                 snake_case_ :    str     			=						head_logprob.index_select(0 ,       UpperCAmelCase__)
                 snake_case_ :    Dict     			=						hidden.index_select(0 ,       UpperCAmelCase__)
             else:
                 snake_case_ :    Dict     			=						hidden

             if i == 0:
                 if labels is not None:
                     snake_case_ :    int     			=						head_logprob_i.gather(1 ,       target_i[:, None]).squeeze(1)
                 else:
                     snake_case_ :    str     			=						head_logprob[:, : self.cutoffs[0]]
             else:
                 snake_case_ :    Tuple     			=						weights[i], biases[i], self.out_projs[i]

                 snake_case_ :    Dict     			=						self._compute_logit(UpperCAmelCase__ ,       UpperCAmelCase__ ,       UpperCAmelCase__ ,       UpperCAmelCase__)
                 snake_case_ :    Dict     			=						nn.functional.log_softmax(UpperCAmelCase__ ,       dim=1)
                 snake_case_ :    Any     			=						self.cutoffs[0] + i - 1  # No probability for the head cluster
                 if labels is not None:
                     snake_case_ :    Any     			=						head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                         1 ,       target_i[:, None]).squeeze(1)
                 else:
                     snake_case_ :    Dict     			=						head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                     snake_case_ :    List[Any]     			=						logprob_i

             if labels is not None:
                 if (hasattr(self ,       "keep_order") and self.keep_order) or keep_order:
                     out.index_copy_(0 ,       UpperCAmelCase__ ,       -logprob_i)
                 else:
                     out[offset : offset + logprob_i.size(0)].copy_(-logprob_i)
                 offset += logprob_i.size(0)

     return out





 def        snake_case__		(    self ,       lowercase_):
     if self.n_clusters == 0:
         snake_case_ :    Optional[Any]     			=						self._compute_logit(UpperCAmelCase__ ,       self.out_layers[0].weight ,       self.out_layers[0].bias ,       self.out_projs[0])
         return nn.functional.log_softmax(UpperCAmelCase__ ,       dim=-1)
     else:
         # construct weights and biases
         snake_case_ :    Optional[Any]     			=						[], []
         for i in range(len(self.cutoffs)):
             if self.div_val == 1:
                 snake_case_ :    Optional[Any]     			=						self.cutoff_ends[i], self.cutoff_ends[i + 1]
                 snake_case_ :    List[str]     			=						self.out_layers[0].weight[l_idx:r_idx]
                 snake_case_ :    Optional[Any]     			=						self.out_layers[0].bias[l_idx:r_idx]
             else:
                 snake_case_ :    Any     			=						self.out_layers[i].weight
                 snake_case_ :    Optional[Any]     			=						self.out_layers[i].bias

             if i == 0:
                 snake_case_ :    Optional[int]     			=						torch.cat([weight_i, self.cluster_weight] ,       dim=0)
                 snake_case_ :    Any     			=						torch.cat([bias_i, self.cluster_bias] ,       dim=0)

             weights.append(UpperCAmelCase__)
             biases.append(UpperCAmelCase__)

         snake_case_ :    Optional[Any]     			=						weights[0], biases[0], self.out_projs[0]
         snake_case_ :    int     			=						self._compute_logit(UpperCAmelCase__ ,       UpperCAmelCase__ ,       UpperCAmelCase__ ,       UpperCAmelCase__)

         snake_case_ :    Optional[int]     			=						hidden.new_empty((head_logit.size(0), self.n_token))
         snake_case_ :    List[Any]     			=						nn.functional.log_softmax(UpperCAmelCase__ ,       dim=1)

         snake_case_ :    int     			=						[0] + self.cutoffs
         for i in range(len(UpperCAmelCase__) - 1):
             snake_case_ :    Union[str, Any]     			=						cutoff_values[i], cutoff_values[i + 1]

             if i == 0:
                 snake_case_ :    int     			=						head_logprob[:, : self.cutoffs[0]]
             else:
                 snake_case_ :    Dict     			=						weights[i], biases[i], self.out_projs[i]

                 snake_case_ :    Tuple     			=						self._compute_logit(UpperCAmelCase__ ,       UpperCAmelCase__ ,       UpperCAmelCase__ ,       UpperCAmelCase__)
                 snake_case_ :    List[Any]     			=						nn.functional.log_softmax(UpperCAmelCase__ ,       dim=1)

                 snake_case_ :    Tuple     			=						head_logprob[:, -i] + tail_logprob_i
                 snake_case_ :    Any     			=						logprob_i

         return out




'''simple docstring'''
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding






class    UpperCAmelCase_				(	snake_case__					):
 UpperCAmelCase_		    =						"""ClapFeatureExtractor"""
 UpperCAmelCase_		    =						("""RobertaTokenizer""", """RobertaTokenizerFast""")





 def __init__(    self ,       lowercase_ ,       lowercase_):
     super().__init__(lowercase_ ,       lowercase_)





 def __call__(    self ,       lowercase_=None ,       lowercase_=None ,       lowercase_=None ,       **lowercase_):
     snake_case_ :    Any     			=						kwargs.pop("sampling_rate" ,       lowercase_)

     if text is None and audios is None:
         raise ValueError("You have to specify either text or audios. Both cannot be none.")

     if text is not None:
         snake_case_ :    Optional[int]     			=						self.tokenizer(lowercase_ ,       return_tensors=lowercase_ ,       **lowercase_)

     if audios is not None:
         snake_case_ :    Any     			=						self.feature_extractor(
             lowercase_ ,       sampling_rate=lowercase_ ,       return_tensors=lowercase_ ,       **lowercase_)

     if text is not None and audios is not None:
         snake_case_ :    Dict     			=						audio_features.input_features
         return encoding
     elif text is not None:
         return encoding
     else:
         return BatchEncoding(data=dict(**lowercase_) ,       tensor_type=lowercase_)





 def        snake_case__		(    self ,       *lowercase_ ,       **lowercase_):
     return self.tokenizer.batch_decode(*lowercase_ ,       **lowercase_)





 def        snake_case__		(    self ,       *lowercase_ ,       **lowercase_):
     return self.tokenizer.decode(*lowercase_ ,       **lowercase_)





 @property
 def        snake_case__		(    self):
     snake_case_ :    Union[str, Any]     			=						self.tokenizer.model_input_names
     snake_case_ :    int     			=						self.feature_extractor.model_input_names
     return list(dict.fromkeys(tokenizer_input_names + feature_extractor_input_names))




