onnx-community/jina-embeddings-v4-vllm-code

1"""Eval composed ONNX sub-parts vs full PyTorch — accepts MULTIPLE build dirs at once and reports2each one's detected precision (from its manifest) alongside its pooled-embedding cosine.3 4  uv run eval.py --model vllm-retrieval --onnx-dir onnx/cpu_fp165  uv run eval.py --model vllm-retrieval --onnx-dir onnx/cpu_fp16 onnx/cpu_fp32 onnx/cpu_fp16-int86 7The full PyTorch reference is computed ONCE (model freed before any ORT session loads — they don't8co-fit in RAM), then every dir is scored against it. Precision/quant is read per-dir from manifest,9so you can eye the fidelity vs size trade-off across builds in one run. CPU only.10"""11import argparse12import gc13import json14from pathlib import Path15 16import numpy as np17import torch18 19from common import (cosine, describe_precision, embed_image_onnx, embed_text_onnx, hf_name,20                    image_rope_and_mask, load_model, load_sessions, load_tokenizer, make_image_inputs,21                    make_inputs, mean_pool, npdt_of, quiet, text_position_ids)22 23TEXTS = [("Query", "capital of France?"), ("Passage", "Paris is the capital of France."),24         ("Query", "def add(a,b): return a+b")]25 26 27def pytorch_refs(model_dir, image, size):28    """(text_refs, img_ref) from the full model; the model is freed before this returns.29    Loaded in fp16 to fit RAM (the fp32 model is ~12 GB and won't co-fit with the ORT sessions)."""30    import torch as _t31    model = load_model(model_dir, dtype=_t.float16, attn="eager")32    tok = load_tokenizer(model_dir)33    text_refs = []34    for prefix, t in TEXTS:35        ids, am = make_inputs(tok, [t], prefix=prefix)36        pos = text_position_ids(am)37        with torch.no_grad():38            h = model.model.language_model(inputs_embeds=model.model.language_model.embed_tokens(ids),39                                           attention_mask=am, position_ids=pos, use_cache=False).last_hidden_state40        text_refs.append((prefix, t, mean_pool(h, am).float().numpy()))41    batch = make_image_inputs(model_dir, image, size)42    ipos, vm = image_rope_and_mask(model, batch)43    with torch.no_grad():44        full = model.model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"],45                           position_ids=ipos, pixel_values=batch["pixel_values"],46                           image_grid_thw=batch["image_grid_thw"], use_cache=False).last_hidden_state47    img_ref = mean_pool(full, vm).float().numpy()48    del model, full, h; gc.collect()49    return tok, text_refs, img_ref50 51 52def eval_dir(onnx_dir, model_dir, tok, text_refs, img_ref, image):53    out = Path(onnx_dir)54    man = json.loads((out / "manifest.json").read_text())55    npdt = npdt_of(man); size = man["image_size"]56    prec = describe_precision(man)57    sess = load_sessions(out, need_vision=True)58    meta = dict(np.load(out / "image_meta.npz"))59    worst = 1.060    print(f"\n--- {out}   precision={prec} ---")61    for prefix, t, ref in text_refs:62        c = cosine(embed_text_onnx(sess, tok, t, prefix, npdt), ref); worst = min(worst, c)63        print(f"  [text ] {prefix+': '+t[:34]!r:44} cos={c:.6f}")64    c = cosine(embed_image_onnx(sess, model_dir, image, size, npdt, meta), img_ref); worst = min(worst, c)65    print(f"  [image] {'synthetic' if not image else image:44} cos={c:.6f}")66    return prec, worst67 68 69def main():70    ap = argparse.ArgumentParser(description="eval jina-embeddings-v4 ONNX sub-parts (multi-dir) vs PyTorch")71    ap.add_argument("--model", default="vllm-retrieval")72    ap.add_argument("dirs", nargs="*", help="one or more build dirs (positional)")73    ap.add_argument("--onnx-dir", nargs="+", default=None, help="one or more build dirs (flag form)")74    ap.add_argument("--image", default=None)75    ap.add_argument("--tol", type=float, default=0.999)76    args = ap.parse_args()77    quiet()78    # dedupe by resolved path (keep first occurrence) so the same dir isn't eval'd twice —79    # e.g. "onnx/fp32" and "onnx/fp32/" are one model80    raw = args.dirs or args.onnx_dir or ["onnx/cpu_fp16"]81    seen, onnx_dirs = set(), []82    for d in raw:83        key = Path(d).resolve()84        if key not in seen:85            seen.add(key); onnx_dirs.append(d)86    print(f"=== eval composed ONNX vs full PyTorch | {hf_name(args.model)} (cpu) ===")87 88    tok, text_refs, img_ref = pytorch_refs(args.model, args.image, 224)89    rows = []90    for d in onnx_dirs:91        prec, worst = eval_dir(d, args.model, tok, text_refs, img_ref, args.image)92        rows.append((d, prec, worst, worst >= args.tol))93 94    print(f"\n=== summary (tol {args.tol}) ===")95    w = max(len(d) for d, *_ in rows)96    for d, prec, worst, ok in rows:97        print(f"  {d:<{w}}  {prec:<22} worst cos {worst:.6f}  {'PASS' if ok else 'FAIL'}")98    raise SystemExit(0 if all(ok for *_, ok in rows) else 1)99 100 101if __name__ == "__main__":102    main()103