Give wordless pictures a way to be found

A screenshot with no OCR-able words used to be a memory nothing could reach. It now has a vector, and a question can be asked with it. - picture space: Xenova/chinese-clip-vit-base-patch16 fp16 ONNX, 361MB, sha256 0fad978faa89b85f…, CPU-runnable on the onnxruntime the repo already ships. Text 76ms, image 29ms/image at batch 8. The ONNX contract was measured, not assumed: it is the fused contrastive graph, so the unused tower is fed a placeholder row, and the tokenizer must be BertTokenizer — a CLIP tokenizer emits ids past the embedding table. - visual_index.py keeps picture vectors in their own family beside the text ones, keyed by atom id, and refuses to answer when the map and the array disagree. retrieval.py stays under 500 lines. - a search can be asked with a picture, and when the words come up thin the picture ranking joins the fusion — but only then, so a text query never pays for the encoder. - visual tokens are counted at 32x32 patches (migration v7): 21,325 image atoms went from 5,331 to 34,618 tokens, and the library from 29.9M to 84.2M. The 100M capsule fits this 50,278-file corpus and nothing much larger — that is what the claim actually means. - the vision-language describer is written and gated: with no snapshot an atom records "deferred: no vision-language model snapshot"; a broken model becomes a note on the atom instead of a dead ingest. - forget now voids the picture index and the copy this process holds, so a forgotten picture stops answering from RAM. Measured, six arms, ScreenSpot-v2 (554 icon crops, 408 whole screenshots), bench/reports/VISUAL.md: icon pool picture 0.208 words 0.175 chance 0.018 screenshot pool words 0.384 picture 0.215 screenshot to icon 0.108 — CLIP is a whole-image model, not a grounding model: it cannot point at one small icon inside a screen. within one screen MRR 0.908 (recall@10 is 1.000 there only because chance is 1.000 too, so MRR is the only honest number in that row). ScreenSpot-v2 instructions are all English, so this run does not test the bilingual claim and the report says so. Gates: 199 tests, ruff clean on 17 files, mypy clean on 16 (fresh cache), the repo's own pyright 0 errors on the 12 files touched.
2026-10-07 20:41:19 -04:00
1 parent 7e1fe649f6
commit 03ca35763f
@@ -0,0 +1,120 @@
+"""Image and text encoding for cross-modal retrieval, over one small ONNX graph.
+
+Chinese-CLIP ViT-B/16 in fp16 (377 MB) keeps bilingual text and images in one vector
+space, so a screenshot can be found by a picture or by a sentence. The graph is the
+fused contrastive model: every call needs both towers' inputs, so the unused side is
+fed a placeholder row.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from PIL import Image
+
+IMAGE_SIZE = 224
+TEXT_MAX_LENGTH = 52
+IMAGE_MEAN = (0.48145466, 0.45782750, 0.40821073)
+IMAGE_STD = (0.26862954, 0.26130258, 0.27577711)
+DOWNLOAD_HINT = (
+    "huggingface-cli download Xenova/chinese-clip-vit-base-patch16 "
+    "--include 'onnx/model_fp16.onnx' 'tokenizer*' 'config.json' --local-dir <dir>"
+)
+
+
+def preprocess_image(image: Image.Image) -> np.ndarray:
+    """Turn a PIL image into the (3, 224, 224) float32 tensor the graph expects.
+
+    The reference processor resizes to a square and normalizes with CLIP constants;
+    doing it here keeps torchvision out of the dependency list.
+    """
+    rgb = image.convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE), Image.Resampling.BICUBIC)
+    array = np.asarray(rgb, dtype=np.float32) / 255.0
+    mean = np.array(IMAGE_MEAN, dtype=np.float32)
+    std = np.array(IMAGE_STD, dtype=np.float32)
+    return ((array - mean) / std).transpose(2, 0, 1)
+
+
+class ClipEncoder:
+    """One ONNX session, both towers, batched image encoding."""
+
+    def __init__(self, model_dir: Path) -> None:
+        self.model_dir = Path(model_dir)
+        self.onnx_path = self.model_dir / "onnx" / "model_fp16.onnx"
+        if not self.onnx_path.exists():
+            raise SystemExit(
+                f"{self.onnx_path} is not a Chinese-CLIP export; fetch the fp16 graph: {DOWNLOAD_HINT}"
+            )
+        self._session: Any | None = None
+        self._tokenizer: Any | None = None
+        self._placeholder_text: dict[str, np.ndarray] | None = None
+        self._placeholder_image: np.ndarray | None = None
+
+    @property
+    def session(self) -> Any:
+        if self._session is None:
+            import onnxruntime as ort
+
+            self._session = ort.InferenceSession(str(self.onnx_path), providers=["CPUExecutionProvider"])
+        return self._session
+
+    @property
+    def tokenizer(self) -> Any:
+        if self._tokenizer is None:
+            from transformers import BertTokenizer
+
+            # The export is BERT-style (vocab.txt, [CLS]=101); transformers 5 unified the fast and
+            # slow tokenizers, so the name the CLIP docs use is an alias for this one.
+            self._tokenizer = BertTokenizer.from_pretrained(str(self.model_dir))
+        return self._tokenizer
+
+    def encode_images(self, images: list[Image.Image], batch_size: int = 8) -> np.ndarray:
+        if not images:
+            return np.zeros((0, 512), dtype=np.float32)
+        rows: list[np.ndarray] = []
+        for start in range(0, len(images), batch_size):
+            chunk = images[start : start + batch_size]
+            pixels = np.stack([preprocess_image(image) for image in chunk])
+            rows.append(self._run("image_embeds", pixels, self._text_feed(1))[0])
+        return np.asarray(np.concatenate(rows), dtype=np.float32)
+
+    def encode_texts(self, texts: list[str], batch_size: int = 16) -> np.ndarray:
+        if not texts:
+            return np.zeros((0, 512), dtype=np.float32)
+        rows: list[np.ndarray] = []
+        for start in range(0, len(texts), batch_size):
+            chunk = texts[start : start + batch_size]
+            feed = self._text_feed(len(chunk), chunk)
+            rows.append(self._run("text_embeds", self._image_feed(len(chunk)), feed)[0])
+        return np.asarray(np.concatenate(rows), dtype=np.float32)
+
+    def _run(self, output: str, pixels: np.ndarray, text: dict[str, np.ndarray]) -> list[np.ndarray]:
+        feed = {"pixel_values": pixels.astype(np.float32), **text}
+        return self.session.run([output], feed)
+
+    def _text_feed(self, batch: int, texts: list[str] | None = None) -> dict[str, np.ndarray]:
+        if texts is None:
+            if self._placeholder_text is None:
+                self._placeholder_text = self._tokenize(["placeholder"])
+            return self._placeholder_text
+        return self._tokenize(texts)
+
+    def _image_feed(self, batch: int) -> np.ndarray:
+        if self._placeholder_image is None:
+            self._placeholder_image = np.zeros((1, 3, IMAGE_SIZE, IMAGE_SIZE), dtype=np.float32)
+        return np.repeat(self._placeholder_image, batch, axis=0)
+
+    def _tokenize(self, texts: list[str]) -> dict[str, np.ndarray]:
+        encoded = self.tokenizer(
+            texts,
+            padding="max_length",
+            truncation=True,
+            max_length=TEXT_MAX_LENGTH,
+            return_tensors="np",
+        )
+        return {
+            "input_ids": encoded["input_ids"].astype(np.int64),
+            "attention_mask": encoded["attention_mask"].astype(np.int64),
+        }