build.py

"""Build jina-embeddings-v4 (vLLM merged, per-task) → ONNX sub-parts (vision / embeddings / backbone).

  uv run build.py --model vllm-retrieval --output onnx/cpu_fp16              # fp16 (default)
  uv run build.py --model vllm-retrieval --output onnx/cpu_fp32 --precision fp32
  uv run build.py --model vllm-retrieval --output onnx/cpu_int8 --precision int8
  uv run build.py --model vllm-retrieval --output onnx/cpu_int4 --precision int4

--precision {fp16,fp32,int8,int4}. int8/int4 build a fp16 graph then weight-quantize the backbone
in place (block-wise MatMulNBits; vision/embeddings stay fp16). int8 keeps pooled-cosine ≥0.999;
int4 does NOT (embedding drift ~5-8%) so it only warns — measure with eval.py before use.
CPU only (no device flag).
"""
import argparse
import gc
import json
import shutil
import tempfile
from pathlib import Path

import numpy as np
import onnx
import torch

from common import (HIDDEN, MATRYOSHKA, BackboneSub, EmbeddingsSub, VisionSub, cosine, hf_name,
                    image_rope_and_mask, load_model, load_tokenizer, make_image_inputs, make_session,
                    mean_pool, quiet)


def _export(enc, sample, names_in, names_out, dyn, raw):
    with torch.no_grad():
        torch.onnx.export(enc, sample, str(raw), input_names=names_in, output_names=names_out,
                          dynamic_axes=dyn, opset_version=19, do_constant_folding=True, dynamo=False)


def _consolidate(out, stem, raw):
    """temp graph+loose externals (in raw's dir) → one {out}/{stem}.onnx + {stem}.onnx.data.
    Only the two final files land in `out`; sibling sub-parts are never touched."""
    final = out / f"{stem}.onnx"
    for p in (final, out / f"{stem}.onnx.data"):
        if p.exists(): p.unlink()
    m = onnx.load(str(raw))
    onnx.save_model(m, str(final), save_as_external_data=True, all_tensors_to_one_file=True,
                    location=f"{stem}.onnx.data", size_threshold=1024, convert_attribute=False)
    del m; gc.collect()
    shutil.rmtree(raw.parent, ignore_errors=True)
    return sum(p.stat().st_size for p in out.glob(f"{stem}.onnx*")) / 1e9


def _tmp(stem):
    return Path(tempfile.mkdtemp()) / f"{stem}.onnx"


TOKENIZER_FILES = ("tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt",
                   "special_tokens_map.json", "added_tokens.json", "preprocessor_config.json",
                   "chat_template.json")


def task_prompts(hf_model):
    """Per-task text prompt convention (from jina's model cards). text-matching is symmetric (both
    sides use 'Query:'); retrieval/code are asymmetric (query 'Query:', document 'Passage:')."""
    if "text-matching" in hf_model:
        return {"query": "Query:", "document": "Query:", "symmetric": True}
    return {"query": "Query:", "document": "Passage:", "symmetric": False}


QUANT_BITS = {"int8": 8, "int4": 4}     # precision → MatMulNBits bit-width (int* → fp16 graph + quant)


def _quant_backbone_inplace(out, bits, block_size):
    """Block-wise n-bit (MatMulNBits) on backbone.onnx MatMul weights, overwriting it in `out`.
    Compute stays float — only the weight matrices are quantized. vision/embeddings stay fp16
    (small, and more precision-sensitive). Returns the new backbone size in GB."""
    from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer
    m = onnx.load(str(out / "backbone.onnx"))
    q = MatMulNBitsQuantizer(m, block_size=block_size, is_symmetric=True, bits=bits)
    q.process()
    for p in (out / "backbone.onnx", out / "backbone.onnx.data"):
        if p.exists(): p.unlink()
    onnx.save_model(q.model.model, str(out / "backbone.onnx"), save_as_external_data=True,
                    all_tensors_to_one_file=True, location="backbone.onnx.data",
                    size_threshold=1024, convert_attribute=False)
    del m, q; gc.collect()
    return sum(p.stat().st_size for p in out.glob("backbone.onnx*")) / 1e9


