whisper_eval.py

import os
from pathlib import Path
from typing import Optional, Tuple

import numpy as np
import onnxruntime as onnxrt
import torch
from datasets import load_dataset
from transformers import (
    AutoConfig,
    AutoProcessor,
    GenerationConfig,
    WhisperForConditionalGeneration,
)
from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput


os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"


model_name = "openai/whisper-tiny.en"
config = AutoConfig.from_pretrained(model_name)
processor = AutoProcessor.from_pretrained(model_name)

batch_size = 1
encoder_num_attention_heads = 6
decoder_num_attention_heads = 6
hidden_size = 384
encoder_sequence_length = 1500
decoder_max_length = 448
num_hidden_layers = 4

encoder_shape = (
    batch_size,
    encoder_num_attention_heads,
    encoder_sequence_length,
    hidden_size // encoder_num_attention_heads,
)
decoder_shape = (
    batch_size,
    decoder_num_attention_heads,
    decoder_max_length,
    hidden_size // decoder_num_attention_heads,
)


# load dataset
ds = load_dataset(
    "hf-internal-testing/librispeech_asr_dummy", "clean", split="validation"
)
idx = 4
inputs = processor.feature_extractor(ds[idx]["audio"]["array"], return_tensors="pt")
input_features = inputs.input_features


# onnx_model_path = "/home/ubuntu/optimum/output_whisper_smooth_quant_4_oct_static_testing"
onnx_model_path = ".\\whisper-tiny-static-shape-quantized-SL-448"
config_file = ".\\other_libs_qdq\\vaip_config_gemm_asr_decoder.json"
encoder_model_path = ".\\whisper-tiny-static-shape-quantized-SL-448\\encoder_model.onnx"
decoder_model_path = ".\\whisper-tiny-static-shape-quantized-SL-448\\decoder_model_quantized.onnx"

print(decoder_model_path)


class ORTEncoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.main_input_name = "input_features"
        self.session = onnxrt.InferenceSession(
            encoder_model_path, providers=["CPUExecutionProvider"]
        )
        self.output_names = {
            output_key.name: idx
            for idx, output_key in enumerate(self.session.get_outputs())
        }

    def forward(
        self,
        input_features: torch.FloatTensor,
        **kwargs,
    ) -> BaseModelOutput:
        onnx_inputs = {"input_features": input_features.cpu().detach().numpy()}

        # Run inference
        outputs = self.session.run(None, onnx_inputs)
        last_hidden_state = torch.from_numpy(
            outputs[self.output_names["last_hidden_state"]]
        )

        return BaseModelOutput(last_hidden_state=last_hidden_state)


class ORTDecoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        sess_options = onnxrt.SessionOptions()
        self.provider = "VitisAIExecutionProvider"
        self.provider_options = {"config_file": config_file}
        sess_options.graph_optimization_level = (
            onnxrt.GraphOptimizationLevel.ORT_DISABLE_ALL
        )
        sess_options.add_session_config_entry("session.disable_quant_qdq", "1")
        self.session = onnxrt.InferenceSession(
            decoder_model_path,
            providers=[self.provider],
            sess_options=sess_options,
            provider_options=[self.provider_options],
        )

        self.generation_config = GenerationConfig.from_model_config(config)
        self.max_length = decoder_max_length

        self.input_names = {
            input_key.name: idx
            for idx, input_key in enumerate(self.session.get_inputs())
        }
        self.output_names = {
            output_key.name: idx
            for idx, output_key in enumerate(self.session.get_outputs())
        }
        self.key_value_input_names = [
            key for key in self.input_names if (".key" in key) or (".value" in key)
        ]
        self.key_value_output_names = [
            key for key in self.output_names if (".key" in key) or (".value" in key)
        ]

        self.reset()

    def reset(self):
        # Set the start model inputs
        self.decoder_attention_mask = np.zeros((batch_size, self.max_length)).astype(
            np.int64
        )
        self.decoder_attention_mask[0, 0] = 1
        self.position_ids = np.array([[0]]).astype(np.int64)

        # Set the input / output names
        self.num_pkv = 4

    def prepare_pkv(self):
        decoder_key_value = torch.rand(*decoder_shape).to(torch.float32)
        encoder_key_value = torch.rand(*encoder_shape).to(torch.float32)

        past_key_values = []
        repeat_count = len(self.key_value_input_names) // 4
        past_key_values = tuple(
            (decoder_key_value, decoder_key_value, encoder_key_value, encoder_key_value)
            for _ in range(repeat_count)
        )

        return tuple(past_key_values)

    def forward(
        self,
        input_ids: torch.LongTensor,
        encoder_hidden_states: torch.FloatTensor,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
    ) -> Seq2SeqLMOutput:
        if past_key_values is None:
            self.reset()

        if self.position_ids[0][0] == self.max_length:
            logits = torch.zeros((len(input_ids), 1, config.vocab_size))
            logits[:, :, config.eos_token_id] = 1

            return Seq2SeqLMOutput(logits=logits, past_key_values=past_key_values)

        onnx_inputs = {"input_ids": input_ids.cpu().detach().numpy()}

        onnx_inputs["position_ids"] = self.position_ids
        onnx_inputs["decoder_attention_mask"] = self.decoder_attention_mask
        onnx_inputs["encoder_hidden_states"] = (
            encoder_hidden_states.cpu().detach().numpy()
        )

        if self.position_ids[0][0] == 0:
            past_key_values = self.prepare_pkv()

        past_key_values = tuple(
            past_key_value
            for pkv_per_layer in past_key_values
            for past_key_value in pkv_per_layer
        )

        for input_name, past_key_value in zip(
            self.key_value_input_names, past_key_values
        ):
            onnx_inputs[input_name] = past_key_value.cpu().detach().numpy()

        # Run inference
        outputs = self.session.run(None, onnx_inputs)

        logits = torch.from_numpy(outputs[self.output_names["logits"]])

        out_past_key_values = tuple(
            torch.from_numpy(outputs[self.output_names[key]])
            for key in self.key_value_output_names
        )

        if self.position_ids[0][0] == 0:
            out_past_key_values = tuple(
                out_past_key_values[i : i + self.num_pkv]
                for i in range(0, len(out_past_key_values), self.num_pkv)
            )
        else:
            out_past_key_values = tuple(
                out_past_key_values[i : i + 2] + past_key_values[i + 2 : i + 4]
                for i in range(0, len(out_past_key_values), self.num_pkv)
            )

        if self.position_ids[0][0] < self.max_length - 1:
            self.decoder_attention_mask[:, self.position_ids[0][0] + 1] = 1
        self.position_ids += 1

        return Seq2SeqLMOutput(logits=logits, past_key_values=out_past_key_values)


class ORTModelForWhisper(WhisperForConditionalGeneration):
    def __init__(self, *args, **kwargs):
        config = AutoConfig.from_pretrained(model_name)
        super().__init__(config)

        self.encoder = ORTEncoder()
        self.decoder = ORTDecoder()

    def get_encoder(self):
        return self.encoder

    def forward(
        self,
        input_features: Optional[torch.FloatTensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        **kwargs,
    ) -> Seq2SeqLMOutput:
        if encoder_outputs is None:
            encoder_outputs = self.encoder(input_features=input_features)

        # Decode
        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids[:, -1:],
            encoder_hidden_states=encoder_outputs.last_hidden_state,
            past_key_values=past_key_values,
        )

        return Seq2SeqLMOutput(
            logits=decoder_outputs.logits,
            past_key_values=decoder_outputs.past_key_values,
        )

    def can_generate(self):
        return True

    def reset(self):
        self.decoder.reset()


model_ort = ORTModelForWhisper()
model = WhisperForConditionalGeneration.from_pretrained(model_name)


def test_ort():
    model = ORTModelForWhisper()

    generated_ids = model.generate(input_features)
    model_output = processor.tokenizer.batch_decode(
        generated_ids, skip_special_tokens=True
    )[0]

    print("ORT: ", model_output, generated_ids)


def test_original():
    model = WhisperForConditionalGeneration.from_pretrained(model_name)

    generated_ids = model.generate(input_features)
    model_output = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

    print("Torch: ", model_output, generated_ids)


test_ort()