Give wordless pictures a way to be found

A screenshot with no OCR-able words used to be a memory nothing could
reach. It now has a vector, and a question can be asked with it.

- picture space: Xenova/chinese-clip-vit-base-patch16 fp16 ONNX, 361MB,
  sha256 0fad978faa89b85f…, CPU-runnable on the onnxruntime the repo
  already ships. Text 76ms, image 29ms/image at batch 8. The ONNX contract
  was measured, not assumed: it is the fused contrastive graph, so the
  unused tower is fed a placeholder row, and the tokenizer must be
  BertTokenizer — a CLIP tokenizer emits ids past the embedding table.
- visual_index.py keeps picture vectors in their own family beside the
  text ones, keyed by atom id, and refuses to answer when the map and the
  array disagree. retrieval.py stays under 500 lines.
- a search can be asked with a picture, and when the words come up thin
  the picture ranking joins the fusion — but only then, so a text query
  never pays for the encoder.
- visual tokens are counted at 32x32 patches (migration v7): 21,325 image
  atoms went from 5,331 to 34,618 tokens, and the library from 29.9M to
  84.2M. The 100M capsule fits this 50,278-file corpus and nothing much
  larger — that is what the claim actually means.
- the vision-language describer is written and gated: with no snapshot an
  atom records "deferred: no vision-language model snapshot"; a broken
  model becomes a note on the atom instead of a dead ingest.
- forget now voids the picture index and the copy this process holds, so
  a forgotten picture stops answering from RAM.

Measured, six arms, ScreenSpot-v2 (554 icon crops, 408 whole screenshots),
bench/reports/VISUAL.md:
  icon pool      picture 0.208  words 0.175  chance 0.018
  screenshot pool words 0.384   picture 0.215
  screenshot to icon 0.108  — CLIP is a whole-image model, not a grounding
  model: it cannot point at one small icon inside a screen.
  within one screen MRR 0.908 (recall@10 is 1.000 there only because
  chance is 1.000 too, so MRR is the only honest number in that row).
ScreenSpot-v2 instructions are all English, so this run does not test the
bilingual claim and the report says so.

Gates: 199 tests, ruff clean on 17 files, mypy clean on 16 (fresh cache),
the repo's own pyright 0 errors on the 12 files touched.
This commit is contained in:
ZacharyZhang-NY committed 2026-10-07 20:41:19 -04:00
1 parent 7e1fe649f6
commit 03ca35763f
20 files changed
+1702 -22

No files matched your search

