Multilingual models converted to TFLite doesn't work

[shadowofsoul]

Hello everyone!

I'm here asking for some guidance. I'm trying to convert the models to TFLite to be used on mobile.
I based my work so far in this discussion and repo: #11 created by the awesome @nyadla-sys

*.en based models work, but not the multilingual ones.
I'm converting the models with hugging face transformers using the following script:

import tensorflow as tf
from transformers import TFWhisperModel,TFWhisperForConditionalGeneration, WhisperFeatureExtractor
from datasets import load_dataset

model = TFWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-tiny")
ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
inputs = feature_extractor(
    ds[0]["audio"]["array"], sampling_rate=ds[0]["audio"]["sampling_rate"], return_tensors="tf"
)
input_features = inputs.input_features
print(input_features)
decoder_input_ids = tf.convert_to_tensor([[1, 1]]) * model.config.decoder_start_token_id
model.save('tf_whisper_saved')

class GenerateModel(tf.Module):
  def __init__(self, model):
    super(GenerateModel, self).__init__()
    self.model = model

  @tf.function(
    # shouldn't need static batch size, but throws exception without it (needs to be fixed)
    input_signature=[
      tf.TensorSpec((1, 80, 3000), tf.float32, name="input_features"), 
    ],
  )
  def serving(self, input_features):
    outputs = self.model.generate(
      input_features,
      max_new_tokens=223, #change as needed
      return_dict_in_generate=True,
    )
    return {"sequences": outputs["sequences"]}

saved_model_dir = 'tf_whisper_saved'
tflite_model_path = 'whisper-tiny.tflite'

generate_model = GenerateModel(model=model)
tf.saved_model.save(generate_model, saved_model_dir, signatures={"serving_default": generate_model.serving})

# Convert the model
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
  tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.
]
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

# Save the model
with open(tflite_model_path, 'wb') as f:
    f.write(tflite_model)

Testing the model on python with TFLite (using interpreter) gives me a bad translation (gibberbish). Testing it on android with TFLite it dies with:

 E  gather index out of bounds
 E  Node number 35 (GATHER) failed to invoke.
 E  Node number 694 (WHILE) failed to invoke.

It looks like the model conversion is not correct with multilingual models. Any hint/tip what i'm doing wrong?

Thanks in advance!

[nyadla-sys]

@shadowofsoul Could you please change the line in the above notebook that reads 'max_new_tokens=223, #change as needed' to 'max_new_tokens=447, #change as needed'?

[dkakaie]

I have the same problem. Tried this and problem still persists with gibberish output running TFLite model eventhough TFWhisperForConditionalGeneration seems OK.

[shadowofsoul]

I can confirm. I think that the problem is in the conversion to PB model from the .save() function of the transformer.

If anyone has some link to know how to parse the FLOAT32 tensors that you get from the encoder to the decoder that is expecting INT64 as the decoder_ids after you run the inference with invoke(), i could try to test tflite models generated from the ONNX conversion maybe get a encoder-decoder code working? that was my last try on this.

[dkakaie]

If it helps, I used ORTModelForSpeechSeq2Seq from optimum, converted Whisper models previously finetuned by Transformers to ONNX. It creates an encoder_model.onnx, decoder_model.onnx and decoder_with_past_model.onnx and used this as a port in C#. Audio gets padded and is fed into encoder to obtain last_hidden_state which is in turn fed into decoder_model with start or previous tokens recursively until reaching eot or a defined max generation.
model = ORTModelForSpeechSeq2Seq.from_pretrained(SOME_MODEL_ID, from_transformers=True, use_cache=True) model.save_pretrained(SOME_FOLDER)

[mingryon]

Hi researcher. I'm very interested in multilingual model. I tried to test whisper-tiny.tflite model, but it doesn't work. Do we have any method to convert specified language to tflite model like english model? I've tried decode & encode method, but speed is very slow.
I hope the answer for converting specified language to tflite model. For example, Italian.

[nyadla-sys]

@dkakaie Could you provide a notebook that demonstrates how to generate the encoder.onnx and decoder.onnx files and perform inference with them?

[dkakaie]

Sure, I will although it may take some time as I'm on the move at the moment.

[dkakaie]

A quick and dirty run available here:

[dkakaie]

@nyadla-sys Any updates? Is the notebook running correctly?

[nyadla-sys]

@dkakaie NOtebook works correctly.

[mingryon]

Hi @nyadla-sys . I'm very interested in multilingual model. I tried to test whisper-tiny.tflite model, but it doesn't work. Do we have any method to convert specified language to tflite model like english model? I've tried decode & encode method, but speed is very slow.
I hope the answer for converting specified language to tflite model. For example, Italian.

