Add Colab notebook showing Hugging Face import via Turbine. (#17093)
This shows a minimal import from Hugging Face using the style of code
from
https://github.com/nod-ai/SHARK-TestSuite/tree/main/e2eshark/pytorch/models.
I'd like to expand on this with other advanced export APIs (e.g. using
parameter files like Safetensors/GGUF/IRPA, Turbine-LLM, and custom
kernels), but figured I'd start with something fairly well tested and
simple.
The test case I forked from also used a text prompt for an audio -> text
model, so this just shows some numbers as output... I'd like to show
text or images as output, though that needs some interfacing
(WhisperProcessor wraps WhisperFeatureExtractor and WhisperTokenizer,
for example).
diff --git a/docs/website/docs/guides/ml-frameworks/pytorch.md b/docs/website/docs/guides/ml-frameworks/pytorch.md
index a9afdb6..14fb2bd 100644
--- a/docs/website/docs/guides/ml-frameworks/pytorch.md
+++ b/docs/website/docs/guides/ml-frameworks/pytorch.md
@@ -231,6 +231,7 @@
| Code samples | |
| -- | -- |
Simple AOT export notebook | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/pytorch_aot_simple.ipynb)
+Import [Whisper](https://huggingface.co/openai/whisper-small) from [:hugging: Hugging Face](https://huggingface.co/) notebook | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/pytorch_huggingface_whisper.ipynb)
Simple MLP export | [`core/examples/aot_mlp/mlp_export_simple.py`](https://github.com/nod-ai/SHARK-Turbine/blob/main/core/examples/aot_mlp/mlp_export_simple.py)
### :octicons-tools-16: Advanced API
diff --git a/samples/colab/README.md b/samples/colab/README.md
index 8851928..8241713 100644
--- a/samples/colab/README.md
+++ b/samples/colab/README.md
@@ -8,14 +8,15 @@
Framework | Notebook file | Description | Link
-------- | ------------- | ----------- | ----
Generic | [low_level_invoke_function](low_level_invoke_function.ipynb) | Shows low level IREE python binding concepts | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/low_level_invoke_function.ipynb)
-PyTorch | [pytorch_aot_advanced](pytorch_aot_advanced.ipynb) | Uses [SHARK-Turbine](https://github.com/nod-ai/SHARK-Turbine) to export a PyTorch program | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/pytorch_aot_advanced.ipynb)
-PyTorch | [pytorch_aot_simple](pytorch_aot_simple.ipynb) | Uses [SHARK-Turbine](https://github.com/nod-ai/SHARK-Turbine) to export a simple PyTorch program | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/pytorch_aot_simple.ipynb)
-PyTorch | [pytorch_jit](pytorch_jit.ipynb) | Uses [SHARK-Turbine](https://github.com/nod-ai/SHARK-Turbine) for eager execution in a PyTorch session | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/pytorch_jit.ipynb)
+PyTorch | [pytorch_aot_advanced](pytorch_aot_advanced.ipynb) | Uses [Turbine](https://github.com/nod-ai/SHARK-Turbine) to export a PyTorch program | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/pytorch_aot_advanced.ipynb)
+PyTorch | [pytorch_aot_simple](pytorch_aot_simple.ipynb) | Uses [Turbine](https://github.com/nod-ai/SHARK-Turbine) to export a simple PyTorch program | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/pytorch_aot_simple.ipynb)
+PyTorch | [pytorch_jit](pytorch_jit.ipynb) | Uses [Turbine](https://github.com/nod-ai/SHARK-Turbine) for eager execution in a PyTorch session | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/pytorch_jit.ipynb)
+PyTorch | [pytorch_huggingface_whisper](pytorch_huggingface_whisper.ipynb) | Imports the [Whisper](https://huggingface.co/openai/whisper-small) model from [Hugging Face](https://huggingface.co/) | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/pytorch_huggingface_whisper.ipynb)
TensorFlow | [tensorflow_edge_detection](tensorflow_edge_detection.ipynb) |Performs image edge detection | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/tensorflow_edge_detection.ipynb)
TensorFlow | [tensorflow_hub_import](tensorflow_hub_import.ipynb) | Imports a [MobileNet V2](https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification) model from [TensorFlow Hub](https://tfhub.dev/) | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/tensorflow_hub_import.ipynb)
TensorFlow | [tensorflow_mnist_training](tensorflow_mnist_training.ipynb) | Compiles, trains, and executes a neural network | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/tensorflow_mnist_training.ipynb)
-TensorFlow | [tensorflow_resnet](tensorflow_resnet.ipynb) | Compiles and runs a pretrained [ResNet50](https://www.tensorflow.org/api_docs/python/tf/keras/applications/ResNet50) model | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/tensorflow_resnet.ipynb)
-TFLite | [tflite_text_classification](tflite_text_classification.ipynb) | Compiles and runs a pretrained [text classification](https://www.tensorflow.org/lite/examples/text_classification/overview) model | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/tflite_text_classification.ipynb)
+TensorFlow | [tensorflow_resnet](tensorflow_resnet.ipynb) | Runs a pretrained [ResNet50](https://www.tensorflow.org/api_docs/python/tf/keras/applications/ResNet50) model | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/tensorflow_resnet.ipynb)
+TFLite | [tflite_text_classification](tflite_text_classification.ipynb) | Runs a pretrained [text classification](https://www.tensorflow.org/lite/examples/text_classification/overview) model | [](https://colab.research.google.com/github/iree-org/iree/blob/main/samples/colab/tflite_text_classification.ipynb)
## Working with GitHub
diff --git a/samples/colab/pytorch_huggingface_whisper.ipynb b/samples/colab/pytorch_huggingface_whisper.ipynb
new file mode 100644
index 0000000..bf8577a
--- /dev/null
+++ b/samples/colab/pytorch_huggingface_whisper.ipynb
@@ -0,0 +1,408 @@
+{
+ "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/ghost.svg\" height=\"20px\"> IREE\n",
+ "\n",
+ "This notebook uses [SHARK-Turbine](https://github.com/nod-ai/SHARK-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"
+ ]
+ }
+ ]
+ }
+ ]
+}
\ No newline at end of file
diff --git a/samples/colab/test_notebooks.py b/samples/colab/test_notebooks.py
index 97decb0..a31d43f 100755
--- a/samples/colab/test_notebooks.py
+++ b/samples/colab/test_notebooks.py
@@ -22,6 +22,7 @@
# PyTorch notebooks using SHARK-Turbine require Python 3.10+ in Docker.
"pytorch_aot_advanced.ipynb",
"pytorch_aot_simple.ipynb",
+ "pytorch_huggingface_whisper.ipynb",
"pytorch_jit.ipynb",
]
