mlboydaisuke's picture
gen-cards: regenerate Use-it block
c3866b2 verified
|
Raw
History Blame Contribute Delete
7.66 kB
---
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.
<!-- gen-cards:use-it begin id=colmodernvbert (managed by scripts/gen-cards β€” edit cards.json / QuickStart.swift, not this block) -->
## Use it
▢️ **Run it (source)** β€” the [DocSearch runner](https://github.com/john-rocky/coreai-kit/tree/main/Examples/DocSearch)
(visual page search over bundled sample pages; the GUI (iPhone) adds tiled where-it-matched highlights):
```bash
git clone https://github.com/john-rocky/coreai-kit
open coreai-kit/Examples/DocSearch/DocSearch.xcodeproj
# β†’ Run, then pick "ColModernVBERT" in the model picker
# agents / headless (macOS):
cd coreai-kit/Examples/DocSearch
swift run docsearch-cli --model colmodernvbert --query "monthly revenue trend"
```
πŸ’» **Build with it** β€” complete; the glue is kit API, copy-paste runs:
```swift
import CoreAIKitEmbeddings
let retriever = try await VisualDocumentRetriever(
catalog: "colmodernvbert")
var corpus: [VisualDocumentRetriever.PageEmbedding] = []
for url in pages {
corpus.append(try await retriever.encode(page: ImageFile.load(url).cgImage))
}
let hits = try await retriever.retrieve(query: query, over: corpus, topK: pages.count)
// hits: pages ranked by MaxSim, best match first β€” no OCR, pages are matched as pictures
```
The take-home is [`Examples/DocSearch/Sources/QuickStart.swift`](https://github.com/john-rocky/coreai-kit/blob/main/Examples/DocSearch/Sources/QuickStart.swift)
β€” this exact code as one typed function, no UI; the CLI is an argument shell over it, and
the GUI drives the same `VisualDocumentRetriever(catalog:)` with tiled per-page encoding.
Encode your corpus once and keep the `PageEmbedding`s β€” scoring a query is then host-side
MaxSim, no model call per page. `encodeTiled(page:)` localizes *where* a query matched.
**Integration checklist**
- SPM: `https://github.com/john-rocky/coreai-kit` β†’ product **CoreAIKitEmbeddings**
- Info.plist: `NSPhotoLibraryUsageDescription` β€” only if you use PhotosPicker to import pages
- Entitlements: none needed
- First run downloads the model β€” 0.7 GB (Mac) / 0.7 GB (iPhone) β€” then it loads from the
local cache (Application Support; progress via the `downloadProgress` callback)
- Measure in Release β€” Debug is ~3Γ— slower on per-token host work
<!-- gen-cards:use-it end -->
## 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)).