{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "31cd209d-39a4-4e0b-b34a-ced26db5b7d3",
   "metadata": {},
   "source": [
    "# ONNX to TF Lite"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c06c3b4b-01c8-4b99-9b30-7c8bd7a676cb",
   "metadata": {},
   "source": [
    "**eIQ AI Toolkit** uses the **Olive** (https://microsoft.github.io/Olive/) framework for model quantization and conversion.\n",
    "This framework applies **passes**, where each pass represents a single transformation of a model.\n",
    "Passes can also be chained together to perform multiple transformations in sequence.\n",
    "\n",
    "This guide explains how to use the **TFLiteConversion** pass to convert a quantized ONNX model into a quantized TFLite model.\n",
    "It demonstrates how to use the eIQ AI Toolkit API to:\n",
    "- Load a quantized ONNX model into application\n",
    "- Run the conversion to a TFLite model\n",
    "- Retrieve the converted model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a5ab8e4eefeb01a",
   "metadata": {},
   "source": [
    "This guide requires the eIQ AI Toolkit backend to be running.\n",
    "If you haven't set it up yet, please refer to the following tutorial:\n",
    "[eIQ AI Toolkit setup & launch](../tools/aiToolkit/installRun.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3457cd50aa3a7ca9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from pathlib import Path\n",
    "\n",
    "# Set your eIQ AI Toolkit url:\n",
    "AI_TOOLKIT_BACKEND_URL = \"http://localhost:8000\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f729522-5f3b-4e11-97dd-5082bdbc75ac",
   "metadata": {},
   "source": [
    "## Load ONNX model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41d7c6375ac9ffb2",
   "metadata": {},
   "source": [
    "Loading any model into the application involves two steps:\n",
    "1. Specify the model metadata\n",
    "2. Upload the model file\n",
    "\n",
    "The metadata must include the model type you are uploading. For example, use `pytorch` for a PyTorch model or `onnx` for an ONNX model.\n",
    "\n",
    "For ONNX models, no additional metadata is required."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ba647b90ad1cd61",
   "metadata": {},
   "source": [
    "### 1. Prepare ONNX model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12604756f4de003a",
   "metadata": {},
   "source": [
    "If you already have a model, update the path to point to its location.\n",
    "If you don’t have a model yet, set the path to the location where the model will be saved.\n",
    "(See the following sections for instructions on downloading a sample model.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdfc8298a855516a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify the path to your ONNX model\n",
    "model_path = Path(\"path_to_quantized_onnx_model.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbcc0c14e9e2b574",
   "metadata": {},
   "source": [
    "Use the following script to download the example model:\n",
    "\n",
    "*Note: Skip this step if you already have your own model.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88827937eb11ec37",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_model_url = \"https://eiq.nxp.com/training-materials/_misc/models/quantized_model.onnx\"\n",
    "with open(model_path, \"wb\") as f:\n",
    "    response = requests.get(\n",
    "        url=example_model_url\n",
    "    )\n",
    "    f.write(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "745b91a140295896",
   "metadata": {},
   "source": [
    "### 2. Specify metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2061ba0d0c5dd66",
   "metadata": {},
   "source": [
    "In this section, we’ll upload the model metadata to eIQ AI Toolkit. This step only involves specifying the type of model you are uploading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "211159d26cb91b17",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODELS_API_URL = f\"{AI_TOOLKIT_BACKEND_URL}/models\"\n",
    "\n",
    "model_metadata = {\n",
    "        \"model_type\": \"onnx\",\n",
    "    }\n",
    "\n",
    "response = requests.post(MODELS_API_URL, json=model_metadata)\n",
    "response_data = response.json()\n",
    "model_uuid = response_data[\"data\"][\"model\"][\"uuid\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ece3806d029dec9c",
   "metadata": {},
   "source": [
    "### 3. Upload model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29b3b7b36a6b7aa5",
   "metadata": {},
   "source": [
    "Now we can upload the model file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edf599661c911faa",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(model_path, \"rb\") as model_file:\n",
    "    response = requests.post(\n",
    "        url=f\"{AI_TOOLKIT_BACKEND_URL}/models/{model_uuid}\", # Model identifier is part of the request URL\n",
    "        files={\n",
    "            \"model_file\": model_file,\n",
    "        }\n",
    "    )\n",
    "\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ea5a807c46c5ca",
   "metadata": {},
   "source": [
    "After uploading the model metadata and file, you can verify its registration and readiness status using the following endpoint. If the status remains `in_progress`, run the check multiple times until it changes to `ready`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbc620d84311bd69",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = requests.get(f\"{AI_TOOLKIT_BACKEND_URL}/models/{model_uuid}\")\n",
    "data = response.json()\n",
    "print(f'Model status: {data[\"data\"][\"model\"][\"status\"]}')\n",
    "print(f'Model status description: {data[\"data\"][\"model\"][\"status_description\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b250b222-12a3-4c9e-a5d6-1b9ecba625d1",
   "metadata": {},
   "source": [
    "## Convert to quantized TFLite"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "136cd7f1c7e6ece4",
   "metadata": {},
   "source": [
    "In the code below, you’ll notice that no parameters are set for the `TFLiteConversion` pass. All values use their defaults. If you want to view the list of available parameters, check this endpoint\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d70b4a79b78a6f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "available_passes_response = requests.get(f\"{AI_TOOLKIT_BACKEND_URL}/optimizations/passes\")\n",
    "available_passes = available_passes_response.json()\n",
    "\n",
    "# This prints configuration parameters only for TFLiteConversion pass. Feel free to change it and explore\n",
    "# other passes as well.\n",
    "tflite_conversion_pass_config = next(_pass for _pass in available_passes[\"data\"][\"passes\"] if _pass[\"type\"] == \"TFLiteConversion\")\n",
    "print(tflite_conversion_pass_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b20d102fb6bbf48",
   "metadata": {},
   "source": [
    "Let's run the conversion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b4773396b515cfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "OPTIMIZATIONS_API_URL = f\"{AI_TOOLKIT_BACKEND_URL}/optimizations\"\n",
    "RUN_OPTIMIZATION_API_URL = f\"{OPTIMIZATIONS_API_URL}/run\"\n",
    "\n",
    "pass_config = {\n",
    "    \"model_uuid\": model_uuid,\n",
    "    \"passes\": [\n",
    "        {\n",
    "            \"type\": \"TFLiteConversion\",\n",
    "            \"config\": {\n",
    "                # Feel free to add specific TFLite conversion parameters here if needed\n",
    "            }\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "\n",
    "optimization_response = requests.post(RUN_OPTIMIZATION_API_URL, json=pass_config)\n",
    "optimization_response_data = optimization_response.json()\n",
    "optimization_uuid = optimization_response_data[\"data\"][\"optimization\"][\"uuid\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7892106885be65e3",
   "metadata": {},
   "source": [
    "The conversion is now running. You can check its status by calling this endpoint. Repeat the check until the status changes to `success`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a60bf1ca3fa1bbd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = requests.get(f\"{OPTIMIZATIONS_API_URL}/{optimization_uuid}\")\n",
    "data = response.json()\n",
    "status = data[\"data\"][\"optimization\"][\"status\"]\n",
    "print(f\"Conversion status: {status}\")\n",
    "\n",
    "if status == \"success\":\n",
    "    artifact_id = data[\"data\"][\"optimization\"][\"artifacts\"][0][\"artifact_id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c40dd296325f063b",
   "metadata": {},
   "source": [
    "### Download quantized TFLite model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa0b2447eced631",
   "metadata": {},
   "source": [
    "Once the conversion is complete and successful, you can download the resulting model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7aab43fc6002e62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change model path to your location\n",
    "dest_model_path = Path(\"quantized_tflite_model.tflite\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56321060-8f85-4616-86f4-b9847a0c6365",
   "metadata": {},
   "outputs": [],
   "source": [
    "download_response = requests.get(f\"{AI_TOOLKIT_BACKEND_URL}/optimizations/{optimization_uuid}/resources/{artifact_id}\")\n",
    "with dest_model_path.open(\"wb\") as f:\n",
    "    f.write(download_response.content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}