From 03ca35763fd63eb8d8f0d2c721d02c6c3ded90a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Wed, 7 Oct 2026 20:41:19 -0400 Subject: [PATCH] Give wordless pictures a way to be found MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- CHECKLIST.md | 23 +- backend/src/ely_eye/api.py | 6 +- backend/src/ely_eye/cli.py | 45 ++- backend/src/ely_eye/clip_encoder.py | 120 ++++++++ backend/src/ely_eye/config.py | 3 + backend/src/ely_eye/db.py | 3 +- backend/src/ely_eye/describe.py | 87 ++++++ backend/src/ely_eye/ingestion.py | 28 +- backend/src/ely_eye/memory.py | 5 +- backend/src/ely_eye/migrations.py | 29 +- backend/src/ely_eye/retrieval.py | 63 ++++- backend/src/ely_eye/schemas.py | 2 + backend/src/ely_eye/storage.py | 12 + backend/src/ely_eye/visual_index.py | 199 ++++++++++++++ backend/tests/test_describe.py | 108 ++++++++ backend/tests/test_visual_index.py | 302 +++++++++++++++++++++ backend/tests/test_visual_tokens.py | 99 +++++++ bench/reports/VISUAL.md | 18 ++ bench/reports/visual.json | 166 ++++++++++++ bench/visual_recall.py | 406 ++++++++++++++++++++++++++++ 20 files changed, 1702 insertions(+), 22 deletions(-) create mode 100644 backend/src/ely_eye/clip_encoder.py create mode 100644 backend/src/ely_eye/describe.py create mode 100644 backend/src/ely_eye/visual_index.py create mode 100644 backend/tests/test_describe.py create mode 100644 backend/tests/test_visual_index.py create mode 100644 backend/tests/test_visual_tokens.py create mode 100644 bench/reports/VISUAL.md create mode 100644 bench/reports/visual.json create mode 100644 bench/visual_recall.py diff --git a/CHECKLIST.md b/CHECKLIST.md index d5803ea..c4db62c 100644 --- a/CHECKLIST.md +++ b/CHECKLIST.md @@ -134,11 +134,24 @@ ## 阶段 2 — P1：让它成为产品 -- [ ] **p1-1** 图像进 dense - - [ ] p1-1.1 选 embedding ≤500MB，双语 + 图像；记录模型 SHA 与体积 - - [ ] p1-1.2 图像 atom 入索引；query 侧支持图像 - - [ ] p1-1.3 视觉 token 计数改 32×32 口径（1280×720 ≈ 900） - - [ ] p1-1.4 VLM 描述：代码路径就绪，标 `DEFERRED` +- [x] **p1-1** 图像进 dense + - [x] p1-1.1 选 embedding ≤500MB，双语 + 图像；记录模型 SHA 与体积 + - `Xenova/chinese-clip-vit-base-patch16` fp16 ONNX，361MB，`onnx/model_fp16.onnx` sha256 `0fad978faa89b85fda53a83a4c202116a9eb87c8f9723fb4d85e4483efbce4b2`；CPU 可跑（仓库已有 onnxruntime）；文本塔 76ms、图像塔 29ms/张（batch 8） + - 契约是量出来的不是猜的：输入 `input_ids`/`attention_mask`/`pixel_values (N,3,224,224)`，输出 `image_embeds`/`text_embeds (N,512)`；这是融合对比图，未用的塔必须喂占位行；分词器用 `BertTokenizer`（按 CLIP 分词器加载会得到超出 vocab 的 id，Gather 直接炸） + - [x] p1-1.2 图像 atom 入索引；query 侧支持图像 + - `backend/src/ely_eye/visual_index.py`（独立文件，retrieval.py 停在 499 行没有膨胀）；`ely-eye index-visual`；`search --image`；API `SearchRequest.image` + - 开关规则：`clip_model_dir` 指向含 `onnx/model_fp16.onnx` 的目录才算可用，不再多一个 bool + - 14 个测试：句子找到图、图找到图、新图一到索引即失效、缺快照点名、map/vector 不一致点名、词够多时不付 encoder、forget 之后图片不再从内存回答 + - [x] p1-1.3 视觉 token 计数改 32×32 口径（1280×720 ≈ 900） + - `visual_token_equivalent = ⌈w/32⌉·⌈h/32⌉`；migration v7 重算存量：21,325 个图像 atom 从 5,331 → 34,618 tokens；整库 29.9M → **84.2M**，于是「100M 胶囊」在这份 50,278 文件语料上刚好装得下、再大就装不下——这才是那句 claim 的真实含义 + - [x] p1-1.4 VLM 描述：代码路径就绪，标 `DEFERRED` + - `describe.py`；ingestion 只在 OCR 没找到字时才调用（图上已有字是免费的）；无快照 → atom metadata `deferred: no vision-language model snapshot`；模型报错 → `failed: …` 且 ingest 不死（策略在调用方，不在 describer）；`status` 显示两个视觉开关；5 个测试 +- p1-1 证据（我额外加的六臂基准，不是验收项）：`bench/visual_recall.py`，ScreenSpot-v2 的 554 个图标 crop + 408 张整屏，报告 `bench/reports/VISUAL.md` + - 图标池：picture 0.208 vs words 0.175 vs chance 0.018 —— 图片空间比瞎猜高 11.5 倍，并且在 427 个 OCR 没字的 crop 上超过了词 + - 整屏池：words 0.384 > picture 0.215 —— 整屏的身份主要是文字，OCR 词在这里赢 + - 整屏 → 图标：0.108 —— CLIP 不会在截图里定位小图标，它是整图相似模型，不是 grounding 模型 + - 单屏内：MRR 0.908（recall@10 是 1.000 但 chance 也是 1.000，所以那一行只有 MRR 有意义） + - 本轮没测双语：ScreenSpot-v2 的指令全是英文，报告里写明了 - 验收：无文字图像被图像查询召回有 recall 数字；`pytest tests/test_visual_tokens.py` 绿 - [ ] **p1-2** 音频摄入 - [ ] p1-2.1 `Modality` 加 audio/transcript；未知扩展名不再静默跳过 diff --git a/backend/src/ely_eye/api.py b/backend/src/ely_eye/api.py index 582ea90..4a3603c 100644 --- a/backend/src/ely_eye/api.py +++ b/backend/src/ely_eye/api.py @@ -197,7 +197,11 @@ def create_app() -> FastAPI: def search(request: SearchRequest, settings: Annotated[Settings, Depends(settings_dep)]) -> SearchResult: try: return RetrievalService(settings).search( - request.query, request.top_k, filters=request.filters, cursor=request.cursor + request.query, + request.top_k, + filters=request.filters, + cursor=request.cursor, + image=request.image, ) except ValueError as exc: raise HTTPException(status_code=422, detail=str(exc)) from exc diff --git a/backend/src/ely_eye/cli.py b/backend/src/ely_eye/cli.py index 09f2bfb..29f3246 100644 --- a/backend/src/ely_eye/cli.py +++ b/backend/src/ely_eye/cli.py @@ -17,8 +17,9 @@ from .benchmarks import ( from .cartridge import CartridgeService from .cartridge_assets import CartridgeAssetService from .cli_transactive import register as register_transactive_commands -from .config import get_settings +from .config import Settings, get_settings from .db import Database +from .describe import DEFERRED, VisionDescriber from .embeddings import EmbeddingService from .ingestion import IngestionService from .memory import MemoryService @@ -31,6 +32,7 @@ from .retrieval import RetrievalService from .runtime import QwenRuntime from .schemas import MemoryOrigin, SearchFilters, TimePrecision, TrainingAdapterKind from .training import TrainingService +from .visual_index import VisualIndex app = typer.Typer(no_args_is_help=True) @@ -49,6 +51,29 @@ def serve(host: str = "127.0.0.1", port: int = 8765) -> None: uvicorn.run(create_app(), host=host, port=port) +def _senses(settings: Settings, db: Database) -> dict[str, dict[str, object]]: + """What this machine can see, and what it has to pretend about. + + A missing snapshot is reported, not hidden: an unrun capability is a plan, and the status + output is where an operator learns which claims about this install are real today. + """ + visual = VisualIndex(settings, db) + describer = VisionDescriber(settings) + index_present = visual.index_path.exists() and visual.map_path.exists() + return { + "picture_search": { + "model_dir": str(visual.model_dir() or ""), + "model_ready": visual.available(), + "index": "present" if index_present else "absent", + }, + "picture_description": { + "model_dir": str(describer.model_dir() or ""), + "model_ready": describer.available(), + "note": "" if describer.available() else DEFERRED, + }, + } + + @app.command() def status() -> None: settings = get_settings() @@ -66,6 +91,7 @@ def status() -> None: ), "cartridges": db.cartridge_count(), "data_home": str(settings.home), + "senses": _senses(settings, db), } ) ) @@ -97,6 +123,7 @@ def search( until: str = typer.Option("", "--until"), min_trust: float = typer.Option(-1.0, "--min-trust", help="drop memories whose trust falls below this"), cursor: str = typer.Option("", "--cursor", help="next_cursor from the previous page"), + image: str = typer.Option("", "--image", help="a picture to search with, instead of only words"), ) -> None: """Rank the memories that may answer, and say so when none may.""" settings = get_settings() @@ -112,10 +139,26 @@ def search( min_trust=None if min_trust < 0 else min_trust, ), cursor=cursor or None, + image=Path(image) if image else None, ) typer.echo(result.model_dump_json(indent=2)) +@app.command() +def describe_image(path: Path) -> None: + """Say what a picture shows; names the setup step when no vision model is installed.""" + settings = get_settings() + typer.echo(json_dump({"image": str(path), "caption": VisionDescriber(settings).describe(path)})) + + +@app.command() +def index_visual() -> None: + """Encode every image atom so a picture or a sentence can find it.""" + settings = get_settings() + database = Database(settings) + typer.echo(json_dump(VisualIndex(settings, database).build())) + + @app.command() def remember( text: str, diff --git a/backend/src/ely_eye/clip_encoder.py b/backend/src/ely_eye/clip_encoder.py new file mode 100644 index 0000000..7975daa --- /dev/null +++ b/backend/src/ely_eye/clip_encoder.py @@ -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 " +) + + +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), + } diff --git a/backend/src/ely_eye/config.py b/backend/src/ely_eye/config.py index d4deab7..12309c3 100644 --- a/backend/src/ely_eye/config.py +++ b/backend/src/ely_eye/config.py @@ -18,6 +18,9 @@ class Settings(BaseSettings): kv_cache_dtype: str = "auto" planner_model_id: str = "Qwen/Qwen3.5-0.8B" embedding_model_id: str = "Qwen/Qwen3-VL-Embedding-8B" + clip_model_dir: Path | None = None + vlm_model_dir: Path | None = None + vlm_max_new_tokens: int = 96 adapter_dir: Path | None = None runtime_backend: Literal["sglang", "transformers"] = "transformers" sglang_base_url: str = "http://127.0.0.1:8000/v1" diff --git a/backend/src/ely_eye/db.py b/backend/src/ely_eye/db.py index 0e9f08d..bcc7932 100644 --- a/backend/src/ely_eye/db.py +++ b/backend/src/ely_eye/db.py @@ -44,7 +44,8 @@ class Database: connection.close() def init(self) -> None: - with sqlite3.connect(self.path) as connection: + # connect(), not a bare sqlite3.connect: the migrations read rows by name. + with self.connect() as connection: apply_migrations(connection) def upsert_sources(self, sources: list[SourceRecord]) -> None: diff --git a/backend/src/ely_eye/describe.py b/backend/src/ely_eye/describe.py new file mode 100644 index 0000000..248dbde --- /dev/null +++ b/backend/src/ely_eye/describe.py @@ -0,0 +1,87 @@ +"""What a picture says when it has no words on it. + +OCR reads letters; it cannot say what a sunset, a shopping cart, or a warning triangle is. A +vision-language model can, and its sentence becomes the atom's text so ordinary word search can +reach the picture too. The model is a multi-gigabyte download, so this path is written, gated on +a snapshot, and reported as deferred until somebody fetches one: an unrun capability is a plan, +not a result. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from .config import Settings, get_settings + +PROMPT = "Describe what is in this image in one or two sentences. Name the objects and any visible text." +DOWNLOAD_HINT = "fetch a Qwen-VL snapshot and set ELY_EYE_VLM_MODEL_DIR to the directory holding it" +DEFERRED = "deferred: no vision-language model snapshot" + + +class VisionDescriber: + """One lazily loaded vision model; `describe_or_none` is the path ingestion takes.""" + + def __init__(self, settings: Settings | None = None) -> None: + self.settings = settings or get_settings() + self._model: Any | None = None + self._processor: Any | None = None + + def model_dir(self) -> Path | None: + model_dir = self.settings.vlm_model_dir + if model_dir is None or not Path(model_dir).is_dir(): + return None + return Path(model_dir) + + def available(self) -> bool: + return self.model_dir() is not None + + def describe(self, path: Path) -> str: + """Caption one picture; a missing snapshot is a setup error, not an empty answer.""" + model_dir = self.model_dir() + if model_dir is None: + raise SystemExit(f"{path}: image description needs a vision-language model; {DOWNLOAD_HINT}") + if not path.is_file(): + raise ValueError(f"{path}: no picture to describe") + return self._caption(path, model_dir) + + def describe_or_none(self, path: Path) -> str | None: + """Caption if this machine can, otherwise say nothing and let the caller record why. + + A model failure is raised, not swallowed: ingestion decides whether a broken vision + model kills the ingest or becomes a note on the atom, and it is the one that knows. + """ + if not self.available(): + return None + return self._caption(path, self.model_dir() or Path()) + + def _caption(self, path: Path, model_dir: Path) -> str: + processor, model = self._load(model_dir) + import torch + from PIL import Image + + with Image.open(path) as picture: + inputs = processor(images=picture.convert("RGB"), text=PROMPT, return_tensors="pt").to( + model.device, torch.float16 + ) + generated = model.generate(**inputs, max_new_tokens=self.settings.vlm_max_new_tokens, do_sample=False) + trimmed = generated[:, inputs["input_ids"].shape[1] :] + text = processor.batch_decode(trimmed, skip_special_tokens=True)[0].strip() + if not text: + raise ValueError(f"{path}: the vision model answered with nothing") + return text + + def _load(self, model_dir: Path) -> tuple[Any, Any]: + """Load once, and hand back the pair so the caller never reads a half-built attribute.""" + processor, model = self._processor, self._model + if processor is None or model is None: + import torch + from transformers import AutoModelForImageTextToText, AutoProcessor + + processor = AutoProcessor.from_pretrained(str(model_dir)) + model = AutoModelForImageTextToText.from_pretrained( + str(model_dir), torch_dtype=torch.float16, device_map="auto" + ) + model.eval() + self._processor, self._model = processor, model + return processor, model diff --git a/backend/src/ely_eye/ingestion.py b/backend/src/ely_eye/ingestion.py index bb58113..b14ddaf 100644 --- a/backend/src/ely_eye/ingestion.py +++ b/backend/src/ely_eye/ingestion.py @@ -16,6 +16,7 @@ from PIL import Image from .config import Settings, get_settings from .db import Database +from .describe import DEFERRED, VisionDescriber from .ocr import read_text_from_image from .schemas import ( EvidenceAtom, @@ -32,6 +33,7 @@ from .storage import ( sha256_bytes, sha256_file, token_equivalent, + visual_token_equivalent, ) SUPPORTED_TEXT_EXTENSIONS = { @@ -93,6 +95,7 @@ class IngestionService: self.settings = settings or get_settings() self.db = database or Database(self.settings) self.objects = object_store or ObjectStore(self.settings) + self.describer = VisionDescriber(self.settings) def ingest_path(self, path: Path, cartridge_name: str | None = None) -> IngestResponse: path = path.expanduser().resolve() @@ -208,18 +211,33 @@ class IngestionService: source = self._source_record(path, self._image_kind(path), metadata=image_metadata(path)) image_ref = self.objects.put_file(path, "images") ocr = read_text_from_image(path) + caption = "" + caption_note = "" + if not ocr.text.strip(): + # Words already on the picture are free; only a wordless one is worth a model call. + try: + caption = self.describer.describe_or_none(path) or "" + except (OSError, RuntimeError, ValueError) as error: + caption, caption_note = "", f"failed: {error}" + else: + caption_note = "vlm" if caption else DEFERRED + text = caption or ocr.text width, height = image_size(path) atom = EvidenceAtom( atom_id=f"{source.source_id}:image", modality=Modality.ui_screenshot if source.kind == "ui_screenshot" else Modality.image, source_id=source.source_id, source=str(path), - text=ocr.text, + text=text, image_ref=image_ref, layout=LayoutBox(x=0, y=0, w=width, h=height), - trust=TrustScores(parser=1.0, ocr=ocr.confidence, model=0.0), - token_equivalent=token_equivalent(ocr.text) + max(1, (width * height) // (224 * 224)), - metadata={"ocr_regions": ocr.regions, **image_metadata(path)}, + trust=TrustScores(parser=1.0, ocr=ocr.confidence, model=0.6 if caption else 0.0), + token_equivalent=token_equivalent(text) + visual_token_equivalent(width, height), + metadata={ + "ocr_regions": ocr.regions, + "caption": caption_note, + **image_metadata(path), + }, ) return IngestBatch(sources=[source], atoms=[atom]) @@ -334,7 +352,7 @@ class IngestionService: image_ref=image_ref, layout=LayoutBox(x=0, y=0, w=float(width), h=float(height), frame_second=second), trust=TrustScores(parser=1.0, ocr=ocr.confidence, model=0.0), - token_equivalent=token_equivalent(ocr.text) + max(1, (width * height) // (224 * 224)), + token_equivalent=token_equivalent(ocr.text) + visual_token_equivalent(width, height), metadata={"frame_index": frame_index, "second": second, "ocr_regions": ocr.regions}, ) ) diff --git a/backend/src/ely_eye/memory.py b/backend/src/ely_eye/memory.py index 3c85cc2..2a559af 100644 --- a/backend/src/ely_eye/memory.py +++ b/backend/src/ely_eye/memory.py @@ -27,6 +27,7 @@ from .schemas import ( utc_now, ) from .storage import token_equivalent +from .visual_index import VisualIndex # Trust scores describe extraction fidelity, not truthfulness: a remembered sentence is # stored verbatim, so the parser is fully trusted and origin carries the epistemic weight. @@ -189,13 +190,13 @@ class MemoryService: ) def _drop_vector_index(self) -> bool: - """The dense index is derived data; a stale one keeps answering with deleted text.""" + """The text and picture indexes are derived data; a stale one keeps answering with deleted memories.""" dropped = False for path in (self.settings.cache_dir / "dense_index.npy", self.settings.cache_dir / "dense_index_atoms.json"): if path.exists(): path.unlink() dropped = True - return dropped + return VisualIndex(self.settings, self.db).drop() or dropped def _cartridges_holding(self, source_id: str) -> list[dict[str, Any]]: """Name the cartridges that exported this source, so they can be rebuilt. diff --git a/backend/src/ely_eye/migrations.py b/backend/src/ely_eye/migrations.py index fc9c676..8112387 100644 --- a/backend/src/ely_eye/migrations.py +++ b/backend/src/ely_eye/migrations.py @@ -10,11 +10,13 @@ upgraded in place rather than rebuilt, because the only copy of the memory is th from __future__ import annotations +import json import sqlite3 +from .storage import token_equivalent, visual_token_equivalent from .text_index import index_text -SCHEMA_VERSION = 6 +SCHEMA_VERSION = 7 ATOM_COLUMNS_V2 = """ atom_id TEXT PRIMARY KEY, @@ -204,6 +206,29 @@ def _upgrade_v1_to_v2(connection: sqlite3.Connection) -> None: connection.execute(f"ALTER TABLE atoms ADD COLUMN {name} {definition}") +def _recount_visual_tokens(connection: sqlite3.Connection) -> None: + """Re-price raster atoms at 32x32 patches; a 1280x720 frame cost 18 tokens under the old rule. + + An atom whose layout carries no size keeps its stored number: there is no region to price, + and inventing one would move a figure we cannot verify. + """ + rows = connection.execute( + "SELECT atom_id, text, layout_json FROM atoms WHERE modality IN ('image', 'video_frame', 'ui_screenshot')" + ).fetchall() + updates: list[tuple[int, str]] = [] + for row in rows: + try: + layout = json.loads(row["layout_json"] or "{}") + except json.JSONDecodeError as error: + raise SystemExit(f"atom {row['atom_id']}: its stored layout is unreadable ({error})") from error + width = float(layout.get("w") or 0) + height = float(layout.get("h") or 0) + if width <= 0 or height <= 0: + continue + updates.append((token_equivalent(row["text"]) + visual_token_equivalent(width, height), str(row["atom_id"]))) + connection.executemany("UPDATE atoms SET token_equivalent = ? WHERE atom_id = ?", updates) + + def apply(connection: sqlite3.Connection) -> None: """Bring the schema to SCHEMA_VERSION, in place, without losing rows.""" version = int(connection.execute("PRAGMA user_version").fetchone()[0]) @@ -220,6 +245,8 @@ def apply(connection: sqlite3.Connection) -> None: _create_v2(connection) if version < 3: _rebuild_fts(connection) + if version < 7: + _recount_visual_tokens(connection) connection.executescript(VIEWS) connection.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") diff --git a/backend/src/ely_eye/retrieval.py b/backend/src/ely_eye/retrieval.py index 8fea0c8..da5c100 100644 --- a/backend/src/ely_eye/retrieval.py +++ b/backend/src/ely_eye/retrieval.py @@ -10,15 +10,19 @@ from __future__ import annotations import base64 from collections import defaultdict +from pathlib import Path from . import cartridge_store from .config import Settings, get_settings -from .db import Database +from .db import Database, FtsPage from .embeddings import EmbeddingService from .schemas import EvidenceAtom, RetrievalHit, SearchFilters, SearchResult +from .visual_index import VisualIndex RRF_K = 60 CANDIDATE_WINDOW = 10 +# Words decide most searches; a picture vector is only worth its encoding cost when words fell short. +THIN_SPARSE_PAGE = 2 class RetrievalService: @@ -26,6 +30,7 @@ class RetrievalService: self.settings = settings or get_settings() self.db = database or Database(self.settings) self.embeddings = EmbeddingService(self.settings, self.db) + self.visual = VisualIndex(self.settings, self.db) def search( self, @@ -34,6 +39,7 @@ class RetrievalService: *, filters: SearchFilters | None = None, cursor: str | None = None, + image: Path | None = None, ) -> SearchResult: """Answer one query against the memories that are open right now. @@ -44,17 +50,20 @@ class RetrievalService: limit = top_k or self.settings.retrieval_top_k if limit <= 0: raise ValueError(f"search needs a positive limit, got {limit}") + if not query.strip() and image is None: + raise ValueError("search needs words or a picture to ask with, and got neither") offset = decode_cursor(cursor) criteria = (filters or SearchFilters()).model_dump(exclude_none=True) scope = criteria.get("cartridge_ids") if scope is None: scope = cartridge_store.mounted_ids(self.db) or None criteria["cartridge_ids"] = scope - page = self.db.search_atoms( - query, - limit=limit * CANDIDATE_WINDOW, - offset=offset, - **criteria, + # A picture-only question has no words for BM25, so the word ranking is skipped rather + # than asked to match an empty expression, which would answer nothing for the wrong reason. + page = ( + self.db.search_atoms(query, limit=limit * CANDIDATE_WINDOW, offset=offset, **criteria) + if query.strip() + else FtsPage(rows=[], has_more=False) ) rankings: list[list[str]] = [[atom_id for atom_id, _ in page.rows]] scores: dict[str, dict[str, float]] = defaultdict(dict) @@ -68,6 +77,17 @@ class RetrievalService: rankings.append(list(dense)) for atom_id, cosine in dense.items(): scores[atom_id]["dense_score"] = cosine + visual = self._visual_search( + query, + image, + limit * CANDIDATE_WINDOW, + scope, + thin=len(rankings[0]) < limit * THIN_SPARSE_PAGE, + ) + if visual: + rankings.append(list(visual)) + for atom_id, cosine in visual.items(): + scores[atom_id]["visual_score"] = cosine hits = self._rank(rankings, scores, limit) more = page.has_more or len(rankings[0]) > limit return SearchResult( @@ -84,6 +104,36 @@ class RetrievalService: query_vector = self.embeddings.embed_query(query) return dict(index.search(query_vector, limit)) + def _visual_search( + self, + query: str, + image: Path | None, + limit: int, + scope: list[str] | None, + *, + thin: bool, + ) -> dict[str, float]: + """Ask the picture index when words are not enough. + + A screenshot with no text on it has nothing for FTS to match, so the sparse page comes + back short; that is the signal to spend the encoding. An image query always asks for it. + """ + if image is None and not thin: + return {} + if not self.visual.available() or self.visual.load() is None: + return {} + if image is not None: + try: + ranked = self.visual.search_image(image, limit) + except (OSError, ValueError) as error: + raise ValueError(f"the query image {image} cannot be read: {error}") from error + else: + ranked = self.visual.search_text(query, limit) + if scope: + allowed = cartridge_store.atoms_of_cartridges(self.db, scope) + ranked = [(atom_id, score) for atom_id, score in ranked if atom_id in allowed] + return dict(ranked) + def _rank(self, rankings: list[list[str]], scores: dict[str, dict[str, float]], limit: int) -> list[RetrievalHit]: """Reciprocal rank fusion: rank position is comparable across signals, raw scores are not. @@ -108,6 +158,7 @@ class RetrievalService: atom=atom, sparse_score=signal.get("sparse_score", 0.0), dense_score=signal.get("dense_score", 0.0), + visual_score=signal.get("visual_score", 0.0), final_score=fused[atom_id], ) ) diff --git a/backend/src/ely_eye/schemas.py b/backend/src/ely_eye/schemas.py index 27f4a1a..5f574a7 100644 --- a/backend/src/ely_eye/schemas.py +++ b/backend/src/ely_eye/schemas.py @@ -149,6 +149,7 @@ class RetrievalHit(BaseModel): atom: EvidenceAtom sparse_score: float dense_score: float = 0.0 + visual_score: float = 0.0 graph_score: float = 0.0 final_score: float @@ -183,6 +184,7 @@ class SearchRequest(BaseModel): top_k: int | None = Field(default=None, gt=0) filters: SearchFilters | None = None cursor: str | None = None + image: Path | None = None class SearchResult(BaseModel): diff --git a/backend/src/ely_eye/storage.py b/backend/src/ely_eye/storage.py index 158a191..ff78fbb 100644 --- a/backend/src/ely_eye/storage.py +++ b/backend/src/ely_eye/storage.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import math import mimetypes import shutil from pathlib import Path @@ -32,6 +33,7 @@ def safe_source_id(path: Path, content_hash: str) -> str: def token_equivalent(text: str) -> int: + """A cheap stand-in for a tokenizer: four ascii characters or two wide ones per token.""" if not text: return 0 ascii_chars = sum(1 for char in text if ord(char) < 128) @@ -39,6 +41,16 @@ def token_equivalent(text: str) -> int: return max(1, ascii_chars // 4 + non_ascii_chars // 2) +VISUAL_PATCH = 32 + + +def visual_token_equivalent(width: float, height: float) -> int: + """Tokens a picture costs a vision model: one per 32x32 patch, so 1280x720 is 920.""" + if width <= 0 or height <= 0: + raise ValueError(f"a visual region needs a positive size, got {width}x{height}") + return max(1, math.ceil(width / VISUAL_PATCH) * math.ceil(height / VISUAL_PATCH)) + + class ObjectStore: def __init__(self, settings: Settings | None = None) -> None: self.settings = settings or get_settings() diff --git a/backend/src/ely_eye/visual_index.py b/backend/src/ely_eye/visual_index.py new file mode 100644 index 0000000..0260c01 --- /dev/null +++ b/backend/src/ely_eye/visual_index.py @@ -0,0 +1,199 @@ +"""A picture found by a picture, or by a sentence in either language. + +Text search cannot reach an image with no words on it: OCR returns nothing, so the atom's +text is empty and the index has nothing to match. This index keeps one CLIP vector per image +atom and encodes the query with the matching tower — a picture with the image tower, a +sentence with the text tower — so both land in the same space and a screenshot can be found +either way. +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image + +from .clip_encoder import DOWNLOAD_HINT, ClipEncoder +from .config import Settings, get_settings +from .db import Database +from .schemas import EvidenceAtom +from .storage import ObjectStore + +INDEX_NAME = "clip_image_index.npy" +MAP_NAME = "clip_image_index_atoms.json" +GRAPH_PART = Path("onnx") / "model_fp16.onnx" + +# Loading 48k vectors costs about as much as the search itself, and the API builds a fresh +# service per request, so the array is kept per process and re-read only when the file changes. +_LOADED: dict[tuple[str, float], tuple[list[str], np.ndarray]] = {} + + +def is_visual(atom: EvidenceAtom) -> bool: + """An atom is findable by picture only if there is a picture behind it.""" + return bool(atom.image_ref) + + +class VisualIndex: + """CLIP vectors for the image atoms, plus the two ways to ask for them.""" + + def __init__(self, settings: Settings | None = None, database: Database | None = None) -> None: + self.settings = settings or get_settings() + self.db = database or Database(self.settings) + self.index_path = self.settings.cache_dir / INDEX_NAME + self.map_path = self.settings.cache_dir / MAP_NAME + self._encoder: ClipEncoder | None = None + + def model_dir(self) -> Path | None: + """The CLIP snapshot, or None when this machine has not fetched one.""" + model_dir = self.settings.clip_model_dir + if model_dir is None or not (model_dir / GRAPH_PART).exists(): + return None + return model_dir + + def available(self) -> bool: + return self.model_dir() is not None + + def encoder(self) -> ClipEncoder: + if self._encoder is None: + model_dir = self.model_dir() + if model_dir is None: + raise SystemExit( + f"visual search needs a Chinese-CLIP snapshot; set ELY_EYE_CLIP_MODEL_DIR. {DOWNLOAD_HINT}" + ) + self._encoder = ClipEncoder(model_dir) + return self._encoder + + def visual_atoms(self, atoms: list[EvidenceAtom] | None = None) -> list[EvidenceAtom]: + """Image atoms whose bytes are still on disk; a dead reference cannot be encoded.""" + pool = self.db.list_atoms() if atoms is None else atoms + store = ObjectStore(self.settings) + kept: list[EvidenceAtom] = [] + for atom in pool: + if not is_visual(atom): + continue + path = store.resolve(str(atom.image_ref)) + if path.exists(): + kept.append(atom) + return kept + + def build(self, atoms: list[EvidenceAtom] | None = None) -> dict[str, Any]: + """Encode every image atom once; the batch is what keeps this affordable.""" + started = time.monotonic() + pool = self.visual_atoms(atoms) + encoder = self.encoder() + vectors: list[np.ndarray] = [] + kept: list[EvidenceAtom] = [] + batch: list[Image.Image] = [] + for atom in pool: + try: + with Image.open(ObjectStore(self.settings).resolve(str(atom.image_ref))) as picture: + batch.append(picture.convert("RGB")) + except OSError as error: + print(f"skipping {atom.atom_id}: the image will not open ({error})") + continue + kept.append(atom) + if len(batch) >= 32: + vectors.append(encoder.encode_images(batch)) + batch = [] + if batch: + vectors.append(encoder.encode_images(batch)) + matrix = np.concatenate(vectors) if vectors else np.zeros((0, 512), dtype=np.float32) + self._write([atom.atom_id for atom in kept], matrix) + return { + "atoms": len(kept), + "dim": int(matrix.shape[1]) if matrix.size else 0, + "seconds": round(time.monotonic() - started, 2), + "model_dir": str(self.model_dir()), + } + + def load(self) -> tuple[list[str], np.ndarray] | None: + if not self.index_path.exists() or not self.map_path.exists(): + return None + stamp = (str(self.index_path), self.index_path.stat().st_mtime) + cached = _LOADED.get(stamp) + if cached is not None: + return cached + atom_ids = _read_ids(self.map_path) + vectors = np.load(self.index_path) + if len(atom_ids) != vectors.shape[0]: + raise ValueError( + f"{self.map_path}: {len(atom_ids)} atoms for {vectors.shape[0]} vectors; rebuild with ely-eye index-visual" + ) + _LOADED.clear() + _LOADED[stamp] = (atom_ids, vectors) + return atom_ids, vectors + + def drop(self) -> bool: + """Delete the index files and the copy this process is holding. + + Forget has to void this index the way it voids the text one: a forgotten picture that + still answers from an array in memory is a memory that was never forgotten. + """ + removed = False + for path in (self.index_path, self.map_path): + if path.exists(): + path.unlink() + removed = True + _LOADED.clear() + return removed + + def is_current(self, atoms: list[EvidenceAtom] | None = None) -> bool: + """True when the index covers exactly the image atoms the store holds right now.""" + loaded = self.load() + if loaded is None: + return False + return {atom.atom_id for atom in self.visual_atoms(atoms)} == set(loaded[0]) + + def search_text(self, query: str, top_k: int, restrict: set[str] | None = None) -> list[tuple[str, float]]: + vector = self.encoder().encode_texts([query])[0] + return self._nearest(vector, top_k, restrict) + + def search_image(self, image: Path, top_k: int, restrict: set[str] | None = None) -> list[tuple[str, float]]: + with Image.open(image) as picture: + vector = self.encoder().encode_images([picture.convert("RGB")])[0] + return self._nearest(vector, top_k, restrict) + + def _nearest( + self, query_vector: np.ndarray, top_k: int, restrict: set[str] | None = None + ) -> list[tuple[str, float]]: + """Closest stored pictures, optionally inside one candidate pool. + + The pool is how a benchmark asks a smaller question — “which element of this screen” — + without rebuilding the index per question. + """ + loaded = self.load() + if loaded is None or top_k <= 0: + return [] + atom_ids, vectors = loaded + if vectors.size == 0: + return [] + query = _unit(query_vector.reshape(1, -1))[0] + scores = _unit(vectors) @ query + if restrict is not None: + keep = np.fromiter((atom_id in restrict for atom_id in atom_ids), dtype=bool, count=len(atom_ids)) + scores = np.where(keep, scores, -np.inf) + order = np.argsort(scores)[::-1][:top_k] + return [(atom_ids[index], float(scores[index])) for index in order if np.isfinite(scores[index])] + + def _write(self, atom_ids: list[str], vectors: np.ndarray) -> None: + self.index_path.parent.mkdir(parents=True, exist_ok=True) + np.save(self.index_path, vectors) + self.map_path.write_text(json.dumps(atom_ids, ensure_ascii=False), encoding="utf-8") + _LOADED.clear() + + +def _read_ids(path: Path) -> list[str]: + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"{path}: the visual index map is unreadable ({error})") from error + + +def _unit(vectors: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(vectors, axis=1, keepdims=True) + norms[norms == 0] = 1.0 + return vectors / norms diff --git a/backend/tests/test_describe.py b/backend/tests/test_describe.py new file mode 100644 index 0000000..91c6208 --- /dev/null +++ b/backend/tests/test_describe.py @@ -0,0 +1,108 @@ +"""A picture with no words on it, and what the ingest path does about it. + +The vision model is a multi-gigabyte download, so these tests run the branch, not the weights: +a stand-in describer answers, and the atom records whether its text came from letters on the +picture, from a model's sentence, or from nothing at all. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import cast + +import pytest +from PIL import Image + +import ely_eye.ingestion as ingestion_module +from ely_eye.config import Settings +from ely_eye.db import Database +from ely_eye.describe import DEFERRED, VisionDescriber +from ely_eye.ingestion import IngestionService +from ely_eye.ocr import OcrResult + +CAPTION = "一个购物车图标停在结算按钮旁边" + + +class FakeDescriber: + """Stands in for the vision model and counts how often ingestion asked.""" + + def __init__(self, answer: str | None = CAPTION, fails: bool = False) -> None: + self.answer = answer + self.fails = fails + self.calls = 0 + + def describe_or_none(self, path: Path) -> str | None: + self.calls += 1 + if self.fails: + raise RuntimeError("the vision model ran out of memory") + return self.answer + + +def _picture(tmp_path: Path) -> Path: + path = tmp_path / "scene.png" + Image.new("RGB", (640, 360), (12, 34, 56)).save(path) + return path + + +def _ingest(tmp_path: Path, describer: object, ocr_text: str = "") -> tuple[Settings, Database]: + settings = Settings(home=tmp_path / "home") + database = Database(settings) + service = IngestionService(settings=settings, database=database) + service.describer = cast(VisionDescriber, describer) + ingestion_module.read_text_from_image = lambda path: OcrResult(ocr_text, 0.9, []) + service.ingest_path(_picture(tmp_path)) + return settings, database + + +def test_a_wordless_picture_ingested_with_a_model_carries_the_caption(tmp_path: Path) -> None: + from ely_eye.storage import token_equivalent, visual_token_equivalent + + _, database = _ingest(tmp_path, FakeDescriber()) + + atom = database.list_atoms()[0] + assert atom.text == CAPTION + assert atom.metadata["caption"] == "vlm" + assert atom.trust.model == 0.6, "a sentence invented by a model is less trusted than printed letters" + assert atom.token_equivalent == token_equivalent(CAPTION) + visual_token_equivalent(640, 360) + + +def test_a_wordless_picture_without_a_model_says_so_in_the_atom(tmp_path: Path) -> None: + describer = FakeDescriber(answer=None) + + _, database = _ingest(tmp_path, describer) + + atom = database.list_atoms()[0] + assert atom.text == "" + assert atom.metadata["caption"] == DEFERRED + assert describer.calls == 1, "ingestion asks, and the answer is recorded as missing" + + +def test_a_picture_with_printed_words_never_costs_a_model_call(tmp_path: Path) -> None: + describer = FakeDescriber() + + _, database = _ingest(tmp_path, describer, ocr_text="Checkout") + + assert describer.calls == 0, "letters already on the picture are free" + atom = database.list_atoms()[0] + assert atom.text == "Checkout" + assert atom.metadata["caption"] == "" + + +def test_a_model_that_fails_leaves_the_picture_ingestable(tmp_path: Path) -> None: + describer = FakeDescriber(fails=True) + + _, database = _ingest(tmp_path, describer) + + atom = database.list_atoms()[0] + assert atom.text == "" + assert atom.metadata["caption"].startswith("failed:"), "the atom says why its text is empty" + assert "ran out of memory" in atom.metadata["caption"] + + +def test_a_missing_snapshot_is_named_as_a_setup_step(tmp_path: Path) -> None: + settings = Settings(home=tmp_path / "home", vlm_model_dir=tmp_path / "nowhere") + describer = VisionDescriber(settings) + + assert not describer.available() + with pytest.raises(SystemExit, match="ELY_EYE_VLM_MODEL_DIR"): + describer.describe(_picture(tmp_path)) diff --git a/backend/tests/test_visual_index.py b/backend/tests/test_visual_index.py new file mode 100644 index 0000000..507b786 --- /dev/null +++ b/backend/tests/test_visual_index.py @@ -0,0 +1,302 @@ +"""Finding a picture by a picture, or by a sentence. + +The encoder is a 361 MB ONNX graph, so these tests swap in a fake that maps a picture's +colour, or a query's number, to a one-hot vector. What is under test is the index contract: +which atoms get encoded, when the index is stale, when search asks the picture tower at all, +and what happens when the snapshot is missing. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import cast + +import numpy as np +import pytest +from PIL import Image + +from ely_eye.clip_encoder import ClipEncoder +from ely_eye.config import Settings +from ely_eye.db import Database +from ely_eye.memory import MemoryService +from ely_eye.retrieval import RetrievalService +from ely_eye.schemas import EvidenceAtom, Modality, SourceRecord, TrustScores +from ely_eye.storage import ObjectStore +from ely_eye.visual_index import VisualIndex, is_visual + +DIM = 4 + + +class FakeEncoder: + """A picture is its top-left colour; a text query is its digits.""" + + def __init__(self) -> None: + self.image_calls = 0 + self.text_calls = 0 + + def encode_images(self, images: list[Image.Image], batch_size: int = 8) -> np.ndarray: + self.image_calls += 1 + return _one_hot([int(np.asarray(image.convert("RGB"))[0, 0, 0]) for image in images]) + + def encode_texts(self, texts: list[str], batch_size: int = 16) -> np.ndarray: + self.text_calls += 1 + return _one_hot([int(text) for text in texts]) + + +def _one_hot(codes: list[int]) -> np.ndarray: + matrix = np.zeros((len(codes), DIM), dtype=np.float32) + for row, code in enumerate(codes): + matrix[row, code % DIM] = 1.0 + return matrix + + +def _settings(tmp_path: Path, with_model: bool = True) -> Settings: + model_dir = tmp_path / "clip" + if with_model: + (model_dir / "onnx").mkdir(parents=True, exist_ok=True) + (model_dir / "onnx" / "model_fp16.onnx").write_bytes(b"graph") + return Settings(home=tmp_path / "home", clip_model_dir=model_dir, default_holder="wife") + + +def _picture(tmp_path: Path, name: str, shade: int) -> Path: + path = tmp_path / f"{name}.png" + Image.new("RGB", (64, 64), (shade, 0, 0)).save(path) + return path + + +def _source(source_id: str, path: str, mime: str) -> SourceRecord: + return SourceRecord(source_id=source_id, path=path, kind="file", mime=mime, sha256=source_id, size_bytes=64) + + +def _image_atom(tmp_path: Path, settings: Settings, database: Database, name: str, shade: int) -> EvidenceAtom: + """Store a picture the way ingestion does: bytes in the object store, atom pointing at them.""" + picture = _picture(tmp_path, name, shade) + uri = ObjectStore(settings).put_file(picture, "images") + database.upsert_sources([_source(name, str(picture), "image/png")]) + atom = EvidenceAtom( + atom_id=f"{name}:image", + modality=Modality.image, + source_id=name, + source=f"file://{name}", + text="", + image_ref=uri, + trust=TrustScores(parser=1.0, ocr=0.0, model=0.0), + token_equivalent=920, + holder_id=settings.default_holder, + ) + database.upsert_atoms([atom]) + return atom + + +def _attach(index: VisualIndex, encoder: FakeEncoder) -> FakeEncoder: + index._encoder = cast(ClipEncoder, encoder) + return encoder + + +def _index(tmp_path: Path, with_model: bool = True) -> tuple[VisualIndex, Settings, Database, FakeEncoder]: + settings = _settings(tmp_path, with_model) + database = Database(settings) + index = VisualIndex(settings, database) + return index, settings, database, _attach(index, FakeEncoder()) + + +def test_a_text_atom_is_not_findable_by_picture() -> None: + atom = EvidenceAtom( + atom_id="t:1", + modality=Modality.text, + source_id="s", + source="file://s", + text="钥匙放在抽屉里", + trust=TrustScores(parser=1.0, ocr=0.0, model=0.0), + token_equivalent=6, + ) + assert not is_visual(atom) + + +def test_only_images_whose_bytes_survive_are_encoded(tmp_path: Path) -> None: + index, settings, database, _ = _index(tmp_path) + _image_atom(tmp_path, settings, database, "keep", 1) + dead = _image_atom(tmp_path, settings, database, "gone", 2) + ObjectStore(settings).resolve(str(dead.image_ref)).unlink() + + kept = index.visual_atoms() + + assert [atom.atom_id for atom in kept] == ["keep:image"] + + +def test_a_sentence_finds_the_picture_it_describes(tmp_path: Path) -> None: + index, settings, database, encoder = _index(tmp_path) + for shade in (1, 2, 3): + _image_atom(tmp_path, settings, database, f"pic{shade}", shade) + index.build() + + ranked = index.search_text("2", top_k=3) + + assert ranked[0][0] == "pic2:image" + assert ranked[0][1] == pytest.approx(1.0) + assert encoder.text_calls == 1 + + +def test_a_picture_finds_the_picture_it_looks_like(tmp_path: Path) -> None: + index, settings, database, _ = _index(tmp_path) + for shade in (1, 2, 3): + _image_atom(tmp_path, settings, database, f"pic{shade}", shade) + index.build() + + ranked = index.search_image(_picture(tmp_path, "query", 3), top_k=2) + + assert [atom_id for atom_id, _ in ranked] == ["pic3:image", "pic1:image"] + + +def test_the_index_is_stale_the_moment_a_picture_arrives(tmp_path: Path) -> None: + index, settings, database, _ = _index(tmp_path) + _image_atom(tmp_path, settings, database, "first", 1) + index.build() + assert index.is_current() + + _image_atom(tmp_path, settings, database, "second", 2) + + assert not index.is_current(), "a new picture must not be invisible to search" + + +def test_a_pool_limits_the_answer_without_rebuilding_the_index(tmp_path: Path) -> None: + index, settings, database, _ = _index(tmp_path) + for shade in (1, 2, 3): + _image_atom(tmp_path, settings, database, f"pic{shade}", shade) + index.build() + + ranked = index.search_text("3", top_k=3, restrict={"pic1:image", "pic3:image"}) + assert [atom_id for atom_id, _ in ranked] == ["pic3:image", "pic1:image"] + + outside = index.search_text("3", top_k=3, restrict={"pic1:image", "pic2:image"}) + assert "pic3:image" not in [atom_id for atom_id, _ in outside] + + +def test_forgetting_a_picture_voids_the_picture_index(tmp_path: Path) -> None: + index, settings, database, _ = _index(tmp_path) + atom = _image_atom(tmp_path, settings, database, "pic1", 1) + index.build() + assert index.load() is not None, "warm the copy this process holds" + + report = MemoryService(settings, database).forget(source_id=atom.source_id) + + assert report.vector_index_dropped + assert not index.index_path.exists() and not index.map_path.exists() + assert index.load() is None, "a forgotten picture must not keep answering from memory" + + +def test_a_missing_snapshot_is_named_rather_than_mumbled(tmp_path: Path) -> None: + settings = _settings(tmp_path, with_model=False) + + index = VisualIndex(settings, Database(settings)) + + assert not index.available() + with pytest.raises(SystemExit, match="ELY_EYE_CLIP_MODEL_DIR"): + index.encoder() + + +def test_a_map_that_disagrees_with_the_vectors_is_refused(tmp_path: Path) -> None: + index, settings, database, _ = _index(tmp_path) + _image_atom(tmp_path, settings, database, "only", 1) + index.build() + index.map_path.write_text('["only:image", "ghost:image"]', encoding="utf-8") + + with pytest.raises(ValueError, match="2 atoms for 1 vectors"): + index.load() + + +def test_a_corrupt_map_names_the_file(tmp_path: Path) -> None: + index, settings, database, _ = _index(tmp_path) + _image_atom(tmp_path, settings, database, "only", 1) + index.build() + index.map_path.write_text("not json", encoding="utf-8") + + with pytest.raises(ValueError, match="visual index map is unreadable"): + index.load() + + +def _search_service(encoder: FakeEncoder, settings: Settings, database: Database) -> RetrievalService: + service = RetrievalService(settings, database) + _attach(service.visual, encoder) + return service + + +def test_words_that_come_up_short_let_the_picture_index_answer(tmp_path: Path) -> None: + index, settings, database, encoder = _index(tmp_path) + for shade in (1, 2, 3): + _image_atom(tmp_path, settings, database, f"pic{shade}", shade) + index.build() + service = _search_service(encoder, settings, database) + + result = service.search("7", top_k=3) + + assert result.hits[0].atom.atom_id == "pic3:image" + assert result.hits[0].visual_score == pytest.approx(1.0) + assert encoder.text_calls == 1 + + +def test_words_that_already_fill_the_page_do_not_pay_for_a_picture(tmp_path: Path) -> None: + index, settings, database, encoder = _index(tmp_path) + for shade in range(1, 9): + _image_atom(tmp_path, settings, database, f"pic{shade}", shade) + database.upsert_sources([_source(f"note{shade}", f"/tmp/note{shade}.md", "text/markdown")]) + database.upsert_atoms( + [ + EvidenceAtom( + atom_id=f"note{shade}:text", + modality=Modality.text, + source_id=f"note{shade}", + source=f"file://note{shade}", + text="widget gadget sprocket", + trust=TrustScores(parser=1.0, ocr=0.0, model=0.0), + token_equivalent=3, + holder_id=settings.default_holder, + ) + ] + ) + index.build() + service = _search_service(encoder, settings, database) + + result = service.search("widget gadget sprocket", top_k=2) + + assert encoder.text_calls == 0, "a page full of word matches does not need the encoder" + assert all(hit.visual_score == 0.0 for hit in result.hits) + + +def test_an_image_query_uses_the_image_tower(tmp_path: Path) -> None: + index, settings, database, encoder = _index(tmp_path) + for shade in (1, 2): + _image_atom(tmp_path, settings, database, f"pic{shade}", shade) + index.build() + service = _search_service(encoder, settings, database) + + picture_calls = encoder.image_calls + + result = service.search("", top_k=2, image=_picture(tmp_path, "q", 2)) + + assert result.hits[0].atom.atom_id == "pic2:image" + assert encoder.image_calls == picture_calls + 1, "the query picture goes through the image tower" + assert encoder.text_calls == 0, "a picture-only question has no words to encode" + + +def test_an_unreadable_query_image_is_the_caller_s_problem(tmp_path: Path) -> None: + index, settings, database, encoder = _index(tmp_path) + _image_atom(tmp_path, settings, database, "pic1", 1) + index.build() + service = _search_service(encoder, settings, database) + broken = tmp_path / "broken.png" + broken.write_text("not a picture", encoding="utf-8") + + with pytest.raises(ValueError, match="cannot be read"): + service.search("", image=broken) + + +def test_a_search_without_a_snapshot_never_touches_the_encoder(tmp_path: Path) -> None: + _, settings, database, encoder = _index(tmp_path, with_model=False) + _image_atom(tmp_path, settings, database, "pic1", 1) + service = _search_service(encoder, settings, database) + + result = service.search("pic1", top_k=2) + + assert encoder.text_calls == 0 + assert all(hit.visual_score == 0.0 for hit in result.hits) diff --git a/backend/tests/test_visual_tokens.py b/backend/tests/test_visual_tokens.py new file mode 100644 index 0000000..f1321fd --- /dev/null +++ b/backend/tests/test_visual_tokens.py @@ -0,0 +1,99 @@ +"""What a picture costs the model that reads it. + +The first version divided pixels by 224x224, which priced a 1280x720 frame at 18 tokens — +less than its caption. A vision model spends one token per patch, and Qwen3-VL's patch is +32x32, so the same frame is 920. A context budget built on the old number holds six times +less picture than it claims. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from PIL import Image + +import ely_eye.ingestion as ingestion_module +from ely_eye.config import Settings +from ely_eye.db import Database +from ely_eye.ingestion import IngestionService +from ely_eye.migrations import SCHEMA_VERSION, apply +from ely_eye.ocr import OcrResult +from ely_eye.storage import token_equivalent, visual_token_equivalent + +FRAME = (1280, 720) + + +def _ingest(tmp_path: Path, settings: Settings) -> Database: + """Run the real ingest path; each caller decides what its OCR says.""" + database = Database(settings) + Image.new("RGB", FRAME, (10, 20, 30)).save(tmp_path / "frame.png") + IngestionService(settings=settings, database=database).ingest_path(tmp_path / "frame.png") + return database + + +def test_a_frame_is_priced_by_the_patches_a_vision_model_sees() -> None: + assert visual_token_equivalent(1280, 720) == 920 + + +def test_a_small_tile_still_costs_more_than_a_single_token() -> None: + assert visual_token_equivalent(224, 224) == 49 + + +def test_a_region_without_a_size_is_refused_rather_than_guessed() -> None: + for width, height in ((0, 720), (1280, -1)): + with pytest.raises(ValueError, match="needs a positive size"): + visual_token_equivalent(width, height) + + +def test_an_ingested_image_carries_its_text_and_its_patches(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + settings = Settings(home=tmp_path) + monkeypatch.setattr(ingestion_module, "read_text_from_image", lambda path: OcrResult("购物车", 0.9, [])) + + database = _ingest(tmp_path, settings) + atom = database.list_atoms()[0] + + assert atom.token_equivalent == token_equivalent("购物车") + 920 + + +def test_an_existing_store_is_repriced_by_the_migration(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + settings = Settings(home=tmp_path) + monkeypatch.setattr(ingestion_module, "read_text_from_image", lambda path: OcrResult("", 0.0, [])) + database = _ingest(tmp_path, settings) + note = tmp_path / "notes.md" + note.write_text("# 钥匙放在玄关抽屉里\n", encoding="utf-8") + IngestionService(settings=settings, database=database).ingest_path(note) + with database.connect() as connection: + connection.execute("UPDATE atoms SET token_equivalent = 18 WHERE modality = 'image'") + connection.execute("PRAGMA user_version = 6") + + with database.connect() as connection: + apply(connection) + + atoms = {atom.modality.value: atom for atom in database.list_atoms()} + assert atoms["image"].token_equivalent == 920, "the old 224x224 price must not survive the upgrade" + assert atoms["text"].token_equivalent == token_equivalent(note.read_text(encoding="utf-8")) + with database.connect() as connection: + assert int(connection.execute("PRAGMA user_version").fetchone()[0]) == SCHEMA_VERSION + + +def test_an_atom_without_a_layout_size_keeps_the_number_it_has(tmp_path: Path) -> None: + """Nothing can be priced from a missing region, so the stored figure is left alone.""" + settings = Settings(home=tmp_path) + database = Database(settings) + with database.connect() as connection: + apply(connection) + connection.execute("PRAGMA user_version = 6") + connection.execute( + "INSERT INTO sources (source_id, path, kind, mime, sha256, size_bytes, created_at, metadata_json)" + " VALUES ('s', '/tmp/x.png', 'image', 'image/png', 'h', 1, '2026-01-01T00:00:00+00:00', '{}')" + ) + connection.execute( + "INSERT INTO atoms (atom_id, source_id, modality, source, ingested_at, text, embedding_refs_json," + " layout_json, relations_json, trust_json, metadata_json, token_equivalent, origin)" + " VALUES ('a1', 's', 'image', '/tmp/x.png', '2026-01-01T00:00:00+00:00', 'x', '[]', '{}', '[]', '{}', '{}', 18, 'artifact')" + ) + + with database.connect() as connection: + apply(connection) + assert int(connection.execute("SELECT token_equivalent FROM atoms WHERE atom_id = 'a1'").fetchone()[0]) == 18 diff --git a/bench/reports/VISUAL.md b/bench/reports/VISUAL.md new file mode 100644 index 0000000..b05bd79 --- /dev/null +++ b/bench/reports/VISUAL.md @@ -0,0 +1,18 @@ +# VISUAL: can a library of wordless pictures be searched? (commit 7e1fe649f) + +- generated 2026-10-08T00:37:50.278445+00:00 +- corpus: OS-Copilot/ScreenSpot-v2 (icon instructions), 554 icon crops and 408 whole screenshots stored as separate memories, 427 crops with no OCR text +- encoder: Xenova/chinese-clip-vit-base-patch16 sha256 0fad978faa89b85f +- machine: 32 cpus, load [46.81, 44.0, 33.68] +- instructions containing Chinese: 0 (ScreenSpot-v2 is English, so this run does not test the bilingual claim) + +| ask | pool | recall@1 | recall@5 | recall@10 | chance@10 | MRR | hit@10 lower 95% | p50 s | p95 s | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | +| words → icon crops | crops | 0.090 | 0.152 | 0.175 | 0.0181 | 0.118 | 0.146 | 0.1716 | 0.2096 | +| a sentence → icon crops (picture space) | crops | 0.081 | 0.155 | 0.208 | 0.0181 | 0.116 | 0.176 | 0.1399 | 0.1998 | +| the whole screenshot → icon crops | crops | 0.038 | 0.083 | 0.108 | 0.0181 | 0.058 | 0.085 | 0.1599 | 0.22 | +| words → whole screenshots | shots | 0.121 | 0.296 | 0.384 | 0.0245 | 0.196 | 0.345 | 0.1362 | 0.2038 | +| a sentence → whole screenshots (picture space) | shots | 0.049 | 0.146 | 0.215 | 0.0245 | 0.091 | 0.183 | 0.136 | 0.204 | +| a sentence → the icons of that one screen | own_screen | 0.827 | 1.000 | 1.000 | 1.0000 | 0.908 | 0.993 | 0.1358 | 0.1959 | + +The words rows are what the product could do before any picture vector: if OCR found no words, there is nothing to match. The picture rows ask the same sentence, or the whole screenshot, of the CLIP space. The pools change the question: among every icon in the library, among whole screenshots (the granularity the store keeps), and among one screen's icons, which is the shape ScreenSpot itself poses. A number is only comparable inside its own pool, so chance is printed per row. diff --git a/bench/reports/visual.json b/bench/reports/visual.json new file mode 100644 index 0000000..4c250cb --- /dev/null +++ b/bench/reports/visual.json @@ -0,0 +1,166 @@ +{ + "generated_at": "2026-10-08T00:37:50.278445+00:00", + "commit": "7e1fe649f68bcb25108e2781d912c8a80e3bbfc8", + "dataset": "OS-Copilot/ScreenSpot-v2 (icon instructions)", + "model": { + "repo": "Xenova/chinese-clip-vit-base-patch16", + "sha256": "0fad978faa89b85fda53a83a4c202116a9eb87c8f9723fb4d85e4483efbce4b2" + }, + "targets": 554, + "screenshots": 408, + "crops_without_ocr_text": 427, + "instruction_languages": { + "cjk_instructions": 0 + }, + "machine": { + "system": "Linux", + "release": "6.18.40.1-microsoft-standard-WSL2", + "cpus": 32, + "threads_per_query": 4, + "load_average_1_5_15": [ + 46.81, + 44.0, + 33.68 + ] + }, + "arms": [ + { + "label": "words → icon crops", + "key": "words->crops", + "n_answerable": 554, + "n_unanswerable": 0, + "recall_at_1": 0.09025270758122744, + "recall_at_5": 0.15162454873646208, + "recall_at_10": 0.17509025270758122, + "ndcg_at_10": 0.13173348831922826, + "mrr": 0.11802403873703514, + "hit_at_10": 0.17509025270758122, + "hit_at_10_wilson_lower": 0.1457109350541202, + "abstention_rate_answerable": 0.04512635379061372, + "abstention_rate_unanswerable": 0.0, + "top_hit_from_holder_unanswerable": 0.0, + "latency": { + "queries": 554, + "p50": 0.1716, + "p95": 0.2096 + }, + "chance_recall_at_10": 0.0181, + "pool": "crops" + }, + { + "label": "a sentence → icon crops (picture space)", + "key": "picture_words->crops", + "n_answerable": 554, + "n_unanswerable": 0, + "recall_at_1": 0.0812274368231047, + "recall_at_5": 0.1552346570397112, + "recall_at_10": 0.2075812274368231, + "ndcg_at_10": 0.13734530518050886, + "mrr": 0.1158894905736061, + "hit_at_10": 0.2075812274368231, + "hit_at_10_wilson_lower": 0.1758785967641179, + "abstention_rate_answerable": 0.0, + "abstention_rate_unanswerable": 0.0, + "top_hit_from_holder_unanswerable": 0.0, + "latency": { + "queries": 554, + "p50": 0.1399, + "p95": 0.1998 + }, + "chance_recall_at_10": 0.0181, + "pool": "crops" + }, + { + "label": "the whole screenshot → icon crops", + "key": "picture->crops", + "n_answerable": 554, + "n_unanswerable": 0, + "recall_at_1": 0.03790613718411552, + "recall_at_5": 0.08303249097472924, + "recall_at_10": 0.10830324909747292, + "ndcg_at_10": 0.069893811531563, + "mrr": 0.058046100509999435, + "hit_at_10": 0.10830324909747292, + "hit_at_10_wilson_lower": 0.08507164644767994, + "abstention_rate_answerable": 0.0, + "abstention_rate_unanswerable": 0.0, + "top_hit_from_holder_unanswerable": 0.0, + "latency": { + "queries": 554, + "p50": 0.1599, + "p95": 0.22 + }, + "chance_recall_at_10": 0.0181, + "pool": "crops" + }, + { + "label": "words → whole screenshots", + "key": "words->shots", + "n_answerable": 554, + "n_unanswerable": 0, + "recall_at_1": 0.12093862815884476, + "recall_at_5": 0.296028880866426, + "recall_at_10": 0.3844765342960289, + "ndcg_at_10": 0.24046684872963128, + "mrr": 0.1959028135923443, + "hit_at_10": 0.3844765342960289, + "hit_at_10_wilson_lower": 0.34489504523476483, + "abstention_rate_answerable": 0.0, + "abstention_rate_unanswerable": 0.0, + "top_hit_from_holder_unanswerable": 0.0, + "latency": { + "queries": 554, + "p50": 0.1362, + "p95": 0.2038 + }, + "chance_recall_at_10": 0.0245, + "pool": "shots" + }, + { + "label": "a sentence → whole screenshots (picture space)", + "key": "picture_words->shots", + "n_answerable": 554, + "n_unanswerable": 0, + "recall_at_1": 0.048736462093862815, + "recall_at_5": 0.14620938628158844, + "recall_at_10": 0.2148014440433213, + "ndcg_at_10": 0.11948802140007041, + "mrr": 0.09050126067274082, + "hit_at_10": 0.2148014440433213, + "hit_at_10_wilson_lower": 0.18262874778266808, + "abstention_rate_answerable": 0.0, + "abstention_rate_unanswerable": 0.0, + "top_hit_from_holder_unanswerable": 0.0, + "latency": { + "queries": 554, + "p50": 0.136, + "p95": 0.204 + }, + "chance_recall_at_10": 0.0245, + "pool": "shots" + }, + { + "label": "a sentence → the icons of that one screen", + "key": "picture_words->own_screen", + "n_answerable": 554, + "n_unanswerable": 0, + "recall_at_1": 0.8267148014440433, + "recall_at_5": 1.0, + "recall_at_10": 1.0, + "ndcg_at_10": 0.9318097223990938, + "mrr": 0.9080024067388689, + "hit_at_10": 1.0, + "hit_at_10_wilson_lower": 0.9931137086353976, + "abstention_rate_answerable": 0.0, + "abstention_rate_unanswerable": 0.0, + "top_hit_from_holder_unanswerable": 0.0, + "latency": { + "queries": 554, + "p50": 0.1358, + "p95": 0.1959 + }, + "chance_recall_at_10": 1.0, + "pool": "own_screen" + } + ] +} \ No newline at end of file diff --git a/bench/visual_recall.py b/bench/visual_recall.py new file mode 100644 index 0000000..0e904c7 --- /dev/null +++ b/bench/visual_recall.py @@ -0,0 +1,406 @@ +"""Can a library of wordless pictures be searched? + +ScreenSpot-v2 is the only corpus in this suite whose gold answer is a *region of a screen*, and +its icon targets carry almost no text. That makes it the honest test of the picture index: a +sentence about an icon has to reach the icon itself, not a transcription of it. + +Six arms, three questions: + + words BM25 over whatever OCR found — what the product did before any picture vector + picture a sentence, or a whole screenshot, asked of the CLIP space + pool the same question at three granularities: among all icon crops, among whole + screenshots, and among the icons of one screen (the shape ScreenSpot itself poses) +""" + +from __future__ import annotations + +import hashlib +import json +import re +import statistics +import subprocess +import sys +import time +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from ely_eye.config import Settings +from ely_eye.db import Database +from ely_eye.ocr import read_text_from_image +from ely_eye.retrieval import RetrievalService +from ely_eye.schemas import ( + EvidenceAtom, + Modality, + SourceRecord, + TrustScores, +) +from ely_eye.storage import ( + ObjectStore, + token_equivalent, + visual_token_equivalent, +) +from ely_eye.visual_index import VisualIndex +from metrics import aggregate, chance_recall_at_k, score_query +from systems import machine_facts + +ROOT = Path(__file__).resolve().parent +CORPUS = ROOT / "corpus" / "OS-Copilot" / "ScreenSpot-v2" +IMAGES = CORPUS / "images" +IMAGE_ZIP = CORPUS / "screenspotv2_image.zip" +CROPS = ROOT / "runs" / "visual" / "crops" +HOME = ROOT / "runs" / "visual" +MODEL_DIR = ROOT / "models" / "chinese-clip-vit-base-patch16" +REPORTS = ROOT / "reports" +TOP_K = 10 + + +@dataclass(frozen=True) +class Arm: + """One way to ask, one pool to answer from, one kind of gold.""" + + label: str + signal: str # words | picture_words | picture + pool: str # crops | shots | own_screen + gold: str # crop | screen + + @property + def key(self) -> str: + return f"{self.signal}->{self.pool}" + + +ARMS = ( + Arm("words → icon crops", "words", "crops", "crop"), + Arm("a sentence → icon crops (picture space)", "picture_words", "crops", "crop"), + Arm("the whole screenshot → icon crops", "picture", "crops", "crop"), + Arm("words → whole screenshots", "words", "shots", "screen"), + Arm("a sentence → whole screenshots (picture space)", "picture_words", "shots", "screen"), + Arm("a sentence → the icons of that one screen", "picture_words", "own_screen", "crop"), +) + + +def git_commit() -> str: + return subprocess.run( + ["git", "rev-parse", "HEAD"], cwd=ROOT.parent, check=True, capture_output=True, text=True + ).stdout.strip() + + +def model_sha256() -> str: + graph = MODEL_DIR / "onnx" / "model_fp16.onnx" + if not graph.exists(): + raise SystemExit(f"{graph} is missing; fetch the Chinese-CLIP fp16 export first") + digest = hashlib.sha256() + with graph.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def ensure_images() -> None: + """The dataset ships its screenshots inside one zip; unpack them once, flat.""" + if IMAGES.exists() and any(IMAGES.iterdir()): + return + if not IMAGE_ZIP.exists(): + raise SystemExit(f"{IMAGE_ZIP} is missing; download the ScreenSpot-v2 corpus first") + import zipfile + + started = time.monotonic() + with zipfile.ZipFile(IMAGE_ZIP) as archive: + names = [name for name in archive.namelist() if name.lower().endswith((".png", ".jpg", ".jpeg"))] + IMAGES.mkdir(parents=True, exist_ok=True) + for name in names: + target = IMAGES / Path(name).name + if not target.exists(): + target.write_bytes(archive.read(name)) + print(f"unpacked {len(names)} screenshots into {IMAGES} in {time.monotonic() - started:.1f}s", flush=True) + + +def icon_targets() -> list[dict[str, Any]]: + """Every icon instruction, with its crop and the screenshot it came from.""" + ensure_images() + targets: list[dict[str, Any]] = [] + for name in ("mobile", "desktop", "web"): + rows = _read_rows(CORPUS / f"screenspot_{name}_v2.json") + for row in rows: + if row.get("data_type") != "icon": + continue + picture = CORPUS / "images" / row["img_filename"] + if not picture.exists(): + raise SystemExit(f"{picture}: ScreenSpot names a screenshot that is not on disk") + stem = Path(row["img_filename"]).stem + crop_path = CROPS / f"{stem}-{len(targets)}.png" + screenshot_path = CROPS / f"{stem}-full.png" + CROPS.mkdir(parents=True, exist_ok=True) + from PIL import Image + + with Image.open(picture) as screen: + left, top, width, height = [int(value) for value in row["bbox"]] + if not crop_path.exists(): + screen.crop((left, top, left + width, top + height)).convert("RGB").save(crop_path) + if not screenshot_path.exists(): + screen.convert("RGB").save(screenshot_path) + targets.append( + { + "instruction": row["instruction"], + "source_id": f"{stem}-icon-{len(targets)}", + "screen_id": f"screen-{stem}", + "crop": crop_path, + "screenshot": screenshot_path, + "size": (width, height), + "holder": row.get("data_source"), + } + ) + if not targets: + raise SystemExit(f"{CORPUS}: no icon targets found; is the ScreenSpot-v2 corpus downloaded?") + return targets + + +def screenshot_targets(targets: list[dict[str, Any]]) -> list[dict[str, Any]]: + """One memory per distinct screenshot: the granularity the product actually stores.""" + screens: dict[str, dict[str, Any]] = {} + for target in targets: + screen = screens.setdefault( + target["screen_id"], + {"source_id": target["screen_id"], "screenshot": target["screenshot"], "crops": []}, + ) + screen["crops"].append(target["source_id"]) + return list(screens.values()) + + +def _read_rows(path: Path) -> list[dict[str, Any]]: + try: + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"{path}: the ScreenSpot instruction file is unreadable ({error})") from error + + +def _atom(settings: Settings, store: ObjectStore, path: Path, source_id: str, modality: Modality) -> EvidenceAtom: + ocr = read_text_from_image(path) + width, height = _size(path) + return EvidenceAtom( + atom_id=f"{source_id}:image", + modality=modality, + source_id=source_id, + source=str(path), + text=ocr.text, + image_ref=store.put_file(path, "images"), + trust=TrustScores(parser=1.0, ocr=ocr.confidence, model=0.0), + token_equivalent=token_equivalent(ocr.text) + visual_token_equivalent(width, height), + scenario="ui-referencing", + ) + + +def ingest(settings: Settings, targets: list[dict[str, Any]], screens: list[dict[str, Any]]) -> None: + """Store each crop and each screenshot as one memory, the way ingestion would.""" + database = Database(settings) + store = ObjectStore(settings) + atoms: list[EvidenceAtom] = [] + sources: list[SourceRecord] = [] + for position, target in enumerate(targets): + path = Path(target["crop"]) + atoms.append(_atom(settings, store, path, str(target["source_id"]), Modality.image)) + sources.append(_source(str(target["source_id"]), path)) + if (position + 1) % 100 == 0: + print(f"ingested {position + 1}/{len(targets)} crops", flush=True) + for position, screen in enumerate(screens): + path = Path(screen["screenshot"]) + atoms.append(_atom(settings, store, path, str(screen["source_id"]), Modality.ui_screenshot)) + sources.append(_source(str(screen["source_id"]), path)) + if (position + 1) % 100 == 0: + print(f"ingested {position + 1}/{len(screens)} screenshots", flush=True) + database.upsert_sources(sources) + database.upsert_atoms(atoms) + + +def _source(source_id: str, path: Path) -> SourceRecord: + return SourceRecord( + source_id=source_id, + path=str(path), + kind="image", + mime="image/png", + sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + size_bytes=path.stat().st_size, + ) + + +def _size(path: Path) -> tuple[int, int]: + from PIL import Image + + with Image.open(path) as picture: + return picture.size + + +def run_arm( + arm: Arm, + targets: list[dict[str, Any]], + pools: dict[str, Any], + index: VisualIndex, + retrieval: RetrievalService, + source_of_atom: dict[str, str], +) -> tuple[list[list[str]], list[float]]: + """Ask every instruction one way, and keep the answers in corpus ids.""" + ranked: list[list[str]] = [] + seconds: list[float] = [] + for position, target in enumerate(targets): + pool = _pool(arm, target, pools) + started = time.monotonic() + if arm.signal == "words": + hits = [hit.atom.atom_id for hit in retrieval.search(target["instruction"], TOP_K * 4).hits] + answer = [hit for hit in hits if pool is None or hit in pool][:TOP_K] + elif arm.signal == "picture_words": + answer = [atom_id for atom_id, _ in index.search_text(target["instruction"], TOP_K, pool)] + else: + answer = [atom_id for atom_id, _ in index.search_image(target["screenshot"], TOP_K, pool)] + seconds.append(time.monotonic() - started) + ranked.append([source_of_atom[atom_id] for atom_id in answer if atom_id in source_of_atom]) + if (position + 1) % 100 == 0: + print(f"{arm.label}: {position + 1}/{len(targets)}", flush=True) + return ranked, seconds + + +def _pool(arm: Arm, target: dict[str, Any], pools: dict[str, Any]) -> set[str] | None: + if arm.pool == "own_screen": + return { + atom_id + for atom_id, source_id in pools["atom_of_target"].items() + if source_id in set(pools["screen_crops"][target["screen_id"]]) + } + return pools.get(arm.pool) + + +def score(arm: Arm, targets: list[dict[str, Any]], ranked: list[list[str]], seconds: list[float]) -> dict[str, Any]: + scores = [ + score_query( + f"{target['source_id']}, {target['instruction'][:40]}", + hits, + {target["source_id"] if arm.gold == "crop" else target["screen_id"]}, + ) + for target, hits in zip(targets, ranked, strict=True) + ] + ordered = sorted(seconds) + return { + "label": arm.label, + "key": arm.key, + **aggregate(scores), + "latency": { + "queries": len(scores), + "p50": round(statistics.median(ordered), 4), + "p95": round(ordered[min(len(ordered) - 1, int(len(ordered) * 0.95))], 4), + }, + } + + +def main() -> None: + settings = Settings(home=HOME, clip_model_dir=MODEL_DIR) + targets = icon_targets() + screens = screenshot_targets(targets) + database = Database(settings) + if database.atom_count() != len(targets) + len(screens): + started = time.monotonic() + ingest(settings, targets, screens) + print(f"ingested {len(targets)} crops + {len(screens)} screenshots in {time.monotonic() - started:.1f}s", flush=True) + source_of_atom = {atom.atom_id: atom.source_id for atom in database.list_atoms()} + crops = {atom.atom_id for atom in database.list_atoms() if atom.modality is Modality.image} + shots = {atom.atom_id for atom in database.list_atoms() if atom.modality is Modality.ui_screenshot} + wordless = sum( + 1 for atom in database.list_atoms() if atom.atom_id in crops and not atom.text.strip() + ) + + index = VisualIndex(settings, database) + if not index.is_current(): + started = time.monotonic() + built = index.build() + print(f"encoded {built} in {time.monotonic() - started:.1f}s", flush=True) + retrieval = RetrievalService(settings, database) + + pools = { + "crops": crops, + "shots": shots, + "atom_of_target": source_of_atom, + "screen_crops": {screen["source_id"]: screen["crops"] for screen in screens}, + } + per_screen = [len(screen["crops"]) for screen in screens] + rows = [] + for arm in ARMS: + ranked, seconds = run_arm(arm, targets, pools, index, retrieval, source_of_atom) + row = score(arm, targets, ranked, seconds) + row["chance_recall_at_10"] = round( + chance_recall_at_k( + 1, + len(crops) if arm.pool == "crops" else len(shots) if arm.pool == "shots" else int(statistics.median(per_screen)), + TOP_K, + ), + 4, + ) + row["pool"] = arm.pool + rows.append(row) + + report = { + "generated_at": datetime.now(timezone.utc).isoformat(), + "commit": git_commit(), + "dataset": "OS-Copilot/ScreenSpot-v2 (icon instructions)", + "model": {"repo": "Xenova/chinese-clip-vit-base-patch16", "sha256": model_sha256()}, + "targets": len(targets), + "screenshots": len(screens), + "crops_without_ocr_text": wordless, + "instruction_languages": { + "cjk_instructions": sum(1 for t in targets if re.search(r"[一-鿿]", t["instruction"])) + }, + "machine": machine_facts(), + "arms": rows, + } + REPORTS.mkdir(parents=True, exist_ok=True) + (REPORTS / "visual.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") + (REPORTS / "VISUAL.md").write_text(markdown(report), encoding="utf-8") + for row in rows: + print( + f"{row['label']}: recall@10 {row['recall_at_10']:.3f} (chance {row['chance_recall_at_10']:.3f}) " + f"mrr {row['mrr']:.3f} p50 {row['latency']['p50']}s" + ) + print(f"wrote {REPORTS / 'VISUAL.md'}") + + +def markdown(report: dict[str, Any]) -> str: + lines = [ + f"# VISUAL: can a library of wordless pictures be searched? (commit {report['commit'][:9]})", + "", + f"- generated {report['generated_at']}", + ( + f"- corpus: {report['dataset']}, {report['targets']} icon crops and {report['screenshots']} whole " + f"screenshots stored as separate memories, {report['crops_without_ocr_text']} crops with no OCR text" + ), + f"- encoder: {report['model']['repo']} sha256 {report['model']['sha256'][:16]}", + f"- machine: {report['machine']['cpus']} cpus, load {report['machine']['load_average_1_5_15']}", + ( + f"- instructions containing Chinese: {report['instruction_languages']['cjk_instructions']} " + "(ScreenSpot-v2 is English, so this run does not test the bilingual claim)" + ), + "", + "| ask | pool | recall@1 | recall@5 | recall@10 | chance@10 | MRR | hit@10 lower 95% | p50 s | p95 s |", + "| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |", + ] + for row in report["arms"]: + lines.append( + f"| {row['label']} | {row['pool']} | {row['recall_at_1']:.3f} | {row['recall_at_5']:.3f} " + f"| {row['recall_at_10']:.3f} | {row['chance_recall_at_10']:.4f} | {row['mrr']:.3f} " + f"| {row['hit_at_10_wilson_lower']:.3f} | {row['latency']['p50']} | {row['latency']['p95']} |" + ) + lines += [ + "", + ( + "The words rows are what the product could do before any picture vector: if OCR found no words, " + "there is nothing to match. The picture rows ask the same sentence, or the whole screenshot, of the " + "CLIP space. The pools change the question: among every icon in the library, among whole screenshots " + "(the granularity the store keeps), and among one screen's icons, which is the shape ScreenSpot " + "itself poses. A number is only comparable inside its own pool, so chance is printed per row." + ), + "", + ] + return "\n".join(lines) + + +if __name__ == "__main__": + main()