def build(model_dir, output, precision, image, image_size, block_size=32):
    out = Path(output); out.mkdir(parents=True, exist_ok=True)
    bits = QUANT_BITS.get(precision)                       # None for fp16/fp32
    base = "fp32" if precision == "fp32" else "fp16"       # int8/int4 quantize a fp16 graph
    dtype = torch.float16 if base == "fp16" else torch.float32
    npdt = np.float16 if base == "fp16" else np.float32
    print(f"=== build SUB-PARTS | {hf_name(model_dir)} → {out} ({precision}, cpu) ===")
    model = load_model(model_dir, dtype=dtype, attn="eager")

    # image sample (fixed resolution) drives vision + exercises the fusion in embeddings/backbone
    batch = make_image_inputs(model_dir, image, image_size)
    ipos, vmask = image_rope_and_mask(model, batch)
    px = batch["pixel_values"].to(dtype)
    grid = batch["image_grid_thw"]

    sizes = {}
    vsub = VisionSub(model, grid).eval()
    with torch.no_grad():
        vfeat = vsub(px)
    print(f"  vision: pixel_values{tuple(px.shape)} grid{grid[0].tolist()} → image_features{tuple(vfeat.shape)}")
    raw = _tmp("vision")
    _export(vsub, (px,), ["pixel_values"], ["image_features"], {"pixel_values": {0: "num_patches"}}, raw)
    sizes["vision"] = _consolidate(out, "vision", raw)

    esub = EmbeddingsSub(model).eval()
    iid = batch["input_ids"]
    with torch.no_grad():
        emb_ref = esub(iid, vfeat)
    raw = _tmp("embeddings")
    _export(esub, (iid, vfeat), ["input_ids", "image_features"], ["inputs_embeds"],
            {"input_ids": {0: "b", 1: "s"}, "image_features": {0: "n"}, "inputs_embeds": {0: "b", 1: "s"}}, raw)
    sizes["embeddings"] = _consolidate(out, "embeddings", raw)

    bsub = BackboneSub(model).eval()
    iam = batch["attention_mask"]
    with torch.no_grad():
        hid_ref = bsub(emb_ref, iam, ipos)
    raw = _tmp("backbone")
    _export(bsub, (emb_ref, iam, ipos), ["inputs_embeds", "attention_mask", "position_ids"],
            ["last_hidden"], {"inputs_embeds": {0: "b", 1: "s"}, "attention_mask": {0: "b", 1: "s"},
                              "position_ids": {2: "s"}, "last_hidden": {0: "b", 1: "s"}}, raw)
    sizes["backbone"] = _consolidate(out, "backbone", raw)

    quantized = {}
    if bits:  # int8/int4: quantize the shipped backbone in place, then sanity-check the real artifact
        sizes["backbone"] = _quant_backbone_inplace(out, bits, block_size)
        quantized["backbone"] = precision

    img_ref = mean_pool(hid_ref, vmask).float().cpu().numpy()
    # the image-prompt layout is FIXED for a given resolution (only pixels vary) → save the host
    # tensors so image inference needs NO model load (just the processor for pixel_values).
    np.savez(out / "image_meta.npz", input_ids=iid.cpu().numpy(), attention_mask=iam.cpu().numpy(),
             position_ids=ipos.cpu().numpy(), vision_mask=vmask.cpu().numpy())
    del model, vsub, esub, bsub; gc.collect()

    for fn in TOKENIZER_FILES:
        s = Path(model_dir) / fn
        if s.exists(): shutil.copy(s, out / fn)
    for stem, gb in sizes.items():
        print(f"  {stem}.onnx (+.data): {gb:.1f} GB")

    # composed sanity: vision → embeddings → backbone → pool  vs  the PyTorch image embedding
    vs = make_session(out / "vision.onnx"); es = make_session(out / "embeddings.onnx")
    bs = make_session(out / "backbone.onnx")
    f = vs.run(None, {"pixel_values": px.float().cpu().numpy().astype(npdt)})[0]
    e = es.run(None, {"input_ids": iid.cpu().numpy(), "image_features": f})[0]
    h = bs.run(None, {"inputs_embeds": e, "attention_mask": iam.cpu().numpy(),
                      "position_ids": ipos.cpu().numpy()})[0]
    c = cosine(mean_pool(h, vmask.cpu().numpy()), img_ref)
    print(f"  composed sanity (image, ORT chain vs PyTorch): cos={c:.6f}")

    man = {"hf_model": hf_name(model_dir), "local_dir": str(Path(model_dir).name),
           "precision": base, "embedding_dim": HIDDEN, "matryoshka_dims": MATRYOSHKA,
           "image_size": image_size, "image_grid_thw": grid[0].tolist(),
           "sub_models": {
               "vision": {"file": "vision.onnx", "in": "pixel_values[N,1176]", "out": "image_features[N,2048]",
                          "task_agnostic": True},
               "embeddings": {"file": "embeddings.onnx", "in": "input_ids[B,S] + image_features[N,2048]",
                              "out": "inputs_embeds[B,S,2048]"},
               "backbone": {"file": "backbone.onnx", "in": "inputs_embeds + attention_mask + position_ids[3,B,S]",
                            "out": "last_hidden[B,S,2048]"}},
           "compose_text": "embeddings(ids,empty) -> backbone -> mean-pool(attn_mask) -> L2norm",
           "compose_image": "vision(px) -> embeddings(ids,feats) -> backbone -> mean-pool(vision-span) -> L2norm",
           "prompts": task_prompts(hf_name(model_dir))}
    if quantized:
        man["quantized"] = quantized     # I/O dtype stays fp16; precision field holds the base
    (out / "manifest.json").write_text(json.dumps(man, indent=2))
    # fp16/fp32 must hit parity; int4 is knowingly lossy (embedding drift), so warn instead of fail.
    if c < 0.999:
        if bits and bits <= 4:
            print(f"  [WARN] {precision} composed cos {c:.4f} < 0.999 — expected for int4 (lossy); "
                  f"verify with eval.py before use")
        else:
            raise SystemExit(f"  [FAIL] composed sanity cos {c:.4f} < 0.999")
    print("  build OK")
    return out


def main():
    ap = argparse.ArgumentParser(description="jina-embeddings-v4 → ONNX sub-parts")
    ap.add_argument("--model", default="vllm-retrieval")
    ap.add_argument("--output", default="onnx/cpu_fp16")
    ap.add_argument("--precision", default="fp16", choices=["fp16", "fp32", "int8", "int4"],
                    help="int8/int4 build a fp16 graph then quantize the backbone in place (block-wise)")
    ap.add_argument("--image", default=None)
    ap.add_argument("--image-size", type=int, default=224)
    ap.add_argument("--block-size", type=int, default=32, help="int8/int4 quantization block size")
    args = ap.parse_args()
    quiet()
    build(args.model, args.output, args.precision, args.image, args.image_size, args.block_size)


if __name__ == "__main__":
    main()