llm_conversion.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "e812825502c9458891c62c1010a6fe54": {
          "model_module": "@jupyter-widgets/output",
          "model_name": "OutputModel",
          "model_module_version": "1.0.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/output",
            "_model_module_version": "1.0.0",
            "_model_name": "OutputModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/output",
            "_view_module_version": "1.0.0",
            "_view_name": "OutputView",
            "layout": "IPY_MODEL_f103a5b69a4042ecaa8c2b5692f68c10",
            "msg_id": "",
            "outputs": [
              {
                "output_type": "stream",
                "name": "stdout",
                "text": [
                  "Setup done.\n"
                ]
              }
            ]
          }
        },
        "f103a5b69a4042ecaa8c2b5692f68c10": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# LLM Conversion\n",
        "\n",
        "<link rel=\"stylesheet\" href=\"/mediapipe/site.css\">\n",
        "\n",
        "<table align=\"left\" class=\"buttons\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/googlesamples/mediapipe/blob/main/examples/llm_inference/conversion/llm_conversion.ipynb\" target=\"_blank\">\n",
        "      <img src=\"https://developers.google.com/static/mediapipe/solutions/customization/colab-logo-32px_1920.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://github.com/googlesamples/mediapipe/blob/main/examples/llm_inference/conversion/llm_conversion.ipynb\" target=\"_blank\">\n",
        "      <img src=\"https://developers.google.com/static/mediapipe/solutions/customization/github-logo-32px_1920.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ],
      "metadata": {
        "id": "iJyYBXrg79-N"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title License information { display-mode: \"form\" }\n",
        "#@markdown Copyright 2024 The MediaPipe Authors.\n",
        "#@markdown Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "#@markdown\n",
        "#@markdown you may not use this file except in compliance with the License.\n",
        "#@markdown You may obtain a copy of the License at\n",
        "#@markdown\n",
        "#@markdown https://www.apache.org/licenses/LICENSE-2.0\n",
        "#@markdown\n",
        "#@markdown Unless required by applicable law or agreed to in writing, software\n",
        "#@markdown distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "#@markdown WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "#@markdown See the License for the specific language governing permissions and\n",
        "#@markdown limitations under the License."
      ],
      "metadata": {
        "id": "Gy8YJJJgx_4G"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "5yYF61P0xseL",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34,
          "referenced_widgets": [
            "e812825502c9458891c62c1010a6fe54",
            "f103a5b69a4042ecaa8c2b5692f68c10"
          ]
        },
        "outputId": "9781c22f-34b7-44e8-d5e4-f025ed1befc1"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Output()"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "e812825502c9458891c62c1010a6fe54"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "# @title Setup { display-mode: \"form\" }\n",
        "# @markdown import ipywidgets\\\n",
        "# @markdown import IPython.display\\\n",
        "# @markdown import os\\\n",
        "# @markdown import huggingface downloader\\\n",
        "# @markdown import mediapipe genai converter\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import display\n",
        "install_out = widgets.Output()\n",
        "display(install_out)\n",
        "with install_out:\n",
        "  !pip install mediapipe\n",
        "  !pip install torch\n",
        "  !pip install huggingface_hub\n",
        "  import os\n",
        "  from huggingface_hub import hf_hub_download\n",
        "  from mediapipe.tasks.python.genai import converter\n",
        "\n",
        "install_out.clear_output()\n",
        "with install_out:\n",
        "  print(\"Setup done.\")\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Run { display-mode: \"form\"}\n",
        "\n",
        "model = widgets.Dropdown(\n",
        "    options=[\"Gemma 2B\", \"Falcon 1B\", \"StableLM 3B\", \"Phi 2\"],\n",
        "    value='Gemma 2B',\n",
        "    description='model',\n",
        "    disabled=False,\n",
        ")\n",
        "\n",
        "backend = widgets.Dropdown(\n",
        "    options=[\"cpu\", \"gpu\"],\n",
        "    value='cpu',\n",
        "    description='backend',\n",
        "    disabled=False,\n",
        ")\n",
        "\n",
        "token = widgets.Password(\n",
        "    value='',\n",
        "    placeholder='huggingface token',\n",
        "    description='HF token:',\n",
        "    disabled=False\n",
        ")\n",
        "\n",
        "def on_change_model(change):\n",
        "  if change[\"new\"] != 'Gemma 2b':\n",
        "    token_description.layout.display = \"none\"\n",
        "    token.layout.display = \"none\"\n",
        "  else:\n",
        "    token_description.layout.display = \"flex\"\n",
        "    token.layout.display = \"flex\"\n",
        "\n",
        "model.observe(on_change_model, names=['value'])\n",
        "\n",
        "\n",
        "\n",
        "display(model)\n",
        "display(backend)\n",
        "\n",
        "token_description = widgets.Output()\n",
        "with token_description:\n",
        "  print(\"Huggingface token needed for gated model (e.g. Gemma)\")\n",
        "  print(\"You can get it from https://huggingface.co/settings/tokens\")\n",
        "\n",
        "display(token_description)\n",
        "display(token)\n",
        "\n",
        "\n",
        "\n",
        "def gemma_download(token):\n",
        "  REPO_ID = \"google/gemma-2b-it\"\n",
        "  FILENAMES = [\"tokenizer.json\", \"tokenizer_config.json\", \"model-00001-of-00002.safetensors\", \"model-00002-of-00002.safetensors\"]\n",
        "  os.environ['HF_TOKEN'] = token\n",
        "  with out:\n",
        "    for filename in FILENAMES:\n",
        "      hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=\"./gemma-2b-it\")\n",
        "\n",
        "def falcon_download():\n",
        "  REPO_ID = \"tiiuae/falcon-rw-1b\"\n",
        "  FILENAMES = [\"tokenizer.json\", \"tokenizer_config.json\", \"pytorch_model.bin\"]\n",
        "  with out:\n",
        "    for filename in FILENAMES:\n",
        "      hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=\"./falcon-rw-1b\")\n",
        "\n",
        "def stablelm_download():\n",
        "  REPO_ID = \"stabilityai/stablelm-3b-4e1t\"\n",
        "  FILENAMES = [\"tokenizer.json\", \"tokenizer_config.json\", \"model.safetensors\"]\n",
        "  with out:\n",
        "    for filename in FILENAMES:\n",
        "      hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=\"./stablelm-3b-4e1t\")\n",
        "\n",
        "def phi2_download():\n",
        "  REPO_ID = \"microsoft/phi-2\"\n",
        "  FILENAMES = [\"tokenizer.json\", \"tokenizer_config.json\", \"model-00001-of-00002.safetensors\", \"model-00002-of-00002.safetensors\"]\n",
        "  with out:\n",
        "    for filename in FILENAMES:\n",
        "      hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=\"./phi-2\")\n",
        "\n",
        "\n",
        "def gemma_convert_config(backend):\n",
        "  input_ckpt = '/content/gemma-2b-it/'\n",
        "  vocab_model_file = '/content/gemma-2b-it/'\n",
        "  output_dir = '/content/intermediate/gemma-2b-it/'\n",
        "  output_tflite_file = f'/content/converted_models/gemma_{backend}.bin'\n",
        "  return converter.ConversionConfig(input_ckpt=input_ckpt, ckpt_format='safetensors', model_type='GEMMA_2B', backend=backend, output_dir=output_dir, combine_file_only=False, vocab_model_file=vocab_model_file, output_tflite_file=output_tflite_file)\n",
        "\n",
        "def falcon_convert_config(backend):\n",
        "  input_ckpt = '/content/falcon-rw-1b/pytorch_model.bin'\n",
        "  vocab_model_file = '/content/falcon-rw-1b/'\n",
        "  output_dir = '/content/intermediate/falcon-rw-1b/'\n",
        "  output_tflite_file = f'/content/converted_models/falcon_{backend}.bin'\n",
        "  return converter.ConversionConfig(input_ckpt=input_ckpt, ckpt_format='pytorch', model_type='FALCON_RW_1B', backend=backend, output_dir=output_dir, combine_file_only=False, vocab_model_file=vocab_model_file, output_tflite_file=output_tflite_file)\n",
        "\n",
        "def stablelm_convert_config(backend):\n",
        "  input_ckpt = '/content/stablelm-3b-4e1t/'\n",
        "  vocab_model_file = '/content/stablelm-3b-4e1t/'\n",
        "  output_dir = '/content/intermediate/stablelm-3b-4e1t/'\n",
        "  output_tflite_file = f'/content/converted_models/stablelm_{backend}.bin'\n",
        "  return converter.ConversionConfig(input_ckpt=input_ckpt, ckpt_format='safetensors', model_type='STABLELM_4E1T_3B', backend=backend, output_dir=output_dir, combine_file_only=False, vocab_model_file=vocab_model_file, output_tflite_file=output_tflite_file)\n",
        "\n",
        "\n",
        "def phi2_convert_config(backend):\n",
        "  input_ckpt = '/content/phi-2'\n",
        "  vocab_model_file = '/content/phi-2/'\n",
        "  output_dir = '/content/intermediate/phi-2/'\n",
        "  output_tflite_file = f'/content/converted_models/phi2_{backend}.bin'\n",
        "\n",
        "  return converter.ConversionConfig(input_ckpt=input_ckpt, ckpt_format='safetensors', model_type='PHI_2', backend=backend, output_dir=output_dir, combine_file_only=False, vocab_model_file=vocab_model_file, output_tflite_file=output_tflite_file)\n",
        "\n",
        "\n",
        "def on_button_clicked(b):\n",
        "  try:\n",
        "    out.clear_output()\n",
        "    with out:\n",
        "      print(\"Downloading model ...\")\n",
        "    button.description = \"Downloading ...\"\n",
        "    button.disabled = True\n",
        "    model.disabled = True\n",
        "    backend.disabled = True\n",
        "\n",
        "    if model.value == 'Gemma 2B':\n",
        "      gemma_download(token.value)\n",
        "    elif model.value == 'Falcon 1B':\n",
        "      falcon_download()\n",
        "    elif model.value == 'StableLM 3B':\n",
        "      stablelm_download()\n",
        "    elif model.value == 'Phi 2':\n",
        "      phi2_download()\n",
        "    else:\n",
        "      raise Exception(\"Invalid model\")\n",
        "\n",
        "    with out:\n",
        "      print(\"Done\")\n",
        "      print(\"Converting model ...\")\n",
        "\n",
        "    button.description = \"Converting ...\"\n",
        "\n",
        "    if model.value == 'Gemma 2B':\n",
        "      config = gemma_convert_config(backend.value)\n",
        "    elif model.value == 'Falcon 1B':\n",
        "      config = falcon_convert_config(backend.value)\n",
        "    elif model.value == 'StableLM 3B':\n",
        "      config = stablelm_convert_config(backend.value)\n",
        "    elif model.value == 'Phi 2':\n",
        "      config = phi2_convert_config(backend.value)\n",
        "    else:\n",
        "      with out:\n",
        "        raise Exception(\"Invalid model\")\n",
        "      return\n",
        "\n",
        "    with out:\n",
        "      converter.convert_checkpoint(config)\n",
        "      print(\"Done\")\n",
        "\n",
        "    button.description = \"Start Conversion\"\n",
        "    button.disabled = False\n",
        "    model.disabled = False\n",
        "    backend.disabled = False\n",
        "\n",
        "  except Exception as e:\n",
        "    button.description = \"Start Conversion\"\n",
        "    button.disabled = False\n",
        "    model.disabled = False\n",
        "    backend.disabled = False\n",
        "    with out:\n",
        "      print(e)\n",
        "\n",
        "\n",
        "button = widgets.Button(description=\"Start Conversion\")\n",
        "\n",
        "button.on_click(on_button_clicked)\n",
        "display(button)\n",
        "\n",
        "out = widgets.Output(layout={'border': '1px solid black'})\n",
        "display(out)\n",
        "\n",
        "print(\"\\nNotice: Converted models are saved under ./converted_models\")"
      ],
      "metadata": {
        "id": "LSSxrLyQPofw",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "acd38821-87b6-41e5-ebd2-3a77ee5d5a46"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Dropdown(description='model', options=('Gemma 2B', 'Falcon 1B', 'StableLM 3B', 'Phi 2'), value='Gemma 2B')"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "2abb456d33314c8f99083e0a95d4741a"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Dropdown(description='backend', options=('cpu', 'gpu'), value='cpu')"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "1707c4ee793a40c988711595ae1e61a5"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Output()"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "c0cd062d121a4a50bd2b7a24f1708c97"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Password(description='HF token:', placeholder='huggingface token')"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "392fc422ab2541b698ff500e10f46814"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Button(description='Start Conversion', style=ButtonStyle())"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "d878aa74559145a2b848526779181011"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Output(layout=Layout(border='1px solid black'))"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f3f94f86a6444147a2b7140bf2cbcaae"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Notice: Converted models are saved under ./converted_models\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Save to google drive { display-mode: \"form\"}\n",
        "google_drive_directory = \"converted_models\" #@param {type:\"string\"}\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "print(\"Copying models ...\")\n",
        "!mkdir -p /content/drive/MyDrive/$google_drive_directory\n",
        "!cp -r -f /content/converted_models/* /content/drive/MyDrive/$google_drive_directory\n",
        "print(\"Done\")"
      ],
      "metadata": {
        "id": "Am3idu-FRkqf",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9ac502e5-dd6a-4e37-c535-567b9acd12ef"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n",
            "Copying models ...\n",
            "cp: cannot stat '/content/converted_models/*': No such file or directory\n",
            "Done\n"
          ]
        }
      ]
    }
  ]
}