Add 🤗 Transformers support

README.md

@@ -27,6 +27,7 @@ tags:
 - pytorch
 - NeMo
 - hf-asr-leaderboard
+- transformers
 license: cc-by-4.0
 widget:
 - example_title: Librispeech sample 1
@@ -197,25 +198,27 @@ NVIDIA Developer [Nemotron](https://developer.nvidia.com/nemotron)<br>
 [NVIDIA Riva Speech](https://developer.nvidia.com/riva?sortBy=developer_learning_library%2Fsort%2Ffeatured_in.riva%3Adesc%2Ctitle%3Aasc#demos)<br>
 [NeMo Documentation](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/models.html)<br>
 
-## NVIDIA NeMo: Training
+## How to Use this Model
+
+The model is available for use in the NeMo toolkit [3], and can be used as a pre-trained checkpoint for inference or for fine-tuning on another dataset.
+
+You can also run Parakeet RNNT with [Transformers](https://github.com/huggingface/transformers) 🤗 (more below).
+
+### 1) NeMo usage
 
 To train, fine-tune or play with the model you will need to install [NVIDIA NeMo](https://github.com/NVIDIA/NeMo). We recommend you install it after you've installed latest PyTorch version.
 ```
 pip install nemo_toolkit['all']
 ``` 
 
-## How to Use this Model
-
-The model is available for use in the NeMo toolkit [3], and can be used as a pre-trained checkpoint for inference or for fine-tuning on another dataset.
-
-### Automatically instantiate the model
+#### Automatically instantiate the model
 
 ```python
 import nemo.collections.asr as nemo_asr
 asr_model = nemo_asr.models.EncDecRNNTBPEModel.from_pretrained(model_name="nvidia/parakeet-rnnt-0.6b")
 ```
 
-### Transcribing using Python
+#### Transcribing using Python
 First, let's get a sample
 ```
 wget https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav
@@ -226,7 +229,7 @@ output = asr_model.transcribe(['2086-149220-0033.wav'])
 print(output[0].text)
 ```
 
-### Transcribing many audio files
+#### Transcribing many audio files
 
 ```shell
 python [NEMO_GIT_FOLDER]/examples/asr/transcribe_speech.py 
@@ -234,6 +237,114 @@ python [NEMO_GIT_FOLDER]/examples/asr/transcribe_speech.py
  audio_dir="<DIRECTORY CONTAINING AUDIO FILES>"
 ```
 
+### 2) [Transformers](https://github.com/huggingface/transformers) 🤗 usage 
+
+
+Until Parakeet RNNT is part of an official Transformers release, you can use it by installing from source.
+
+```bash
+pip install git+https://github.com/huggingface/transformers
+```
+
+<details>
+  <summary>➡️ Pipeline usage</summary>
+
+```python
+from transformers import pipeline
+
+pipe = pipeline("automatic-speech-recognition", model="eustlb/parakeet-rnnt-0.6b")
+out = pipe("https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/bcn_weather.mp3")
+print(out)
+```
+</details>
+
+<details>
+  <summary>➡️ AutoModel</summary>
+
+```python
+from transformers import AutoModelForRNNT, AutoProcessor
+from datasets import load_dataset, Audio
+
+num_samples = 3
+
+model_id = "eustlb/parakeet-rnnt-0.6b"
+processor = AutoProcessor.from_pretrained(model_id)
+model = AutoModelForRNNT.from_pretrained(model_id, dtype="auto", device_map="auto")
+
+ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+speech_samples = [el["array"] for el in ds["audio"][:num_samples]]
+
+inputs = processor(speech_samples, sampling_rate=processor.feature_extractor.sampling_rate)
+inputs.to(model.device, dtype=model.dtype)
+output = model.generate(**inputs, return_dict_in_generate=True)
+print(processor.decode(output.sequences, skip_special_tokens=True))
+```
+</details>
+
+
+<details>
+  <summary>➡️ Timestamping</summary>
+
+```python
+from datasets import Audio, load_dataset
+from transformers import AutoModelForRNNT, AutoProcessor
+
+num_samples = 3
+
+model_id = "eustlb/parakeet-rnnt-0.6b"
+processor = AutoProcessor.from_pretrained(model_id)
+model = AutoModelForRNNT.from_pretrained(model_id, dtype="auto", device_map="auto")
+
+ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+speech_samples = [el["array"] for el in ds["audio"][:num_samples]]
+
+inputs = processor(speech_samples, sampling_rate=processor.feature_extractor.sampling_rate)
+inputs.to(model.device, dtype=model.dtype)
+output = model.generate(**inputs, return_dict_in_generate=True)
+decoded_output, decoded_timestamps = processor.decode(
+    output.sequences,
+    durations=output.durations,
+    skip_special_tokens=True,
+)
+print("Transcription:", decoded_output)
+print("Timestamped tokens:", decoded_timestamps)
+```
+</details>
+
+<details>
+  <summary>➡️ Training</summary>
+
+```python
+from transformers import AutoModelForRNNT, AutoProcessor
+from datasets import load_dataset, Audio
+import torch
+
+model_id = "eustlb/parakeet-rnnt-0.6b"
+NUM_SAMPLES = 4
+
+processor = AutoProcessor.from_pretrained(model_id)
+model = AutoModelForRNNT.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
+model.train()
+
+ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+ds = ds.cast_column("audio", Audio(sampling_rate=processor.feature_extractor.sampling_rate))
+speech_samples = [el["array"] for el in ds["audio"][:NUM_SAMPLES]]
+text_samples = ds["text"][:NUM_SAMPLES]
+
+# passing `text` to the processor will prepare inputs' `labels` key
+inputs = processor(audio=speech_samples, text=text_samples, sampling_rate=processor.feature_extractor.sampling_rate)
+inputs.to(device=model.device, dtype=model.dtype)
+
+outputs = model(**inputs)
+print("Loss:", outputs.loss.item())
+outputs.loss.backward()
+```
+</details>
+
+For more details about usage, please refer to the [Transformers' documentation](https://huggingface.co/docs/transformers/en/model_doc/parakeet).
+
 ### Input
 
 This model accepts 16000 Hz mono-channel audio (wav files) as input.
@@ -303,6 +414,4 @@ Check out [Riva live demo](https://developer.nvidia.com/riva#demos).
 
 [4] [Suno.ai](https://suno.ai/)
 
-[5] [HuggingFace ASR Leaderboard](https://huggingface.co/spaces/hf-audio/open_asr_leaderboard)
-
-
+[5] [HuggingFace ASR Leaderboard](https://huggingface.co/spaces/hf-audio/open_asr_leaderboard)

config.json

@@ -0,0 +1,42 @@
+{
+  "architectures": [
+    "ParakeetForRNNT"
+  ],
+  "blank_token_id": 1024,
+  "decoder_hidden_size": 640,
+  "dtype": "float32",
+  "encoder_config": {
+    "activation_dropout": 0.1,
+    "attention_bias": true,
+    "attention_dropout": 0.1,
+    "conv_kernel_size": 9,
+    "convolution_bias": true,
+    "dropout": 0.1,
+    "dropout_positions": 0.0,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "layerdrop": 0.1,
+    "max_position_embeddings": 5000,
+    "model_type": "parakeet_encoder",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 8,
+    "num_mel_bins": 80,
+    "scale_input": true,
+    "subsampling_conv_channels": 256,
+    "subsampling_conv_kernel_size": 3,
+    "subsampling_conv_stride": 2,
+    "subsampling_factor": 8
+  },
+  "hidden_act": "relu",
+  "initializer_range": 0.02,
+  "is_encoder_decoder": true,
+  "max_symbols_per_step": 10,
+  "model_type": "parakeet_rnnt",
+  "num_decoder_layers": 2,
+  "pad_token_id": 0,
+  "transformers_version": "5.10.0.dev0",
+  "vocab_size": 1025
+}

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "_from_model_config": true,
+  "decoder_start_token_id": 1024,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 0,
+  "transformers_version": "5.10.0.dev0"
+}

model.safetensors

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

processor_config.json

@@ -0,0 +1,17 @@
+{
+  "blank_token": "<blank>",
+  "decoder_type": "rnnt",
+  "feature_extractor": {
+    "feature_extractor_type": "ParakeetFeatureExtractor",
+    "feature_size": 80,
+    "hop_length": 160,
+    "n_fft": 512,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "preemphasis": 0.97,
+    "return_attention_mask": true,
+    "sampling_rate": 16000,
+    "win_length": 400
+  },
+  "processor_class": "ParakeetProcessor"
+}

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<blank>",
+  "processor_class": "ParakeetProcessor",
+  "tokenizer_class": "ParakeetTokenizer",
+  "unk_token": "<unk>"
+}