tokimoa/sarashina2.2-ocr-mlx-4bit

Sarashina2.2-OCR — MLX 4bit

sbintuitions/sarashina2.2-ocr（SB Intuitions・MIT）の Apple Silicon（MLX）向け4bit量子化版です。tokimoaが変換・検証して公開しています。

• 2.8GB（元は約7.5GB）。LLM側のみ4bit量子化、vision tower・mergerはbf16のまま（6.107 bits/weight）
• M4 Max実測: 広告画像1枚あたり平均2.8秒（transformers/MPS bf16版と同速・メモリ約1/3）
• 検証済み: 日本語広告画像15枚（横書き・縦書き・バナー混在）で、決定論の業務指標 （禁止語の抽出網羅・誤検出ゼロ・商品名保持）がtransformers bf16版と15/15で同値

Sarashina2.2-OCRは日英文書解析に特化した3BのOCRモデルで、縦書き・複雑レイアウト・表・数式に 対応します（詳細は元モデルカード参照）。

使い方

このモデルのアーキテクチャ（sarashina2_vision）はmlx-vlm本体に未収録のため、 リポジトリ同梱のモデルクラスを先にimportしてください。

pip install "mlx-vlm==0.6.6" torch torchvision pillow protobuf sentencepiece

import sys
from huggingface_hub import snapshot_download

path = snapshot_download("tokimoa/sarashina2.2-ocr-mlx-4bit")
sys.path.insert(0, path)
import sarashina2_vision  # noqa: F401  mlx_vlm.models.sarashina2_vision として登録

import mlx.core as mx
from pathlib import Path
from PIL import Image
from transformers import AutoProcessor
from mlx_vlm.utils import load_model
from mlx_vlm.models.cache import KVCache

model = load_model(Path(path))
processor = AutoProcessor.from_pretrained(path, trust_remote_code=True)

image = Image.open("document.png").convert("RGB")
message = [{"role": "user", "content": [{"type": "image", "image": image}]}]
inputs = processor.apply_chat_template(
    message, tokenize=True, add_generation_prompt=True,
    return_dict=True, return_tensors="np")

# greedy + repetition_penalty 1.2（元モデル推奨設定）
cache = [KVCache() for _ in model.layers]
logits = model(mx.array(inputs["input_ids"]), mx.array(inputs["pixel_values"]),
               cache=cache, image_grid_thw=mx.array(inputs["image_grid_thw"])).logits
seen = {int(t) for t in inputs["input_ids"][0]}
tokens = []
for _ in range(3000):
    logit = logits[0, -1, :].astype(mx.float32)
    idx = mx.array(sorted(seen))
    vals = logit[idx]
    logit[idx] = mx.where(vals < 0, vals * 1.2, vals / 1.2)
    tok = int(mx.argmax(logit).item())
    if tok == 2:  # </s>
        break
    tokens.append(tok)
    seen.add(tok)
    logits = model.language_model(mx.array([[tok]]), cache=cache).logits

print(processor.decode(tokens, skip_special_tokens=True))

変換メモ

• 変換: mlx_vlm.convert（q_bits=4, group_size=64）＋独自モデルクラス （Qwen2-VL ViT＋deepstack merger＋Llama-3B＋interleaved M-RoPE spatial_resetバリアント）
• configuration_sarashina2_vision.py はtransformers v5の型検証 （Qwen2VLVisionConfig.mlp_ratio: int と実値3.7362の衝突）を回避したパッチ版
• ライセンス: モデル重みはMIT（元モデル準拠）。同梱pythonコードのうちSB Intuitions由来の ファイルはApache-2.0（ヘッダ表記のとおり）

関連

• bf16版: tokimoa/sarashina2.2-ocr-mlx-bf16
• 変換・検証: tokimoa（日本語実務でのローカルLLM実測をやっています）