[nyadla-sys]

@dkakaie i managed to generate encoder and decoder tflite models and just pending to complete the decoder post processing to generate tokens with text
https://colab.research.google.com/github/usefulsensors/openai-whisper/blob/main/notebooks/whisper_encoder_decoder_tflite.ipynb

If you could help in post-processing that's appreciated.

[dkakaie]

OK, I exported decoder to have its shape dynamic. Will inform you of any updates after testing it.

[nyadla-sys]

@dkakaie Great ,thank you so much

[nyadla-sys]

I would greatly appreciate it if you could share the code for exporting it with dynamic shape, as I would like to try it out as well.

[shadowofsoul]

@dkakaie i ask for the same. I'm looking at the tokenizer right now to make some port to JS / mobile and testing it with the onnx models you generated.

[nyadla-sys]

@dkakaie I am not sure if something like below

import whisper
import torch
tiny_model = whisper.load_model("tiny")
input_names=["serving_default_input.1","serving_default_onnx_tf__tf_MatMul_1_85dc0e00"]
output_names=["StatefulPartitionedCall"]
dynamic_axes={'serving_default_input.1': {0: 'batch_size', 1: 'sequence_len'}, 'StatefulPartitionedCall': {0: 'batch_size', 1: 'sequence_len'}}
torch.onnx.export(tiny_model.decoder, (torch.tensor([[50258, 50259, 50359, 50363]]).to(device), torch.randn(1, 1500, 384).to(device)), "/content/whisper-decoder-language.onnx", input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes)

[dkakaie]

@nyadla-sys @shadowofsoul here it is:

device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.onnx.export(model.decoder,
                  (torch.tensor([[50258, 50259, 50359, 50363]]).to(device), #sample input ids
                   torch.randn(1, 1500, 384).to(device)), #encoder outputs
                  "./decoder.onnx", #final onnx model name
                  opset_version=10, #opset 13 fails for me with unsupported squeeze sth
                  input_names=['tokens', 'hidden_states'], # the model's input names, 
                  output_names=['output'], #self-set output node name
                  dynamic_axes={'tokens': {1: 'toks'}, # variable length axes, inputs ids, tokens are index=1 and we want that dimension
                                'output': {1: 'toks'}}) #variable output axes

[dkakaie]

TFLite from previous ONNX export not promising. I am not even sure if TFLite is OK with variable length inputs.

[nyadla-sys]

i guess we need to use interpreter.resize_tensor_input() to change the shape we wish and then use settensor()..anyway i am also trying will update here

[nyadla-sys]

with the below code i am able to give dynamic input to tflite model
'''
import numpy as np
tflite_model_path='/content/whisper-decoder-language-hybrid.tflite'
#tflite_model_path='/content/whisper-decoder_main-int8.tflite'
print(tflite_model_path)

interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
interpreter.allocate_tensors()

decoder_input_ids = torch.tensor([50258, 50259, 50359, 50363])
decoder_input_ids = tf.expand_dims(decoder_input_ids, 0)
print(decoder_input_ids.shape)
print(encoder_output_data.shape)

input_tensor_1 = interpreter.get_input_details()[0]['index']
interpreter.set_tensor(input_tensor_1, encoder_output_data)

input_tensor_2 = interpreter.get_input_details()[1]['index']
interpreter.resize_tensor_input(input_tensor_2, decoder_input_ids.shape)
interpreter.allocate_tensors()
interpreter.set_tensor(input_tensor_2, decoder_input_ids)

output_tensor = interpreter.get_output_details()[0]['index']

interpreter.invoke()

output_data = interpreter.get_tensor(output_tensor)

cleaned = np.argmax(output_data, axis=-1)
last_token = cleaned[0,-1]
print(last_token)
'''

[nyadla-sys]

@dkakaie I managed to run decoder tflite model
https://colab.research.google.com/github/usefulsensors/openai-whisper/blob/main/notebooks/whisper_encoder_decoder_tflite.ipynb

[shadowofsoul]

@dkakaie if you are looking speed improvement, maybe it worth to take a look at the openvinno implementation: #208

[mingryon]

Hi, @nyadla-sys I've looked your encoder & decoder codebase and looks good. Can I ask one thing separately? If I want decode for Italian language, what is language code?
Seems 50259 is English.

[nyadla-sys]

<|it|>": 50274,

[Ayaa17]

Hi @nyadla-sys,
I want try whisper_encoder_decoder_tflite.ipynb,
but it throw an error.

