{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "UUXnh11hA75x",
        "jbcW5jMLK8gK"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2024 The IREE Authors"
      ],
      "metadata": {
        "id": "UUXnh11hA75x"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Licensed under the Apache License v2.0 with LLVM Exceptions.\n",
        "# See https://llvm.org/LICENSE.txt for license information.\n",
        "# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception"
      ],
      "metadata": {
        "cellView": "form",
        "id": "FqsvmKpjBJO2"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# <img src=\"https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.png\" height=\"20px\"> Hugging Face to <img src=\"https://upload.wikimedia.org/wikipedia/commons/thumb/1/10/PyTorch_logo_icon.svg/640px-PyTorch_logo_icon.svg.png\" height=\"20px\"> PyTorch to <img src=\"https://raw.githubusercontent.com/iree-org/iree/main/docs/website/docs/assets/images/IREE_Logo_Icon_Color.svg\" height=\"20px\"> IREE\n",
        "\n",
        "This notebook uses [iree-turbine](https://github.com/iree-org/iree-turbine) to export a pretrained [Hugging Face Transformers](https://huggingface.co/docs/transformers/) model to [IREE](https://github.com/iree-org/iree), leveraging [torch-mlir](https://github.com/llvm/torch-mlir) under the covers.\n",
        "\n",
        "* The pretrained [whisper-small](https://huggingface.co/openai/whisper-small)\n",
        "  model is showcased here as it is small enough to fit comfortably into a Colab\n",
        "  notebook. Other pretrained models can be found at\n",
        "  https://huggingface.co/docs/transformers/index."
      ],
      "metadata": {
        "id": "38UDc27KBPD1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup"
      ],
      "metadata": {
        "id": "jbcW5jMLK8gK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "#@title Uninstall existing packages\n",
        "#   This avoids some warnings when installing specific PyTorch packages below.\n",
        "!python -m pip uninstall -y fastai torchaudio torchdata torchtext torchvision"
      ],
      "metadata": {
        "id": "KsPubQSvCbXd",
        "cellView": "form"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!python -m pip install --pre --index-url https://download.pytorch.org/whl/test/cpu --upgrade torch==2.3.0"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oO1tirq2ggmO",
        "outputId": "c3390361-9f40-4a49-b5c7-898a62614143"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://download.pytorch.org/whl/test/cpu\n",
            "Collecting torch==2.3.0\n",
            "  Downloading https://download.pytorch.org/whl/test/cpu/torch-2.3.0%2Bcpu-cp310-cp310-linux_x86_64.whl (190.4 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m190.4/190.4 MB\u001b[0m \u001b[31m4.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (3.13.4)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (4.11.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (3.3)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch==2.3.0) (2023.6.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch==2.3.0) (2.1.5)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch==2.3.0) (1.3.0)\n",
            "Installing collected packages: torch\n",
            "  Attempting uninstall: torch\n",
            "    Found existing installation: torch 2.2.1+cu121\n",
            "    Uninstalling torch-2.2.1+cu121:\n",
            "      Successfully uninstalled torch-2.2.1+cu121\n",
            "Successfully installed torch-2.3.0+cpu\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4iJFDHbsAzo4",
        "outputId": "94721ae8-e222-4203-c356-888b42bc20b9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting iree-turbine\n",
            "  Downloading iree_turbine-2.3.0rc20240410-py3-none-any.whl (150 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/150.4 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m\u001b[90m━\u001b[0m \u001b[32m143.4/150.4 kB\u001b[0m \u001b[31m4.4 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m150.4/150.4 kB\u001b[0m \u001b[31m3.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from iree-turbine) (1.25.2)\n",
            "Collecting iree-compiler>=20240410.859 (from iree-turbine)\n",
            "  Downloading iree_compiler-20240410.859-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (64.4 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m64.4/64.4 MB\u001b[0m \u001b[31m10.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting iree-runtime>=20240410.859 (from iree-turbine)\n",
            "  Downloading iree_runtime-20240410.859-cp310-cp310-manylinux_2_28_x86_64.whl (7.4 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.4/7.4 MB\u001b[0m \u001b[31m26.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: torch>=2.1.0 in /usr/local/lib/python3.10/dist-packages (from iree-turbine) (2.3.0+cpu)\n",
            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from iree-compiler>=20240410.859->iree-turbine) (6.0.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->iree-turbine) (3.13.4)\n",
            "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->iree-turbine) (4.11.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->iree-turbine) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->iree-turbine) (3.3)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->iree-turbine) (3.1.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->iree-turbine) (2023.6.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=2.1.0->iree-turbine) (2.1.5)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=2.1.0->iree-turbine) (1.3.0)\n",
            "Installing collected packages: iree-runtime, iree-compiler, iree-turbine\n",
            "Successfully installed iree-compiler-20240410.859 iree-runtime-20240410.859 iree-turbine-2.3.0rc20240410\n"
          ]
        }
      ],
      "source": [
        "!python -m pip install iree-turbine"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Report version information\n",
        "!echo \"Installed iree-turbine, $(python -m pip show iree_turbine | grep Version)\"\n",
        "\n",
        "!echo -e \"\\nInstalled IREE, compiler version information:\"\n",
        "!iree-compile --version\n",
        "\n",
        "import torch\n",
        "print(\"\\nInstalled PyTorch, version:\", torch.__version__)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nkVLzRpcDnVL",
        "outputId": "ee4e956f-ca7d-45ac-9913-672ad444d89f"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Installed iree-turbine, Version: 2.3.0rc20240410\n",
            "\n",
            "Installed IREE, compiler version information:\n",
            "IREE (https://iree.dev):\n",
            "  IREE compiler version 20240410.859 @ b4273a4bfc66ba6dd8f62f6483d74d42a7b936f1\n",
            "  LLVM version 19.0.0git\n",
            "  Optimized build\n",
            "\n",
            "Installed PyTorch, version: 2.3.0+cpu\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load and run whisper-small"
      ],
      "metadata": {
        "id": "I0OfTFxwOud1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load the pretrained model from https://huggingface.co/openai/whisper-small.\n",
        "\n",
        "See also:\n",
        "\n",
        "* Model card: https://huggingface.co/docs/transformers/model_doc/whisper\n",
        "* Test case in [SHARK-TestSuite](https://github.com/nod-ai/SHARK-TestSuite/): [`pytorch/models/whisper-small/model.py`](https://github.com/nod-ai/SHARK-TestSuite/blob/main/e2eshark/pytorch/models/whisper-small/model.py)"
      ],
      "metadata": {
        "id": "94Ji4URLT_xM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "\n",
        "# https://huggingface.co/docs/transformers/model_doc/auto\n",
        "# AutoModelForCausalLM -> WhisperForCausalLM\n",
        "# AutoTokenizer        -> WhisperTokenizerFast\n",
        "\n",
        "modelname = \"openai/whisper-small\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(modelname)\n",
        "\n",
        "# Some of the options here affect how the model is exported. See the test cases\n",
        "# at https://github.com/nod-ai/SHARK-TestSuite/tree/main/e2eshark/pytorch/models\n",
        "# for other options that may be useful to set.\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    modelname,\n",
        "    output_attentions=False,\n",
        "    output_hidden_states=False,\n",
        "    attn_implementation=\"eager\",\n",
        "    torchscript=True,\n",
        ")\n",
        "\n",
        "# This is just a simple demo to get some data flowing through the model.\n",
        "# Depending on this model and what input it expects (text, image, audio, etc.)\n",
        "# this might instead use a specific Processor class. For Whisper,\n",
        "# WhisperProcessor runs audio input pre-processing and output post-processing.\n",
        "example_prompt = \"Hello world!\"\n",
        "example_encoding = tokenizer(example_prompt, return_tensors=\"pt\")\n",
        "example_input = example_encoding[\"input_ids\"].cpu()\n",
        "example_args = (example_input,)"
      ],
      "metadata": {
        "id": "HLbfUuoBPHgH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Test exporting using [`torch.export()`](https://pytorch.org/docs/stable/export.html#torch.export.export). If `torch.export` works, `aot.export()` from Turbine should work as well."
      ],
      "metadata": {
        "id": "vQlF_ua3UNvo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "exported_program = torch.export.export(model, example_args)"
      ],
      "metadata": {
        "id": "-4LykgffY9uH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Export using the simple [`aot.export()`](https://iree.dev/guides/ml-frameworks/pytorch/#simple-api) API from Turbine."
      ],
      "metadata": {
        "id": "wXZI4GliUazA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import shark_turbine.aot as aot\n",
        "# Note: aot.export() wants the example args to be unpacked.\n",
        "whisper_compiled_module = aot.export(model, *example_args)"
      ],
      "metadata": {
        "id": "R7-rN_z2Y_5z"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Compile using Turbine/IREE then run the program."
      ],
      "metadata": {
        "id": "YK3hjpTpUdhc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "binary = whisper_compiled_module.compile(save_to=None)\n",
        "\n",
        "import iree.runtime as ireert\n",
        "config = ireert.Config(\"local-task\")\n",
        "vm_module = ireert.load_vm_module(\n",
        "    ireert.VmModule.wrap_buffer(config.vm_instance, binary.map_memory()),\n",
        "    config,\n",
        ")\n",
        "\n",
        "iree_outputs = vm_module.main(example_args[0])\n",
        "print(iree_outputs[0].to_host())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FctBxxEXZBan",
        "outputId": "12c042d0-f740-4de1-e246-36ed9b4b357b"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[[  5.8126216   3.9667568   4.5749426 ...   2.7658575   2.6436937\n",
            "     1.5479789]\n",
            "  [  7.563438    6.0299625   5.1000338 ...   6.4327035   6.101554\n",
            "     6.434801 ]\n",
            "  [  0.9380368  -4.4696164  -4.012759  ...  -6.24863    -7.791795\n",
            "    -6.84537  ]\n",
            "  [  0.7450911  -3.7631674  -7.4870267 ...  -6.7348223  -6.966235\n",
            "   -10.022385 ]\n",
            "  [ -0.9628638  -3.5101964  -6.0158615 ...  -7.116393   -6.7086525\n",
            "   -10.225711 ]\n",
            "  [  3.3470955   2.4927258  -3.3042645 ...  -1.5709444  -1.8455245\n",
            "    -2.9991858]]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the program using native PyTorch to compare outputs."
      ],
      "metadata": {
        "id": "5WuFpyFfUjh8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch_outputs = model(example_args[0])\n",
        "print(torch_outputs[0].detach().numpy())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IxPYkcPycG4r",
        "outputId": "d1a3b111-4a6f-4e2a-f80a-645c192e57e3"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[[  5.8126183    3.9667587    4.5749483  ...   2.7658575    2.643694\n",
            "     1.5479784 ]\n",
            "  [  7.563436     6.029952     5.100036   ...   6.4327083    6.101557\n",
            "     6.4348083 ]\n",
            "  [  0.93802685  -4.469646    -4.012787   ...  -6.2486415   -7.7918167\n",
            "    -6.8453975 ]\n",
            "  [  0.74507916  -3.763197    -7.487034   ...  -6.734877    -6.966276\n",
            "   -10.022424  ]\n",
            "  [ -0.96288276  -3.510221    -6.0158725  ...  -7.1164136   -6.708687\n",
            "   -10.225745  ]\n",
            "  [  3.3470666    2.492654    -3.304323   ...  -1.5709934   -1.8455791\n",
            "    -2.9992423 ]]]\n"
          ]
        }
      ]
    }
  ]
}