+5 -1
View File
@@ -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
+44 -1
View File
@@ -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,
+120
View File
@@ -0,0 +1,120 @@
"""Image and text encoding for cross-modal retrieval, over one small ONNX graph.
Chinese-CLIP ViT-B/16 in fp16 (377 MB) keeps bilingual text and images in one vector
space, so a screenshot can be found by a picture or by a sentence. The graph is the
fused contrastive model: every call needs both towers' inputs, so the unused side is
fed a placeholder row.
"""
from __future__ import annotations
from pathlib import Path
from typing import Any
import numpy as np
from PIL import Image
IMAGE_SIZE = 224
TEXT_MAX_LENGTH = 52
IMAGE_MEAN = (0.48145466, 0.45782750, 0.40821073)
IMAGE_STD = (0.26862954, 0.26130258, 0.27577711)
DOWNLOAD_HINT = (
"huggingface-cli download Xenova/chinese-clip-vit-base-patch16 "
"--include 'onnx/model_fp16.onnx' 'tokenizer*' 'config.json' --local-dir <dir>"
)
def preprocess_image(image: Image.Image) -> np.ndarray:
"""Turn a PIL image into the (3, 224, 224) float32 tensor the graph expects.
The reference processor resizes to a square and normalizes with CLIP constants;
doing it here keeps torchvision out of the dependency list.
"""
rgb = image.convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE), Image.Resampling.BICUBIC)
array = np.asarray(rgb, dtype=np.float32) / 255.0
mean = np.array(IMAGE_MEAN, dtype=np.float32)
std = np.array(IMAGE_STD, dtype=np.float32)
return ((array - mean) / std).transpose(2, 0, 1)
class ClipEncoder:
"""One ONNX session, both towers, batched image encoding."""
def __init__(self, model_dir: Path) -> None:
self.model_dir = Path(model_dir)
self.onnx_path = self.model_dir / "onnx" / "model_fp16.onnx"
if not self.onnx_path.exists():
raise SystemExit(
f"{self.onnx_path} is not a Chinese-CLIP export; fetch the fp16 graph: {DOWNLOAD_HINT}"
)
self._session: Any | None = None
self._tokenizer: Any | None = None
self._placeholder_text: dict[str, np.ndarray] | None = None
self._placeholder_image: np.ndarray | None = None
@property
def session(self) -> Any:
if self._session is None:
import onnxruntime as ort
self._session = ort.InferenceSession(str(self.onnx_path), providers=["CPUExecutionProvider"])
return self._session
@property
def tokenizer(self) -> Any:
if self._tokenizer is None:
from transformers import BertTokenizer
# The export is BERT-style (vocab.txt, [CLS]=101); transformers 5 unified the fast and
# slow tokenizers, so the name the CLIP docs use is an alias for this one.
self._tokenizer = BertTokenizer.from_pretrained(str(self.model_dir))
return self._tokenizer
def encode_images(self, images: list[Image.Image], batch_size: int = 8) -> np.ndarray:
if not images:
return np.zeros((0, 512), dtype=np.float32)
rows: list[np.ndarray] = []
for start in range(0, len(images), batch_size):
chunk = images[start : start + batch_size]
pixels = np.stack([preprocess_image(image) for image in chunk])
rows.append(self._run("image_embeds", pixels, self._text_feed(1))[0])
return np.asarray(np.concatenate(rows), dtype=np.float32)
def encode_texts(self, texts: list[str], batch_size: int = 16) -> np.ndarray:
if not texts:
return np.zeros((0, 512), dtype=np.float32)
rows: list[np.ndarray] = []
for start in range(0, len(texts), batch_size):
chunk = texts[start : start + batch_size]
feed = self._text_feed(len(chunk), chunk)
rows.append(self._run("text_embeds", self._image_feed(len(chunk)), feed)[0])
return np.asarray(np.concatenate(rows), dtype=np.float32)
def _run(self, output: str, pixels: np.ndarray, text: dict[str, np.ndarray]) -> list[np.ndarray]:
feed = {"pixel_values": pixels.astype(np.float32), **text}
return self.session.run([output], feed)
def _text_feed(self, batch: int, texts: list[str] | None = None) -> dict[str, np.ndarray]:
if texts is None:
if self._placeholder_text is None:
self._placeholder_text = self._tokenize(["placeholder"])
return self._placeholder_text
return self._tokenize(texts)
def _image_feed(self, batch: int) -> np.ndarray:
if self._placeholder_image is None:
self._placeholder_image = np.zeros((1, 3, IMAGE_SIZE, IMAGE_SIZE), dtype=np.float32)
return np.repeat(self._placeholder_image, batch, axis=0)
def _tokenize(self, texts: list[str]) -> dict[str, np.ndarray]:
encoded = self.tokenizer(
texts,
padding="max_length",
truncation=True,
max_length=TEXT_MAX_LENGTH,
return_tensors="np",
)
return {
"input_ids": encoded["input_ids"].astype(np.int64),
"attention_mask": encoded["attention_mask"].astype(np.int64),
}
+3
View File
@@ -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"
+2 -1
View File
@@ -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:
+87
View File
@@ -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
+23 -5
View File
@@ -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},
)
)
+3 -2
View File
@@ -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.
+28 -1
View File
@@ -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}")
+57 -6
View File
@@ -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],
)
)
+2
View File
@@ -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):
+12
View File
@@ -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()
+199
View File
@@ -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