KeyError: in user code:

  File "/usr/local/lib/python3.10/dist-packages/onnx_tf/backend_tf_module.py", line 99, in __call__  *
      output_ops = self.backend._onnx_node_to_tensorflow_op(onnx_node,
  File "/usr/local/lib/python3.10/dist-packages/onnx_tf/backend.py", line 347, in _onnx_node_to_tensorflow_op  *
      return handler.handle(node, tensor_dict=tensor_dict, strict=strict)
  File "/usr/local/lib/python3.10/dist-packages/onnx_tf/handlers/handler.py", line 59, in handle  *
      return ver_handle(node, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/onnx_tf/handlers/backend/conv.py", line 15, in version_11  *
      return cls.conv(node, kwargs["tensor_dict"])
  File "/usr/local/lib/python3.10/dist-packages/onnx_tf/handlers/backend/conv_mixin.py", line 29, in conv  *
      x = input_dict[node.inputs[0]]

  KeyError: 'x.1'

If you could help that's appreciated.

[yangjundy]

Hi @nyadla-sys,

I have a follow-up question: how can I create TFLite models for both the encoder and decoder from transformers TFWhisperModel? I noticed your excellent example in “whisper_encoder_decoder_tflite.ipynb” uses torch model “whisper.load_model” conversion to ONNX, and the conversion to tflite which is no longer actively developed for TFLite conversion support.

Thanks!

[nyadla-sys]

this will open up multilingual support for whisper tflite models

[nyadla-sys]

@dkakaie Can you assist me in running inference on the below Whisper ONNX model from HuggingFace? This will help me understand the input and output processing for creating a full int8 model.

!python -m transformers.onnx --model=openai/whisper-tiny.en onnx/

[dkakaie]

Sure, I'm currently using encoder and decoder models produced using Optimum but will check this one out too.

[Sriramraja05]

Hi, I too also on the same page, I am looking for the word timestamp after converting the tflite model, Is there any possibility should we implement the word timestamp of tflite model?

[nyadla-sys]

use this whisper tflite model for timestamp generation https://colab.research.google.com/drive/1qXcgILcA-HPEYqAYPrxQQ1TRwXerErDk?usp=sharing

[Sriramraja05]

Thanks @nyadla-sys , I have checked on it ,It basically gives the audio chunk time. Which is 30 second for each phrase, I am basically expecting each word start and end timestamp. I see that transformer supports timestamp. Or should I have to work on signature while converting tflite to support the word timestamp?

[mingryon]

Hi, NiranJan Yadla. Thanks for your reply. Can please explain where or how I can get total language code list? Thanks.

…

On Fri, Oct 27, 2023 at 10:57 PM Niranjan Yadla ***@***.***> wrote: <|it|>": 50274, — Reply to this email directly, view it on GitHub <#778 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AZUX3PHQBDY6SN27TNQZDTDYBQGZTAVCNFSM6AAAAAATM3MIXKVHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM3TIMBXGE3TA> . You are receiving this because you commented.Message ID: ***@***.***>

[dkakaie]

Here's the list

[mingryon]

Thanks a lot

…

On Sat, Oct 28, 2023 at 7:34 PM Davud Kakaie ***@***.***> wrote: Here <https://huggingface.co/openai/whisper-large-v2/blob/main/added_tokens.json>'s the list — Reply to this email directly, view it on GitHub <#778 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AZUX3PB2PWOWLYYJWFLREPLYBUXZ7AVCNFSM6AAAAAATM3MIXKVHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM3TIMJSGE3DS> . You are receiving this because you commented.Message ID: ***@***.***>

[mingryon]

Hi researcher I'm very interested in multilingual model. I tried to test whisper-tiny.tflite model, but it doesn't work. Do we have any method to convert specified language to tflite model like english model? I've tried decode & encode method, but speed is very slow. I hope the answer for converting specified language to tflite model. For example, Italian.

…

On Sun, Oct 29, 2023 at 7:13 PM Li Ming ***@***.***> wrote: Thanks a lot On Sat, Oct 28, 2023 at 7:34 PM Davud Kakaie ***@***.***> wrote: > Here > <https://huggingface.co/openai/whisper-large-v2/blob/main/added_tokens.json>'s > the list > > — > Reply to this email directly, view it on GitHub > <#778 (reply in thread)>, > or unsubscribe > <https://github.com/notifications/unsubscribe-auth/AZUX3PB2PWOWLYYJWFLREPLYBUXZ7AVCNFSM6AAAAAATM3MIXKVHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM3TIMJSGE3DS> > . > You are receiving this because you commented.Message ID: > ***@***.***> >