ColModernVBERT -> Core AI: query + doc encoders (fp16/fp32), visual document retrieval (late-interaction / MaxSim)

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+doc/colmodernvbert-doc_float16_s89_static.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
+fp32/doc/colmodernvbert-doc_float32_s89_static.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
+fp32/query/colmodernvbert-query_float32_s32_static.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text
+query/colmodernvbert-query_float16_s32_static.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,108 @@
+---
+license: mit
+library_name: coreai
+pipeline_tag: visual-document-retrieval
+tags:
+  - core-ai
+  - apple
+  - on-device
+  - visual-document-retrieval
+  - late-interaction
+  - colbert
+  - colpali
+  - retrieval
+base_model: ModernVBERT/colmodernvbert
+---
+
+# ColModernVBERT — Core AI
+
+**The zoo's first visual document retriever and first late-interaction (ColBERT / MaxSim)
+multi-vector model**, running as static `.aimodel` graphs on Apple Silicon (Mac GPU / iPhone).
+A Core AI port of [`ModernVBERT/colmodernvbert`](https://huggingface.co/ModernVBERT/colmodernvbert)
+(MIT) — a compact 250M visual document retriever: a **ModernBERT-150M bidirectional text
+encoder** + **SigLIP2 vision encoder** (pixel-shuffle ×4) with a `custom_text_proj` head that
+emits a **per-token L2-normalized 128-d multi-vector**. Retrieval is **late interaction**: you
+encode a text query and a page *image* into token-level vectors and score them with **MaxSim**
+(`score = Σ_q max_d ⟨E_q, E_d⟩`). No OCR — the page is matched as a picture, so tables, charts
+and complex layouts are first-class.
+
+This completes the on-device RAG trifecta alongside the text
+[Qwen3-Embedding](https://huggingface.co/mlboydaisuke/Qwen3-Embedding-0.6B-CoreAI) (text→text
+dense) and [Qwen3-Reranker](https://huggingface.co/mlboydaisuke/Qwen3-Reranker-0.6B-CoreAI)
+(cross-encoder): **embed → rerank → visual-retrieval**, all on device.
+
+## Two encoders (two graphs)
+
+| graph | input | output | fp16 size |
+|---|---|---|---|
+| **query** | `input_ids [1,32] i32`, `attention_mask [1,32] i32` | `query_embeddings [1,32,128]` | 298 MB |
+| **doc** | `pixel_values [1,1,3,512,512]`, `pixel_attention_mask [1,1,512,512] i32` | `doc_embeddings [1,89,128]` | 407 MB |
+
+Both are single bidirectional forwards — no KV cache, no generation. The per-token L2-norm and
+the `attention_mask` masking are baked in-graph; **MaxSim runs on the host** (a tiny matmul +
+max + sum). Each bundle directory holds one `*.aimodel` plus a `tokenizer/` folder.
+
+- **query**: right-pad the tokenized query to the 32-token grid (queries are short; ModernBERT's
+  sliding-window(128) sees the full sequence → full attention). Slice to the real token count
+  before MaxSim.
+- **doc**: a **single 512×512 tile** ("global image") layout — the text template (CLS + image
+  markers + 64 `<image>` placeholders + SEP) is baked as a graph constant, so the only runtime
+  inputs are the pixels. Preprocess the page like Idefics3: resize so the longest edge ≤ 512,
+  pad to 512×512, rescale ×1/255, normalize with mean/std = 0.5, and build the
+  `pixel_attention_mask` (1 for real pixels, 0 for padding).
+
+> **Single-tile v1.** This release ships the single 512px global-image document path: lightweight,
+> iPhone-friendly, and accurate on typical pages. The model's full high-resolution mode (split a
+> page into multiple 512px tiles + the global image, 800+ doc tokens) is a planned follow-up for
+> dense small-print documents.
+
+## Repo layout
+
+```
+query/   colmodernvbert-query_float16_s32_static.aimodel + tokenizer/   (298 MB, fp16 — iPhone)
+doc/     colmodernvbert-doc_float16_s89_static.aimodel                  (407 MB, fp16 — iPhone)
+fp32/query/  colmodernvbert-query_float32_s32_static.aimodel + tokenizer/  (595 MB — Mac)
+fp32/doc/    colmodernvbert-doc_float32_s89_static.aimodel                 (813 MB — Mac)
+README.md · reference_query.json · reference_doc.json · test_doc.png
+```
+
+Each `query/` and `doc/` directory is a complete bundle root (one `.aimodel`, plus `tokenizer/`
+on the query side). fp16 ships for iPhone (~705 MB for both encoders); fp32 is for Mac / max
+precision.
+
+## On-device (CoreAIKit)
+
+```swift
+import CoreAIKitEmbeddings
+
+// Downloads query/ + doc/ (fp16) from this repo, or uses a sideloaded copy if present.
+let retriever = try await VisualDocumentRetriever()   // .colModernVBERTQuery / .colModernVBERTDoc
+
+// Encode a page as tiles (reliable spatial grounding), rank queries, and locate the match.
+let page = try await retriever.encodeTiled(page: cgImage, rows: 6, cols: 4)
+let q = try await retriever.encode(query: "total revenue in the third quarter")
+let score = retriever.score(query: q, tiledPage: page)     // MaxSim, page ranking
+let rect  = retriever.bestTile(query: q, tiledPage: page)  // normalized region to highlight
+```
+
+See [`Examples/DocSearch`](https://github.com/john-rocky/coreai-kit/tree/main/Examples/DocSearch)
+for a full iPhone demo (bundled + imported documents, query → ranked pages → highlighted region).
+
+## Parity (Core AI engine vs. PyTorch reference, M4 Max GPU)
+
+Per-token cosine of the 128-d multi-vectors against the `colpali_engine` PyTorch model:
+
+| encoder | float32 | float16 |
+|---|---|---|
+| query | min/mean **1.000000** | min 0.999997 / mean 0.999999 |
+| doc | min/mean **1.000000** | min 0.999994 / mean 0.999998 |
+
+End-to-end retrieval: the host **MaxSim reproduces `processor.score` exactly** (max |Δ| = 0.0000),
+the engine ranking matches the PyTorch ranking on every clear-margin query, and the single-tile
+engine retrieves the intended page **3/3** on a rendered-text corpus.
+
+## License
+
+MIT, inherited from [`ModernVBERT/colmodernvbert`](https://huggingface.co/ModernVBERT/colmodernvbert).
+See the upstream model and paper *ModernVBERT: Towards Smaller Visual Document Retrievers*
+([arXiv:2510.01149](https://arxiv.org/abs/2510.01149)).

doc/colmodernvbert-doc_float16_s89_static.aimodel/main.hash

@@ -0,0 +1 @@
+ZE��?��5�PT!�DIv�k�+��3��wK

doc/colmodernvbert-doc_float16_s89_static.aimodel/main.mlirb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a45f482e93fcee4bd3595501d5421f7444976916bb92beea6d033c1ed19774b
+size 426703966

doc/colmodernvbert-doc_float16_s89_static.aimodel/metadata.json

@@ -0,0 +1,7 @@
+{
+  "description" : "ColModernVBERT visual document retriever (document\/image encoder, single 512px tile): SigLIP2 vision + pixel-shuffle x4 -> ModernBERT-150M -> custom_text_proj(768->128) -> per-token L2-normalized 128-d multi-vector for ColBERT-style late-interaction (MaxSim). https:\/\/huggingface.co\/ModernVBERT\/colmodernvbert",
+  "creationDate" : "20260624T025647Z",
+  "license" : "MIT",
+  "assetVersion" : "2.0",
+  "author" : "ModernVBERT (Illuin Technology \/ ETH)"
+}

fp32/doc/colmodernvbert-doc_float32_s89_static.aimodel/main.hash

@@ -0,0 +1 @@
+sYL�vJ�^��c��H\|�g��<;?�l

fp32/doc/colmodernvbert-doc_float32_s89_static.aimodel/main.mlirb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7359114cfd76184a835ea911b163adee98485c7ca667bfaa3c120f3b3fbc1e6c
+size 852898790

fp32/doc/colmodernvbert-doc_float32_s89_static.aimodel/metadata.json

@@ -0,0 +1,7 @@
+{
+  "description" : "ColModernVBERT visual document retriever (document\/image encoder, single 512px tile): SigLIP2 vision + pixel-shuffle x4 -> ModernBERT-150M -> custom_text_proj(768->128) -> per-token L2-normalized 128-d multi-vector for ColBERT-style late-interaction (MaxSim). https:\/\/huggingface.co\/ModernVBERT\/colmodernvbert",
+  "license" : "MIT",
+  "assetVersion" : "2.0",
+  "author" : "ModernVBERT (Illuin Technology \/ ETH)",
+  "creationDate" : "20260624T025601Z"
+}

fp32/query/colmodernvbert-query_float32_s32_static.aimodel/main.hash

@@ -0,0 +1 @@
+V�z��BL%�/9B�,0�J���Ľme��Y�

fp32/query/colmodernvbert-query_float32_s32_static.aimodel/main.mlirb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1156c67a95a0424c25972f143942ba2c309d4ac1be9f1cc4bd6d65e40cdf59c8
+size 624058345

fp32/query/colmodernvbert-query_float32_s32_static.aimodel/metadata.json

@@ -0,0 +1,7 @@
+{
+  "assetVersion" : "2.0",
+  "license" : "MIT",
+  "author" : "ModernVBERT (Illuin Technology \/ ETH)",
+  "description" : "ColModernVBERT visual document retriever (query\/text encoder): ModernBERT-150M bidirectional backbone -> custom_text_proj Linear(768->128) -> per-token L2-normalized 128-d multi-vector for ColBERT-style late-interaction (MaxSim) retrieval. Source: https:\/\/huggingface.co\/ModernVBERT\/colmodernvbert",
+  "creationDate" : "20260624T024133Z"
+}

fp32/query/tokenizer/chat_template.jinja

@@ -0,0 +1,2 @@
+<|begin_of_text|>{% for message in messages %}{{message['role'] | capitalize}}{% if message['content'][0]['type'] == 'image' %}{{':'}}{% else %}{{': '}}{% endif %}{% for line in message['content'] %}{% if line['type'] == 'text' %}{{line['text']}}{% elif line['type'] == 'image' %}{{ '<image>' }}{% endif %}{% endfor %}<end_of_utterance>
+{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}

fp32/query/tokenizer/tokenizer_config.json

@@ -0,0 +1,33 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": [
+    "<fake_token_around_image>",
+    "<image>",
+    "<end_of_utterance>"
+  ],
+  "is_local": false,
+  "legacy": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "max_length": 8192,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask",
+    "pixel_values",
+    "pixel_attention_mask"
+  ],
+  "model_max_length": 8192,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "left",
+  "processor_class": "ColModernVBertProcessor",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "tokenizer_class": "TokenizersBackend",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}

query/colmodernvbert-query_float16_s32_static.aimodel/main.hash

@@ -0,0 +1,2 @@
+I
+9���kc�_��}��6�b�&����e�t�n�

query/colmodernvbert-query_float16_s32_static.aimodel/main.mlirb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:490a399f909e6b6306cc5fc8d87ddecf3618c462c226b191a2a1659f74fc6e98
+size 312212955

query/colmodernvbert-query_float16_s32_static.aimodel/metadata.json

@@ -0,0 +1,7 @@
+{
+  "assetVersion" : "2.0",
+  "creationDate" : "20260624T024253Z",
+  "license" : "MIT",
+  "author" : "ModernVBERT (Illuin Technology \/ ETH)",
+  "description" : "ColModernVBERT visual document retriever (query\/text encoder): ModernBERT-150M bidirectional backbone -> custom_text_proj Linear(768->128) -> per-token L2-normalized 128-d multi-vector for ColBERT-style late-interaction (MaxSim) retrieval. Source: https:\/\/huggingface.co\/ModernVBERT\/colmodernvbert"
+}

query/tokenizer/chat_template.jinja

@@ -0,0 +1,2 @@
+<|begin_of_text|>{% for message in messages %}{{message['role'] | capitalize}}{% if message['content'][0]['type'] == 'image' %}{{':'}}{% else %}{{': '}}{% endif %}{% for line in message['content'] %}{% if line['type'] == 'text' %}{{line['text']}}{% elif line['type'] == 'image' %}{{ '<image>' }}{% endif %}{% endfor %}<end_of_utterance>
+{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}

query/tokenizer/tokenizer_config.json

@@ -0,0 +1,33 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": [
+    "<fake_token_around_image>",
+    "<image>",
+    "<end_of_utterance>"
+  ],
+  "is_local": false,
+  "legacy": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "max_length": 8192,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask",
+    "pixel_values",
+    "pixel_attention_mask"
+  ],
+  "model_max_length": 8192,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "left",
+  "processor_class": "ColModernVBertProcessor",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "tokenizer_class": "TokenizersBackend",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}