HomeAgent/scripts/export_chineseclip_onnx.py

#!/usr/bin/env python3
"""把 Chinese-CLIP ViT-B/16 导出成 HomeAgent 的规范 ONNX 产物。

为什么是 Chinese-CLIP：
  text+image 的默认向量空间要同时满足「小、可商用、中文原生」。
  Chinese-CLIP ViT-B/16 = 188M 参数 / 721MB ONNX / 实测常驻 1.15GB，
  许可是 Apache-2.0（可随发行版分发），且原生在 2 亿中文图文对上训练。
  对比：jina-v5-omni-nano 2.23GB 但 CC BY-NC（不可商用）；Qwen3-VL-Emb-2B
  9.4GB（质量最好，保留为可选 provider）。

产物（--out 目录，会被清空重建）：
  TextEncoder.onnx     input_ids[·,52] + attention_mask[·,52] → text_features[·,512]
  VisionEncoder.onnx   pixel_values[·,3,224,224]              → image_features[·,512]
  embed_config.json    维度/预处理/分词超参/文件名（provider 侧的唯一权威）
  vocab.txt            分词器词表（来自官方模型目录）
  reference.json       冻结参考：逐文本 token id + 逐样本参考向量（Go 侧回归用）
  SHA256SUMS

自检纪律（对齐 export_qwen3vl_embedding_onnx.py）：
  1. 每个用例都真跑一次 ONNX 并与 PyTorch 对比，打印逐用例 cos；
  2. 计划用例集合与实际执行集合必须相等，否则非零退出（防「先跳过再校验」的假通过）；
  3. 不采信退出码，失败一律非零退出并说明原因。

用法：
  scripts/export_chineseclip_onnx.py --out DIR [--model-dir DIR|--model-id REPO]
                                     [--verify-only] [--no-reference] [--skip-verify]
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import sys
import time

import numpy as np

DEFAULT_MODEL_ID = "OFA-Sys/chinese-clip-vit-base-patch16"
TEXT_FILE = "TextEncoder.onnx"
VISION_FILE = "VisionEncoder.onnx"
CONFIG_FILE = "embed_config.json"
REFERENCE_FILE = "reference.json"
OPSET = 17
MAX_LENGTH = 52
IMAGE_SIZE = 224

# 冻结用例：文本覆盖纯中文/中英混/长文本/标点，图像覆盖纯色与渐变。
FIXTURE_TEXTS = [
    "一张红色方块的图片",
    "一只猫在草地上",
    "蓝色的天空",
    "HomeAgent 是一个本地 AI 管家",
    "这是一段比较长的中文文本，用来验证分词器在超过五十个 token 时的截断行为是否正确，"
    "同时检查标点符号、数字 12345 和英文单词 embedding 的处理。",
]
FIXTURE_IMAGES = [
    ("red", (220, 30, 30)),
    ("green", (60, 120, 60)),
    ("blue", (30, 30, 220)),
    ("gray", (128, 128, 128)),
]


def solid(color: tuple[int, int, int]) -> np.ndarray:
    from PIL import Image

    img = Image.new("RGB", (320, 320), color)
    return np.asarray(img, dtype=np.uint8)


def gradient() -> np.ndarray:
    """确定性的横向渐变，避免只有纯色导致区分度不足。"""
    row = np.linspace(0, 255, 320, dtype=np.uint8)
    img = np.zeros((320, 320, 3), dtype=np.uint8)
    img[:, :, 0] = row[None, :]
    img[:, :, 1] = row[:, None]
    img[:, :, 2] = 64
    return img


def norm_cos(a: np.ndarray, b: np.ndarray) -> float:
    a = a.reshape(-1).astype(np.float64)
    b = b.reshape(-1).astype(np.float64)
    na, nb = np.linalg.norm(a), np.linalg.norm(b)
    if na == 0 or nb == 0:
        return 0.0
    return float(np.dot(a, b) / (na * nb))


def sha256_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def load_model(model_dir: str | None, model_id: str, local_only: bool):
    import torch
    from transformers import ChineseCLIPModel, ChineseCLIPProcessor

    src = model_dir or model_id
    kw = {"local_files_only": True} if local_only else {}
    print(f"[load] {src}")
    t0 = time.time()
    processor = ChineseCLIPProcessor.from_pretrained(src, **kw)
    model = ChineseCLIPModel.from_pretrained(src, **kw).eval()
    print(f"[load] 用时 {time.time() - t0:.1f}s")
    return model, processor


def tower_forward_text(model, input_ids, attention_mask):
    import torch

    with torch.inference_mode():
        out = model.text_model(input_ids=input_ids, attention_mask=attention_mask)
        pooled = out.pooler_output
        if pooled is None:
            pooled = out.last_hidden_state[:, 0]
        return model.text_projection(pooled)


def tower_forward_vision(model, pixel_values):
    import torch

    with torch.inference_mode():
        out = model.vision_model(pixel_values=pixel_values)
        pooled = out.pooler_output
        if pooled is None:
            pooled = out.last_hidden_state[:, 0]
        return model.visual_projection(pooled)


class TextTowerWrapper:
    """torch.onnx.export 需要 nn.Module，这里在函数内构造以避免顶层 import torch。"""


def make_wrappers(model):
    import torch

    class TextTower(torch.nn.Module):
        def __init__(self, m):
            super().__init__()
            self.m = m

        def forward(self, input_ids, attention_mask):
            out = self.m.text_model(input_ids=input_ids, attention_mask=attention_mask)
            pooled = out.pooler_output
            if pooled is None:
                pooled = out.last_hidden_state[:, 0]
            return self.m.text_projection(pooled)

    class VisionTower(torch.nn.Module):
        def __init__(self, m):
            super().__init__()
            self.m = m

        def forward(self, pixel_values):
            out = self.m.vision_model(pixel_values=pixel_values)
            pooled = out.pooler_output
            if pooled is None:
                pooled = out.last_hidden_state[:, 0]
            return self.m.visual_projection(pooled)

    return TextTower(model).eval(), VisionTower(model).eval()


def export_onnx(model, processor, out_dir: str) -> None:
    import torch

    text_tower, vision_tower = make_wrappers(model)
    tok = processor.tokenizer

    enc = tok(["占位"], padding="max_length", truncation=True,
              max_length=MAX_LENGTH, return_tensors="pt")
    pixel = torch.zeros(1, 3, IMAGE_SIZE, IMAGE_SIZE, dtype=torch.float32)

    print(f"[export] {TEXT_FILE}")
    torch.onnx.export(
        text_tower,
        (enc["input_ids"], enc["attention_mask"]),
        os.path.join(out_dir, TEXT_FILE),
        input_names=["input_ids", "attention_mask"],
        output_names=["text_features"],
        dynamic_axes={"input_ids": {0: "batch"}, "attention_mask": {0: "batch"},
                      "text_features": {0: "batch"}},
        opset_version=OPSET,
        do_constant_folding=True,
        dynamo=False,
    )
    print(f"[export] {VISION_FILE}")
    torch.onnx.export(
        vision_tower,
        (pixel,),
        os.path.join(out_dir, VISION_FILE),
        input_names=["pixel_values"],
        output_names=["image_features"],
        dynamic_axes={"pixel_values": {0: "batch"}, "image_features": {0: "batch"}},
        opset_version=OPSET,
        do_constant_folding=True,
        dynamo=False,
    )


def preprocess_images(processor, images: list[np.ndarray]):
    """用官方 processor 做图像预处理，得到与 PyTorch 完全一致的像素张量。"""
    from PIL import Image

    pil = [Image.fromarray(a) for a in images]
    enc = processor(images=pil, return_tensors="pt")
    return enc["pixel_values"]


def run_verification(model, processor, out_dir: str, plan: list[str]) -> dict:
    """逐个用例真跑 ONNX 并与 PyTorch 比对；返回参考数据。"""
    import onnxruntime as ort

    tok = processor.tokenizer
    text_sess = ort.InferenceSession(os.path.join(out_dir, TEXT_FILE),
                                     providers=["CPUExecutionProvider"])
    vision_sess = ort.InferenceSession(os.path.join(out_dir, VISION_FILE),
                                       providers=["CPUExecutionProvider"])

    executed: list[str] = []
    reference: dict = {"texts": [], "images": []}

    print("\n[verify] 文本塔")
    for text in FIXTURE_TEXTS:
        enc = tok([text], padding="max_length", truncation=True,
                  max_length=MAX_LENGTH, return_tensors="pt")
        ids = enc["input_ids"].numpy().astype(np.int64)
        mask = enc["attention_mask"].numpy().astype(np.int64)
        pt = tower_forward_text(model, enc["input_ids"], enc["attention_mask"]).numpy()
        ox = text_sess.run(["text_features"], {"input_ids": ids, "attention_mask": mask})[0]
        cos = norm_cos(pt, ox)
        name = f"text:{text[:24]}"
        executed.append(name)
        print(f"  cos={cos:.9f}  ids[:8]={ids[0][:8].tolist()}  {text[:28]}")
        if cos < 0.9999:
            raise SystemExit(f"文本塔导出不一致: {name} cos={cos}")
        reference["texts"].append({"text": text, "input_ids": ids[0].tolist(),
                                   "attention_mask": mask[0].tolist(),
                                   "vector": [float(v) for v in ox.reshape(-1)]})

    print("\n[verify] 视觉塔")
    images = [solid(c) for _, c in FIXTURE_IMAGES] + [gradient()]
    names = [n for n, _ in FIXTURE_IMAGES] + ["gradient"]
    pixel = preprocess_images(processor, images)
    import torch

    for i, (nm, _) in enumerate(zip(names, images)):
        px = pixel[i:i + 1]
        pt = tower_forward_vision(model, px).numpy()
        ox = vision_sess.run(["image_features"],
                             {"pixel_values": px.numpy().astype(np.float32)})[0]
        cos = norm_cos(pt, ox)
        executed.append(f"image:{nm}")
        print(f"  cos={cos:.9f}  {nm}")
        if cos < 0.9999:
            raise SystemExit(f"视觉塔导出不一致: {nm} cos={cos}")
        # 参考向量直接存像素张量的 sha256，Go 侧用同一预处理即可复算
        reference["images"].append({
            "name": nm,
            "pixels_sha256": hashlib.sha256(px.numpy().astype(np.float32).tobytes()).hexdigest(),
            "vector": [float(v) for v in ox.reshape(-1)],
        })

    missing = [c for c in plan if c not in executed]
    extra = [c for c in executed if c not in plan]
    if missing or extra:
        raise SystemExit(f"用例覆盖不一致: 缺 {missing} 多 {extra}")
    print(f"\n[verify] 覆盖度 OK（{len(executed)} 个用例，计划 {len(plan)}）")
    return reference


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--model-dir", default="")
    ap.add_argument("--model-id", default=DEFAULT_MODEL_ID)
    ap.add_argument("--verify-only", action="store_true")
    ap.add_argument("--no-reference", action="store_true")
    ap.add_argument("--skip-verify", action="store_true")
    args = ap.parse_args()

    plan = [f"text:{t[:24]}" for t in FIXTURE_TEXTS] + \
           [f"image:{n}" for n, _ in FIXTURE_IMAGES] + ["image:gradient"]

    if not args.verify_only:
        # 清空重建，避免旧产物被当成这次的成果
        if os.path.isdir(args.out):
            shutil.rmtree(args.out)
        os.makedirs(args.out, exist_ok=True)

    model, processor = load_model(args.model_dir or None, args.model_id,
                                  local_only=bool(args.model_dir))
    if not os.path.isdir(args.out):
        os.makedirs(args.out, exist_ok=True)

    if not args.verify_only:
        export_onnx(model, processor, args.out)

    # 词表随产物一起放：provider 只依赖这个目录
    src_vocab = os.path.join(args.model_dir, "vocab.txt") if args.model_dir else None
    if src_vocab and os.path.exists(src_vocab):
        shutil.copy2(src_vocab, os.path.join(args.out, "vocab.txt"))

    reference = None
    if not args.skip_verify:
        reference = run_verification(model, processor, args.out, plan)
    else:
        print("[verify] 已按 --skip-verify 跳过（不据此宣布成功）")

    if not args.verify_only:
        cfg = {
            "arch": "chinese-clip-vit-base-patch16",
            "dim": 512,
            "text_onnx": TEXT_FILE,
            "vision_onnx": VISION_FILE,
            "max_length": MAX_LENGTH,
            "image_size": IMAGE_SIZE,
            "resample": "bicubic",
            "rescale": 1.0 / 255.0,
            "image_mean": [0.48145466, 0.4578275, 0.40821073],
            "image_std": [0.26862954, 0.26130258, 0.27577711],
            "normalize_vector": True,   # provider 必须 L2 归一化后再入库
            "tokenizer": {
                "type": "bert-wordpiece",
                "vocab": "vocab.txt",
                "do_lower_case": True,
                "tokenize_chinese_chars": True,
                "cls_id": 101, "sep_id": 102, "pad_id": 0, "unk_id": 100,
            },
            "modalities": ["text", "image"],
            "unsupported_modalities": ["audio", "video"],
            "notes": "Chinese-CLIP ViT-B/16：视觉 ViT-B/16 + 文本 RoBERTa-wwm-base，"
                     "输出 512 维共享空间。文本塔取 CLS（pooler）后过 text_projection，"
                     "视觉塔取 CLS 后过 visual_projection；两者均未在图中归一化，"
                     "归一化由 provider 负责。",
        }
        with open(os.path.join(args.out, CONFIG_FILE), "w") as f:
            json.dump(cfg, f, ensure_ascii=False, indent=2)

        if reference is not None and not args.no_reference:
            reference["source"] = {"model_id": args.model_id,
                                   "model_dir": args.model_dir or "(hub)"}
            reference["artifacts"] = {n: sha256_file(os.path.join(args.out, n))
                                      for n in (TEXT_FILE, VISION_FILE)}
            with open(os.path.join(args.out, REFERENCE_FILE), "w") as f:
                json.dump(reference, f, ensure_ascii=False, indent=2)

        with open(os.path.join(args.out, "SHA256SUMS"), "w") as f:
            for n in sorted(os.listdir(args.out)):
                if n == "SHA256SUMS":
                    continue
                p = os.path.join(args.out, n)
                if os.path.isfile(p):
                    f.write(f"{sha256_file(p)}  {n}\n")

    print(f"\n[out] {args.out}")
    for n in sorted(os.listdir(args.out)):
        p = os.path.join(args.out, n)
        if os.path.isfile(p):
            print(f"  {n:<20} {os.path.getsize(p) / 1e6:9.1f} MB")
    return 0


if __name__ == "__main__":
    sys.exit(main())