| --- |
| 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. |
|
|
