Add new model for top world languages supported by Whisper

Generate_tflite_for_whisper_base_TOP_WORLD_version.ipynb

@@ -0,0 +1,764 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c5g9NTF_Ixad"
+      },
+      "source": [
+        "##Install Tranformers and datasets"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "w4VPaSlnHUvT",
+        "collapsed": true
+      },
+      "outputs": [],
+      "source": [
+        "!pip install transformers==4.33.0\n",
+        "!pip install tensorflow==2.14.0\n",
+        "!pip install numpy==1.26.4"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ClniiYCWHK4b",
+        "collapsed": true
+      },
+      "outputs": [],
+      "source": [
+        "! pip install datasets"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "pljpioLsJOtb"
+      },
+      "source": [
+        "##Load pre trained TF Whisper Base model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BJNOxn5vHaGi"
+      },
+      "outputs": [],
+      "source": [
+        "import tensorflow as tf\n",
+        "from transformers import TFWhisperModel, WhisperFeatureExtractor\n",
+        "from datasets import load_dataset\n",
+        "\n",
+        "model = TFWhisperModel.from_pretrained(\"openai/whisper-base\")\n",
+        "feature_extractor = WhisperFeatureExtractor.from_pretrained(\"openai/whisper-base\")\n",
+        "\n",
+        "ds = load_dataset(\"google/fleurs\", \"fr_fr\", split=\"test\")\n",
+        "inputs = feature_extractor(\n",
+        "    ds[0][\"audio\"][\"array\"], sampling_rate=ds[0][\"audio\"][\"sampling_rate\"], return_tensors=\"tf\"\n",
+        ")\n",
+        "input_features = inputs.input_features\n",
+        "print(input_features)\n",
+        "decoder_input_ids = tf.convert_to_tensor([[1, 1]]) * model.config.decoder_start_token_id\n",
+        "last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state\n",
+        "list(last_hidden_state.shape)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "W9XP25uhJl44"
+      },
+      "source": [
+        "##Generate Saved model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vpYwMmgyHf0B"
+      },
+      "outputs": [],
+      "source": [
+        "model.save('/content/tf_whisper_saved')"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "TY_79jFEJYyJ"
+      },
+      "source": [
+        "##Convert saved model to TFLite model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "owez2zvzHl-p"
+      },
+      "outputs": [],
+      "source": [
+        "import tensorflow as tf\n",
+        "\n",
+        "saved_model_dir = '/content/tf_whisper_saved'\n",
+        "tflite_model_path = 'whisper.tflite'\n",
+        "\n",
+        "# Convert the model\n",
+        "converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)\n",
+        "converter.target_spec.supported_ops = [\n",
+        "  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.\n",
+        "  tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.\n",
+        "]\n",
+        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
+        "tflite_model = converter.convert()\n",
+        "\n",
+        "# Save the model\n",
+        "with open(tflite_model_path, 'wb') as f:\n",
+        "    f.write(tflite_model)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tFkzUrjIbNcH"
+      },
+      "outputs": [],
+      "source": [
+        "%ls -la"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fpEnWZt7iQJK"
+      },
+      "source": [
+        "##Evaluate TF model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "-RuFFohHg2ho"
+      },
+      "outputs": [],
+      "source": [
+        "import tensorflow as tf\n",
+        "from transformers import WhisperProcessor, TFWhisperForConditionalGeneration\n",
+        "from datasets import load_dataset\n",
+        "\n",
+        "processor = WhisperProcessor.from_pretrained(\"openai/whisper-base\")\n",
+        "model = TFWhisperForConditionalGeneration.from_pretrained(\"openai/whisper-base\")\n",
+        "\n",
+        "ds = load_dataset(\"google/fleurs\", \"fr_fr\", split=\"test\")\n",
+        "\n",
+        "inputs = processor(ds[0][\"audio\"][\"array\"], return_tensors=\"tf\")\n",
+        "input_features = inputs.input_features\n",
+        "\n",
+        "generated_ids = model.generate(input_features)\n",
+        "\n",
+        "transcription = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]\n",
+        "transcription"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "U-eKuy_cG4u0"
+      },
+      "source": [
+        "## Evaluate TF Lite model (naive)\n",
+        "\n",
+        "We can load the model as defined above... but the model is useless on its own. Generation is much more complex that a model forward pass."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "wnfHirgyG0W4"
+      },
+      "outputs": [],
+      "source": [
+        "tflite_model_path = 'whisper.tflite'\n",
+        "interpreter = tf.lite.Interpreter(tflite_model_path)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a8VJQuHJKzl4"
+      },
+      "source": [
+        "## Create generation-enabled TF Lite model\n",
+        "\n",
+        "The solution consists in defining a model whose serving function is the generation call. Here's an example of how to do it:"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JmIgqWVgVBZN"
+      },
+      "source": [
+        "Now with monkey-patch for fixing NaN errors with -inf values"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "id": "e5P8s66yU7Kv"
+      },
+      "outputs": [],
+      "source": [
+        "import tensorflow as tf\n",
+        "import numpy as np\n",
+        "from transformers import TFForceTokensLogitsProcessor, TFLogitsProcessor\n",
+        "from typing import List, Optional, Union, Any\n",
+        "\n",
+        "# Patching methods of class TFForceTokensLogitsProcessor(TFLogitsProcessor):\n",
+        "\n",
+        "def my__init__(self, force_token_map: List[List[int]]):\n",
+        "    force_token_map = dict(force_token_map)\n",
+        "    # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the\n",
+        "    # index of the array corresponds to the index of the token to be forced, for XLA compatibility.\n",
+        "    # Indexes without forced tokens will have an negative value.\n",
+        "    force_token_array = np.ones((max(force_token_map.keys()) + 1), dtype=np.int32) * -1\n",
+        "    for index, token in force_token_map.items():\n",
+        "        if token is not None:\n",
+        "            force_token_array[index] = token\n",
+        "    self.force_token_array = tf.convert_to_tensor(force_token_array, dtype=tf.int32)\n",
+        "\n",
+        "def my__call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:\n",
+        "    def _force_token(generation_idx):\n",
+        "        batch_size = scores.shape[0]\n",
+        "        current_token = self.force_token_array[generation_idx]\n",
+        "\n",
+        "        # Original code below generates NaN values when the model is exported to tflite\n",
+        "        # it just needs to be a negative number so that the forced token's value of 0 is the largest\n",
+        "        # so it will get chosen\n",
+        "        #new_scores = tf.ones_like(scores, dtype=scores.dtype) * -float(\"inf\")\n",
+        "        new_scores = tf.ones_like(scores, dtype=scores.dtype) * -float(1)\n",
+        "        indices = tf.stack((tf.range(batch_size), tf.tile([current_token], [batch_size])), axis=1)\n",
+        "        updates = tf.zeros((batch_size,), dtype=scores.dtype)\n",
+        "        new_scores = tf.tensor_scatter_nd_update(new_scores, indices, updates)\n",
+        "        return new_scores\n",
+        "\n",
+        "    scores = tf.cond(\n",
+        "        tf.greater_equal(cur_len, tf.shape(self.force_token_array)[0]),\n",
+        "        # If the current length is geq than the length of force_token_array, the processor does nothing.\n",
+        "        lambda: tf.identity(scores),\n",
+        "        # Otherwise, it may force a certain token.\n",
+        "        lambda: tf.cond(\n",
+        "            tf.greater_equal(self.force_token_array[cur_len], 0),\n",
+        "            # Only valid (positive) tokens are forced\n",
+        "            lambda: _force_token(cur_len),\n",
+        "            # Otherwise, the processor does nothing.\n",
+        "            lambda: scores,\n",
+        "        ),\n",
+        "    )\n",
+        "    return scores\n",
+        "\n",
+        "TFForceTokensLogitsProcessor.__init__ = my__init__\n",
+        "TFForceTokensLogitsProcessor.__call__ = my__call__"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rIkUCdiyU7ZT"
+      },
+      "outputs": [],
+      "source": [
+        "import tensorflow as tf\n",
+        "\n",
+        "class GenerateModel(tf.Module):\n",
+        "    def __init__(self, model):\n",
+        "        super(GenerateModel, self).__init__()\n",
+        "        self.model = model\n",
+        "\n",
+        "    @tf.function(\n",
+        "        input_signature=[\n",
+        "            tf.TensorSpec((1, 80, 3000), tf.float32, name=\"input_features\"),\n",
+        "            tf.TensorSpec((1,), tf.int32, name=\"lang_token\"),\n",
+        "        ],\n",
+        "    )\n",
+        "    def transcribe_lang(self, input_features, lang_token):\n",
+        "        if lang_token == 50259:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50259], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50260:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50260], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50261:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50261], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50262:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50262], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50263:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50263], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50264:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50264], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50265:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50265], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50266:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50266], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50267:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50267], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50268:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50268], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50269:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50269], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50270:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50270], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50271:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50271], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50272:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50272], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50273:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50273], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50274:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50274], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50275:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50275], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50277:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50277], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50278:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50278], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50279:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50279], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50280:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50280], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50281:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50281], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50282:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50282], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50283:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50283], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50284:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50284], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50285:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50285], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50286:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50286], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50287:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50287], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50288:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50288], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50289:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50289], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50290:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50290], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50291:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50291], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50292:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50292], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50293:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50293], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50298:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50298], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50301:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50301], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50305:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50305], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50307:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50307], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50343:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50343], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        elif lang_token == 50345:\n",
+        "            outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[1, 50345], [2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "\n",
+        "        else:\n",
+        "             outputs = self.model.generate(\n",
+        "                input_features,\n",
+        "                max_new_tokens=450,  # change as needed\n",
+        "                return_dict_in_generate=True,\n",
+        "                forced_decoder_ids=[[2, 50359], [3, 50363]],\n",
+        "            )\n",
+        "        return {\"sequences\": outputs[\"sequences\"]}\n",
+        "\n",
+        "\n",
+        "    @tf.function(\n",
+        "        input_signature=[\n",
+        "            tf.TensorSpec((1, 80, 3000), tf.float32, name=\"input_features\"),\n",
+        "        ],\n",
+        "    )\n",
+        "    def transcribe(self, input_features):\n",
+        "        outputs = self.model.generate(\n",
+        "            input_features,\n",
+        "            max_new_tokens=450,  # change as needed\n",
+        "            return_dict_in_generate=True,\n",
+        "            forced_decoder_ids=[[2, 50359], [3, 50363]],\n",
+        "        )\n",
+        "        return {\"sequences\": outputs[\"sequences\"]}\n",
+        "\n",
+        "    @tf.function(\n",
+        "        input_signature=[\n",
+        "            tf.TensorSpec((1, 80, 3000), tf.float32, name=\"input_features\"),\n",
+        "        ],\n",
+        "    )\n",
+        "    def translate(self, input_features):\n",
+        "        outputs = self.model.generate(\n",
+        "            input_features,\n",
+        "            max_new_tokens=450,  # change as needed\n",
+        "            return_dict_in_generate=True,\n",
+        "            forced_decoder_ids=[[2, 50358], [3, 50363]],\n",
+        "        )\n",
+        "        return {\"sequences\": outputs[\"sequences\"]}\n",
+        "\n",
+        "# Assuming `model` is already defined and loaded\n",
+        "saved_model_dir = '/content/tf_whisper_saved'\n",
+        "tflite_model_path = 'whisper.tflite'\n",
+        "\n",
+        "generate_model = GenerateModel(model=model)\n",
+        "tf.saved_model.save(generate_model, saved_model_dir, signatures={\n",
+        "    \"serving_default\": generate_model.transcribe,\n",
+        "    \"serving_transcribe\": generate_model.transcribe,\n",
+        "    \"serving_translate\": generate_model.translate,\n",
+        "    \"serving_transcribe_lang\": generate_model.transcribe_lang,\n",
+        "\n",
+        "})\n",
+        "\n",
+        "# Convert the model\n",
+        "converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)\n",
+        "converter.target_spec.supported_ops = [\n",
+        "  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.\n",
+        "  tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.\n",
+        "]\n",
+        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
+        "tflite_model = converter.convert()\n",
+        "\n",
+        "# Save the model\n",
+        "with open(tflite_model_path, 'wb') as f:\n",
+        "    f.write(tflite_model)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "pwd"
+      ],
+      "metadata": {
+        "id": "llf-5421rZ-G"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!zip -r /content/tf_whisper_saved.zip /content/tf_whisper_saved/"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "collapsed": true,
+        "id": "7pnAWtGZp6MJ",
+        "outputId": "42d6c775-1af9-4482-837a-eb3537d5e2c0"
+      },
+      "execution_count": null,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "  adding: content/tf_whisper_saved/ (stored 0%)\n",
+            "  adding: content/tf_whisper_saved/assets/ (stored 0%)\n",
+            "  adding: content/tf_whisper_saved/variables/ (stored 0%)\n",
+            "  adding: content/tf_whisper_saved/variables/variables.data-00000-of-00001 (deflated 41%)\n",
+            "  adding: content/tf_whisper_saved/variables/variables.index (deflated 79%)\n",
+            "  adding: content/tf_whisper_saved/fingerprint.pb (stored 0%)\n",
+            "  adding: content/tf_whisper_saved/keras_metadata.pb (deflated 96%)\n",
+            "  adding: content/tf_whisper_saved/saved_model.pb (deflated 93%)\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "u9MustgMU7oI"
+      },
+      "outputs": [],
+      "source": [
+        "# loaded model... now with generate!\n",
+        "tflite_model_path = 'whisper.tflite'\n",
+        "interpreter = tf.lite.Interpreter(tflite_model_path)\n",
+        "\n",
+        "tflite_generate = interpreter.get_signature_runner('serving_translate')\n",
+        "lang_token = tf.constant([50261], dtype=tf.int32)\n",
+        "generated_ids = tflite_generate(input_features=input_features)[\"sequences\"]\n",
+        "transcription = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]\n",
+        "transcription\n",
+        "\n",
+        "\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "machine_shape": "hm",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

whisper-base.TOP_WORLD.tflite

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

whisper-base.TOP_WORLD.tokens

@@ -0,0 +1,3 @@
+[50259, 50260, 50261, 50262,50263, 50264, 50265, 50266, 50267, 50268, 50269,50270, 50271,50272, 50273, 50274,50275, 50277,50278,50279,50280, 50281,50282, 50283, 50284, 50285, 50286,50287, 50288, 50289,50290, 50291, 50292, 50293, 50298, 50301, 50305, 50307, 50343, 50345]
+
+