{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "gpuType": "T4"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert Huihui-gemma-4-E2B-it-abliterated \u2192 ONNX (q4f16) for Transformers.js\n",
    "\n",
    "**Runtime \u2192 Change runtime type \u2192 T4 GPU**, then run the cells in order (~20 min total).\n",
    "\n",
    "This notebook:\n",
    "1. Exports [huihui-ai/Huihui-gemma-4-E2B-it-abliterated](https://huggingface.co/huihui-ai/Huihui-gemma-4-E2B-it-abliterated) to ONNX (`image-text-to-text`, CUDA)\n",
    "2. Quantizes weights to 4-bit / activations fp16 (`q4f16`) with ONNX Runtime\n",
    "3. Uploads the web-ready subset to **`kermelp/Huihui-gemma-4-E2B-it-abliterated-ONNX`**\n",
    "\n",
    "The result powers the WebGPU demo Space (runs fully in the visitor's browser)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title 1. Install dependencies (~2 min)\n",
    "%pip install -q -U \"transformers>=5.5\" \"optimum[onnx]\" onnxruntime onnx onnxconverter-common huggingface_hub sentencepiece pillow\n",
    "import transformers\n",
    "print(\"transformers:\", transformers.__version__)\n",
    "assert transformers.__version__.startswith(\"5.\"), \"gemma4 needs transformers >= 5.5 - Runtime > Restart session, then rerun\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title 2. Export to ONNX (~10-15 min, downloads ~10 GB)\n",
    "import subprocess\n",
    "\n",
    "MODEL_ID = \"huihui-ai/Huihui-gemma-4-E2B-it-abliterated\"\n",
    "OUT_DIR = \"/content/onnx_export\"\n",
    "import os, shutil\n",
    "shutil.rmtree(OUT_DIR, ignore_errors=True)\n",
    "\n",
    "h = subprocess.run([\"optimum-cli\", \"export\", \"onnx\", \"--help\"], capture_output=True, text=True)\n",
    "flags = h.stdout + h.stderr\n",
    "\n",
    "attempts = [\n",
    "    [\"--task\", \"image-text-to-text\", \"--device\", \"cuda\", \"--dtype\", \"float16\"],\n",
    "    [\"--task\", \"image-text-to-text\", \"--device\", \"cuda\"],\n",
    "    [\"--task\", \"image-text-to-text\"],\n",
    "    [],\n",
    "]\n",
    "for extra in attempts:\n",
    "    # drop flags this optimum version does not know\n",
    "    cmd_extra = []\n",
    "    skip = False\n",
    "    for a in extra:\n",
    "        if skip:\n",
    "            skip = False\n",
    "            continue\n",
    "        if a.startswith(\"--\") and (a + \" \") not in flags and a not in flags:\n",
    "            print(\"flag not supported, skipping:\", a)\n",
    "            if a == \"--device\":\n",
    "                continue\n",
    "            skip = True  # also skip its value\n",
    "            continue\n",
    "        cmd_extra.append(a)\n",
    "    cmd = [\"optimum-cli\", \"export\", \"onnx\", \"--model\", MODEL_ID] + cmd_extra + [OUT_DIR]\n",
    "    print(\"\\nTRY:\", \" \".join(cmd), flush=True)\n",
    "    r = subprocess.run(cmd, env={**__import__(\"os\").environ, \"HF_XET_HIGH_PERFORMANCE\": \"1\"})\n",
    "    if r.returncode == 0:\n",
    "        print(\"\\nEXPORT OK\")\n",
    "        break\n",
    "    print(\"attempt failed, trying next configuration...\")\n",
    "else:\n",
    "    raise RuntimeError(\"all export attempts failed - see errors above\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title 3. Quantize to q4f16 (+ fp16 fallback) (~5 min)\n",
    "import glob, os\n",
    "import onnx\n",
    "from onnxconverter_common import float16 as onnx_float16\n",
    "from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer\n",
    "\n",
    "ONNX_DIR = os.path.join(OUT_DIR, \"onnx\") if os.path.isdir(os.path.join(OUT_DIR, \"onnx\")) else OUT_DIR\n",
    "\n",
    "for base_file in sorted(glob.glob(os.path.join(ONNX_DIR, \"*.onnx\"))):\n",
    "    if base_file.endswith((\"_fp16.onnx\", \"_q4.onnx\", \"_q4f16.onnx\")):\n",
    "        continue\n",
    "    stem = base_file[:-5]\n",
    "    print(\"quantizing\", os.path.basename(base_file), flush=True)\n",
    "    model = onnx.load(base_file)\n",
    "    try:\n",
    "        q = MatMulNBitsQuantizer(model, bits=4, block_size=32, is_symmetric=True)\n",
    "        q.process()\n",
    "        q4_path = stem + \"_q4.onnx\"\n",
    "        onnx.save(q.model.model, q4_path)\n",
    "        try:\n",
    "            fp16_q4 = onnx_float16.convert_float_to_float16(onnx.load(q4_path), keep_io_types=True)\n",
    "            onnx.save(fp16_q4, stem + \"_q4f16.onnx\")\n",
    "            print(\"  -> q4f16 ok\")\n",
    "        except Exception as e:\n",
    "            print(\"  -> q4f16 failed (kept q4):\", e)\n",
    "    except Exception as e:\n",
    "        print(\"  -> q4 quantization failed, falling back to fp16:\", e)\n",
    "        m = onnx.load(base_file)\n",
    "        onnx.save(onnx_float16.convert_float_to_float16(m, keep_io_types=True), stem + \"_fp16.onnx\")\n",
    "\n",
    "for f in sorted(glob.glob(os.path.join(ONNX_DIR, \"*\"))):\n",
    "    print(f\"{os.path.getsize(f)/1e9:6.2f} GB  {os.path.basename(f)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title 4. Log in to Hugging Face\n",
    "# Get a WRITE token at https://huggingface.co/settings/tokens - paste it when prompted.\n",
    "from huggingface_hub import notebook_login\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title 5. Upload web-ready files (~5 min)\n",
    "import glob, os, shutil\n",
    "from huggingface_hub import upload_folder, list_repo_files, whoami\n",
    "\n",
    "REPO_ID = f\"{whoami()['name']}/Huihui-gemma-4-E2B-it-abliterated-ONNX\"\n",
    "print(\"uploading to\", REPO_ID)\n",
    "\n",
    "STAGE = \"/content/stage\"\n",
    "!rm -rf {STAGE}\n",
    "os.makedirs(os.path.join(STAGE, \"onnx\"), exist_ok=True)\n",
    "\n",
    "# 1) root config/tokenizer files produced by the export\n",
    "KEEP_ROOT_EXT = (\".json\", \".jinja\", \".model\", \".txt\")\n",
    "for f in glob.glob(os.path.join(OUT_DIR, \"*\")):\n",
    "    name = os.path.basename(f)\n",
    "    if os.path.isfile(f) and name.lower().endswith(KEEP_ROOT_EXT):\n",
    "        shutil.copy2(f, os.path.join(STAGE, name))\n",
    "\n",
    "# 2) browser-ready ONNX variants: *_fp16* and *_q4f16* graphs + their data blobs,\n",
    "#    plus small (<200 MB) unsuffixed encoder/embed graphs as fallbacks\n",
    "def wanted(base, size):\n",
    "    is_variant_data = (\"_fp16.\" in base or \"_q4f16.\" in base) and (base.endswith(\".onnx\") or \".onnx_data\" in base)\n",
    "    small_base_graph = base.endswith(\".onnx\") and size < 200e6 and not any(s in base for s in (\"decoder_model_merged\",))\n",
    "    return is_variant_data or small_base_graph\n",
    "\n",
    "for f in glob.glob(os.path.join(ONNX_DIR, \"*\")):\n",
    "    base = os.path.basename(f)\n",
    "    if os.path.isfile(f) and wanted(base, os.path.getsize(f)):\n",
    "        shutil.copy2(f, os.path.join(STAGE, \"onnx\", base))\n",
    "\n",
    "total = sum(os.path.getsize(f) for f in glob.glob(STAGE + \"/**/*\", recursive=True) if os.path.isfile(f))\n",
    "print(f\"staging {total/1e9:.2f} GB\")\n",
    "for f in sorted(glob.glob(STAGE + \"/**/*\", recursive=True)):\n",
    "    if os.path.isfile(f):\n",
    "        print(f\"  {os.path.getsize(f)/1e9:6.2f} GB  {f.replace(STAGE + '/', '')}\")\n",
    "\n",
    "upload_folder(folder_path=STAGE, repo_id=REPO_ID, repo_type=\"model\")\n",
    "print(\"\\ndone! files now in repo:\")\n",
    "for f in sorted(list_repo_files(REPO_ID)):\n",
    "    print(\" \", f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finished!\n",
    "Open the demo: **https://huggingface.co/spaces/kermelp/Huihui-gemma4-abliterated-webgpu**\n",
    "(first load downloads ~4 GB into your browser cache; needs a WebGPU-capable browser - Chrome/Edge).\n",
    "\n",
    "If cell 3 shows `q4f16` for `decoder_model_merged`, `embed_tokens`, `vision_encoder`, and `audio_encoder`, everything the demo needs is up."
   ]
  }
 ]
}