upload hubert for seq classification

config.json

@@ -0,0 +1,75 @@
+{
+  "activation_dropout": 0.0,
+  "apply_spec_augment": true,
+  "architectures": [
+    "HubertModel"
+  ],
+  "attention_dropout": 0.1,
+  "bos_token_id": 1,
+  "conv_bias": true,
+  "conv_dim": [
+    512,
+    512,
+    512,
+    512,
+    512,
+    512,
+    512
+  ],
+  "conv_kernel": [
+    10,
+    3,
+    3,
+    3,
+    3,
+    2,
+    2
+  ],
+  "conv_stride": [
+    5,
+    2,
+    2,
+    2,
+    2,
+    2,
+    2
+  ],
+  "ctc_loss_reduction": "sum",
+  "ctc_zero_infinity": false,
+  "do_stable_layer_norm": true,
+  "eos_token_id": 2,
+  "feat_extract_activation": "gelu",
+  "feat_extract_dropout": 0.0,
+  "feat_extract_norm": "layer",
+  "feat_proj_dropout": 0.1,
+  "final_dropout": 0.0,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "layerdrop": 0.1,
+  "mask_channel_length": 10,
+  "mask_channel_min_space": 1,
+  "mask_channel_other": 0.0,
+  "mask_channel_prob": 0.0,
+  "mask_channel_selection": "static",
+  "mask_feature_length": 10,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_space": 1,
+  "mask_time_other": 0.0,
+  "mask_time_prob": 0.075,
+  "mask_time_selection": "static",
+  "model_type": "hubert",
+  "num_attention_heads": 16,
+  "num_conv_pos_embedding_groups": 16,
+  "num_conv_pos_embeddings": 128,
+  "num_feat_extract_layers": 7,
+  "num_hidden_layers": 24,
+  "pad_token_id": 0,
+  "transformers_version": "4.9.0.dev0",
+  "vocab_size": 32
+}

flax_model.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9182ceb19f342e64595b56f8cefbe3aad87fc906c46824798741c3221f55956
+size 417861048

hubert_for_sequence_classification.py

@@ -0,0 +1,89 @@
+#!/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

run_hubert_classifier.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+from hubert_for_sequence_classification import FlaxHubertForSequenceClassification, FlaxHubertModel
+import numpy as np
+
+# need to do some ugly save/reload because of a bug
+model = FlaxHubertModel.from_pretrained("facebook/hubert-large-ll60k", from_pt=True)
+model.save_pretrained("./")
+model = FlaxHubertForSequenceClassification.from_pretrained("./")
+
+dummy_input = np.array(2 * [1024 * [1.0]], dtype=np.float32)
+
+logits = model(dummy_input).logits
+
+# output shape is (batch_size, 2)
+print("output shape", logits.shape)