Upload Speech-JEPA weights

README.md

@@ -0,0 +1,92 @@
+---
+language: en
+license: bsd-3-clause
+library_name: transformers
+pipeline_tag: feature-extraction
+datasets:
+- openslr/librispeech_asr
+tags:
+- speech
+- audio
+- self-supervised-learning
+- jepa
+- speech-jepa
+---
+
+# Speech-JEPA Base
+
+[Code and recipes](https://github.com/labhamlet/SpeechJEPA)
+
+The base model pretrained on 16 kHz sampled speech audio (LibriSpeech, 960 h). When using the model make
+sure that your speech input is also sampled at 16 kHz.
+
+**Note**: This model does not have a tokenizer as it was pretrained on audio alone. In order to use this
+model for speech recognition, a tokenizer should be created and the model should be fine-tuned on labeled
+text data. Fine-tuned checkpoints for the Libri-light splits are listed below.
+
+## Model description
+
+Speech-JEPA is a Joint-Embedding Predictive Architecture for speech. A context encoder sees a partially
+masked sequence of convolutional speech features and a light-weight convolutional decoder predicts the
+latent representations of the masked spans produced by an exponential-moving-average teacher
+(instance-normalised average of the top-8 teacher layers, as in data2vec 2.0). There is no
+reconstruction of audio and no quantised targets.
+
+| | |
+|---|---|
+| Parameters | 89.7 M |
+| Feature encoder | 7-layer 1-D CNN (wav2vec 2.0 layout), 320× downsampling, 20 ms frames |
+| Transformer | 12 layers, 768-d, 12 heads, rotary position embeddings, post-norm |
+| Pretraining data | LibriSpeech 960 h |
+| Pretraining updates | 370000 |
+
+`jepa_pretraining_extras.safetensors` additionally holds the JEPA decoder, mask token and the EMA teacher encoder so that pre-training can be resumed with the original code.
+
+## Usage
+
+```python
+import torch
+from datasets import load_dataset
+from transformers import AutoModel, Wav2Vec2FeatureExtractor
+
+feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("labhamlet/speech-jepa-base")
+model = AutoModel.from_pretrained("labhamlet/speech-jepa-base", trust_remote_code=True)
+
+ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+inputs = feature_extractor(ds[0]["audio"]["array"], sampling_rate=16000, return_tensors="pt")
+
+with torch.no_grad():
+    hidden_states = model(**inputs).last_hidden_state   # (1, frames, 768), one frame every 20 ms
+```
+
+For batched input pass `padding=True` to the feature extractor; the returned `attention_mask` is used
+by the model to ignore padded samples. `output_hidden_states=True` returns the input embedding and the
+output of every Transformer layer.
+
+The original training code can also load this repository:
+
+```python
+from wavjepa.hub import load_pretrained      # github.com/labhamlet/SpeechJEPA
+model = load_pretrained("labhamlet/speech-jepa-base")          # wavjepa.jepa_d2v2.JEPA, incl. decoder and EMA teacher
+```
+
+## Fine-tuned checkpoints (CTC, LibriSpeech character vocabulary)
+
+| Labelled data | Repository |
+|---|---|
+| Libri-light 10 min | [labhamlet/speech-jepa-base-asr-10min](https://huggingface.co/labhamlet/speech-jepa-base-asr-10min) |
+| Libri-light 1 h | [labhamlet/speech-jepa-base-asr-1h](https://huggingface.co/labhamlet/speech-jepa-base-asr-1h) |
+| Libri-light 10 h | [labhamlet/speech-jepa-base-asr-10h](https://huggingface.co/labhamlet/speech-jepa-base-asr-10h) |
+| LibriSpeech train-clean-100 | [labhamlet/speech-jepa-base-asr-100h](https://huggingface.co/labhamlet/speech-jepa-base-asr-100h) |
+
+## Files
+
+* `model.safetensors`, `config.json` — `SpeechJEPAModel` weights and configuration
+* `modeling_speech_jepa.py`, `configuration_speech_jepa.py` — model code (`trust_remote_code=True`)
+* `preprocessor_config.json` — `Wav2Vec2FeatureExtractor` (16 kHz, per-utterance normalisation)
+* `jepa_pretraining_extras.safetensors` — JEPA decoder, mask token and EMA teacher (original key names)
+* `training_config.json` — hyper-parameters of the pretraining run
+
+## License
+
+BSD-3-Clause.

config.json

@@ -0,0 +1,68 @@
+{
+  "activation_dropout": 0.1,
+  "apply_spec_augment": true,
+  "architectures": [
+    "SpeechJEPAModel"
+  ],
+  "attention_dropout": 0.1,
+  "auto_map": {
+    "AutoConfig": "configuration_speech_jepa.SpeechJEPAConfig",
+    "AutoModel": "modeling_speech_jepa.SpeechJEPAModel",
+    "AutoModelForCTC": "modeling_speech_jepa.SpeechJEPAForCTC"
+  },
+  "conv_bias": false,
+  "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": "mean",
+  "ctc_zero_infinity": true,
+  "dtype": "float32",
+  "feat_extract_norm": "group",
+  "final_dropout": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-06,
+  "layerdrop": 0.05,
+  "mask_feature_length": 64,
+  "mask_feature_min_masks": 0,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_masks": 2,
+  "mask_time_prob": 0.065,
+  "max_position_embeddings": 8192,
+  "model_type": "speech_jepa",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "rope_theta": 10000.0,
+  "transformers_version": "5.2.0",
+  "use_weighted_layer_sum": false,
+  "vocab_size": 29
+}

configuration_speech_jepa.py

@@ -0,0 +1,86 @@
+"""Speech-JEPA model configuration (Hugging Face `transformers` format)."""
+from transformers import PretrainedConfig
+
+
+class SpeechJEPAConfig(PretrainedConfig):
+    r"""
+    Configuration of a Speech-JEPA model (a wav2vec 2.0 style convolutional feature encoder followed by a
+    post-norm Transformer encoder with rotary position embeddings). Field names follow
+    `Wav2Vec2Config` wherever the semantics are identical.
+    """
+
+    model_type = "speech_jepa"
+
+    def __init__(
+        self,
+        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),
+        conv_bias=False,
+        feat_extract_norm="group",
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        layer_norm_eps=1e-6,
+        rope_theta=10000.0,
+        max_position_embeddings=8192,
+        hidden_dropout=0.1,
+        activation_dropout=0.1,
+        attention_dropout=0.1,
+        final_dropout=0.1,
+        layerdrop=0.05,
+        apply_spec_augment=True,
+        mask_time_prob=0.065,
+        mask_time_length=10,
+        mask_time_min_masks=2,
+        mask_feature_prob=0.0,
+        mask_feature_length=64,
+        mask_feature_min_masks=0,
+        use_weighted_layer_sum=False,
+        vocab_size=29,
+        ctc_loss_reduction="mean",
+        ctc_zero_infinity=True,
+        pad_token_id=0,
+        initializer_range=0.02,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+        self.conv_dim = list(conv_dim)
+        self.conv_kernel = list(conv_kernel)
+        self.conv_stride = list(conv_stride)
+        self.conv_bias = conv_bias
+        self.feat_extract_norm = feat_extract_norm
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.layer_norm_eps = layer_norm_eps
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_dropout = hidden_dropout
+        self.activation_dropout = activation_dropout
+        self.attention_dropout = attention_dropout
+        self.final_dropout = final_dropout
+        self.layerdrop = layerdrop
+        self.apply_spec_augment = apply_spec_augment
+        self.mask_time_prob = mask_time_prob
+        self.mask_time_length = mask_time_length
+        self.mask_time_min_masks = mask_time_min_masks
+        self.mask_feature_prob = mask_feature_prob
+        self.mask_feature_length = mask_feature_length
+        self.mask_feature_min_masks = mask_feature_min_masks
+        self.use_weighted_layer_sum = use_weighted_layer_sum
+        self.vocab_size = vocab_size
+        self.ctc_loss_reduction = ctc_loss_reduction
+        self.ctc_zero_infinity = ctc_zero_infinity
+        self.initializer_range = initializer_range
+
+        if not (len(self.conv_dim) == len(self.conv_kernel) == len(self.conv_stride)):
+            raise ValueError("conv_dim, conv_kernel and conv_stride must have the same length")
+
+    @property
+    def num_feat_extract_layers(self):
+        return len(self.conv_dim)

jepa_pretraining_extras.safetensors

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

model.safetensors

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

modeling_speech_jepa.py

@@ -0,0 +1,401 @@
+"""PyTorch Speech-JEPA model (Hugging Face `transformers` format).
+
+Self-contained port of the encoder trained in https://github.com/labhamlet/SpeechJEPA
+(`wavjepa/jepa_d2v2.py`, `wavjepa/modules.py`). Only `torch` and `transformers` are required.
+"""
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
+
+from .configuration_speech_jepa import SpeechJEPAConfig
+
+
+# --------------------------------------------------------------------------- #
+# SpecAugment helper (identical to transformers' wav2vec2 implementation)
+# --------------------------------------------------------------------------- #
+def _compute_mask_indices(shape, mask_prob, mask_length, attention_mask=None, min_masks=0) -> np.ndarray:
+    batch_size, sequence_length = shape
+    if mask_length < 1:
+        raise ValueError("`mask_length` has to be bigger than 0.")
+    if mask_length > sequence_length:
+        raise ValueError(f"`mask_length` ({mask_length}) must be <= `sequence_length` ({sequence_length})")
+
+    epsilon = np.random.rand(1).item()
+
+    def compute_num_masked_span(input_length):
+        num_masked_span = int(mask_prob * input_length / mask_length + epsilon)
+        num_masked_span = max(num_masked_span, min_masks)
+        if num_masked_span * mask_length > sequence_length:
+            num_masked_span = sequence_length // mask_length
+        if input_length - (mask_length - 1) < num_masked_span:
+            num_masked_span = max(input_length - (mask_length - 1), 0)
+        return num_masked_span
+
+    input_lengths = (
+        attention_mask.detach().sum(-1).tolist() if attention_mask is not None else [sequence_length] * batch_size
+    )
+    spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)
+    spec_aug_mask_idxs = []
+    max_num_masked_span = compute_num_masked_span(sequence_length)
+    if max_num_masked_span == 0:
+        return spec_aug_mask
+
+    for input_length in input_lengths:
+        num_masked_span = compute_num_masked_span(input_length)
+        idx = np.random.choice(np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False)
+        dummy = sequence_length - 1 if len(idx) == 0 else idx[0]
+        idx = np.concatenate([idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy])
+        spec_aug_mask_idxs.append(idx)
+
+    spec_aug_mask_idxs = np.array(spec_aug_mask_idxs)
+    spec_aug_mask_idxs = np.broadcast_to(spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length))
+    spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length)
+    offsets = np.broadcast_to(np.arange(mask_length)[None, None, :], (batch_size, max_num_masked_span, mask_length))
+    spec_aug_mask_idxs = spec_aug_mask_idxs + offsets.reshape(batch_size, max_num_masked_span * mask_length)
+    spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1
+    np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)
+    return spec_aug_mask
+
+
+# --------------------------------------------------------------------------- #
+# convolutional feature encoder
+# --------------------------------------------------------------------------- #
+class SpeechJEPAConvLayer(nn.Module):
+    def __init__(self, config: SpeechJEPAConfig, layer_id: int):
+        super().__init__()
+        in_dim = config.conv_dim[layer_id - 1] if layer_id > 0 else 1
+        out_dim = config.conv_dim[layer_id]
+        self.conv = nn.Conv1d(in_dim, out_dim, config.conv_kernel[layer_id],
+                              stride=config.conv_stride[layer_id], bias=config.conv_bias)
+        self.layer_norm = None
+        if config.feat_extract_norm == "group" and layer_id == 0:
+            self.layer_norm = nn.GroupNorm(out_dim, out_dim, affine=True)
+        self.activation = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_states):
+        hidden_states = self.conv(hidden_states)
+        if self.layer_norm is not None:
+            hidden_states = self.layer_norm(hidden_states)
+        return self.activation(hidden_states)
+
+
+class SpeechJEPAFeatureEncoder(nn.Module):
+    def __init__(self, config: SpeechJEPAConfig):
+        super().__init__()
+        self.conv_layers = nn.ModuleList([SpeechJEPAConvLayer(config, i) for i in range(config.num_feat_extract_layers)])
+        self._requires_grad = True
+
+    def _freeze_parameters(self):
+        for p in self.parameters():
+            p.requires_grad = False
+        self._requires_grad = False
+
+    def forward(self, input_values):
+        hidden_states = input_values[:, None]  # (B, 1, T)
+        for layer in self.conv_layers:
+            hidden_states = layer(hidden_states)
+        return hidden_states  # (B, C, T')
+
+
+class SpeechJEPAFeatureProjection(nn.Module):
+    """LayerNorm(conv features) -> Linear -> LayerNorm(hidden)."""
+
+    def __init__(self, config: SpeechJEPAConfig):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm(config.conv_dim[-1])
+        self.projection = nn.Linear(config.conv_dim[-1], config.hidden_size)
+        self.post_layer_norm = nn.LayerNorm(config.hidden_size)
+
+    def forward(self, hidden_states):
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states = self.projection(hidden_states)
+        return self.post_layer_norm(hidden_states)
+
+
+# --------------------------------------------------------------------------- #
+# Transformer encoder with rotary position embeddings
+# --------------------------------------------------------------------------- #
+def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Interleaved-pair RoPE (same convention as torchtune's RotaryPositionalEmbeddings).
+
+    x: (B, S, H, D); cos/sin: (S, D // 2).
+    """
+    dtype = x.dtype
+    x = x.float().reshape(*x.shape[:-1], -1, 2)
+    cos = cos[None, :, None, :]
+    sin = sin[None, :, None, :]
+    out = torch.stack([x[..., 0] * cos - x[..., 1] * sin, x[..., 1] * cos + x[..., 0] * sin], dim=-1)
+    return out.flatten(3).to(dtype)
+
+
+class SpeechJEPARotaryEmbedding(nn.Module):
+    """cos/sin tables for interleaved RoPE. Computed on the fly (and cached as plain attributes, not buffers,
+    so that meta-device instantiation in `from_pretrained` cannot leave them uninitialised)."""
+
+    def __init__(self, dim: int, max_seq_len: int, base: float):
+        super().__init__()
+        self.dim, self.base, self.max_seq_len = dim, base, max_seq_len
+        self._cache = {}
+
+    def forward(self, seq_len: int, device):
+        key = (seq_len, str(device))
+        if key not in self._cache:
+            inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device).float() / self.dim))
+            pos = torch.arange(seq_len, dtype=torch.float32, device=device)
+            freqs = torch.outer(pos, inv_freq)
+            self._cache = {key: (freqs.cos(), freqs.sin())}
+        return self._cache[key]
+
+
+class SpeechJEPAAttention(nn.Module):
+    def __init__(self, config: SpeechJEPAConfig):
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.dropout = config.attention_dropout
+        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size)
+        self.k_proj = nn.Linear(config.hidden_size, config.hidden_size)
+        self.v_proj = nn.Linear(config.hidden_size, config.hidden_size)
+        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)
+
+    def forward(self, hidden_states, cos, sin, attention_mask=None):
+        B, S, _ = hidden_states.shape
+        q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
+        k = self.k_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
+        v = self.v_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
+        q = apply_rotary(q, cos, sin).transpose(1, 2)
+        k = apply_rotary(k, cos, sin).transpose(1, 2)
+        v = v.transpose(1, 2)
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask,
+                                             dropout_p=self.dropout if self.training else 0.0, is_causal=False)
+        return self.out_proj(out.transpose(1, 2).reshape(B, S, -1))
+
+
+class SpeechJEPAFeedForward(nn.Module):
+    def __init__(self, config: SpeechJEPAConfig):
+        super().__init__()
+        self.intermediate_dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        self.intermediate_dropout = nn.Dropout(config.activation_dropout)
+        self.output_dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.output_dropout = nn.Dropout(config.hidden_dropout)
+
+    def forward(self, hidden_states):
+        hidden_states = self.intermediate_dropout(self.intermediate_act_fn(self.intermediate_dense(hidden_states)))
+        return self.output_dropout(self.output_dense(hidden_states))
+
+
+class SpeechJEPAEncoderLayer(nn.Module):
+    """Post-norm Transformer layer: LN(x + attn(x)) -> LN(h + ffn(h))."""
+
+    def __init__(self, config: SpeechJEPAConfig):
+        super().__init__()
+        self.attention = SpeechJEPAAttention(config)
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = SpeechJEPAFeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states, cos, sin, attention_mask=None):
+        hidden_states = self.layer_norm(hidden_states + self.dropout(self.attention(hidden_states, cos, sin, attention_mask)))
+        return self.final_layer_norm(hidden_states + self.feed_forward(hidden_states))
+
+
+class SpeechJEPAEncoder(nn.Module):
+    def __init__(self, config: SpeechJEPAConfig):
+        super().__init__()
+        self.config = config
+        self.rotary = SpeechJEPARotaryEmbedding(config.hidden_size // config.num_attention_heads,
+                                                config.max_position_embeddings, config.rope_theta)
+        self.layers = nn.ModuleList([SpeechJEPAEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+
+    def forward(self, hidden_states, attention_mask=None, output_hidden_states=False, return_dict=True):
+        B, S, _ = hidden_states.shape
+        sdpa_mask = None
+        if attention_mask is not None:
+            # (B, S) 1 = real frame  ->  (B, 1, 1, S) bool, True = may attend to key
+            sdpa_mask = attention_mask.bool()[:, None, None, :]
+        cos, sin = self.rotary(S, hidden_states.device)
+
+        all_hidden_states = (hidden_states,) if output_hidden_states else None
+        for layer in self.layers:
+            if self.training and torch.rand(1).item() < self.config.layerdrop:
+                continue
+            hidden_states = layer(hidden_states, cos, sin, sdpa_mask)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in (hidden_states, all_hidden_states) if v is not None)
+        return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states)
+
+
+# --------------------------------------------------------------------------- #
+# pretrained-model wrappers
+# --------------------------------------------------------------------------- #
+class SpeechJEPAPreTrainedModel(PreTrainedModel):
+    config_class = SpeechJEPAConfig
+    base_model_prefix = "speech_jepa"
+    main_input_name = "input_values"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.Conv1d):
+            nn.init.kaiming_normal_(module.weight)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, SpeechJEPAModel):
+            module.masked_spec_embed.data.uniform_()
+
+    def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
+        for kernel, stride in zip(self.config.conv_kernel, self.config.conv_stride):
+            input_lengths = torch.div(input_lengths - kernel, stride, rounding_mode="floor") + 1
+        return input_lengths
+
+    def _get_feature_vector_attention_mask(self, feature_vector_length: int, attention_mask: torch.LongTensor):
+        output_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long)
+        batch_size = attention_mask.shape[0]
+        mask = torch.zeros((batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device)
+        mask[(torch.arange(batch_size, device=mask.device), output_lengths - 1)] = 1
+        return mask.flip([-1]).cumsum(-1).flip([-1]).bool()
+
+
+class SpeechJEPAModel(SpeechJEPAPreTrainedModel):
+    def __init__(self, config: SpeechJEPAConfig):
+        super().__init__(config)
+        self.config = config
+        self.feature_extractor = SpeechJEPAFeatureEncoder(config)
+        self.feature_projection = SpeechJEPAFeatureProjection(config)
+        self.masked_spec_embed = nn.Parameter(torch.Tensor(config.hidden_size).uniform_())
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.encoder = SpeechJEPAEncoder(config)
+        self.post_init()
+
+    def freeze_feature_encoder(self):
+        self.feature_extractor._freeze_parameters()
+
+    def _mask_hidden_states(self, hidden_states, mask_time_indices=None, attention_mask=None):
+        if not getattr(self.config, "apply_spec_augment", True):
+            return hidden_states
+        batch_size, sequence_length, hidden_size = hidden_states.size()
+        if mask_time_indices is not None:
+            hidden_states[mask_time_indices] = self.masked_spec_embed.to(hidden_states.dtype)
+        elif self.config.mask_time_prob > 0 and self.training:
+            idx = _compute_mask_indices((batch_size, sequence_length), self.config.mask_time_prob,
+                                        self.config.mask_time_length, attention_mask=attention_mask,
+                                        min_masks=self.config.mask_time_min_masks)
+            idx = torch.tensor(idx, device=hidden_states.device, dtype=torch.bool)
+            hidden_states[idx] = self.masked_spec_embed.to(hidden_states.dtype)
+        if self.config.mask_feature_prob > 0 and self.training:
+            idx = _compute_mask_indices((batch_size, hidden_size), self.config.mask_feature_prob,
+                                        self.config.mask_feature_length, min_masks=self.config.mask_feature_min_masks)
+            idx = torch.tensor(idx, device=hidden_states.device, dtype=torch.bool)[:, None].expand(-1, sequence_length, -1)
+            hidden_states[idx] = 0
+        return hidden_states
+
+    def forward(
+        self,
+        input_values: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        mask_time_indices: Optional[torch.BoolTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        """``input_values``: (B, T) waveform at 16 kHz, normalised to zero mean / unit variance per utterance
+        (use ``Wav2Vec2FeatureExtractor`` with ``do_normalize=True``). ``attention_mask``: (B, T), 1 = real sample."""
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        features = self.feature_extractor(input_values).transpose(1, 2)  # (B, T', C)
+        hidden_states = self.feature_projection(features)
+
+        frame_mask = None
+        if attention_mask is not None:
+            frame_mask = self._get_feature_vector_attention_mask(hidden_states.shape[1], attention_mask)
+
+        hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices, attention_mask=frame_mask)
+        hidden_states = self.dropout(hidden_states)
+
+        encoder_outputs = self.encoder(hidden_states, attention_mask=frame_mask,
+                                       output_hidden_states=output_hidden_states, return_dict=return_dict)
+        if not return_dict:
+            return encoder_outputs
+        return BaseModelOutput(last_hidden_state=encoder_outputs.last_hidden_state, hidden_states=encoder_outputs.hidden_states)
+
+
+class SpeechJEPAForCTC(SpeechJEPAPreTrainedModel):
+    def __init__(self, config: SpeechJEPAConfig):
+        super().__init__(config)
+        self.speech_jepa = SpeechJEPAModel(config)
+        self.dropout = nn.Dropout(config.final_dropout)
+        if config.vocab_size is None:
+            raise ValueError("`config.vocab_size` must be set for SpeechJEPAForCTC")
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
+        if config.use_weighted_layer_sum:
+            self.layer_weights = nn.Parameter(torch.ones(config.num_hidden_layers + 1) / (config.num_hidden_layers + 1))
+        self.post_init()
+
+    def freeze_feature_encoder(self):
+        self.speech_jepa.feature_extractor._freeze_parameters()
+
+    def freeze_base_model(self):
+        for p in self.speech_jepa.parameters():
+            p.requires_grad = False
+
+    def forward(
+        self,
+        input_values: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutput]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_hidden_states = True if self.config.use_weighted_layer_sum else output_hidden_states
+
+        outputs = self.speech_jepa(input_values, attention_mask=attention_mask,
+                                   output_hidden_states=output_hidden_states, return_dict=True)
+        if self.config.use_weighted_layer_sum:
+            hidden_states = torch.stack(outputs.hidden_states, dim=1)
+            norm_weights = F.softmax(self.layer_weights, dim=-1)
+            hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1)
+        else:
+            hidden_states = outputs.last_hidden_state
+
+        logits = self.lm_head(self.dropout(hidden_states))
+
+        loss = None
+        if labels is not None:
+            if labels.max() >= self.config.vocab_size:
+                raise ValueError(f"Label values must be <= vocab_size: {self.config.vocab_size}")
+            attention_mask = attention_mask if attention_mask is not None else torch.ones_like(input_values, dtype=torch.long)
+            input_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long)
+            labels_mask = labels >= 0
+            target_lengths = labels_mask.sum(-1)
+            flattened_targets = labels.masked_select(labels_mask)
+            log_probs = F.log_softmax(logits, dim=-1, dtype=torch.float32).transpose(0, 1)
+            with torch.backends.cudnn.flags(enabled=False):
+                loss = F.ctc_loss(log_probs, flattened_targets, input_lengths, target_lengths,
+                                  blank=self.config.pad_token_id, reduction=self.config.ctc_loss_reduction,
+                                  zero_infinity=self.config.ctc_zero_infinity)
+
+        if not return_dict:
+            output = (logits,) + ((outputs.hidden_states,) if output_hidden_states else ())
+            return ((loss,) + output) if loss is not None else output
+        return CausalLMOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states if output_hidden_states else None)
+
+
+__all__ = ["SpeechJEPAConfig", "SpeechJEPAPreTrainedModel", "SpeechJEPAModel", "SpeechJEPAForCTC"]

preprocessor_config.json

@@ -0,0 +1,9 @@
+{
+  "do_normalize": true,
+  "feature_extractor_type": "Wav2Vec2FeatureExtractor",
+  "feature_size": 1,
+  "padding_side": "right",
+  "padding_value": 0.0,
+  "return_attention_mask": true,
+  "sampling_rate": 16000
+}

training_config.json

@@ -0,0 +1,36 @@
+{
+  "transformer_encoder_layers_cfg": {
+    "d_model": 768,
+    "nhead": 12,
+    "batch_first": true,
+    "norm_first": false,
+    "bias": true,
+    "dim_feedforward": 3072,
+    "dropout": 0.0,
+    "activation": "GELU(approximate='none')",
+    "layer_norm_eps": 1e-06
+  },
+  "transformer_encoder_cfg": {
+    "num_layers": 12,
+    "enable_nested_tensor": false,
+    "mask_check": true
+  },
+  "lr": 0.0004,
+  "adam_betas": [
+    0.9,
+    0.98
+  ],
+  "adam_eps": 1e-06,
+  "adam_weight_decay": 0.01,
+  "ema_decay": 0.999,
+  "ema_end_decay": 0.99999,
+  "ema_anneal_end_step": 100000,
+  "warmup_steps": 100000,
+  "average_top_k_layers": 8,
+  "resample_sr": 16000,
+  "original_sr": 16000,
+  "use_gradient_checkpointing": false,
+  "compile_modules": false,
+  "size": "base",
+  "clean_audio_ratio": 0.0
+}