{ "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 }