File size: 3,081 Bytes
004e84f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
#!/usr/bin/env python3
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"
# binary classification
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):
# 1D convolutional layer output length formula taken
# from https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html
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
|