Upload notebook for model generation

Generate_tflite_for_whisper_base_with_transcribe_and_translate_signatures.ipynb

@@ -0,0 +1,380 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c5g9NTF_Ixad"
+      },
+      "source": [
+        "##Install Tranformers and datasets"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "w4VPaSlnHUvT"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install transformers==4.33.0\n",
+        "!pip install tensorflow==2.14.0"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ClniiYCWHK4b"
+      },
+      "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": null,
+      "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",
+        "        ],\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]],  # forced to transcribe any language with no timestamps\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]],  # different forced_decoder_ids\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",
+        "\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",
+      "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_default')\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
+}