When converting Wav2Vec2 CTC models, it is recommended to bake feature normalization (do_normalize) directly into the ONNX graph. This allows the model to accept raw waveforms without requiring a separate feature-extractor asset.
Key Conversion Details:
- Vocabulary Mapping: The pad token must be mapped to
<blk> (the CTC blank token). The word-delimiter token (|) should be mapped to a literal space ( ). - Subsampling Factor: This must be included in
config.json. It is calculated as the product of the strides in the feature-encoder convolutional layers (e.g., 320 for standard XLS-R stacks). - Input/Output Names: The graph should use
input_values and input_lengths as inputs, and logprobs as the output.
import json
from pathlib import Path
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
hf_model_id = "proxectonos/Nos_ASR-wav2vec2-xls-r-300m-gl"
onnx_dir = Path("wav2vec2-onnx")
onnx_dir.mkdir(exist_ok=True)
model = Wav2Vec2ForCTC.from_pretrained(hf_model_id).eval()
processor = Wav2Vec2Processor.from_pretrained(hf_model_id)
do_normalize = bool(processor.feature_extractor.do_normalize)
class NormalizedWav2Vec2(torch.nn.Module):
def forward(self, input_values, input_lengths):
if do_normalize:
mask = torch.arange(input_values.shape[1])[None, :] < input_lengths[:, None]
count = input_lengths.to(input_values.dtype).clamp(min=1)[:, None]
mean = (input_values * mask).sum(dim=1, keepdim=True) / count
var = (((input_values - mean) * mask) ** 2).sum(dim=1, keepdim=True) / count
input_values = torch.where(mask, (input_values - mean) / torch.sqrt(var + 1e-5), input_values)
logits = model(input_values).logits
return torch.nn.functional.log_softmax(logits, dim=-1)
torch.onnx.export(
NormalizedWav2Vec2(),
(torch.randn(1, 16000), torch.tensor([16000], dtype=torch.int64)),
str(onnx_dir / "model.onnx"),
input_names=["input_values", "input_lengths"],
output_names=["logprobs"],
dynamic_axes={
"input_values": {0: "batch", 1: "time"},
"input_lengths": {0: "batch"},
"logprobs": {0: "batch", 1: "frames"},
},
opset_version=18,
)
vocab = processor.tokenizer.get_vocab()
pad_token = processor.tokenizer.pad_token
word_delimiter = processor.tokenizer.word_delimiter_token
with (onnx_dir / "vocab.txt").open("wt") as f:
for token, idx in sorted(vocab.items(), key=lambda kv: kv[1]):
token = "<blk>" if token == pad_token else " " if token == word_delimiter else token
f.write(f"{token} {idx}\n")
subsampling_factor = 1
for layer in model.wav2vec2.feature_extractor.conv_layers:
stride = layer.conv.stride
subsampling_factor *= stride[0] if isinstance(stride, tuple) else stride
with (onnx_dir / "config.json").open("wt") as f:
json.dump({"model_type": "wav2vec2-ctc", "subsampling_factor": subsampling_factor}, f, indent=2)
