{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##### Copyright 2024 The IREE Authors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!python -m pip install --pre --index-url https://download.pytorch.org/whl/cpu --upgrade torch==2.11.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!python -m pip install --pre iree-turbine -f https://iree.dev/pip-release-links.html"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "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__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load and run whisper-small"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "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 [AMD-SHARK-TestSuite](https://github.com/nod-ai/AMD-SHARK-TestSuite/): [`pytorch/models/whisper-small/model.py`](https://github.com/nod-ai/AMD-SHARK-TestSuite/blob/main/e2eamdshark/pytorch/models/whisper-small/model.py)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "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/AMD-SHARK-TestSuite/tree/main/e2eamdshark/pytorch/models\n",
        "# for other options that may be useful to set.\n",
        "# use_cache=False keeps DynamicCache objects out of the forward outputs, which\n",
        "# torch.export cannot trace through.\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    modelname,\n",
        "    output_attentions=False,\n",
        "    output_hidden_states=False,\n",
        "    attn_implementation=\"eager\",\n",
        "    use_cache=False,\n",
        ")\n",
        "model.eval()\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,)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "# run_decompositions() lowers composite ops such as aten.diff (emitted by the\n",
        "# transformers attention mask code) that torch-mlir has no lowering for.\n",
        "exported_program = torch.export.export(model, example_args).run_decompositions()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Export using the simple [`aot.export()`](https://iree.dev/guides/ml-frameworks/pytorch/#simple-api) API from Turbine."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import iree.turbine.aot as aot\n",
        "whisper_compiled_module = aot.export(exported_program)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compile using Turbine/IREE then run the program."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "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.to_host())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Run the program using native PyTorch to compare outputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "torch_outputs = model(example_args[0])\n",
        "print(torch_outputs[0].detach().numpy())"
      ]
    }
  ],
  "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 0
}