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()