ZacharyZhang-NY
cd9851a12c
The dense leg used to be a pair of files next to the database: a full .npy re-encoded from every atom, re-normalized on every query, and stale the moment anything was ingested. It is now a vec0 table inside the same SQLite file, written in the same transaction as the atom it describes, so a memory is queryable the instant it is remembered. - vector_index.py maintains the table; the writer inserts, the reader never rebuilds. A repair scan (ely-eye build-dense-index) catches what a crash hid, and it is not on the search path. - Ranking stays in the engine; every boundary (closed, retired, not yet true, another holder's private memory, the mounted cartridge) is checked once against the atoms table afterwards. A visibility subquery inside the KNN cost 48.9s on 153,857 vectors; the same check over 200 candidates costs 0.2ms. When the boundaries reject enough candidates the ask widens once, never twice. - The Go mirror and the web panel carry the new status fields. - bench/dense_scale.py and bench/quant_probe.py measure the leg at corpus size; bench/scale_check.py now works on a copy so the frozen baseline stays frozen. Measured on the frozen 165,973-atom corpus. bench/reports/DENSE.md, QUANT.md and SCALE.md carry the numbers and the machine they were taken on. Gates: 285 pytest, ruff clean, pyright 0 on the touched files, go build/vet/test and gofmt clean, frontend build and eslint clean. mypy is 25 findings against the parent's 24: the new one is "Skipping analyzing sqlite_vec: missing library stubs", the same class as the two existing qwen_vl_utils lines, and the repo keeps no type: ignore anywhere.
303 lines
11 KiB
Python
303 lines
11 KiB
Python
"""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.visual_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)
|