Blog

import pathlib
from itertools import groupby

from datasets import load_dataset
import librosa
import numpy as np
import onnx
import onnxruntime as ort
import torch
import transformers


def get_model():
    """
    Get the Wav2Vec2.0 model.

    The processor will be a Wav2Vec2Processor object.

    Source code: https://github.com/huggingface/transformers/tree/main/src/transformers/models/wav2vec2
    Looking at the code, the processor for audio may only be a feature extractor.

    :return tuple: Model and processor.
    """
    checkpoint = "bookbot/wav2vec2-ljspeech-gruut"

    model = transformers.AutoModelForCTC.from_pretrained(checkpoint)
    processor = transformers.AutoProcessor.from_pretrained(checkpoint)
    return model, processor

def get_audio_data(local_file=None):
    """Read an audio file as an audio array."""
    if local_file is None:
        # load dummy dataset and read soundfiles
        ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
        audio_array = ds[0]["audio"]["array"]
    else:
        # or, read a single audio file
        audio_array, _ = librosa.load(local_file, sr=16000)
    return audio_array

def get_logits(audio_array, processor, model, sampling_rate=16000):
    """
    Get the logits for the audio.

    For Unity: https://discussions.unity.com/t/how-to-use-logits-output-type/1540746

    :param audio_array: Input audio
    :param processor: Processor or feature extractor for the input.
    :param model: Actual model.
    :param sampling_rate: Samplig rate
    :return:
    """
    # On easy cases, inputs = (audio_array - audio_array.mean()) / sqrt(audio_array.var())
    inputs = processor(
        audio_array,
        return_tensors="pt",
        padding=True,
        sampling_rate=sampling_rate
    )

    with torch.no_grad():
        logits = model(inputs["input_values"]).logits
    return logits


def decode_phonemes(
    ids: torch.Tensor,
    processor: transformers.Wav2Vec2Processor,
    ignore_stress: bool = False
) -> str:
    """
    CTC-like decoding.
    First removes consecutive duplicates, then removes special tokens.
    """
    # removes consecutive duplicates
    ids = [id_ for id_, _ in groupby(ids)]

    special_token_ids = processor.tokenizer.all_special_ids + [
        processor.tokenizer.word_delimiter_token_id
    ]
    # converts id to token, skipping special tokens
    phonemes = [processor.decode(id_) for id_ in ids if id_ not in special_token_ids]

    # joins phonemes
    prediction = " ".join(phonemes)

    # whether to ignore IPA stress marks
    if ignore_stress:
        prediction = prediction.replace("ˈ", "").replace("ˌ", "")

    return prediction


def predicted_answer(logits, processor):
    """From a matrix of logits, get the best prediction match."""
    predicted_ids = torch.argmax(logits, dim=-1)
    my_prediction = decode_phonemes(predicted_ids[0], processor, ignore_stress=True)
    return my_prediction


def run_analysis():
    """Run the model."""
    model, processor = get_model()
    sampling_rate = processor.feature_extractor.sampling_rate
    audio_array = get_audio_data()
    logits = get_logits(
        audio_array,
        processor,
        model,
        sampling_rate
    )
    prediction = predicted_answer(logits, processor)
    return prediction

 
def convert_to_onnx(onnx_model_name="wav2vec2.onnx"):
    """
    Convert a Wav2Vec2 model using a low-level method.

    https://github.com/ccoreilly/wav2vec2-service/blob/master/convert_torch_to_onnx.py

    :param string onnx_model_name: Name of the output file.
    """
    model, _ = get_model()
    audio_len = 250000

    model_input = torch.randn(1, audio_len, requires_grad=True)

    torch.onnx.export(
        model,                          # model being run
        model_input,                    # model input (or a tuple for multiple inputs)
        onnx_model_name,                # where to save the model (can be a file or file-like object)
        export_params=True,             # store the trained parameter weights inside the model file
        opset_version=15,               # the ONNX version to export the model to, 15 is recommended version for Unity
        do_constant_folding=True,       # whether to execute constant folding for optimization
        input_names=['input'],          # the model's input names
        output_names=['output'],        # the model's output names
        dynamic_axes={
            'input': {1: 'audio_len'},    # variable length axes
            'output': {1: 'audio_len'}
        }
    )
    print(f"Model saved to {onnx_model_name} successfully!")

def run_from_onnx(audio_array, model_file="wav2vec2.onnx", check_model=False):
    """
    Run the Wav2Vec2.0 model from an ONNX file.
     
    :param list[float] audio_array: Input data, at 16 kHz.
    :param string model_file: Path to the ONNX file.  
    :param bool check_model: To apply supplementary check to the model file. 
    :return string: Predicted phonemes. 
    """
    audio_array = np.array(audio_array, dtype=np.float32).reshape((1, -1))
    if check_model:
        onnx_model = onnx.load(model_file)
        onnx.checker.check_model(onnx_model)

    ort_sess = ort.InferenceSession(model_file)
    logits = ort_sess.run(None, {'input': audio_array})[0]
    
    # The following is to pass from the logits to the prediction, not the hardest part
    _, processor = get_model()
    prediction = predicted_answer(torch.as_tensor(logits), processor)
    return prediction