Upload glm_ocr_coreml.ipynb with huggingface_hub

glm_ocr_coreml.ipynb

@@ -0,0 +1,566 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# GLM-OCR to CoreML Conversion\n",
+        "\n",
+        "This notebook converts the [GLM-OCR](https://huggingface.co/aoiandroid/GLM-OCR) model (image-to-text OCR) to CoreML for use on iOS/macOS.\n",
+        "\n",
+        "**Model**: Multimodal OCR (CogViT visual encoder + cross-modal connector + GLM-0.5B decoder).  \n",
+        "**Output**: Vision encoder as CoreML (`vision_encoder.mlpackage`), plus tokenizer/config for app-side use.\n",
+        "\n",
+        "**Requirements**: Python 3.10+, PyTorch, transformers (main branch for GLM-OCR support), coremltools. Colab or local GPU recommended."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Install dependencies (uncomment in Colab or fresh env).\n",
+        "# For reproducible builds: pip install -r glm_ocr_coreml_requirements.txt\n",
+        "# Or with versions:\n",
+        "# !pip install -q torch==2.3.0 torchvision==0.18.0\n",
+        "# !pip install -q \"git+https://github.com/huggingface/transformers.git@main\"\n",
+        "# !pip install -q coremltools==7.2\n",
+        "# !pip install -q huggingface_hub>=0.23.0 pillow>=10.3.0"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "from pathlib import Path\n",
+        "\n",
+        "import numpy as np\n",
+        "import torch\n",
+        "import coremltools as ct\n",
+        "from PIL import Image"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 1. Load model and processor\n",
+        "\n",
+        "Using `aoiandroid/GLM-OCR` (duplicate of `zai-org/GLM-OCR`). Ensure transformers supports GLM-OCR (install from main: `pip install git+https://github.com/huggingface/transformers.git`)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "MODEL_ID = \"aoiandroid/GLM-OCR\"  # or \"zai-org/GLM-OCR\"\n",
+        "OUTPUT_DIR = Path(\"./glm_ocr_coreml\")\n",
+        "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
+        "\n",
+        "# Load processor and model (use float32 for tracing; bfloat16 may not trace well)\n",
+        "from transformers import AutoProcessor, AutoModelForImageTextToText\n",
+        "\n",
+        "processor = AutoProcessor.from_pretrained(MODEL_ID)\n",
+        "model = AutoModelForImageTextToText.from_pretrained(\n",
+        "    MODEL_ID,\n",
+        "    torch_dtype=torch.float32,\n",
+        ")\n",
+        "model.eval()\n",
+        "\n",
+        "# Vision config for input shape (default image_size=336)\n",
+        "vision_config = getattr(model.config, \"vision_config\", None)\n",
+        "image_size = 336\n",
+        "if vision_config is not None:\n",
+        "    image_size = getattr(vision_config, \"image_size\", 336)\n",
+        "if isinstance(image_size, (list, tuple)):\n",
+        "    image_size = image_size[0]\n",
+        "hidden_size = getattr(model.config, \"hidden_size\", None) or (getattr(model.config.text_config, \"hidden_size\", 1024) if getattr(model.config, \"text_config\", None) else 1024)\n",
+        "print(f\"Image size: {image_size}, hidden_size: {hidden_size}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### 1.1 Model structure validation\n",
+        "\n",
+        "Verify that the loaded model has the expected attributes (`model.model`, `get_image_features`). Check for a language/decoder submodule for decoder export."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Model structure validation (required for decoder export)\n",
+        "print(\"=== Model structure ===\")\n",
+        "print(f\"Model class: {type(model).__name__}\")\n",
+        "print(f\"Public attributes: {[a for a in dir(model) if not a.startswith('_')]}\")\n",
+        "\n",
+        "inner = getattr(model, \"model\", None)\n",
+        "if inner is None:\n",
+        "    raise RuntimeError(\"model.model not found. Inspect the loaded model structure.\")\n",
+        "\n",
+        "if not hasattr(inner, \"get_image_features\"):\n",
+        "    raise RuntimeError(\n",
+        "        \"get_image_features not found. Install transformers from main: \"\n",
+        "        \"pip install git+https://github.com/huggingface/transformers.git\"\n",
+        "    )\n",
+        "\n",
+        "print(f\"vision_config: {getattr(model.config, 'vision_config', 'N/A')}\")\n",
+        "print(f\"hidden_size: {getattr(model.config, 'hidden_size', 'N/A')}\")\n",
+        "\n",
+        "# For decoder: look for language/text/decoder submodule on model or model.model\n",
+        "decoder_candidates = [\"language_model\", \"text_model\", \"decoder\", \"model\"]\n",
+        "for name in decoder_candidates:\n",
+        "    obj = getattr(model, name, None) or getattr(inner, name, None)\n",
+        "    if obj is not None and hasattr(obj, \"forward\"):\n",
+        "        print(f\"Decoder candidate: {name} (on model or model.model)\")\n",
+        "print(\"Structure validation OK\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 2. Export vision encoder to CoreML\n",
+        "\n",
+        "The vision part of GLM-OCR turns `pixel_values` into hidden states consumed by the language model. We trace `get_image_features(pixel_values)` to obtain a CoreML vision encoder. The app can then run this and feed the outputs into a separate decoder or use the rest of the pipeline in Swift."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Wrapper: pixel_values -> last_hidden_state\n",
+        "# GlmOcrForConditionalGeneration has .model (GlmOcrModel) with get_image_features\n",
+        "class VisionEncoderWrapper(torch.nn.Module):\n",
+        "    def __init__(self, parent_model):\n",
+        "        super().__init__()\n",
+        "        self.base = getattr(parent_model, \"model\", parent_model)\n",
+        "        if not hasattr(self.base, \"get_image_features\"):\n",
+        "            raise AttributeError(\"Loaded model has no get_image_features; ensure transformers supports GLM-OCR.\")\n",
+        "\n",
+        "    def forward(self, pixel_values: torch.Tensor):\n",
+        "        out = self.base.get_image_features(pixel_values=pixel_values)\n",
+        "        return out.last_hidden_state\n",
+        "\n",
+        "wrapper = VisionEncoderWrapper(model)\n",
+        "wrapper.eval()\n",
+        "\n",
+        "batch, channels = 1, 3\n",
+        "dummy_pixel = torch.randn(batch, channels, image_size, image_size, dtype=torch.float32)\n",
+        "\n",
+        "with torch.no_grad():\n",
+        "    traced = torch.jit.trace(\n",
+        "        wrapper,\n",
+        "        (dummy_pixel,),\n",
+        "        check_trace=False,\n",
+        "        strict=False,\n",
+        "    )\n",
+        "# Check output shape\n",
+        "with torch.no_grad():\n",
+        "    out = traced(dummy_pixel)\n",
+        "print(f\"Vision encoder output shape: {out.shape}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Convert vision encoder to CoreML\n",
+        "# Output shape (1, vision_seq_len, hidden_size) - use actual shape from trace\n",
+        "vision_seq_len = out.shape[1]\n",
+        "hidden_size = out.shape[2]\n",
+        "\n",
+        "input_types = [\n",
+        "    ct.TensorType(\n",
+        "        name=\"pixel_values\",\n",
+        "        shape=(1, channels, image_size, image_size),\n",
+        "        dtype=np.float32,\n",
+        "    )\n",
+        "]\n",
+        "output_types = [ct.TensorType(name=\"vision_hidden_states\")]\n",
+        "\n",
+        "# Use iOS16 for reliability; set to iOS15 or iOS17 per target device if needed\n",
+        "vision_mlmodel = ct.convert(\n",
+        "    traced,\n",
+        "    inputs=input_types,\n",
+        "    outputs=output_types,\n",
+        "    convert_to=\"mlprogram\",\n",
+        "    minimum_deployment_target=ct.target.iOS16,\n",
+        "    compute_units=ct.ComputeUnit.ALL,\n",
+        ")\n",
+        "\n",
+        "vision_path = OUTPUT_DIR / \"vision_encoder.mlpackage\"\n",
+        "vision_mlmodel.save(str(vision_path))\n",
+        "print(f\"Saved vision encoder to {vision_path}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Save vision encoder spec for Swift (vision_seq_len, hidden_size, image_size)\n",
+        "import json\n",
+        "\n",
+        "model_spec = {\n",
+        "    \"vision_encoder\": {\n",
+        "        \"input\": {\n",
+        "            \"name\": \"pixel_values\",\n",
+        "            \"shape\": [1, 3, int(image_size), int(image_size)],\n",
+        "            \"dtype\": \"float32\",\n",
+        "        },\n",
+        "        \"output\": {\n",
+        "            \"name\": \"vision_hidden_states\",\n",
+        "            \"shape\": [1, int(vision_seq_len), int(hidden_size)],\n",
+        "            \"dtype\": \"float32\",\n",
+        "        },\n",
+        "    },\n",
+        "    \"image_size\": int(image_size),\n",
+        "    \"vision_seq_len\": int(vision_seq_len),\n",
+        "    \"hidden_size\": int(hidden_size),\n",
+        "    \"model_id\": MODEL_ID,\n",
+        "}\n",
+        "\n",
+        "spec_path = OUTPUT_DIR / \"model_spec.json\"\n",
+        "with open(spec_path, \"w\") as f:\n",
+        "    json.dump(model_spec, f, indent=2)\n",
+        "print(f\"Model spec saved: {spec_path}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 3. Save processor and config\n",
+        "\n",
+        "Copy tokenizer and config so the app can run preprocessing and decoding. Full autoregressive decoding (image + prompt -> text) would require either exporting the decoder as a second CoreML model or implementing the generation loop in Swift using the vision encoder output."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Save processor (tokenizer + image processor) and config to output dir\n",
+        "processor.save_pretrained(OUTPUT_DIR)\n",
+        "model.config.save_pretrained(OUTPUT_DIR)\n",
+        "print(f\"Saved processor and config to {OUTPUT_DIR}\")\n",
+        "print(\"Contents:\", list(OUTPUT_DIR.iterdir()))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 4. Verify CoreML I/O (optional)\n",
+        "\n",
+        "Inspect input/output names and shapes for integration in Swift."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "loaded = ct.models.MLModel(str(vision_path))\n",
+        "spec = loaded.get_spec()\n",
+        "print(\"Vision encoder inputs:\", [d.name for d in spec.description.input])\n",
+        "print(\"Vision encoder outputs:\", [d.name for d in spec.description.output])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 5. Optional: Decoder or full-model export\n",
+        "\n",
+        "The full GLM-OCR pipeline (image + prompt -> generated text) uses `model.generate()` with cache and variable sequence length, which is hard to export as a single CoreML model. Options:\n",
+        "\n",
+        "- **Vision encoder only** (done above): Use `vision_encoder.mlpackage` in the app and implement the decoder/generation loop in Swift, or call a separate decoder CoreML if you export it.\n",
+        "- **Decoder export**: Trace the text model with fixed `encoder_hidden_states` (from the vision encoder output) and `input_ids` to get logits; then run autoregressive generation in the app. This requires building a wrapper that takes (input_ids, encoder_hidden_states, attention_mask) and returns logits, similar to T5/encoder-decoder conversion scripts.\n",
+        "- **Quantization**: Use `coremltools.optimize.coreml.palettize_weights` or `linear_quantize_weights` to reduce vision encoder size (e.g. INT8 or 4-bit)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### 2.1 Quantization (FP16 / INT8) and size comparison\n",
+        "\n",
+        "Apply FP16 and INT8 quantization to reduce vision encoder size for iOS. **After INT8 quantization, run the accuracy verification cell (Section 6) below.**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import shutil\n",
+        "from coremltools.optimize.coreml import (\n",
+        "    linear_quantize_weights,\n",
+        "    OptimizationConfig,\n",
+        "    OpLinearQuantizerConfig,\n",
+        ")\n",
+        "\n",
+        "# FP16 (minimal accuracy loss)\n",
+        "vision_fp16 = ct.models.MLModel(str(vision_path))\n",
+        "vision_fp16_path = OUTPUT_DIR / \"vision_encoder_fp16.mlpackage\"\n",
+        "try:\n",
+        "    q16 = ct.models.neural_network.quantization_utils.quantize_weights(vision_fp16, nbits=16)\n",
+        "    q16.save(str(vision_fp16_path))\n",
+        "except Exception as e:\n",
+        "    print(f\"FP16 quantization failed: {e}\")\n",
+        "    vision_fp16_path = None\n",
+        "\n",
+        "# INT8 (smaller; run accuracy verification after)\n",
+        "config = OptimizationConfig(\n",
+        "    global_config=OpLinearQuantizerConfig(mode=\"linear_symmetric\", weight_threshold=512)\n",
+        ")\n",
+        "vision_int8 = linear_quantize_weights(vision_mlmodel, config)\n",
+        "vision_int8_path = OUTPUT_DIR / \"vision_encoder_int8.mlpackage\"\n",
+        "vision_int8.save(str(vision_int8_path))\n",
+        "\n",
+        "# Size comparison (MB)\n",
+        "for label, path in [\n",
+        "    (\"FP32 (original)\", vision_path),\n",
+        "    (\"FP16\", vision_fp16_path),\n",
+        "    (\"INT8\", vision_int8_path),\n",
+        "]:\n",
+        "    if path is not None and path.exists():\n",
+        "        size_mb = sum(f.stat().st_size for f in path.rglob(\"*\") if f.is_file()) / 1e6\n",
+        "        print(f\"{label}: {size_mb:.1f} MB\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 6. Accuracy verification (PyTorch vs CoreML)\n",
+        "\n",
+        "Compare vision encoder outputs: PyTorch traced model vs CoreML. Use a test image (or a dummy image if `test_image.png` is missing). Cosine similarity per token should be close to 1.0."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "from numpy.linalg import norm\n",
+        "\n",
+        "# Test image: use test_image.png if present, else dummy (shape-only check)\n",
+        "test_image_path = Path(\"test_image.png\")\n",
+        "if test_image_path.exists():\n",
+        "    test_image = Image.open(test_image_path).convert(\"RGB\")\n",
+        "    inputs = processor(images=test_image, return_tensors=\"pt\")\n",
+        "    pixel_values = inputs[\"pixel_values\"].to(torch.float32)\n",
+        "    if pixel_values.shape[2] != image_size or pixel_values.shape[3] != image_size:\n",
+        "        pixel_values = torch.nn.functional.interpolate(\n",
+        "            pixel_values, size=(image_size, image_size), mode=\"bilinear\"\n",
+        "        )\n",
+        "else:\n",
+        "    pixel_values = torch.randn(1, 3, image_size, image_size, dtype=torch.float32)\n",
+        "    print(\"No test_image.png; using dummy tensor for shape verification.\")\n",
+        "\n",
+        "# PyTorch output\n",
+        "with torch.no_grad():\n",
+        "    pt_out = traced(pixel_values).numpy()\n",
+        "\n",
+        "# CoreML output (FP32 model)\n",
+        "pv_np = pixel_values.cpu().numpy() if pixel_values.is_cuda else pixel_values.numpy()\n",
+        "coreml_out = vision_mlmodel.predict({\"pixel_values\": pv_np})[\"vision_hidden_states\"]\n",
+        "\n",
+        "# Cosine similarity per token (average and min)\n",
+        "cos_sims = []\n",
+        "for i in range(pt_out.shape[1]):\n",
+        "    a, b = pt_out[0, i], coreml_out[0, i]\n",
+        "    n = norm(a) * norm(b)\n",
+        "    cos_sims.append(np.dot(a, b) / n if n > 0 else 1.0)\n",
+        "print(f\"Cosine similarity (PyTorch vs CoreML FP32) mean: {np.mean(cos_sims):.6f}, min: {np.min(cos_sims):.6f}\")\n",
+        "assert np.mean(cos_sims) > 0.999, \"Accuracy drop too large; check conversion settings.\"\n",
+        "print(\"Accuracy verification OK\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 7. Decoder export (single-step, optional)\n",
+        "\n",
+        "Export a one-step decoder: `(input_ids, encoder_hidden_states, attention_mask) -> logits`, so the app can run an autoregressive loop in Swift. **GLM-OCR may not expose a separate decoder API** (it merges vision and text in one forward). If trace fails, only the vision encoder is used; implement generation in Swift or call the full model in Python."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Decoder export: try single-step (input_ids, encoder_hidden_states, attention_mask) -> logits\n",
+        "# GLM-OCR may merge vision+text in one forward; we try building inputs_embeds from vision + text embeddings.\n",
+        "decoder_exported = False\n",
+        "DECODER_MAX_LEN = max(256, int(vision_seq_len) + 64)  # ensure text segment exists for trace\n",
+        "\n",
+        "try:\n",
+        "    inner = model.model\n",
+        "    embed_fn = getattr(model, \"get_input_embeddings\", None) or getattr(inner, \"get_input_embeddings\", None)\n",
+        "    if embed_fn is None:\n",
+        "        raise AttributeError(\"No get_input_embeddings on model\")\n",
+        "\n",
+        "    class DecoderStepWrapper(torch.nn.Module):\n",
+        "        def __init__(self, parent_model):\n",
+        "            super().__init__()\n",
+        "            self.inner = parent_model.model\n",
+        "            self.lm_head = parent_model.lm_head\n",
+        "            self.embed = parent_model.get_input_embeddings()\n",
+        "\n",
+        "        def forward(\n",
+        "            self,\n",
+        "            input_ids: torch.Tensor,\n",
+        "            encoder_hidden_states: torch.Tensor,\n",
+        "            attention_mask: torch.Tensor,\n",
+        "        ):\n",
+        "            # Assume sequence layout: [image tokens (vision_seq_len), text tokens (rest)]\n",
+        "            seq_len = input_ids.shape[1]\n",
+        "            if encoder_hidden_states.shape[1] != vision_seq_len:\n",
+        "                raise ValueError(\"encoder_hidden_states seq len must match vision_seq_len\")\n",
+        "            text_len = seq_len - vision_seq_len\n",
+        "            if text_len <= 0:\n",
+        "                text_emb = self.embed(input_ids)\n",
+        "                inputs_embeds = encoder_hidden_states\n",
+        "            else:\n",
+        "                text_emb = self.embed(input_ids[:, vision_seq_len:])\n",
+        "                inputs_embeds = torch.cat([encoder_hidden_states, text_emb], dim=1)\n",
+        "            out = self.inner(\n",
+        "                attention_mask=attention_mask,\n",
+        "                inputs_embeds=inputs_embeds,\n",
+        "                use_cache=False,\n",
+        "            )\n",
+        "            return self.lm_head(out.last_hidden_state)\n",
+        "\n",
+        "    dec_wrapper = DecoderStepWrapper(model)\n",
+        "    dec_wrapper.eval()\n",
+        "    dummy_ids = torch.randint(0, 1000, (1, DECODER_MAX_LEN), dtype=torch.long)\n",
+        "    dummy_enc = torch.randn(1, vision_seq_len, hidden_size, dtype=torch.float32)\n",
+        "    dummy_attn = torch.ones(1, DECODER_MAX_LEN, dtype=torch.long)\n",
+        "    with torch.no_grad():\n",
+        "        dec_traced = torch.jit.trace(\n",
+        "            dec_wrapper,\n",
+        "            (dummy_ids, dummy_enc, dummy_attn),\n",
+        "            check_trace=False,\n",
+        "            strict=False,\n",
+        "        )\n",
+        "    print(\"Decoder trace OK; converting to CoreML...\")\n",
+        "except Exception as e:\n",
+        "    print(f\"Decoder export skipped: {e}\")\n",
+        "    print(\"Use vision encoder only; implement autoregressive decoding in Swift or run full model in Python.\")\n",
+        "    dec_traced = None"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "if dec_traced is not None:\n",
+        "    dec_input_types = [\n",
+        "        ct.TensorType(name=\"input_ids\", shape=(1, DECODER_MAX_LEN), dtype=np.int32),\n",
+        "        ct.TensorType(name=\"encoder_hidden_states\", shape=(1, vision_seq_len, hidden_size), dtype=np.float32),\n",
+        "        ct.TensorType(name=\"attention_mask\", shape=(1, DECODER_MAX_LEN), dtype=np.int32),\n",
+        "    ]\n",
+        "    dec_output_types = [ct.TensorType(name=\"logits\")]\n",
+        "    decoder_mlmodel = ct.convert(\n",
+        "        dec_traced,\n",
+        "        inputs=dec_input_types,\n",
+        "        outputs=dec_output_types,\n",
+        "        convert_to=\"mlprogram\",\n",
+        "        minimum_deployment_target=ct.target.iOS16,\n",
+        "        compute_units=ct.ComputeUnit.ALL,\n",
+        "    )\n",
+        "    decoder_path = OUTPUT_DIR / \"decoder.mlpackage\"\n",
+        "    decoder_mlmodel.save(str(decoder_path))\n",
+        "    print(f\"Saved decoder to {decoder_path}\")\n",
+        "    decoder_exported = True\n",
+        "    # Update model_spec with decoder I/O\n",
+        "    model_spec[\"decoder\"] = {\n",
+        "        \"input\": {\"names\": [\"input_ids\", \"encoder_hidden_states\", \"attention_mask\"], \"shapes\": [(1, DECODER_MAX_LEN), (1, vision_seq_len, hidden_size), (1, DECODER_MAX_LEN)]},\n",
+        "        \"output\": {\"name\": \"logits\", \"shape\": [1, DECODER_MAX_LEN, int(getattr(model.config, \"vocab_size\", getattr(model.config.text_config, \"vocab_size\", 59392)))]},\n",
+        "    }\n",
+        "    with open(spec_path, \"w\") as f:\n",
+        "        json.dump(model_spec, f, indent=2)\n",
+        "else:\n",
+        "    print(\"Decoder not exported; model_spec unchanged.\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 8. Swift integration sketch\n",
+        "\n",
+        "Use the vision encoder (and optional decoder) in an iOS app as below. Add `vision_encoder.mlpackage` to the Xcode project; if the decoder was exported, add `decoder.mlpackage` and run an autoregressive loop."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "swift_example = \"\"\"\n",
+        "// Swift: CoreML vision encoder + optional decoder loop\n",
+        "// 1. Add vision_encoder.mlpackage (and decoder.mlpackage if exported) to the Xcode project.\n",
+        "// 2. Preprocess image to 336x336 float32 and run vision encoder.\n",
+        "\n",
+        "import CoreML\n",
+        "import Vision\n",
+        "\n",
+        "let visionModel = try VisionEncoder(configuration: MLModelConfiguration())\n",
+        "let pixelValues = preprocessImage(uiImage)  // shape (1, 3, 336, 336), Float32\n",
+        "\n",
+        "let input = VisionEncoderInput(pixel_values: pixelValues)\n",
+        "let output = try visionModel.prediction(input: input)\n",
+        "let hiddenStates = output.vision_hidden_states  // (1, vision_seq_len, hidden_size)\n",
+        "\n",
+        "// Pass hiddenStates to the decoder for text generation:\n",
+        "// - If decoder.mlpackage was exported: load DecoderStep, then in a loop feed\n",
+        "//   (input_ids, encoder_hidden_states, attention_mask) and take argmax(logits) for next token.\n",
+        "// - Otherwise implement the generation loop in Swift or call the full model elsewhere.\n",
+        "\"\"\"\n",
+        "print(swift_example)"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.10.0"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 4
+}