| |
| from transformers.models.wav2vec2.modeling_flax_wav2vec2 import FlaxWav2Vec2Module, FlaxWav2Vec2PreTrainedModel |
| from typing import Union |
| from transformers import HubertConfig |
| from transformers.modeling_flax_outputs import FlaxSequenceClassifierOutput |
| import flax.linen as nn |
| import jax.numpy as jnp |
| import jax |
|
|
|
|
| class FlaxHubertForSequenceClassificationModule(nn.Module): |
| config: HubertConfig |
| dtype: jnp.dtype = jnp.float32 |
|
|
| def setup(self): |
| self.hubert = FlaxWav2Vec2Module(self.config, dtype=self.dtype) |
| self.dropout = nn.Dropout(rate=self.config.final_dropout) |
| self.reduce = "mean" |
|
|
| |
| self.lm_head = nn.Dense( |
| 2, |
| kernel_init=jax.nn.initializers.normal(self.config.initializer_range, self.dtype), |
| dtype=self.dtype, |
| ) |
|
|
| def __call__( |
| self, |
| input_values, |
| attention_mask=None, |
| mask_time_indices=None, |
| deterministic=True, |
| output_attentions=None, |
| output_hidden_states=None, |
| return_dict=None, |
| ): |
| outputs = self.hubert( |
| input_values, |
| attention_mask=attention_mask, |
| mask_time_indices=mask_time_indices, |
| deterministic=deterministic, |
| output_attentions=output_attentions, |
| output_hidden_states=output_hidden_states, |
| return_dict=return_dict, |
| ) |
|
|
| hidden_states = outputs[0] |
| if self.reduce == "mean": |
| hidden_states = jnp.mean(hidden_states, axis=1) |
|
|
| hidden_states = jax.nn.relu(hidden_states) |
| logits = self.lm_head(hidden_states) |
|
|
| if not return_dict: |
| return (logits,) + outputs[2:] |
|
|
| return FlaxSequenceClassifierOutput(logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions) |
|
|
| def _get_feat_extract_output_lengths(self, input_lengths: Union[jnp.ndarray, int]): |
| """ |
| Computes the output length of the convolutional layers |
| """ |
|
|
| def _conv_out_length(input_length, kernel_size, stride): |
| |
| |
| return (input_length - kernel_size) // stride + 1 |
|
|
| for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride): |
| input_lengths = _conv_out_length(input_lengths, kernel_size, stride) |
|
|
| return input_lengths |
|
|
|
|
| class FlaxHubertPreTrainedModel(FlaxWav2Vec2PreTrainedModel): |
| config_class = HubertConfig |
| base_model_prefix: str = "hubert" |
| module_class: nn.Module = None |
|
|
| def _get_feat_extract_output_lengths(self, input_lengths: Union[jnp.ndarray, int]): |
| return self.module._get_feat_extract_output_lengths(input_lengths) |
|
|
|
|
| class FlaxHubertModel(FlaxHubertPreTrainedModel): |
| module_class = FlaxWav2Vec2Module |
|
|
|
|
| class FlaxHubertForSequenceClassification(FlaxHubertPreTrainedModel): |
| module_class = FlaxHubertForSequenceClassificationModule |
|
|
