examples/getting_started/olive_quickstart.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MEfAinRUdxmG"
   },
   "source": [
    "# 🚀 Olive Quickstart\n",
    "\n",
    "This notebook shows you how to get started with Olive - an AI model optimization toolkit for the ONNX Runtime. In this notebook, you will:\n",
    "\n",
    "1. Use Olive's automatic model optimizer via a CLI command to optimize an SLM model for the ONNX Runtime (for CPU devices).\n",
    "1. Use the ONNX Runtime Python binding to execute a simple chat interface that consumes the optimized model.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PUV-Rwz83L9E"
   },
   "source": [
    "## 🐍 Install Python dependencies\n",
    "\n",
    "First, install the Olive CLI using `pip`:\n",
    "\n",
    "We recommend installing Olive in a [virtual environment](https://docs.python.org/3/library/venv.html) or a [conda environment](https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BmPoU34Vdh6v"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "%pip install olive-ai[ort-genai,auto-opt]\n",
    "%pip install transformers==4.44.2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🤗 Cache model from Hugging Face\n",
    "\n",
    "In this quickstart you'll be optimizing [HuggingFaceTB/SmolLM2-135M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct), which has many model variants in the Hugging Face repo that are not required by Olive. To minimize the download, you can cache the original model files (safetensors and configuration) in the main folder of the Hugging Face repo:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!huggingface-cli download HuggingFaceTB/SmolLM2-135M-Instruct *.json *.safetensors *.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "J7D8rEOx7jnv"
   },
   "source": [
    "## 🪄 Automatic model optimization with Olive\n",
    "\n",
    "Next, you'll execute Olive's automatic optimizer using the `auto-opt` CLI command, which will:\n",
    "\n",
    "1. Acquire the [HuggingFaceTB/SmolLM2-135M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct) model from cache (note: if the model is not cached then it will download from Hugging Face).\n",
    "1. Capture the model into an ONNX graph and convert the weights into the ONNX format.\n",
    "1. Optimize the ONNX graph (e.g. fuse nodes, reshape, etc)\n",
    "1. Quantize the weights into `int4` precision using the RTN method.\n",
    "\n",
    "It takes around 60secs to optimize the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lrAGXFeSd_O5"
   },
   "outputs": [],
   "source": [
    "!olive auto-opt \\\n",
    "    --model_name_or_path HuggingFaceTB/SmolLM2-135M-Instruct \\\n",
    "    --output_path models/smolm2 \\\n",
    "    --device cpu \\\n",
    "    --provider CPUExecutionProvider \\\n",
    "    --use_ort_genai \\\n",
    "    --precision int4 \\\n",
    "    --log_level 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xTEL8Zwe7xVn"
   },
   "source": [
    "With the `auto-opt` command, you can change the input model to one that is available on Hugging Face - for example, to [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) - or a model that resides on local disk. Olive, will go through the same process of *automatically* converting (to ONNX), optimizing the graph and quantizing the weights. The model can be optimized for different providers and devices - for example, you can choose DirectML (for Windows) as the provider and target either the NPU, GPU, or CPU device.\n",
    "\n",
    "If you are using a Hugging Face gated model like Llama-3.2-1B-Instruct then you'll first need to login to Hugging Face using\n",
    "\n",
    "```bash\n",
    "huggingface-cli login --token USER_ACCESS_TOKEN\n",
    "```\n",
    "\n",
    "For more information on user access tokens, [read the Hugging face documentation on user access tokens](https://huggingface.co/docs/hub/security-tokens).\n",
    "\n",
    "## 🧠 Inference model using ONNX Runtime\n",
    "\n",
    "The ONNX Runtime (ORT) is a fast and light-weight package (available in many programming languages) that runs cross-platform. ORT enables you to infuse your AI models into your applications so that inference is handled *on-device*. The following code creates a simple console-based chat interface that inferences your optimized model.\n",
    "\n",
    "### How to use\n",
    "You'll be prompted to enter a message to the SLM - for example, you could ask *what is the golden ratio*, or *def print_hello_world():*. To exit type *exit* in the chat interface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "amGg0c3swnDF"
   },
   "outputs": [],
   "source": [
    "import onnxruntime_genai as og\n",
    "\n",
    "model_folder = \"models/smolm2/model\"\n",
    "\n",
    "# Load the base model and tokenizer\n",
    "model = og.Model(model_folder)\n",
    "tokenizer = og.Tokenizer(model)\n",
    "tokenizer_stream = tokenizer.create_stream()\n",
    "\n",
    "# Set the max length to something sensible by default,\n",
    "# since otherwise it will be set to the entire context length\n",
    "search_options = {}\n",
    "search_options['max_length'] = 200\n",
    "search_options['past_present_share_buffer'] = False\n",
    "\n",
    "chat_template = \"<|im_start|>user\\n{input}<|im_end|>\\n<|im_start|>assistant\\n\"\n",
    "\n",
    "text = input(\"Input: \")\n",
    "\n",
    "# Keep asking for input phrases\n",
    "while text != \"exit\":\n",
    "    if not text:\n",
    "        print(\"Error, input cannot be empty\")\n",
    "        exit\n",
    "\n",
    "    # generate prompt (prompt template + input)\n",
    "    prompt = f'{chat_template.format(input=text)}'\n",
    "\n",
    "    # encode the prompt using the tokenizer\n",
    "    input_tokens = tokenizer.encode(prompt)\n",
    "\n",
    "    params = og.GeneratorParams(model)\n",
    "    params.set_search_options(**search_options)\n",
    "    params.input_ids = input_tokens\n",
    "    generator = og.Generator(model, params)\n",
    "\n",
    "    print(\"Output: \", end='', flush=True)\n",
    "    # stream the output\n",
    "    try:\n",
    "        while not generator.is_done():\n",
    "            generator.compute_logits()\n",
    "            generator.generate_next_token()\n",
    "\n",
    "            new_token = generator.get_next_tokens()[0]\n",
    "            print(tokenizer_stream.decode(new_token), end='', flush=True)\n",
    "    except KeyboardInterrupt:\n",
    "        print(\"  --control+c pressed, aborting generation--\")\n",
    "\n",
    "    print()\n",
    "    text = input(\"Input: \")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}