From 9cb46c18d3d608519b31c5b0a6b825abbc41e277 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Thu, 8 Oct 2026 08:20:21 -0400 Subject: [PATCH] Keep every file inside the rules the repo already wrote down The 500-line rule and the one-line-comment rule were enforced only from memory, so they drifted. bench/filesize_check.py now checks the budget on every tracked source file, and it passes. Split, with behaviour proven unchanged by comparing each moved function body against HEAD byte for byte (53 training bodies, 21 benchmark bodies, 0 differences): db.py 503 -> 396 store_search.py 105 bench/run.py 609 -> 351 query_set.py 80, report.py 211 benchmarks.py 619 -> 424 benchmark_scoring.py 262 training.py 1630 -> 94 training_{common,encode,metrics,datasets,lora,retrieval} TrainingService stays as a thin facade for the CLI; the facade's parameter lists are pinned against the module functions by inspect, which is what catches a leaked self. Deleted the eight top-level functions no file called, after fixing the scan that had reported 201 false positives by skipping the defining file. The v1/v2/v3 sweep found nothing to delete and says so: every remaining version tag is an on-disk format marker, a migration step, a fixture name, or a third-party model name. Collapsed 39 multi-line comment blocks (86 lines) to one line each, keeping the measured facts and error strings that justify the code. 269 tests pass; ruff clean; mypy 24 errors in 9 files against a 24-in-6 baseline, net zero; pyright 4, the same four HEAD measures. --- CHECKLIST.md | 22 +- backend/src/ely_eye/adapters.py | 2 +- backend/src/ely_eye/api.py | 2 +- backend/src/ely_eye/api_cartridges.py | 4 +- backend/src/ely_eye/benchmark_scoring.py | 262 +++ backend/src/ely_eye/benchmarks.py | 227 +-- backend/src/ely_eye/cartridge.py | 3 +- backend/src/ely_eye/clip_encoder.py | 3 +- backend/src/ely_eye/compiler.py | 3 +- backend/src/ely_eye/db.py | 109 +- backend/src/ely_eye/memory.py | 6 +- backend/src/ely_eye/migrations.py | 8 +- backend/src/ely_eye/proofs/arena.py | 3 +- .../src/ely_eye/proofs/checks_cartridge.py | 14 +- backend/src/ely_eye/proofs/checks_training.py | 10 +- backend/src/ely_eye/proofs/suite.py | 4 +- backend/src/ely_eye/proofs/training_trace.py | 4 +- backend/src/ely_eye/retrieval.py | 8 +- backend/src/ely_eye/store_search.py | 106 ++ backend/src/ely_eye/text_index.py | 6 +- backend/src/ely_eye/training.py | 1574 +---------------- backend/src/ely_eye/training_common.py | 108 ++ backend/src/ely_eye/training_datasets.py | 387 ++++ backend/src/ely_eye/training_encode.py | 177 ++ backend/src/ely_eye/training_lora.py | 408 +++++ backend/src/ely_eye/training_metrics.py | 307 ++++ backend/src/ely_eye/training_retrieval.py | 358 ++++ backend/src/ely_eye/transactive.py | 3 +- backend/src/ely_eye/visual_index.py | 3 +- backend/tests/test_host_profiles.py | 2 +- backend/tests/test_memory_schema.py | 5 +- backend/tests/test_proofs.py | 9 +- backend/tests/test_text_index.py | 5 +- backend/tests/test_training_split.py | 152 ++ bench/audit.py | 6 +- bench/build_code.py | 3 +- bench/build_visrag.py | 3 +- bench/ingest_baseline.py | 20 +- bench/query_set.py | 83 + bench/report.py | 211 +++ bench/run.py | 289 +-- bench/run_external.py | 3 +- bench/scale_check.py | 3 +- bench/sets_common.py | 3 +- bench/systems.py | 3 +- bench/tests/test_joint.py | 13 +- bench/tests/test_systems.py | 6 +- 47 files changed, 2678 insertions(+), 2272 deletions(-) create mode 100644 backend/src/ely_eye/benchmark_scoring.py create mode 100644 backend/src/ely_eye/store_search.py create mode 100644 backend/src/ely_eye/training_common.py create mode 100644 backend/src/ely_eye/training_datasets.py create mode 100644 backend/src/ely_eye/training_encode.py create mode 100644 backend/src/ely_eye/training_lora.py create mode 100644 backend/src/ely_eye/training_metrics.py create mode 100644 backend/src/ely_eye/training_retrieval.py create mode 100644 backend/tests/test_training_split.py create mode 100644 bench/query_set.py create mode 100644 bench/report.py diff --git a/CHECKLIST.md b/CHECKLIST.md index f4206ae..2441480 100644 --- a/CHECKLIST.md +++ b/CHECKLIST.md @@ -318,12 +318,22 @@ - 第 2 轮 FAIL 三条：① **翻页其实没被测到**——E2E 等的是「行数变化」，而搜索列表与浏览列表共用 `.atom-row`，那一刻的 13 行里有一条是搜索命中；把 `nextPage()` 里的 `setCursor`+`loadBrowse` 删掉、或把翻页改成追加而不是替换，E2E 都还 PASS。改成 `data-testid="browse-list"` + 每行 `data-atom`，断言第一页正好 12 行、下一页零重叠、末页 = `total−12`、Prev 复原；两个变异现在各自红：`Next page did not swap the page: the pager is not paging` 与 `the next page repeated 12 of 12 atoms from the first page` ② CHECKLIST 写 `App.tsx`「130 行的壳」，`wc -l` 是 131 ③ CHECKLIST 引的拒绝挂载原文被 smoke 的 `.slice(0,120)` 砍掉了后半句 `; set ELY_EYE_MODEL_DIR to the model snapshot`，改成 slice 200 并照原文记。 - 同一轮实测为真：263 passed；ruff 干净；全仓 mypy 24 错/6 文件（净新增 0）；go build/vet/test+gofmt 干净；`npm run build`+`eslint .` 干净；`ui_smoke.sh` PASS（4.0s）；`mcp_write_smoke.sh` PASS（3.2s）；30 条路由与 `HTTP_SURFACE` 一致；senses 单一真相（`api.py:52/:105` 与 `maintenance.py:238` 都只渲染 `sense_facts()`）；26 个 `types.ts` 接口对着 pydantic 模型逐个 diff 零编造字段；8 条被删路由无人引用，它们背后的函数都还从 CLI/proof 可达。pyright 在那 8 个文件：reviewer 的工作树 **0 errors**，我这棵树多两条 `runtime.py:283` 的 tuple 告警——它只在 `qwen_vl_utils` 那个无类型导入被解析时出现，HEAD 的副本同样有，不是本轮引入；两边都不是新增。 - 第 3 轮 **PASS**（`9f91ff7..b0895eb`，~40 条 shell / 14 次工具调用，本轮 34 分钟）。reviewer 是跑过不是读过：`ui_smoke.sh` exit 0（4.26s，拒绝挂载整句可见）；三个变异各自红——删 `nextPage()` 的 setCursor+loadBrowse → `the pager is not paging`，翻页改追加 → `repeated 12 of 12`，它自己加的第三个（删 `previousPage()` 的复原）→ `Timeout 15000ms`；改完 `git checkout` 后回到 exit 0。门禁重跑：263 passed、npm build+lint 干净、go build/vet/test+gofmt 干净、`mcp_write_smoke.sh` PASS 3.26s、ruff 干净、mypy 冷缓存查 8 个文件 Success、全仓 mypy 24 错/6 文件。数字全对：schemas.py 469、status_schemas.py 74（只从 schemas 拿 `TrainingAdapterKind`，反向无导入）、多行 docstring 114、App.tsx 131、拒绝文案与 `cartridge.py:169-173` 逐字一致、30 条路由与 `HTTP_SURFACE` 完全相等且 8 条删掉的都报 gone、`PAGE_SIZE=12` 两处同值、末页 135−12=123 行、`types.ts` 26 个接口、前端最大文件 `styles.css` 492。它同时指出两点不是缺陷但记着：`PAGE_SIZE` 是手拄到 harness 的（组件改了会变红不会变绿，已在记录里写明）；末页的 `total−12` 读的是未过滤的 `GET /api/atoms?limit=1`，只对是因为前一行刚好把 modality 清掉了。 -- [ ] **p1-6** 代码卫生 - - [ ] p1-6.1 `evals.py` 2413 行拆到 ≤500 行/文件 - - [ ] p1-6.2 删 dead code 与未调用函数 - - [ ] p1-6.3 删 v1/v2/v3 废弃端口 - - [ ] p1-6.4 注释压到 ≤1 行 - - 验收：`python -m bench.filesize_check` 报无文件 >500 行；vulture/grep 证据 +- [x] **p1-6** 代码卫生 + - [x] p1-6.1 每个文件 ≤500 行 + - 验收命令 `python -m bench.filesize_check` → `OK: no tracked source file exceeds 500 lines`。这个脚本是本轮加的：500 行规则之前只靠记忆执行 + - `evals.py` 2413 行在 p0-6 已拆成 `proofs/*`（最大 430）。本轮拆的：`db.py` 503→396（BM25 检索搬进 `store_search.py` 105）、`bench/run.py` 609→351（`query_set.py` 80 + `report.py` 211）、`benchmarks.py` 619→424（`benchmark_scoring.py` 262）、`training.py` 1630→94（`training_{common,encode,metrics,datasets,lora,retrieval}.py` = 108/177/307/387/408/358，`TrainingService` 留成薄壳喂 `cli.py` 的 10 个调用点）、`schemas.py` 534→469（p1-5 拆出 `status_schemas.py` 74） + - 「拆分不改行为」的证据：脚本把每个搬走函数的**函数体去缩进后与 HEAD 逐字节比对**——training 侧 53 个、benchmarks 侧 21 个，**全部相同，0 处不同，0 个缺失**。方法转函数只重写 def 头，`self.settings`→`settings` 这类替换是正文里唯一的改动 + - 新增 `backend/tests/test_training_split.py`（6 条）：CPU 能跑的纯函数逐个测；最后一条用 `inspect.signature` 断言薄壳参数表 = 模块函数参数表去掉注入的两个参数——这条就是抓 `self` 泄漏那类 bug 的 + - [x] p1-6.2 删 dead code 与未调用函数 + - 度量口径（第一版误报 201 个，因为只查了别的文件）：每个顶层 `def` 对**所有**跟踪文件计数（含自己），再排除 pytest 用例与 typer 命令（框架按名字调）→ 8 个零引用，全删：`cartridge_store.py:47 cartridges_holding_atom`、`cartridge_store.py:97 clear_mounts`、`training.py:1316 gradient_checkpointing_enabled`、`retrieval.py:172 group_hits_by_source`、`bench/ingest_baseline.py:239 ingested_paths`、`proofs/report.py:83 is_finite_number`、`runtime.py:432 parse_runtime_json`、`ingestion_repository.py:317 write_sources_json`；重扫 **0** + - 两个为什么是死的：`memory.py` 靠读每颗卡带自己的 `sources.json` 回答归属（导出的卡带可以不在库里）；`set_mounts(database, [], stamp)` 换集合本来就能清空挂载表 + - 顺手删 `backend/build/`（712KB 的 `python -m build` 残留，gitignore 内，0 个跟踪文件）——它让 grep 撒谎 + - [x] p1-6.3 删 v1/v2/v3 废弃端口 + - 量完的结论是**没有可删的**，不是没做：每个 `v1/v2/v3` 都承重——磁盘格式标签（`ely-eye-memory-capsule-v1` `proofs/checks_cartridge.py:128`、`ely-eye-{text,visual}-vectors-v1`、`ely-eye-kivi-replay-v1` `cartridge_assets.py:100,127,158`）、迁移阶梯（`migrations.py:185,190,243`）、矛盾语料 fixture 名（`proofs/contradictions.py:114-135`）、第三方名字（MiniLM-L12-v2、ScreenSpot-v2、LongBench v2）、OpenAI 兼容的 `/v1` URL、缓存键里的 `compiler-v2` salt。没有废弃端口、别名或双代码路径 + - [x] p1-6.4 注释压到 ≤1 行 + - 口径：`tokenize` 找连续 `#` 行块。全仓 **39 块 / 86 行 → 0 块**，每块压成一行，压的时候留承重那句，实测数字与错误原文不删（`Unable to allocate 582. MiB for an array with shape (5632, 9024, 3)`、`288,266 rows for 165,973 atoms`、`[^[^\W\s]` 在 3.14 会吃掉 CJK） + - docstring 不动（全仓多行 docstring 114 条，与父提交同数）：docstring 是模块与函数的契约，不是注释；p1-4 review 用的口径也是「新增多行 docstring 数不得增加」 + - 门禁：`pytest -q` **269 passed**；`ruff check backend/src/ely_eye backend/tests bench` All checks passed；mypy 冷缓存 **24 errors / 9 files**（基线 24/6，training.py 那 6 条随拆分落到 4 个新文件，净新增 0）；`pyright` 查 46 个改动文件 → **4 errors**，与 HEAD 单个 `training.py` 在包副本里量到的 4 条同源（qwen_vl_utils 元组返回、两处 `save_pretrained(Path)`、`bitsandbytes` 未解析）；go build/vet/test + gofmt 干净；`npm run build` + eslint 干净 --- diff --git a/backend/src/ely_eye/adapters.py b/backend/src/ely_eye/adapters.py index df49b96..78002d5 100644 --- a/backend/src/ely_eye/adapters.py +++ b/backend/src/ely_eye/adapters.py @@ -194,7 +194,7 @@ def record_load(source: AdapterSource, model: Any) -> LoadRecord: def load_into_model(model: Any, path: Path) -> Any: """Attach one LoRA to a live model; the only place generation imports peft.""" - from .training import bitsandbytes_gaudi_probe_guard + from .training_metrics import bitsandbytes_gaudi_probe_guard with bitsandbytes_gaudi_probe_guard(): from peft import PeftModel diff --git a/backend/src/ely_eye/api.py b/backend/src/ely_eye/api.py index e998eb8..fab8fd6 100644 --- a/backend/src/ely_eye/api.py +++ b/backend/src/ely_eye/api.py @@ -51,7 +51,7 @@ from .schemas import ( ) from .senses import sense_facts from .status_schemas import MemoryStatus, SystemStatus -from .training import safetensor_weight_stats +from .training_metrics import safetensor_weight_stats from .verifier import Verifier, VisualContradictionLens diff --git a/backend/src/ely_eye/api_cartridges.py b/backend/src/ely_eye/api_cartridges.py index 79d3d7a..9b2378a 100644 --- a/backend/src/ely_eye/api_cartridges.py +++ b/backend/src/ely_eye/api_cartridges.py @@ -1,8 +1,6 @@ """Cartridge and mount routes: a cartridge is a retrieval scope, not only an export.""" -# No `from __future__ import annotations` here: FastAPI reads the route signatures, and -# string annotations would be resolved against this module's globals, where the injected -# settings_dep lives only as a parameter of register(). +# No `from __future__ import annotations`: FastAPI resolves route annotations against this module's globals, where settings_dep exists only as a register() parameter. from collections.abc import Callable from typing import Annotated diff --git a/backend/src/ely_eye/benchmark_scoring.py b/backend/src/ely_eye/benchmark_scoring.py new file mode 100644 index 0000000..a0a60c0 --- /dev/null +++ b/backend/src/ely_eye/benchmark_scoring.py @@ -0,0 +1,262 @@ +"""Score a benchmark prediction: the prompt, the answer match, and the sample normalising.""" + +import json +import math +import re +from ast import literal_eval +from pathlib import Path +from typing import Any + +from .evaluation_schemas import ( + BenchmarkSpec, +) + + +def get_nested(record: dict[str, Any], path: str) -> Any: + current: Any = record + for part in path.split("."): + if not isinstance(current, dict) or part not in current: + return None + current = current[part] + return current + + +def exact_or_contains_score(prediction: str, answer: str) -> float: + normalized_prediction = normalize_answer(prediction) + normalized_answer = normalize_answer(answer) + if not normalized_prediction or not normalized_answer: + return 0.0 + if normalized_prediction == normalized_answer or normalized_answer in normalized_prediction: + return 1.0 + prediction_tokens = normalized_prediction.split() + answer_tokens = normalized_answer.split() + if not prediction_tokens or not answer_tokens: + return 0.0 + overlap = len(set(prediction_tokens) & set(answer_tokens)) + precision = overlap / len(set(prediction_tokens)) + recall = overlap / len(set(answer_tokens)) + if precision + recall == 0: + return 0.0 + return 2 * precision * recall / (precision + recall) + + +def benchmark_score( + benchmark: BenchmarkSpec, + row: dict[str, Any], + prediction: str, + answer: str | None, +) -> float | None: + if answer is None: + return None + if "multiple_choice" in benchmark.task_type: + option_text = answer_option_text(row, answer) + if exact_or_contains_score(prediction, answer): + return 1.0 + if option_text and exact_or_contains_score(prediction, option_text): + return 1.0 + return 0.0 + return exact_or_contains_score(prediction, answer) + + +def answer_option_text(row: dict[str, Any], answer: str) -> str | None: + normalized_answer = answer.strip().upper() + if len(normalized_answer) != 1 or not normalized_answer.isalpha(): + return None + options = row.get("options") + if isinstance(options, str): + try: + options = literal_eval(options) + except (SyntaxError, ValueError): + options = None + index = ord(normalized_answer) - ord("A") + if isinstance(options, (list, tuple)) and 0 <= index < len(options): + return str(options[index]) + if isinstance(options, dict): + return canonical_answer(options.get(normalized_answer)) + choice_key = f"choice_{normalized_answer}" + if choice_key in row: + return canonical_answer(row[choice_key]) + return None + + +def normalize_answer(value: str) -> str: + return re.sub(r"\s+", " ", re.sub(r"[^\w\u4e00-\u9fff]+", " ", value.lower())).strip() + + +def benchmark_prompt(benchmark: BenchmarkSpec, row: dict[str, Any]) -> str: + question = first_present_raw(row, ("question", "query", "text")) or "" + options = first_present_raw(row, ("options", "choices")) or "" + context = first_present_raw(row, ("context", "document", "passage")) or "" + context = compact_preview(context, 6000) + media_note = "Use the attached image evidence." if any(is_pil_image(value) for value in row.values()) else "" + if benchmark.name == "OCRBench": + return f"""Use /no_think mode. Return compact JSON only. + +Benchmark: OCRBench +Task type: image text recognition +Metric: exact_or_contains_accuracy +Use the attached image evidence. + +Question: +{question} + +Read the exact text in the image. For handwritten English words, return the standard spelling with the same letters. + +Return exactly this JSON shape: +{{"answer":""}} +""" + return f"""Use /no_think mode. Return compact JSON only. + +Benchmark: {benchmark.name} +Task type: {benchmark.task_type} +Metric: {benchmark.metric} +{media_note} + +Question: +{question} + +Options: +{options} + +Context: +{context} + +Return exactly this JSON shape: +{{"answer":""}} +""" + + +def extract_answer(text: str) -> str: + stripped = text.strip() + try: + payload = json.loads(stripped) + except json.JSONDecodeError: + start = stripped.find("{") + end = stripped.rfind("}") + payload = None + if start >= 0 and end > start: + try: + payload = json.loads(stripped[start : end + 1]) + except json.JSONDecodeError: + payload = None + if isinstance(payload, dict) and payload.get("answer") is not None: + return str(payload["answer"]).strip() + return compact_preview(stripped, 500) + + +def canonical_answer(value: Any) -> str | None: + if value is None: + return None + if isinstance(value, (list, tuple)): + return " | ".join(str(item) for item in value) + return str(value) + + +def first_present_raw(row: dict[str, Any], keys: tuple[str, ...]) -> Any: + for key in keys: + value = row.get(key) + if value is not None: + return value + return None + + +def normalize_sample(row: dict[str, Any], output_root: Path) -> tuple[dict[str, Any], list[str]]: + media_artifacts: list[str] = [] + normalized: dict[str, Any] = {} + for key, value in row.items(): + normalized_value, artifact = normalize_sample_value(key, value, output_root) + normalized[key] = normalized_value + if artifact: + media_artifacts.append(artifact) + return normalized, media_artifacts + + +def prepare_prediction_media(benchmark: BenchmarkSpec, media_artifacts: list[str]) -> list[str]: + if benchmark.name != "OCRBench": + return media_artifacts + return [str(enhance_ocr_benchmark_image(Path(path))) for path in media_artifacts] + + +def enhance_ocr_benchmark_image(path: Path) -> Path: + from PIL import Image, ImageOps + + image = Image.open(path).convert("RGB") + max_side = max(image.size) + scale = max(1, min(10, math.ceil(768 / max_side))) + if scale > 1: + image = image.resize((image.width * scale, image.height * scale), Image.Resampling.BICUBIC) + border = max(8, scale * 2) + image = ImageOps.expand(image, border=border, fill="white") + target = path.with_name(f"{path.stem}_ocr_x{scale}{path.suffix}") + image.save(target) + return target + + +def normalize_sample_value(key: str, value: Any, output_root: Path) -> tuple[Any, str | None]: + if is_pil_image(value): + digest = image_sha256(value) + artifact = output_root / f"{slugify(key)}_{digest[:12]}.png" + value.save(artifact) + return ( + { + "type": "image", + "mode": getattr(value, "mode", None), + "size": list(getattr(value, "size", ())), + "sha256": digest, + }, + str(artifact), + ) + if isinstance(value, (str, int, float, bool)) or value is None: + return value, None + if isinstance(value, (list, tuple)): + return [normalize_sample_value(f"{key}_{index}", item, output_root)[0] for index, item in enumerate(value)], None + if isinstance(value, dict): + return {str(item_key): normalize_sample_value(str(item_key), item_value, output_root)[0] for item_key, item_value in value.items()}, None + return str(value), None + + +def is_pil_image(value: Any) -> bool: + return hasattr(value, "save") and hasattr(value, "mode") and hasattr(value, "size") + + +def image_sha256(image: Any) -> str: + import io + + buffer = io.BytesIO() + image.save(buffer, format="PNG") + return sha256_bytes(buffer.getvalue()) + + +def sample_identifier(row: dict[str, Any]) -> str | None: + for key in ("id", "_id", "question_id", "query-id", "text_id", "video_id", "corpus-id"): + value = row.get(key) + if value is not None: + return str(value) + return None + + +def first_present(row: dict[str, Any], keys: tuple[str, ...]) -> str | None: + for key in keys: + value = row.get(key) + if value is not None: + return compact_preview(value) + return None + + +def compact_preview(value: Any, limit: int = 240) -> str: + text = re.sub(r"\s+", " ", str(value)).strip() + return text[:limit] + + +def sha256_json(value: dict[str, Any]) -> str: + return sha256_bytes(json.dumps(value, ensure_ascii=False, sort_keys=True, default=str).encode("utf-8")) + + +def sha256_bytes(data: bytes) -> str: + import hashlib + + return hashlib.sha256(data).hexdigest() + + +def slugify(value: str) -> str: + return re.sub(r"[^a-zA-Z0-9_.-]+", "-", value).strip("-").lower() or "item" diff --git a/backend/src/ely_eye/benchmarks.py b/backend/src/ely_eye/benchmarks.py index 0ae76ec..09e6607 100644 --- a/backend/src/ely_eye/benchmarks.py +++ b/backend/src/ely_eye/benchmarks.py @@ -1,13 +1,26 @@ from __future__ import annotations import json -import math -import re import tempfile -from ast import literal_eval from pathlib import Path from typing import Any +from .benchmark_scoring import ( + benchmark_prompt, + benchmark_score, + canonical_answer, + exact_or_contains_score, + extract_answer, + first_present, + first_present_raw, + get_nested, + normalize_sample, + prepare_prediction_media, + sample_identifier, + sha256_bytes, + sha256_json, + slugify, +) from .evaluation_schemas import ( BenchmarkPredictionReport, BenchmarkPredictionStatus, @@ -369,251 +382,43 @@ def load_records(path: Path) -> list[dict[str, Any]]: raise ValueError(f"Unsupported benchmark dataset format: {path}") -def get_nested(record: dict[str, Any], path: str) -> Any: - current: Any = record - for part in path.split("."): - if not isinstance(current, dict) or part not in current: - return None - current = current[part] - return current -def exact_or_contains_score(prediction: str, answer: str) -> float: - normalized_prediction = normalize_answer(prediction) - normalized_answer = normalize_answer(answer) - if not normalized_prediction or not normalized_answer: - return 0.0 - if normalized_prediction == normalized_answer or normalized_answer in normalized_prediction: - return 1.0 - prediction_tokens = normalized_prediction.split() - answer_tokens = normalized_answer.split() - if not prediction_tokens or not answer_tokens: - return 0.0 - overlap = len(set(prediction_tokens) & set(answer_tokens)) - precision = overlap / len(set(prediction_tokens)) - recall = overlap / len(set(answer_tokens)) - if precision + recall == 0: - return 0.0 - return 2 * precision * recall / (precision + recall) -def benchmark_score( - benchmark: BenchmarkSpec, - row: dict[str, Any], - prediction: str, - answer: str | None, -) -> float | None: - if answer is None: - return None - if "multiple_choice" in benchmark.task_type: - option_text = answer_option_text(row, answer) - if exact_or_contains_score(prediction, answer): - return 1.0 - if option_text and exact_or_contains_score(prediction, option_text): - return 1.0 - return 0.0 - return exact_or_contains_score(prediction, answer) -def answer_option_text(row: dict[str, Any], answer: str) -> str | None: - normalized_answer = answer.strip().upper() - if len(normalized_answer) != 1 or not normalized_answer.isalpha(): - return None - options = row.get("options") - if isinstance(options, str): - try: - options = literal_eval(options) - except (SyntaxError, ValueError): - options = None - index = ord(normalized_answer) - ord("A") - if isinstance(options, (list, tuple)) and 0 <= index < len(options): - return str(options[index]) - if isinstance(options, dict): - return canonical_answer(options.get(normalized_answer)) - choice_key = f"choice_{normalized_answer}" - if choice_key in row: - return canonical_answer(row[choice_key]) - return None -def normalize_answer(value: str) -> str: - return re.sub(r"\s+", " ", re.sub(r"[^\w\u4e00-\u9fff]+", " ", value.lower())).strip() -def benchmark_prompt(benchmark: BenchmarkSpec, row: dict[str, Any]) -> str: - question = first_present_raw(row, ("question", "query", "text")) or "" - options = first_present_raw(row, ("options", "choices")) or "" - context = first_present_raw(row, ("context", "document", "passage")) or "" - context = compact_preview(context, 6000) - media_note = "Use the attached image evidence." if any(is_pil_image(value) for value in row.values()) else "" - if benchmark.name == "OCRBench": - return f"""Use /no_think mode. Return compact JSON only. - -Benchmark: OCRBench -Task type: image text recognition -Metric: exact_or_contains_accuracy -Use the attached image evidence. - -Question: -{question} - -Read the exact text in the image. For handwritten English words, return the standard spelling with the same letters. - -Return exactly this JSON shape: -{{"answer":""}} -""" - return f"""Use /no_think mode. Return compact JSON only. - -Benchmark: {benchmark.name} -Task type: {benchmark.task_type} -Metric: {benchmark.metric} -{media_note} - -Question: -{question} - -Options: -{options} - -Context: -{context} - -Return exactly this JSON shape: -{{"answer":""}} -""" -def extract_answer(text: str) -> str: - stripped = text.strip() - try: - payload = json.loads(stripped) - except json.JSONDecodeError: - start = stripped.find("{") - end = stripped.rfind("}") - payload = None - if start >= 0 and end > start: - try: - payload = json.loads(stripped[start : end + 1]) - except json.JSONDecodeError: - payload = None - if isinstance(payload, dict) and payload.get("answer") is not None: - return str(payload["answer"]).strip() - return compact_preview(stripped, 500) -def canonical_answer(value: Any) -> str | None: - if value is None: - return None - if isinstance(value, (list, tuple)): - return " | ".join(str(item) for item in value) - return str(value) -def first_present_raw(row: dict[str, Any], keys: tuple[str, ...]) -> Any: - for key in keys: - value = row.get(key) - if value is not None: - return value - return None -def normalize_sample(row: dict[str, Any], output_root: Path) -> tuple[dict[str, Any], list[str]]: - media_artifacts: list[str] = [] - normalized: dict[str, Any] = {} - for key, value in row.items(): - normalized_value, artifact = normalize_sample_value(key, value, output_root) - normalized[key] = normalized_value - if artifact: - media_artifacts.append(artifact) - return normalized, media_artifacts -def prepare_prediction_media(benchmark: BenchmarkSpec, media_artifacts: list[str]) -> list[str]: - if benchmark.name != "OCRBench": - return media_artifacts - return [str(enhance_ocr_benchmark_image(Path(path))) for path in media_artifacts] -def enhance_ocr_benchmark_image(path: Path) -> Path: - from PIL import Image, ImageOps - - image = Image.open(path).convert("RGB") - max_side = max(image.size) - scale = max(1, min(10, math.ceil(768 / max_side))) - if scale > 1: - image = image.resize((image.width * scale, image.height * scale), Image.Resampling.BICUBIC) - border = max(8, scale * 2) - image = ImageOps.expand(image, border=border, fill="white") - target = path.with_name(f"{path.stem}_ocr_x{scale}{path.suffix}") - image.save(target) - return target -def normalize_sample_value(key: str, value: Any, output_root: Path) -> tuple[Any, str | None]: - if is_pil_image(value): - digest = image_sha256(value) - artifact = output_root / f"{slugify(key)}_{digest[:12]}.png" - value.save(artifact) - return ( - { - "type": "image", - "mode": getattr(value, "mode", None), - "size": list(getattr(value, "size", ())), - "sha256": digest, - }, - str(artifact), - ) - if isinstance(value, (str, int, float, bool)) or value is None: - return value, None - if isinstance(value, (list, tuple)): - return [normalize_sample_value(f"{key}_{index}", item, output_root)[0] for index, item in enumerate(value)], None - if isinstance(value, dict): - return {str(item_key): normalize_sample_value(str(item_key), item_value, output_root)[0] for item_key, item_value in value.items()}, None - return str(value), None -def is_pil_image(value: Any) -> bool: - return hasattr(value, "save") and hasattr(value, "mode") and hasattr(value, "size") -def image_sha256(image: Any) -> str: - import io - - buffer = io.BytesIO() - image.save(buffer, format="PNG") - return sha256_bytes(buffer.getvalue()) -def sample_identifier(row: dict[str, Any]) -> str | None: - for key in ("id", "_id", "question_id", "query-id", "text_id", "video_id", "corpus-id"): - value = row.get(key) - if value is not None: - return str(value) - return None -def first_present(row: dict[str, Any], keys: tuple[str, ...]) -> str | None: - for key in keys: - value = row.get(key) - if value is not None: - return compact_preview(value) - return None -def compact_preview(value: Any, limit: int = 240) -> str: - text = re.sub(r"\s+", " ", str(value)).strip() - return text[:limit] -def sha256_json(value: dict[str, Any]) -> str: - return sha256_bytes(json.dumps(value, ensure_ascii=False, sort_keys=True, default=str).encode("utf-8")) -def sha256_bytes(data: bytes) -> str: - import hashlib - - return hashlib.sha256(data).hexdigest() -def slugify(value: str) -> str: - return re.sub(r"[^a-zA-Z0-9_.-]+", "-", value).strip("-").lower() or "item" diff --git a/backend/src/ely_eye/cartridge.py b/backend/src/ely_eye/cartridge.py index cfae38e..5f20545 100644 --- a/backend/src/ely_eye/cartridge.py +++ b/backend/src/ely_eye/cartridge.py @@ -275,8 +275,7 @@ class CartridgeService: for field in JSON_PARQUET_COLUMNS: row[field] = json.dumps(row[field], ensure_ascii=False, sort_keys=True) rows.append(row) - # An empty struct column has no child field, and Parquet cannot write one: a memory - # atom carries metadata={} and would crash the export. + # Parquet cannot write an empty struct column, and a memory atom carries metadata={}. table = pa.Table.from_pylist(rows) if rows else pa.Table.from_pydict({"atom_id": []}) pq.write_table(table, path, compression="zstd") diff --git a/backend/src/ely_eye/clip_encoder.py b/backend/src/ely_eye/clip_encoder.py index 7975daa..1b74684 100644 --- a/backend/src/ely_eye/clip_encoder.py +++ b/backend/src/ely_eye/clip_encoder.py @@ -65,8 +65,7 @@ class ClipEncoder: 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. + # The export is BERT-style (vocab.txt, [CLS]=101); transformers 5 aliases the fast and slow tokenizers, so this is the name the CLIP docs mean. self._tokenizer = BertTokenizer.from_pretrained(str(self.model_dir)) return self._tokenizer diff --git a/backend/src/ely_eye/compiler.py b/backend/src/ely_eye/compiler.py index c58b58e..91746bd 100644 --- a/backend/src/ely_eye/compiler.py +++ b/backend/src/ely_eye/compiler.py @@ -32,8 +32,7 @@ class ContextCompiler: if compiled.hits: self.cache.record_hit(SEMANTIC_LAYER, cache_key) return compiled - # A cached context with no evidence in it is not an answer, so counting it as a hit - # would let an empty cache entry look like a warm cache. + # A cached context with no evidence is not an answer; counting it would make an empty entry look like a warm cache. self.cache.record_miss( SEMANTIC_LAYER, cache_key, metadata={"reason": "cached context has no hits"} ) diff --git a/backend/src/ely_eye/db.py b/backend/src/ely_eye/db.py index bcc7932..3282ad1 100644 --- a/backend/src/ely_eye/db.py +++ b/backend/src/ely_eye/db.py @@ -3,7 +3,6 @@ from __future__ import annotations import sqlite3 from collections.abc import Iterator from contextlib import contextmanager -from dataclasses import dataclass from datetime import datetime from pathlib import Path from typing import Any @@ -12,16 +11,7 @@ from .config import Settings, get_settings from .migrations import apply as apply_migrations from .row_codec import atom_params, from_json, row_to_atom, row_to_source, to_json from .schemas import EvidenceAtom, SourceRecord, utc_now -from .text_index import fts_expression, fts_terms, index_text, rare_terms - - -@dataclass(frozen=True) -class FtsPage: - """One page of BM25 hits plus whether the corpus holds more: the cursor is an offset, and a - ranking this small re-reads rather than keeps a server-side position.""" - - rows: list[tuple[str, float]] - has_more: bool +from .text_index import index_text class Database: @@ -173,103 +163,6 @@ class Database: by_id = {row["atom_id"]: row_to_atom(row) for row in rows} return [by_id[atom_id] for atom_id in atom_ids if atom_id in by_id] - def search_fts(self, query: str, limit: int) -> list[tuple[str, float]]: - """Best BM25 matches over the whole corpus, no filters: the ranking the product needs most days.""" - return self.search_atoms(query, limit=limit).rows - - def search_atoms( - self, - query: str, - *, - limit: int, - offset: int = 0, - holder: str | None = None, - scenario: str | None = None, - source_id: str | None = None, - modality: str | None = None, - since: datetime | None = None, - until: datetime | None = None, - min_trust: float | None = None, - cartridge_ids: list[str] | None = None, - as_holder: str | None = None, - ) -> FtsPage: - """BM25 over the FTS index joined to memories that are open right now. - - Every filter is a `? IS NULL OR` predicate in one literal statement: the ranking never - depends on string-built SQL, and the planner sees one shape. A memory with no event time - cannot answer a time filter, so it drops out when one is asked for. - - `as_holder` applies the sharing rule: a memory is readable by its owner, by anyone named - in its visible_to list, by everyone when it is public, and by anyone when nobody owns it. - A memory with no holder cannot be private to a holder that does not exist. Without an - asker the query runs over the whole library, which is what an operator debugging wants. - """ - with self.connect() as connection: - expression = fts_expression(rare_terms(connection, fts_terms(query))) - rows = connection.execute( - """ - SELECT f.atom_id, bm25(atom_fts) AS score - FROM atom_fts AS f - JOIN v_atoms_current AS a ON a.atom_id = f.atom_id - WHERE atom_fts MATCH ? - AND (? IS NULL OR a.holder_id = ?) - AND (? IS NULL OR a.scenario = ?) - AND (? IS NULL OR a.source_id = ?) - AND (? IS NULL OR a.modality = ?) - AND (? IS NULL OR (a.event_time IS NOT NULL AND a.event_time >= ?)) - AND (? IS NULL OR (a.event_time IS NOT NULL AND a.event_time <= ?)) - AND (? IS NULL OR ( - a.trust_json IS NOT NULL - AND min( - json_extract(a.trust_json, '$.parser'), - json_extract(a.trust_json, '$.ocr'), - json_extract(a.trust_json, '$.model') - ) >= ? - )) - AND (? IS NULL OR EXISTS ( - SELECT 1 FROM cartridge_atoms AS ca - WHERE ca.atom_id = f.atom_id - AND ca.cartridge_id IN (SELECT value FROM json_each(?)) - )) - AND (? IS NULL OR a.holder_id = '' - OR json_extract(a.visibility_json, '$.scope') = 'everyone' - OR a.holder_id = ? - OR EXISTS ( - SELECT 1 FROM json_each(json_extract(a.visibility_json, '$.holders')) - WHERE value = ? - )) - ORDER BY score - LIMIT ? OFFSET ? - """, - ( - expression, - holder, - holder, - scenario, - scenario, - source_id, - source_id, - modality, - modality, - None if since is None else since.isoformat(), - None if since is None else since.isoformat(), - None if until is None else until.isoformat(), - None if until is None else until.isoformat(), - min_trust, - min_trust, - None if cartridge_ids is None else to_json(cartridge_ids), - None if cartridge_ids is None else to_json(cartridge_ids), - as_holder, - as_holder, - as_holder, - limit + 1, - offset, - ), - ).fetchall() - matched = [(str(row["atom_id"]), -float(row["score"])) for row in rows[:limit]] - return FtsPage(rows=matched, has_more=len(rows) > limit) - - def source_id_of_atom(self, atom_id: str) -> str | None: with self.connect() as connection: row = connection.execute("SELECT source_id FROM atoms WHERE atom_id = ?", (atom_id,)).fetchone() diff --git a/backend/src/ely_eye/memory.py b/backend/src/ely_eye/memory.py index 2a559af..06e696d 100644 --- a/backend/src/ely_eye/memory.py +++ b/backend/src/ely_eye/memory.py @@ -29,8 +29,7 @@ from .schemas import ( 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. +# Trust scores describe extraction fidelity, not truthfulness: a remembered sentence is stored verbatim, so origin carries the epistemic weight. MEMORY_TRUST = TrustScores(parser=1.0, ocr=0.0, model=0.0) @@ -135,8 +134,7 @@ class MemoryService: f"{atom_id} was already closed at {previous[0].valid_to.isoformat()}; " "superseding a closed memory hides which replacement is current" ) - # A correction takes effect when it is asserted: the old memory stops being current - # and the new one starts there, so the pair never claims to be true at the same time. + # A correction takes effect when asserted: the old memory stops being current where the new one starts, so the pair is never true at the same instant. moment = event_time or utc_now() replacement = self.remember( text, diff --git a/backend/src/ely_eye/migrations.py b/backend/src/ely_eye/migrations.py index 8112387..f0057fd 100644 --- a/backend/src/ely_eye/migrations.py +++ b/backend/src/ely_eye/migrations.py @@ -135,10 +135,7 @@ tokenize = 'unicode61' -- per CJK character and per identifier fragment, which is what makes those queries match. """ -# A memory is usable now when nothing has closed it, it had begun by now, and its -# retention has not lapsed. Unknown bounds read as visible: absence of evidence about -# when a legacy row was true is not evidence that it is false. The clock is formatted to -# match how Python writes an aware datetime, so the comparison is exact to the millisecond. +# A memory is usable now when nothing closed it, it had begun by now, and its retention has not lapsed; unknown bounds read as visible, because silence about when a legacy row was true is not evidence that it is false. VIEWS = """ CREATE VIEW IF NOT EXISTS v_atoms_current AS SELECT * FROM atoms @@ -153,8 +150,7 @@ VIEWS = """ CREATE INDEX IF NOT EXISTS idx_atoms_valid_to ON atoms(valid_to); """ -# Columns the first version lacked, in the order they are added. Every one is nullable -# or has a default, so an existing row survives the upgrade with its meaning intact. +# Columns the first version lacked, in the order added; each is nullable or defaulted so an existing row keeps its meaning. _V1_TO_V2_COLUMNS = ( ("holder_id", "TEXT NOT NULL DEFAULT ''"), ("origin", "TEXT NOT NULL DEFAULT 'artifact'"), diff --git a/backend/src/ely_eye/proofs/arena.py b/backend/src/ely_eye/proofs/arena.py index a1f7a7a..7c3888d 100644 --- a/backend/src/ely_eye/proofs/arena.py +++ b/backend/src/ely_eye/proofs/arena.py @@ -11,8 +11,7 @@ from typing import Any, Literal from PIL import Image, ImageDraw, ImageFont -# The four arena positions a model may answer with. Naming them once keeps the puzzle, the parser, -# and the scoring keyed on the same values instead of four spellings in three files. +# The four arena positions a model may answer with, named once so the puzzle, parser and scoring share one spelling. SlotName = Literal["top-left", "top-right", "bottom-left", "bottom-right"] SLOTS: tuple[SlotName, ...] = ("top-left", "top-right", "bottom-left", "bottom-right") diff --git a/backend/src/ely_eye/proofs/checks_cartridge.py b/backend/src/ely_eye/proofs/checks_cartridge.py index 055cfba..2d4dd4d 100644 --- a/backend/src/ely_eye/proofs/checks_cartridge.py +++ b/backend/src/ely_eye/proofs/checks_cartridge.py @@ -182,15 +182,11 @@ class CartridgeChecks: question = "Which Ely-Eye PRD requirements prove long context memory, HashHop, and adapters?" compiler = ContextCompiler(self.settings, self.db) artifact = proof_root / "cartridge_replay.json" - # The replay must come out of the cartridge, not out of whatever the live store happens - # to hold: mounting is what makes the cartridge the retrieval scope (retrieval scopes by - # the mounted set), so an unmounted replay would prove only that compiling is stable. + # The replay must come out of the cartridge, so mounting is what makes the cartridge the retrieval scope; an unmounted replay would prove only that compiling is stable. try: self.cartridges.mount(cartridge_id) except ValueError as error: - # The KV fingerprint gate refuses to mount a cartridge whose model snapshot this - # machine cannot read. A missing input is not a failed claim, so the check reports - # the refusal instead of crashing the suite. + # The KV fingerprint gate refuses a cartridge whose model snapshot this machine cannot read; a missing input is not a failed claim, so report the refusal instead of crashing the suite. write_json(artifact, {"verdict": "deferred", "why": str(error)}) return ProofCheck( name="Context Cartridge Replay", @@ -271,8 +267,7 @@ class CartridgeChecks: call_graph = context["code"]["call_graph"] git_context = context["git"] issue_pr_history = context["issue_pr_history"] - # Building a context object from files on disk is not evidence the repository reached the - # store: the claim is that the repository is queryable, so the store has to hold it. + # Building a context from files on disk is not evidence the repository reached the store; the claim is that it is queryable. stored_sources = [source for source in self.db.list_sources() if source.path.startswith(str(root))] stored_atoms = self.db.atom_count() artifact = proof_root / "repository_input.json" @@ -355,8 +350,7 @@ class CartridgeChecks: reason="an empty store answers every question with an empty context, which would make the two compiles agree for the wrong reason", rerun="ely-eye remember or ely-eye ingest first, then re-run this proof", ) - # The nonce keeps the key unique; the quoted text keeps the question answerable, so a pass - # means the cache served a real context rather than two identical empty ones. + # The nonce keeps the key unique and the quoted text keeps the question answerable, so a pass means the cache served a real context. question = f"{probe[0].text[:60]} cache proof {secrets.token_hex(6)}" before = _semantic_counters(fabric) first = compiler.compile(question, RuntimeProfile.library_100m) diff --git a/backend/src/ely_eye/proofs/checks_training.py b/backend/src/ely_eye/proofs/checks_training.py index 8bc9fb6..8914e0c 100644 --- a/backend/src/ely_eye/proofs/checks_training.py +++ b/backend/src/ely_eye/proofs/checks_training.py @@ -18,11 +18,8 @@ from ..schemas import ( AdapterManifest, TrainingAdapterKind, ) -from ..training import ( - safetensor_weight_stats, - sha256_file, - sha256_tree, -) +from ..training_common import sha256_file, sha256_tree +from ..training_metrics import safetensor_weight_stats # A fixed puzzle count per hop depth. The old suite re-rolled until the running rate crossed its from .report import ( @@ -85,8 +82,7 @@ class TrainingChecks: trace_report = validate_training_trace(manifest, adapter_root) proof_report = validate_training_proof(manifest, adapter_root) contract_report = validate_training_contract(manifest) - # Named conditions, so a failure can say which promise broke instead of listing the - # adapters that happened to be present. + # Named conditions, so a failure says which promise broke instead of listing the adapters that happened to be present. conditions: list[tuple[str, bool]] = [ ("sample_count", manifest.sample_count > 0), ("max_steps", manifest.max_steps > 0), diff --git a/backend/src/ely_eye/proofs/suite.py b/backend/src/ely_eye/proofs/suite.py index f61e64f..b5f3b40 100644 --- a/backend/src/ely_eye/proofs/suite.py +++ b/backend/src/ely_eye/proofs/suite.py @@ -71,9 +71,7 @@ class PRDProofSuite(TrainingChecks, CartridgeChecks, RuntimeChecks, VisualChecks self._cache_trace_check(proof_root), self._visual_contradiction_check(proof_root), ] - # A deferred or removed check is not a failure and not a pass: the suite is only as green - # as the checks that actually ran. A claim that never ran holds the verdict, because - # "passed" would say the PRD is proven while part of it was never attempted. + # A deferred or removed check is neither a failure nor a pass: the suite is only as green as the checks that ran, because "passed" would claim a proof that never ran. failed = [check for check in checks if check.status == "failed"] unrun = [check for check in checks if check.status in ("deferred", "removed")] status: Literal["passed", "failed", "deferred"] = ( diff --git a/backend/src/ely_eye/proofs/training_trace.py b/backend/src/ely_eye/proofs/training_trace.py index b1c5c3b..f935d24 100644 --- a/backend/src/ely_eye/proofs/training_trace.py +++ b/backend/src/ely_eye/proofs/training_trace.py @@ -12,16 +12,16 @@ from ..schemas import ( AdapterManifest, TrainingAdapterKind, ) -from ..training import ( +from ..training_common import ( AUTOCAST_DTYPE, GRADIENT_CHECKPOINTING_MODE, TRAINING_FRAMEWORK, TRAINING_METHOD, TRAINING_PRECISION, TRAINING_PROOF_VERSION, - safetensor_weight_stats, sha256_file, ) +from ..training_metrics import safetensor_weight_stats def cuda_bf16_training(environment: dict[str, Any]) -> bool: diff --git a/backend/src/ely_eye/retrieval.py b/backend/src/ely_eye/retrieval.py index 2e57cc4..02f2b87 100644 --- a/backend/src/ely_eye/retrieval.py +++ b/backend/src/ely_eye/retrieval.py @@ -14,9 +14,10 @@ from pathlib import Path from . import cartridge_store from .config import Settings, get_settings -from .db import Database, FtsPage +from .db import Database from .embeddings import EmbeddingService from .schemas import EvidenceAtom, RetrievalHit, SearchFilters, SearchResult +from .store_search import FtsPage, search_atoms from .visual_index import VisualIndex RRF_K = 60 @@ -58,10 +59,9 @@ class RetrievalService: if scope is None: scope = cartridge_store.mounted_ids(self.db) or None criteria["cartridge_ids"] = scope - # 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. + # A picture-only question has no words for BM25, so the word ranking is skipped rather than matching an empty expression. page = ( - self.db.search_atoms(query, limit=limit * CANDIDATE_WINDOW, offset=offset, **criteria) + search_atoms(self.db, query, limit=limit * CANDIDATE_WINDOW, offset=offset, **criteria) if query.strip() else FtsPage(rows=[], has_more=False) ) diff --git a/backend/src/ely_eye/store_search.py b/backend/src/ely_eye/store_search.py new file mode 100644 index 0000000..e61cc1e --- /dev/null +++ b/backend/src/ely_eye/store_search.py @@ -0,0 +1,106 @@ +"""BM25 search over the store, kept out of the query layer by the line budget.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime + +from .db import Database +from .row_codec import to_json +from .text_index import fts_expression, fts_terms, rare_terms + + +@dataclass(frozen=True) +class FtsPage: + """A page of (atom_id, score) plus whether more matches exist.""" + + rows: list[tuple[str, float]] + has_more: bool + + +def search_fts(database: Database, query: str, limit: int) -> list[tuple[str, float]]: + """Best BM25 matches over the whole corpus, no filters.""" + return search_atoms(database, query, limit=limit).rows + + +def search_atoms( + database: Database, + query: str, + *, + limit: int, + offset: int = 0, + holder: str | None = None, + scenario: str | None = None, + source_id: str | None = None, + modality: str | None = None, + since: datetime | None = None, + until: datetime | None = None, + min_trust: float | None = None, + cartridge_ids: list[str] | None = None, + as_holder: str | None = None, +) -> FtsPage: + """BM25 over the FTS index joined to memories that are open right now and readable by the asker.""" + with database.connect() as connection: + expression = fts_expression(rare_terms(connection, fts_terms(query))) + rows = connection.execute( + """ + SELECT f.atom_id, bm25(atom_fts) AS score + FROM atom_fts AS f + JOIN v_atoms_current AS a ON a.atom_id = f.atom_id + WHERE atom_fts MATCH ? + AND (? IS NULL OR a.holder_id = ?) + AND (? IS NULL OR a.scenario = ?) + AND (? IS NULL OR a.source_id = ?) + AND (? IS NULL OR a.modality = ?) + AND (? IS NULL OR (a.event_time IS NOT NULL AND a.event_time >= ?)) + AND (? IS NULL OR (a.event_time IS NOT NULL AND a.event_time <= ?)) + AND (? IS NULL OR ( + a.trust_json IS NOT NULL + AND min( + json_extract(a.trust_json, '$.parser'), + json_extract(a.trust_json, '$.ocr'), + json_extract(a.trust_json, '$.model') + ) >= ? + )) + AND (? IS NULL OR EXISTS ( + SELECT 1 FROM cartridge_atoms AS ca + WHERE ca.atom_id = f.atom_id + AND ca.cartridge_id IN (SELECT value FROM json_each(?)) + )) + AND (? IS NULL OR a.holder_id = '' + OR json_extract(a.visibility_json, '$.scope') = 'everyone' + OR a.holder_id = ? + OR EXISTS ( + SELECT 1 FROM json_each(json_extract(a.visibility_json, '$.holders')) + WHERE value = ? + )) + ORDER BY score + LIMIT ? OFFSET ? + """, + ( + expression, + holder, + holder, + scenario, + scenario, + source_id, + source_id, + modality, + modality, + None if since is None else since.isoformat(), + None if since is None else since.isoformat(), + None if until is None else until.isoformat(), + None if until is None else until.isoformat(), + min_trust, + min_trust, + None if cartridge_ids is None else to_json(cartridge_ids), + None if cartridge_ids is None else to_json(cartridge_ids), + as_holder, + as_holder, + as_holder, + limit + 1, + offset, + ), + ).fetchall() + matched = [(str(row["atom_id"]), -float(row["score"])) for row in rows[:limit]] + return FtsPage(rows=matched, has_more=len(rows) > limit) diff --git a/backend/src/ely_eye/text_index.py b/backend/src/ely_eye/text_index.py index 70bd607..d96f199 100644 --- a/backend/src/ely_eye/text_index.py +++ b/backend/src/ely_eye/text_index.py @@ -12,15 +12,13 @@ import re import sqlite3 from unicodedata import normalize -# Han, its extension A block, the compatibility ideographs, kana, and Hangul syllables: -# scripts written without spaces. One source of truth for both the tokenizer and the query. +# Han, its extension A block, compatibility ideographs, kana and Hangul syllables: scripts written without spaces, shared by tokenizer and query. _BLOCKS = ((0x3400, 0x4DBF), (0x4E00, 0x9FFF), (0xF900, 0xFAFF), (0x3040, 0x30FF), (0xAC00, 0xD7AF)) _CJK_CLASS = "".join(f"{chr(start)}-{chr(end)}" for start, end in _BLOCKS) _CJK_CHARS = frozenset("".join("".join(chr(code) for code in range(start, end + 1)) for start, end in _BLOCKS)) _CAMEL_BOUNDARY = re.compile(r"(?<=[a-z0-9])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|(?<=[A-Za-z])(?=\d)") -# Everything a reader would not type as a word: symbols, punctuation, and the underscore -# that joins identifiers. `[^[^\W\s]` looks equivalent but deletes CJK characters on 3.14. +# What a reader would not type as a word; `[^[^\W\s]` looks equivalent but deletes CJK characters on 3.14. _SEPARATORS = re.compile(r"[^\w\s]|_") _CJK_RUN = re.compile(f"[{_CJK_CLASS}]+") _SPACES = re.compile(r"\s+") diff --git a/backend/src/ely_eye/training.py b/backend/src/ely_eye/training.py index e3dfb21..171eb78 100644 --- a/backend/src/ely_eye/training.py +++ b/backend/src/ely_eye/training.py @@ -1,43 +1,22 @@ +"""Train the adapters: the service is a façade over the dataset, LoRA and retrieval modules.""" + from __future__ import annotations -import hashlib -import json -import math -import platform -import re -import shutil -import subprocess -import time -from contextlib import contextmanager -from datetime import UTC, datetime -from importlib.metadata import PackageNotFoundError, version from pathlib import Path -from typing import Any, TextIO - -from pydantic import BaseModel, Field from .config import Settings, get_settings from .db import Database -from .schemas import AdapterManifest, EvidenceAtom, TrainingAdapterKind - -SYSTEM_PROMPT = ( - "You are Ely-Eye. Answer from supplied evidence, return compact JSON, " - "and cite atom ids exactly." +from .schemas import AdapterManifest, TrainingAdapterKind +from .training_common import TrainingSample +from .training_datasets import ( + build_prd_dataset, + build_retrieval_dataset, + build_router_dataset, + build_visual_dataset, + validate_dataset, ) -TRAINING_PROOF_VERSION = 2 -TRAINING_METHOD = "BF16 LoRA" -TRAINING_PRECISION = "bf16" -TRAINING_FRAMEWORK = "unsloth_patched_transformers_peft" -GRADIENT_CHECKPOINTING_MODE = "unsloth_patched_torch_non_reentrant" -AUTOCAST_DTYPE = "bf16" - - -class TrainingSample(BaseModel): - id: str - task: str - inputs: list[dict[str, Any]] - memory: dict[str, Any] = Field(default_factory=dict) - output: dict[str, Any] +from .training_lora import train_lora +from .training_retrieval import train_retrieval_lora class TrainingService: @@ -46,18 +25,7 @@ class TrainingService: self.db = database or Database(self.settings) def validate_dataset(self, path: Path) -> list[TrainingSample]: - samples: list[TrainingSample] = [] - with path.open("r", encoding="utf-8") as handle: - for line_number, line in enumerate(handle, start=1): - if not line.strip(): - continue - try: - samples.append(TrainingSample.model_validate_json(line)) - except Exception as exc: - raise ValueError(f"Invalid training sample at line {line_number}: {exc}") from exc - if not samples: - raise ValueError(f"Training dataset is empty: {path}") - return samples + return validate_dataset(path) def build_prd_dataset( self, @@ -66,62 +34,7 @@ class TrainingService: cartridge_id: str | None = None, max_sections: int = 48, ) -> Path: - prd_path = prd_path.resolve() - text = prd_path.read_text(encoding="utf-8") - atoms = self._atoms_for_path(prd_path) - atom_id = atoms[0].atom_id if atoms else f"{prd_path.stem}:{sha256_bytes(text.encode('utf-8'))[:12]}" - cartridge_id = cartridge_id or self._latest_cartridge_id() - sections = split_markdown_sections(text) - if not sections: - raise ValueError(f"No markdown sections found in {prd_path}") - output_path = output_path or (self.settings.training_dir / "hme_prd_sft.jsonl") - output_path.parent.mkdir(parents=True, exist_ok=True) - - samples: list[TrainingSample] = [] - for index, section in enumerate(sections[:max_sections], start=1): - title = section["title"] - body = compact_text(section["body"], limit=5200) - if len(body) < 160: - continue - citation = f"{atom_id}#section={slugify(title)}" - answer = compact_text(body, limit=900) - samples.append( - TrainingSample( - id=f"hme_prd_{index:04d}", - task="hme_core_lora_evidence_json", - inputs=[ - { - "type": "text", - "text": ( - f"Section title: {title}\n" - f"Evidence atom: {atom_id}\n" - f"Evidence:\n{body}\n\n" - "Question: State the Ely-Eye requirement in this section. " - "Return JSON with answer, citations, contradiction_notes, confidence." - ), - } - ], - memory={ - "cartridge": cartridge_id, - "required_atoms": [atom_id], - "source": str(prd_path), - }, - output={ - "answer": answer, - "citations": [citation], - "contradiction_notes": [], - "confidence": 1.0, - "reasoning_policy": "evidence-first", - }, - ) - ) - - if not samples: - raise ValueError(f"No usable PRD training sections found in {prd_path}") - with output_path.open("w", encoding="utf-8", newline="\n") as handle: - for sample in samples: - handle.write(sample.model_dump_json() + "\n") - return output_path + return build_prd_dataset(self.settings, self.db, prd_path, output_path, cartridge_id, max_sections) def build_visual_dataset( self, @@ -130,110 +43,7 @@ class TrainingService: image_paths: list[Path] | None = None, cartridge_id: str | None = None, ) -> Path: - prd_path = prd_path.resolve() - text = prd_path.read_text(encoding="utf-8") - atom_id = self._atom_id_for_path(prd_path, text) - cartridge_id = cartridge_id or self._latest_cartridge_id() - media_dir = self.settings.training_dir / "media" - media_dir.mkdir(parents=True, exist_ok=True) - - prepared_images: list[Path] = [] - for image_path in image_paths or []: - if image_path.exists(): - prepared_images.append(prepare_training_image(image_path, media_dir)) - prd_card = render_prd_training_card(text, atom_id, media_dir / "prd_training_card.png") - prepared_images.append(prd_card) - if not prepared_images: - raise ValueError("No visual training images are available") - - output_path = output_path or (self.settings.training_dir / "hme_visual_sft.jsonl") - samples: list[TrainingSample] = [] - visual_tasks = [ - ( - "hme_vision_lora_visual_grounding", - "12.1 训练目标", - ( - "Ely-Eye trains HME-Vision-LoRA on Qwen3.5-9B vision projector " - "and selected vision layers so UI, PDF, OCR, and chart evidence can be " - "grounded into cited multimodal memory." - ), - ), - ( - "hme_ttt_vl_visual_memory", - "11.2 X — TTT-VL：视觉语言测试时学习", - ( - "TTT-VL keeps Qwen3.5-9B frozen while mutable MLP and visual LoRA " - "adapters learn low-step visual memory updates that can be consolidated " - "into cartridge-local adapter deltas." - ), - ), - ( - "hme_visual_hashhop_grounding", - "11.3 X — Visual HashHop", - ( - "Visual HashHop converts random chain ids into non-text geometric image " - "blocks, then evaluates long-range visual addressing without OCR shortcuts." - ), - ), - ( - "hme_dvkv_visual_cache", - "11.4 S/H — DVKV：Diversity-Aware Visual KV Compression", - ( - "DVKV keeps visual KV tokens by combining attention importance, visual " - "diversity, and evidence value so OCR, layout, chart, and coordinate details " - "survive compression." - ), - ), - ( - "hme_visual_contradiction_lens", - "11.15 X — Visual Contradiction Lens", - ( - "The Visual Contradiction Lens compares visual versions and cites layout, " - "design-token, copy, temporal, and visual-code drift as evidence graph facts." - ), - ), - ] - for index, image_path in enumerate(prepared_images, start=1): - for task_index, (task, section_title, answer) in enumerate(visual_tasks, start=1): - requirement = section_by_title(text, section_title) - samples.append( - TrainingSample( - id=f"hme_visual_{index:02d}_{task_index:02d}", - task=task, - inputs=[ - { - "type": "ui_screenshot", - "path": str(image_path.resolve()), - "max_pixels": 262144, - }, - { - "type": "text", - "text": ( - f"Evidence atom: {atom_id}\n" - f"Visual source: {image_path.name}\n" - f"PRD visual requirement:\n" - f"{compact_text(requirement, 2200)}\n\n" - "Question: Map this visual evidence to the Ely-Eye PRD module. " - "Return JSON with answer, citations, contradiction_notes, confidence." - ), - }, - ], - memory={ - "cartridge": cartridge_id, - "required_atoms": [atom_id], - "source": str(prd_path), - }, - output={ - "answer": answer, - "citations": [f"{atom_id}#section={slugify(section_title)}"], - "contradiction_notes": [], - "confidence": 1.0, - "reasoning_policy": "evidence-first", - }, - ) - ) - write_jsonl(output_path, samples) - return output_path + return build_visual_dataset(self.settings, self.db, prd_path, output_path, image_paths, cartridge_id) def build_router_dataset( self, @@ -241,103 +51,7 @@ class TrainingService: output_path: Path | None = None, cartridge_id: str | None = None, ) -> Path: - prd_path = prd_path.resolve() - text = prd_path.read_text(encoding="utf-8") - atom_id = self._atom_id_for_path(prd_path, text) - cartridge_id = cartridge_id or self._latest_cartridge_id() - output_path = output_path or (self.settings.training_dir / "hme_router_sft.jsonl") - scenarios = [ - ( - "Summarize a stable live-demo answer from local PRD evidence.", - "live_demo", - 262144, - ["mixed"], - "Prefix cache + sparse evidence packing", - ), - ( - "Plan a 1.01M extreme context run with YaRN and parked context.", - "extreme_context", - 1010000, - ["mixed"], - "YaRN + V-NSA sparse blocks + DVKV visual KV compression + CPU/NVMe parked context", - ), - ( - "Route a 100M library question through Context Cartridge memory.", - "library_100m", - 100000000, - ["mixed"], - "Context Cartridge retrieval + ColQwen-compatible dense vectors + BM25 + evidence graph pack", - ), - ( - "Audit UI screenshot and PRD differences for the research demo.", - "research_theater", - 262144, - ["image", "ui_screenshot"], - "Visual patch cache + DVKV evidence-biased region retention", - ), - ( - "Find visual token compression methods across a 100M library cartridge.", - "library_100m", - 100000000, - ["pdf_page", "image", "video_frame"], - "ColQwen-compatible retrieval + graph grouping + cited evidence pack", - ), - ( - "Compare UI versions and cite coordinates plus code evidence.", - "research_theater", - 262144, - ["ui_screenshot", "code", "text"], - "Visual Contradiction Lens + evidence graph writeback", - ), - ( - "Run Visual HashHop at 262K with a two-hop random image chain.", - "research_theater", - 262144, - ["image"], - "Visual HashHop arena + verifier contract + no OCR-dependent evidence", - ), - ( - "Answer from a TTT-VL cartridge after original images are absent.", - "library_100m", - 100000000, - ["mixed"], - "Load cartridge adapters + learned memory recall + cited adapter proof", - ), - ] - samples = [ - TrainingSample( - id=f"hme_router_{index:04d}", - task="hme_router_context_plan", - inputs=[ - { - "type": "text", - "text": ( - f"Evidence atom: {atom_id}\n" - f"PRD routing source:\n{compact_text(section_by_title(text, '1. 产品定义'), 1800)}\n\n" - f"User request: {request}\n" - "Return the Ely-Eye routing plan as compact JSON." - ), - } - ], - memory={ - "cartridge": cartridge_id, - "required_atoms": [atom_id], - "source": str(prd_path), - }, - output={ - "profile": profile, - "token_budget": token_budget, - "required_modalities": modalities, - "retrieval_query": request, - "compression_strategy": compression, - "citations": [f"{atom_id}#section=1-产品定义"], - "confidence": 1.0, - }, - ) - for index, (request, profile, token_budget, modalities, compression) in enumerate(scenarios, start=1) - ] - write_jsonl(output_path, samples) - return output_path + return build_router_dataset(self.settings, self.db, prd_path, output_path, cartridge_id) def build_retrieval_dataset( self, @@ -346,69 +60,7 @@ class TrainingService: cartridge_id: str | None = None, max_sections: int = 32, ) -> Path: - prd_path = prd_path.resolve() - text = prd_path.read_text(encoding="utf-8") - atom_id = self._atom_id_for_path(prd_path, text) - cartridge_id = cartridge_id or self._latest_cartridge_id() - sections = [section for section in split_markdown_sections(text) if len(section["body"]) >= 160] - output_path = output_path or (self.settings.training_dir / "hme_retrieval_pairs.jsonl") - samples: list[TrainingSample] = [] - for index, section in enumerate(sections[:max_sections], start=1): - title = section["title"] - positive = compact_text(section["body"], 1600) - samples.append( - TrainingSample( - id=f"hme_retrieval_{index:04d}", - task="hme_retrieval_contrastive_pair", - inputs=[ - { - "type": "text", - "text": f"Find the PRD section that answers: {title}", - } - ], - memory={ - "cartridge": cartridge_id, - "required_atoms": [atom_id], - "source": str(prd_path), - }, - output={ - "positive": positive, - "positive_title": title, - "citations": [f"{atom_id}#section={slugify(title)}"], - "confidence": 1.0, - }, - ) - ) - if len(samples) < 2: - raise ValueError("Retrieval contrastive training needs at least two PRD sections") - write_jsonl(output_path, samples) - return output_path - - def write_lora_config( - self, - output_dir: Path, - kind: TrainingAdapterKind = TrainingAdapterKind.hme_core_lora, - rank: int = 16, - alpha: int = 32, - model_id: str | None = None, - ) -> Path: - output_dir.mkdir(parents=True, exist_ok=True) - config = { - "base_model": model_id or self.settings.model_id, - "adapter_kind": kind.value, - "method": "BF16 LoRA", - "framework": TRAINING_FRAMEWORK, - "rank": rank, - "alpha": alpha, - "gradient_checkpointing": True, - "gradient_checkpointing_mode": GRADIENT_CHECKPOINTING_MODE, - "precision": "bf16", - "optimizer": "paged_adamw_8bit", - "ensure_weight_tying": kind == TrainingAdapterKind.hme_visual_mtp, - } - path = output_dir / "lora_config.json" - path.write_text(json.dumps(config, ensure_ascii=False, indent=2), encoding="utf-8") - return path + return build_retrieval_dataset(self.settings, self.db, prd_path, output_path, cartridge_id, max_sections) def train_lora( self, @@ -424,321 +76,7 @@ class TrainingService: learning_rate: float | None = None, cartridge_id: str | None = None, ) -> AdapterManifest: - samples = self.validate_dataset(dataset_path) - output_dir.mkdir(parents=True, exist_ok=True) - model_id = model_id or self._model_for_kind(kind) - max_length = max_length or self.settings.train_max_length - gradient_accumulation_steps = ( - gradient_accumulation_steps or self.settings.train_gradient_accumulation_steps - ) - learning_rate = learning_rate or self.settings.train_learning_rate - started_at = datetime.now(UTC) - wall_start = time.perf_counter() - trace_path = output_dir / "training_trace.jsonl" - - import torch - - activate_unsloth_training() - from peft import LoraConfig, get_peft_model - from transformers import AutoProcessor - - try: - from transformers import AutoModelForImageTextToText - - model_cls = AutoModelForImageTextToText - except ImportError: - from transformers import AutoModelForCausalLM - - model_cls = AutoModelForCausalLM - - require_cuda_bf16(torch) - processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) - dtype = torch.bfloat16 - device_map: dict[str, int] = {"": 0} - model = model_cls.from_pretrained( - model_id, - dtype=dtype, - device_map=device_map, - trust_remote_code=True, - ) - model.config.use_cache = False - input_require_grads = False - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - input_require_grads = True - - target_modules = infer_lora_targets(model, kind) - gradient_checkpointing = enable_gradient_checkpointing(model) - peft_config_options: dict[str, Any] = { - "r": rank, - "lora_alpha": alpha, - "target_modules": target_modules, - "lora_dropout": 0.0, - "bias": "none", - "task_type": "CAUSAL_LM", - } - if kind == TrainingAdapterKind.hme_visual_mtp: - peft_config_options["ensure_weight_tying"] = True - with bitsandbytes_gaudi_probe_guard(): - model = get_peft_model(model, LoraConfig(**peft_config_options)) - unsloth_model_type = model_type_for_proof(model) - trainable_params, total_params = model.get_nb_trainable_parameters() - optimizer, optimizer_name = make_optimizer(model, learning_rate) - model.train() - - losses: list[float] = [] - optimizer_updates: list[dict[str, Any]] = [] - trainable_tensors = trainable_named_parameters(model) - initial_trainable_sha = trainable_state_sha256(trainable_tensors) - initial_trainable_l2 = trainable_l2_norm(trainable_tensors) - optimizer.zero_grad(set_to_none=True) - environment = training_environment(torch) - contract = training_contract( - gradient_checkpointing=gradient_checkpointing, - input_require_grads=input_require_grads, - optimizer_name=optimizer_name, - unsloth_model_type=unsloth_model_type, - ) - with trace_path.open("w", encoding="utf-8", newline="\n") as trace: - write_trace( - trace, - { - "event": "run_start", - **contract, - "started_at": started_at.isoformat(), - "adapter_kind": kind.value, - "base_model": model_id, - "dataset_path": str(dataset_path.resolve()), - "dataset_sha256": sha256_file(dataset_path), - "sample_count": len(samples), - "max_steps": max_steps, - "max_length": max_length, - "rank": rank, - "alpha": alpha, - "gradient_accumulation_steps": gradient_accumulation_steps, - "learning_rate": learning_rate, - "target_modules": target_modules, - "trainable_params": int(trainable_params), - "total_params": int(total_params), - "optimizer": optimizer_name, - "environment": environment, - }, - ) - for step in range(max_steps): - sample = samples[step % len(samples)] - batch = encode_training_sample(sample, processor, model.device, max_length) - with torch.autocast( - device_type="cuda", - dtype=torch.bfloat16, - enabled=torch.cuda.is_available(), - ): - outputs = model(**batch) - loss = outputs.loss / gradient_accumulation_steps - loss.backward() - optimizer_step = (step + 1) % gradient_accumulation_steps == 0 or (step + 1) == max_steps - grad_l2 = gradient_l2_norm(trainable_tensors) - update_l2 = None - if optimizer_step: - before_update = trainable_state_cpu(trainable_tensors) - optimizer.step() - update_l2 = trainable_update_l2(before_update, trainable_tensors) - optimizer_updates.append( - { - "step": step + 1, - "gradient_l2": grad_l2, - "update_l2": update_l2, - } - ) - optimizer.zero_grad(set_to_none=True) - actual_loss = float(loss.detach().cpu()) * gradient_accumulation_steps - losses.append(actual_loss) - write_trace( - trace, - { - "event": "step", - "step": step + 1, - "sample_id": sample.id, - "task": sample.task, - "loss": actual_loss, - "optimizer_step": optimizer_step, - "gradient_l2": grad_l2, - "update_l2": update_l2, - }, - ) - - final_trainable_sha = trainable_state_sha256(trainable_tensors) - final_trainable_l2 = trainable_l2_norm(trainable_tensors) - max_update_l2 = max((item["update_l2"] for item in optimizer_updates), default=None) - loss_delta = losses[0] - losses[-1] if len(losses) >= 2 else None - model.save_pretrained(output_dir) - processor.save_pretrained(output_dir) - self.write_lora_config(output_dir, kind, rank, alpha, model_id) - weight_stats = safetensor_weight_stats(output_dir / "adapter_model.safetensors") - finished_at = datetime.now(UTC) - wall_seconds = time.perf_counter() - wall_start - summary_path = output_dir / "training_summary.json" - proof_path = output_dir / "training_proof.json" - training_proof = { - **contract, - "adapter_kind": kind.value, - "base_model": model_id, - "dataset_path": str(dataset_path.resolve()), - "dataset_sha256": sha256_file(dataset_path), - "sample_count": len(samples), - "max_steps": max_steps, - "max_length": max_length, - "rank": rank, - "alpha": alpha, - "target_modules": target_modules, - "trainable_params": int(trainable_params), - "total_params": int(total_params), - "optimizer": optimizer_name, - "learning_rate": learning_rate, - "gradient_accumulation_steps": gradient_accumulation_steps, - "bf16": bool(torch.cuda.is_available()), - "environment": environment, - "initial_trainable_sha256": initial_trainable_sha, - "final_trainable_sha256": final_trainable_sha, - "initial_trainable_l2": initial_trainable_l2, - "final_trainable_l2": final_trainable_l2, - "optimizer_update_count": len(optimizer_updates), - "max_optimizer_update_l2": max_update_l2, - "loss_delta": loss_delta, - "loss_history": losses, - "optimizer_updates": optimizer_updates, - "adapter_weight_sha256": weight_stats["sha256"], - "adapter_weight_stats": weight_stats, - } - write_training_summary(proof_path, training_proof) - write_training_summary( - summary_path, - { - **contract, - "adapter_kind": kind.value, - "base_model": model_id, - "dataset_path": str(dataset_path.resolve()), - "dataset_sha256": sha256_file(dataset_path), - "sample_count": len(samples), - "max_steps": max_steps, - "max_length": max_length, - "rank": rank, - "alpha": alpha, - "target_modules": target_modules, - "trainable_params": int(trainable_params), - "total_params": int(total_params), - "optimizer": optimizer_name, - "learning_rate": learning_rate, - "gradient_accumulation_steps": gradient_accumulation_steps, - "bf16": bool(torch.cuda.is_available()), - "loss_history": losses, - "final_loss": losses[-1] if losses else None, - "initial_trainable_sha256": initial_trainable_sha, - "final_trainable_sha256": final_trainable_sha, - "initial_trainable_l2": initial_trainable_l2, - "final_trainable_l2": final_trainable_l2, - "optimizer_update_count": len(optimizer_updates), - "max_optimizer_update_l2": max_update_l2, - "loss_delta": loss_delta, - "started_at": started_at.isoformat(), - "finished_at": finished_at.isoformat(), - "wall_seconds": wall_seconds, - "environment": environment, - "peak_cuda_memory_mb": cuda_peak_memory_mb(torch), - "adapter_weight_sha256": weight_stats["sha256"], - "adapter_weight_stats": weight_stats, - }, - ) - with trace_path.open("a", encoding="utf-8", newline="\n") as trace: - write_trace( - trace, - { - "event": "run_finish", - **contract, - "finished_at": finished_at.isoformat(), - "wall_seconds": wall_seconds, - "final_loss": losses[-1] if losses else None, - "initial_trainable_sha256": initial_trainable_sha, - "final_trainable_sha256": final_trainable_sha, - "optimizer_update_count": len(optimizer_updates), - "max_optimizer_update_l2": max_update_l2, - "loss_delta": loss_delta, - "peak_cuda_memory_mb": cuda_peak_memory_mb(torch), - "adapter_weight_sha256": weight_stats["sha256"], - }, - ) - adapter_sha = sha256_tree(output_dir) - manifest = AdapterManifest( - adapter_id=adapter_id(kind, dataset_path, model_id), - kind=kind, - created_at=datetime.now(UTC), - proof_version=TRAINING_PROOF_VERSION, - training_method=TRAINING_METHOD, - precision=TRAINING_PRECISION, - framework=TRAINING_FRAMEWORK, - base_model=model_id, - dataset_path=str(dataset_path.resolve()), - dataset_sha256=sha256_file(dataset_path), - sample_count=len(samples), - max_steps=max_steps, - max_length=max_length, - rank=rank, - alpha=alpha, - target_modules=target_modules, - trainable_params=int(trainable_params), - total_params=int(total_params), - optimizer=optimizer_name, - optimizer_family=optimizer_family(optimizer_name), - learning_rate=learning_rate, - gradient_accumulation_steps=gradient_accumulation_steps, - gradient_checkpointing=gradient_checkpointing, - gradient_checkpointing_mode=GRADIENT_CHECKPOINTING_MODE if gradient_checkpointing else None, - autocast_dtype=AUTOCAST_DTYPE, - bf16=torch.cuda.is_available(), - unsloth_version=package_version("unsloth"), - triton_version=package_version("triton-windows") or package_version("triton"), - xformers_version=package_version("xformers"), - unsloth_model_type=unsloth_model_type, - final_loss=losses[-1] if losses else None, - loss_history=losses, - initial_trainable_sha256=initial_trainable_sha, - final_trainable_sha256=final_trainable_sha, - initial_trainable_l2=initial_trainable_l2, - final_trainable_l2=final_trainable_l2, - optimizer_update_count=len(optimizer_updates), - max_optimizer_update_l2=max_update_l2, - loss_delta=loss_delta, - training_proof_path=str(proof_path.resolve()), - training_proof_sha256=sha256_file(proof_path), - training_trace_path=str(trace_path.resolve()), - training_trace_sha256=sha256_file(trace_path), - training_summary_path=str(summary_path.resolve()), - training_summary_sha256=sha256_file(summary_path), - training_started_at=started_at, - training_finished_at=finished_at, - training_wall_seconds=wall_seconds, - training_device=str(environment.get("device_name") or environment.get("device") or ""), - torch_version=str(environment.get("torch_version") or ""), - cuda_version=( - str(environment["torch_cuda_version"]) - if environment.get("torch_cuda_version") is not None - else None - ), - adapter_weight_sha256=str(weight_stats["sha256"] or ""), - adapter_tensor_count=int(weight_stats["tensor_count"]), - adapter_lora_tensor_count=int(weight_stats["lora_tensor_count"]), - adapter_total_elements=int(weight_stats["total_elements"]), - adapter_nonzero_elements=int(weight_stats["nonzero_elements"]), - adapter_max_abs=weight_stats["max_abs"], - adapter_weights_finite=bool(weight_stats["all_finite"]), - artifacts=relative_artifacts(output_dir), - sha256=adapter_sha, - cartridge_id=cartridge_id, - ) - (output_dir / "adapter_manifest.json").write_text( - manifest.model_dump_json(indent=2), - encoding="utf-8", - ) - return manifest + return train_lora(self.settings, self.db, dataset_path, output_dir, max_steps, kind, model_id, rank, alpha, max_length, gradient_accumulation_steps, learning_rate, cartridge_id) def train_retrieval_lora( self, @@ -751,880 +89,6 @@ class TrainingService: learning_rate: float | None = None, cartridge_id: str | None = None, ) -> AdapterManifest: - samples = self.validate_dataset(dataset_path) - output_dir.mkdir(parents=True, exist_ok=True) - max_length = max_length or self.settings.train_max_length - learning_rate = learning_rate or self.settings.train_learning_rate - started_at = datetime.now(UTC) - wall_start = time.perf_counter() - trace_path = output_dir / "training_trace.jsonl" + return train_retrieval_lora(self.settings, self.db, dataset_path, output_dir, max_steps, rank, alpha, max_length, learning_rate, cartridge_id) - import torch - import torch.nn.functional as F - activate_unsloth_training() - from peft import LoraConfig, TaskType, get_peft_model - from transformers import AutoModel, AutoProcessor - - model_id = self.settings.embedding_model_id - require_cuda_bf16(torch) - processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) - model = AutoModel.from_pretrained( - model_id, - dtype=torch.bfloat16, - device_map={"": 0}, - trust_remote_code=True, - ) - if hasattr(model.config, "use_cache"): - model.config.use_cache = False - target_modules = infer_lora_targets(model, TrainingAdapterKind.hme_retrieval) - gradient_checkpointing = enable_gradient_checkpointing(model) - with bitsandbytes_gaudi_probe_guard(): - model = get_peft_model( - model, - LoraConfig( - r=rank, - lora_alpha=alpha, - target_modules=target_modules, - lora_dropout=0.0, - bias="none", - task_type=TaskType.FEATURE_EXTRACTION, - ), - ) - unsloth_model_type = model_type_for_proof(model) - trainable_params, total_params = model.get_nb_trainable_parameters() - optimizer, optimizer_name = make_optimizer(model, learning_rate) - model.train() - - losses: list[float] = [] - optimizer_updates: list[dict[str, Any]] = [] - trainable_tensors = trainable_named_parameters(model) - initial_trainable_sha = trainable_state_sha256(trainable_tensors) - initial_trainable_l2 = trainable_l2_norm(trainable_tensors) - batch_size = min(4, len(samples)) - environment = training_environment(torch) - contract = training_contract( - gradient_checkpointing=gradient_checkpointing, - input_require_grads=False, - optimizer_name=optimizer_name, - unsloth_model_type=unsloth_model_type, - ) - with trace_path.open("w", encoding="utf-8", newline="\n") as trace: - write_trace( - trace, - { - "event": "run_start", - **contract, - "started_at": started_at.isoformat(), - "adapter_kind": TrainingAdapterKind.hme_retrieval.value, - "base_model": model_id, - "dataset_path": str(dataset_path.resolve()), - "dataset_sha256": sha256_file(dataset_path), - "sample_count": len(samples), - "max_steps": max_steps, - "max_length": max_length, - "rank": rank, - "alpha": alpha, - "target_modules": target_modules, - "trainable_params": int(trainable_params), - "total_params": int(total_params), - "optimizer": optimizer_name, - "learning_rate": learning_rate, - "environment": environment, - }, - ) - for step in range(max_steps): - window = [samples[(step + offset) % len(samples)] for offset in range(batch_size)] - queries = [str(sample.inputs[0]["text"]) for sample in window] - documents = [str(sample.output["positive"]) for sample in window] - query_batch = encode_retrieval_batch(queries, processor, model.device, max_length) - document_batch = encode_retrieval_batch(documents, processor, model.device, max_length) - query_embeddings = pool_embeddings(model(**query_batch), query_batch["attention_mask"]) - document_embeddings = pool_embeddings(model(**document_batch), document_batch["attention_mask"]) - logits = query_embeddings @ document_embeddings.T / 0.05 - labels = torch.arange(logits.shape[0], device=logits.device) - loss = F.cross_entropy(logits.float(), labels) - loss.backward() - grad_l2 = gradient_l2_norm(trainable_tensors) - before_update = trainable_state_cpu(trainable_tensors) - optimizer.step() - update_l2 = trainable_update_l2(before_update, trainable_tensors) - optimizer_updates.append( - { - "step": step + 1, - "gradient_l2": grad_l2, - "update_l2": update_l2, - } - ) - optimizer.zero_grad(set_to_none=True) - actual_loss = float(loss.detach().cpu()) - losses.append(actual_loss) - write_trace( - trace, - { - "event": "step", - "step": step + 1, - "sample_ids": [sample.id for sample in window], - "task": "hme_retrieval_contrastive_pair", - "loss": actual_loss, - "optimizer_step": True, - "gradient_l2": grad_l2, - "update_l2": update_l2, - }, - ) - - final_trainable_sha = trainable_state_sha256(trainable_tensors) - final_trainable_l2 = trainable_l2_norm(trainable_tensors) - max_update_l2 = max((item["update_l2"] for item in optimizer_updates), default=None) - loss_delta = losses[0] - losses[-1] if len(losses) >= 2 else None - model.save_pretrained(output_dir) - processor.save_pretrained(output_dir) - self.write_lora_config(output_dir, TrainingAdapterKind.hme_retrieval, rank, alpha, model_id) - weight_stats = safetensor_weight_stats(output_dir / "adapter_model.safetensors") - finished_at = datetime.now(UTC) - wall_seconds = time.perf_counter() - wall_start - summary_path = output_dir / "training_summary.json" - proof_path = output_dir / "training_proof.json" - training_proof = { - **contract, - "adapter_kind": TrainingAdapterKind.hme_retrieval.value, - "base_model": model_id, - "dataset_path": str(dataset_path.resolve()), - "dataset_sha256": sha256_file(dataset_path), - "sample_count": len(samples), - "max_steps": max_steps, - "max_length": max_length, - "rank": rank, - "alpha": alpha, - "target_modules": target_modules, - "trainable_params": int(trainable_params), - "total_params": int(total_params), - "optimizer": optimizer_name, - "learning_rate": learning_rate, - "gradient_accumulation_steps": 1, - "bf16": bool(torch.cuda.is_available()), - "environment": environment, - "initial_trainable_sha256": initial_trainable_sha, - "final_trainable_sha256": final_trainable_sha, - "initial_trainable_l2": initial_trainable_l2, - "final_trainable_l2": final_trainable_l2, - "optimizer_update_count": len(optimizer_updates), - "max_optimizer_update_l2": max_update_l2, - "loss_delta": loss_delta, - "loss_history": losses, - "optimizer_updates": optimizer_updates, - "adapter_weight_sha256": weight_stats["sha256"], - "adapter_weight_stats": weight_stats, - } - write_training_summary(proof_path, training_proof) - write_training_summary( - summary_path, - { - **contract, - "adapter_kind": TrainingAdapterKind.hme_retrieval.value, - "base_model": model_id, - "dataset_path": str(dataset_path.resolve()), - "dataset_sha256": sha256_file(dataset_path), - "sample_count": len(samples), - "max_steps": max_steps, - "max_length": max_length, - "rank": rank, - "alpha": alpha, - "target_modules": target_modules, - "trainable_params": int(trainable_params), - "total_params": int(total_params), - "optimizer": optimizer_name, - "learning_rate": learning_rate, - "gradient_accumulation_steps": 1, - "bf16": bool(torch.cuda.is_available()), - "loss_history": losses, - "final_loss": losses[-1] if losses else None, - "initial_trainable_sha256": initial_trainable_sha, - "final_trainable_sha256": final_trainable_sha, - "initial_trainable_l2": initial_trainable_l2, - "final_trainable_l2": final_trainable_l2, - "optimizer_update_count": len(optimizer_updates), - "max_optimizer_update_l2": max_update_l2, - "loss_delta": loss_delta, - "started_at": started_at.isoformat(), - "finished_at": finished_at.isoformat(), - "wall_seconds": wall_seconds, - "environment": environment, - "peak_cuda_memory_mb": cuda_peak_memory_mb(torch), - "adapter_weight_sha256": weight_stats["sha256"], - "adapter_weight_stats": weight_stats, - }, - ) - with trace_path.open("a", encoding="utf-8", newline="\n") as trace: - write_trace( - trace, - { - "event": "run_finish", - **contract, - "finished_at": finished_at.isoformat(), - "wall_seconds": wall_seconds, - "final_loss": losses[-1] if losses else None, - "initial_trainable_sha256": initial_trainable_sha, - "final_trainable_sha256": final_trainable_sha, - "optimizer_update_count": len(optimizer_updates), - "max_optimizer_update_l2": max_update_l2, - "loss_delta": loss_delta, - "peak_cuda_memory_mb": cuda_peak_memory_mb(torch), - "adapter_weight_sha256": weight_stats["sha256"], - }, - ) - manifest = AdapterManifest( - adapter_id=adapter_id(TrainingAdapterKind.hme_retrieval, dataset_path, model_id), - kind=TrainingAdapterKind.hme_retrieval, - created_at=datetime.now(UTC), - proof_version=TRAINING_PROOF_VERSION, - training_method=TRAINING_METHOD, - precision=TRAINING_PRECISION, - framework=TRAINING_FRAMEWORK, - base_model=model_id, - dataset_path=str(dataset_path.resolve()), - dataset_sha256=sha256_file(dataset_path), - sample_count=len(samples), - max_steps=max_steps, - max_length=max_length, - rank=rank, - alpha=alpha, - target_modules=target_modules, - trainable_params=int(trainable_params), - total_params=int(total_params), - optimizer=optimizer_name, - optimizer_family=optimizer_family(optimizer_name), - learning_rate=learning_rate, - gradient_accumulation_steps=1, - gradient_checkpointing=gradient_checkpointing, - gradient_checkpointing_mode=GRADIENT_CHECKPOINTING_MODE if gradient_checkpointing else None, - autocast_dtype=AUTOCAST_DTYPE, - bf16=torch.cuda.is_available(), - unsloth_version=package_version("unsloth"), - triton_version=package_version("triton-windows") or package_version("triton"), - xformers_version=package_version("xformers"), - unsloth_model_type=unsloth_model_type, - final_loss=losses[-1] if losses else None, - loss_history=losses, - initial_trainable_sha256=initial_trainable_sha, - final_trainable_sha256=final_trainable_sha, - initial_trainable_l2=initial_trainable_l2, - final_trainable_l2=final_trainable_l2, - optimizer_update_count=len(optimizer_updates), - max_optimizer_update_l2=max_update_l2, - loss_delta=loss_delta, - training_proof_path=str(proof_path.resolve()), - training_proof_sha256=sha256_file(proof_path), - training_trace_path=str(trace_path.resolve()), - training_trace_sha256=sha256_file(trace_path), - training_summary_path=str(summary_path.resolve()), - training_summary_sha256=sha256_file(summary_path), - training_started_at=started_at, - training_finished_at=finished_at, - training_wall_seconds=wall_seconds, - training_device=str(environment.get("device_name") or environment.get("device") or ""), - torch_version=str(environment.get("torch_version") or ""), - cuda_version=( - str(environment["torch_cuda_version"]) - if environment.get("torch_cuda_version") is not None - else None - ), - adapter_weight_sha256=str(weight_stats["sha256"] or ""), - adapter_tensor_count=int(weight_stats["tensor_count"]), - adapter_lora_tensor_count=int(weight_stats["lora_tensor_count"]), - adapter_total_elements=int(weight_stats["total_elements"]), - adapter_nonzero_elements=int(weight_stats["nonzero_elements"]), - adapter_max_abs=weight_stats["max_abs"], - adapter_weights_finite=bool(weight_stats["all_finite"]), - artifacts=relative_artifacts(output_dir), - sha256=sha256_tree(output_dir), - cartridge_id=cartridge_id, - ) - (output_dir / "adapter_manifest.json").write_text( - manifest.model_dump_json(indent=2), - encoding="utf-8", - ) - return manifest - - def _atoms_for_path(self, path: Path) -> list[EvidenceAtom]: - resolved = path.resolve() - atoms: list[EvidenceAtom] = [] - for atom in self.db.list_atoms(): - try: - if Path(atom.source).resolve() == resolved: - atoms.append(atom) - except OSError: - if atom.source == str(path): - atoms.append(atom) - return atoms - - def _latest_cartridge_id(self) -> str | None: - rows = self.db.list_cartridges() - return str(rows[0]["cartridge_id"]) if rows else None - - def _model_for_kind(self, kind: TrainingAdapterKind) -> str: - if kind == TrainingAdapterKind.hme_router: - return self.settings.planner_model_id - return self.settings.model_id - - def _atom_id_for_path(self, path: Path, text: str) -> str: - atoms = self._atoms_for_path(path) - return atoms[0].atom_id if atoms else f"{path.stem}:{sha256_bytes(text.encode('utf-8'))[:12]}" - - -def encode_training_sample( - sample: TrainingSample, - processor: Any, - device: Any, - max_length: int, -) -> dict[str, Any]: - messages = messages_for_sample(sample, include_answer=True) - prompt_messages = messages_for_sample(sample, include_answer=False) - image_inputs, video_inputs = vision_inputs(messages) - text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) - prompt_text = processor.apply_chat_template( - prompt_messages, - tokenize=False, - add_generation_prompt=True, - enable_thinking=False, - ) - inputs = processor( - text=[text], - images=image_inputs, - videos=video_inputs, - return_tensors="pt", - ) - prompt_inputs = processor( - text=[prompt_text], - images=image_inputs, - videos=video_inputs, - return_tensors="pt", - ) - labels = inputs["input_ids"].clone() - prompt_length = min(int(prompt_inputs["input_ids"].shape[-1]), int(labels.shape[-1])) - labels[:, :prompt_length] = -100 - pad_token_id = getattr(processor, "tokenizer", processor).pad_token_id - if pad_token_id is not None: - labels[labels == pad_token_id] = -100 - if int((labels != -100).sum().item()) == 0: - raise ValueError(f"Sample {sample.id} produced no trainable tokens") - inputs["labels"] = labels - if int(inputs["input_ids"].shape[-1]) > max_length: - for key in ("input_ids", "attention_mask", "labels"): - if key in inputs: - inputs[key] = inputs[key][:, -max_length:] - return {key: value.to(device) if hasattr(value, "to") else value for key, value in inputs.items()} - - -def encode_retrieval_batch( - texts: list[str], - processor: Any, - device: Any, - max_length: int, -) -> dict[str, Any]: - tokenizer = getattr(processor, "tokenizer", processor) - inputs = tokenizer( - texts, - padding=True, - truncation=True, - max_length=max_length, - return_tensors="pt", - ) - return {key: value.to(device) if hasattr(value, "to") else value for key, value in inputs.items()} - - -def pool_embeddings(outputs: Any, attention_mask: Any) -> Any: - import torch.nn.functional as F - - hidden = outputs.last_hidden_state - mask = attention_mask.unsqueeze(-1).to(hidden.dtype) - pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0) - return F.normalize(pooled, dim=-1) - - -def messages_for_sample(sample: TrainingSample, include_answer: bool) -> list[dict[str, Any]]: - user_content: list[dict[str, Any]] = [] - for item in sample.inputs: - item_type = item.get("type") - if item_type in {"image", "pdf_page", "ui_screenshot"} and item.get("path"): - path = str(item["path"]) - if path.startswith("object://"): - user_content.append({"type": "text", "text": f"Image object reference: {path}"}) - else: - image_part: dict[str, Any] = {"type": "image", "image": path} - for key in ("min_pixels", "max_pixels"): - if key in item: - image_part[key] = item[key] - user_content.append(image_part) - elif item_type == "text" and item.get("text"): - user_content.append({"type": "text", "text": str(item["text"])}) - else: - user_content.append({"type": "text", "text": json.dumps(item, ensure_ascii=False)}) - messages: list[dict[str, Any]] = [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": user_content}, - ] - if include_answer: - messages.append( - { - "role": "assistant", - "content": json.dumps(sample.output, ensure_ascii=False, separators=(",", ":")), - } - ) - return messages - - -def vision_inputs(messages: list[dict[str, Any]]) -> tuple[list[Any] | None, list[Any] | None]: - has_media = any( - isinstance(content, list) - and any(part.get("type") in {"image", "video"} for part in content if isinstance(part, dict)) - for content in (message.get("content") for message in messages) - ) - if not has_media: - return None, None - from qwen_vl_utils import process_vision_info - - return process_vision_info(messages) - - -def infer_lora_targets(model: Any, kind: TrainingAdapterKind) -> list[str]: - import torch - - core_leafs = {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"} - mlp_leafs = {"gate_proj", "up_proj", "down_proj"} - found_leafs: set[str] = set() - exact: set[str] = set() - for name, module in model.named_modules(): - if not isinstance(module, torch.nn.Linear): - continue - leaf = name.rsplit(".", 1)[-1] - lower = name.lower() - if kind == TrainingAdapterKind.hme_core_lora and leaf in core_leafs: - found_leafs.add(leaf) - elif kind == TrainingAdapterKind.hme_vision_lora and ("visual" in lower or "vision" in lower): - exact.add(name) - elif kind == TrainingAdapterKind.hme_ttt_vl and (leaf in mlp_leafs or "visual" in lower): - exact.add(name if "visual" in lower else leaf) - elif kind == TrainingAdapterKind.hme_visual_mtp and ( - "mtp" in lower or "draft" in lower or name == "lm_head" - ): - exact.add(name) - elif kind in {TrainingAdapterKind.hme_router, TrainingAdapterKind.hme_retrieval} and leaf in core_leafs: - found_leafs.add(leaf) - targets = sorted(exact or found_leafs) - if not targets: - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear) and not name.endswith("lm_head"): - targets.append(name.rsplit(".", 1)[-1]) - targets = sorted(set(targets)) - if not targets: - raise ValueError(f"No LoRA target modules found for {kind.value}") - return targets - - -def trainable_named_parameters(model: Any) -> list[tuple[str, Any]]: - return [(name, param) for name, param in model.named_parameters() if param.requires_grad] - - -def trainable_state_cpu(named_parameters: list[tuple[str, Any]]) -> dict[str, Any]: - return {name: param.detach().float().cpu().clone() for name, param in named_parameters} - - -def trainable_state_sha256(named_parameters: list[tuple[str, Any]]) -> str: - digest = hashlib.sha256() - for name, param in named_parameters: - tensor = param.detach().float().cpu().contiguous() - digest.update(name.encode("utf-8")) - digest.update(str(tuple(tensor.shape)).encode("utf-8")) - digest.update(tensor.numpy().tobytes()) - return digest.hexdigest() - - -def trainable_l2_norm(named_parameters: list[tuple[str, Any]]) -> float: - total = 0.0 - for _, param in named_parameters: - tensor = param.detach().float() - total += float(tensor.square().sum().detach().cpu()) - return math.sqrt(total) - - -def gradient_l2_norm(named_parameters: list[tuple[str, Any]]) -> float: - total = 0.0 - for _, param in named_parameters: - if param.grad is None: - continue - grad = param.grad.detach().float() - total += float(grad.square().sum().detach().cpu()) - return math.sqrt(total) - - -def trainable_update_l2(before_update: dict[str, Any], named_parameters: list[tuple[str, Any]]) -> float: - total = 0.0 - for name, param in named_parameters: - delta = param.detach().float().cpu() - before_update[name] - total += float(delta.square().sum()) - return math.sqrt(total) - - -def require_cuda_bf16(torch_module: Any) -> None: - if not torch_module.cuda.is_available(): - raise RuntimeError("training needs a CUDA device; none is available.") - bf16_supported = getattr(torch_module.cuda, "is_bf16_supported", lambda: False) - if not bool(bf16_supported()): - raise RuntimeError("PRD 12.2 training requires BF16-capable CUDA hardware.") - - -def activate_unsloth_training() -> None: - required = { - "unsloth": package_version("unsloth"), - "unsloth_zoo": package_version("unsloth_zoo"), - "triton": package_version("triton-windows") or package_version("triton"), - "xformers": package_version("xformers"), - } - missing = [name for name, installed in required.items() if not installed] - if missing: - raise RuntimeError( - "PRD 12.2 training requires installed Unsloth dependencies: " + ", ".join(missing) - ) - with bitsandbytes_gaudi_probe_guard(): - __import__("unsloth") - - -def enable_gradient_checkpointing(model: Any) -> bool: - if not hasattr(model, "gradient_checkpointing_enable"): - return False - try: - model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) - except TypeError: - model.gradient_checkpointing_enable() - return True - - -def model_type_for_proof(model: Any) -> str | None: - current = model - seen: set[int] = set() - while current is not None and id(current) not in seen: - seen.add(id(current)) - config = getattr(current, "config", None) - model_type = getattr(config, "model_type", None) - if isinstance(model_type, str) and model_type: - return model_type - current = getattr(current, "model", None) - return None - - -def make_optimizer(model: Any, learning_rate: float) -> tuple[Any, str]: - params = [param for param in model.parameters() if param.requires_grad] - try: - with bitsandbytes_gaudi_probe_guard(): - import bitsandbytes as bnb - except Exception as exc: - raise RuntimeError("PRD 12.2 training requires bitsandbytes PagedAdamW8bit.") from exc - try: - return bnb.optim.PagedAdamW8bit(params, lr=learning_rate), "paged_adamw_8bit" - except Exception as exc: - raise RuntimeError("bitsandbytes PagedAdamW8bit failed to initialize.") from exc - - -def optimizer_family(optimizer_name: str) -> str: - if optimizer_name == "paged_adamw_8bit": - return "8-bit AdamW / paged optimizer" - return optimizer_name - - -@contextmanager -def bitsandbytes_gaudi_probe_guard() -> Any: - original_run = subprocess.run - - def guarded_run(*args: Any, **kwargs: Any) -> Any: - command = args[0] if args else kwargs.get("args") - if isinstance(command, str) and "habana-torch-plugin" in command: - return subprocess.CompletedProcess(command, 0, "", "") - return original_run(*args, **kwargs) - - subprocess.run = guarded_run - try: - yield - finally: - subprocess.run = original_run - - -def training_contract( - *, - gradient_checkpointing: bool, - input_require_grads: bool, - optimizer_name: str, - unsloth_model_type: str | None, -) -> dict[str, Any]: - return { - "proof_version": TRAINING_PROOF_VERSION, - "training_method": TRAINING_METHOD, - "precision": TRAINING_PRECISION, - "framework": TRAINING_FRAMEWORK, - "optimizer_family": optimizer_family(optimizer_name), - "gradient_checkpointing": gradient_checkpointing, - "gradient_checkpointing_mode": GRADIENT_CHECKPOINTING_MODE if gradient_checkpointing else None, - "input_require_grads": input_require_grads, - "autocast_dtype": AUTOCAST_DTYPE, - "unsloth_version": package_version("unsloth"), - "triton_version": package_version("triton-windows") or package_version("triton"), - "xformers_version": package_version("xformers"), - "unsloth_model_type": unsloth_model_type, - } - - -def training_environment(torch_module: Any) -> dict[str, Any]: - cuda_available = bool(torch_module.cuda.is_available()) - device_name = torch_module.cuda.get_device_name(0) if cuda_available else platform.processor() - capability = torch_module.cuda.get_device_capability(0) if cuda_available else None - bf16_supported = ( - bool(getattr(torch_module.cuda, "is_bf16_supported", lambda: False)()) - if cuda_available - else False - ) - return { - "python": platform.python_version(), - "platform": platform.platform(), - "device": "cuda:0" if cuda_available else "cpu", - "device_name": device_name, - "cuda_available": cuda_available, - "bf16_supported": bf16_supported, - "cuda_capability": list(capability) if capability else None, - "torch_version": getattr(torch_module, "__version__", None), - "torch_cuda_version": getattr(torch_module.version, "cuda", None), - "transformers_version": package_version("transformers"), - "peft_version": package_version("peft"), - "bitsandbytes_version": package_version("bitsandbytes"), - "unsloth_version": package_version("unsloth"), - "unsloth_zoo_version": package_version("unsloth_zoo"), - "triton_version": package_version("triton-windows") or package_version("triton"), - "xformers_version": package_version("xformers"), - } - - -def package_version(package_name: str) -> str | None: - try: - return version(package_name) - except PackageNotFoundError: - return None - - -def cuda_peak_memory_mb(torch_module: Any) -> float | None: - if not torch_module.cuda.is_available(): - return None - return float(torch_module.cuda.max_memory_allocated() / 1024 / 1024) - - -def safetensor_weight_stats(path: Path) -> dict[str, Any]: - report: dict[str, Any] = { - "path": str(path), - "exists": path.exists(), - "sha256": sha256_file(path) if path.exists() else None, - "tensor_count": 0, - "lora_tensor_count": 0, - "total_elements": 0, - "nonzero_elements": 0, - "nonzero_tensor_count": 0, - "finite_tensor_count": 0, - "all_finite": False, - "max_abs": None, - "mean_abs": None, - "dtype_counts": {}, - "key_sample": [], - "valid": False, - } - if not path.exists(): - return report - - import torch - from safetensors import safe_open - - total_abs = 0.0 - max_abs = 0.0 - dtype_counts: dict[str, int] = {} - with safe_open(path, framework="pt", device="cpu") as handle: - keys = list(handle.keys()) - report["tensor_count"] = len(keys) - report["key_sample"] = keys[:12] - for key in keys: - tensor = handle.get_tensor(key) - numel = int(tensor.numel()) - dtype = str(tensor.dtype).replace("torch.", "") - dtype_counts[dtype] = dtype_counts.get(dtype, 0) + 1 - report["total_elements"] += numel - if "lora_" in key: - report["lora_tensor_count"] += 1 - if numel == 0: - continue - tensor_float = tensor.detach().float() - finite = bool(torch.isfinite(tensor_float).all().item()) - if finite: - report["finite_tensor_count"] += 1 - nonzero = int(torch.count_nonzero(tensor).item()) - report["nonzero_elements"] += nonzero - if nonzero > 0: - report["nonzero_tensor_count"] += 1 - abs_tensor = tensor_float.abs() - tensor_max_abs = float(abs_tensor.max().item()) - max_abs = max(max_abs, tensor_max_abs) - total_abs += float(abs_tensor.sum().item()) - - report["dtype_counts"] = dtype_counts - report["all_finite"] = ( - report["tensor_count"] > 0 and report["finite_tensor_count"] == report["tensor_count"] - ) - if int(report["total_elements"]) > 0: - report["max_abs"] = max_abs - report["mean_abs"] = total_abs / int(report["total_elements"]) - report["valid"] = ( - report["exists"] - and int(report["tensor_count"]) > 0 - and int(report["lora_tensor_count"]) > 0 - and int(report["total_elements"]) > 0 - and int(report["nonzero_elements"]) > 0 - and bool(report["all_finite"]) - and report["max_abs"] is not None - and float(report["max_abs"]) > 0.0 - ) - return report - - -def write_trace(handle: TextIO, payload: dict[str, Any]) -> None: - handle.write(json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n") - handle.flush() - - -def write_training_summary(path: Path, payload: dict[str, Any]) -> None: - path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") - - -def split_markdown_sections(text: str) -> list[dict[str, str]]: - matches = list(re.finditer(r"^(#{2,3})\s+(.+?)\s*$", text, flags=re.MULTILINE)) - sections: list[dict[str, str]] = [] - for index, match in enumerate(matches): - start = match.end() - end = matches[index + 1].start() if index + 1 < len(matches) else len(text) - sections.append({"title": match.group(2).strip(), "body": text[start:end].strip()}) - return sections - - -def section_by_title(text: str, title_fragment: str) -> str: - for section in split_markdown_sections(text): - if title_fragment in section["title"]: - return section["body"] - return "" - - -def write_jsonl(path: Path, samples: list[TrainingSample]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("w", encoding="utf-8", newline="\n") as handle: - for sample in samples: - handle.write(sample.model_dump_json() + "\n") - - -def prepare_training_image(source: Path, media_dir: Path, max_side: int = 896) -> Path: - from PIL import Image - - image = Image.open(source).convert("RGB") - image.thumbnail((max_side, max_side)) - target = media_dir / f"{source.stem}_{sha256_file(source)[:10]}.png" - image.save(target) - return target - - -def render_prd_training_card(text: str, atom_id: str, target: Path) -> Path: - from PIL import Image, ImageDraw, ImageFont - - target.parent.mkdir(parents=True, exist_ok=True) - image = Image.new("RGB", (1024, 768), "#f7f5ef") - draw = ImageDraw.Draw(image) - font_path = Path("C:/Windows/Fonts/msyh.ttc") - title_font = ImageFont.truetype(str(font_path), 34) if font_path.exists() else ImageFont.load_default() - body_font = ImageFont.truetype(str(font_path), 21) if font_path.exists() else ImageFont.load_default() - small_font = ImageFont.truetype(str(font_path), 18) if font_path.exists() else ImageFont.load_default() - draw.rectangle((28, 28, 996, 740), outline="#1f5f4a", width=3) - draw.text((52, 48), "Ely-Eye PRD Training Card", fill="#09201b", font=title_font) - draw.text((52, 100), f"Evidence Atom: {atom_id}", fill="#1f5f4a", font=small_font) - source = ( - section_by_title(text, "12.1 训练目标") - + "\n" - + section_by_title(text, "12.2 训练策略") - + "\n" - + section_by_title(text, "12.3 训练数据结构") - ) - y = 148 - for line in wrap_text(compact_text(source, 1500), width=42): - draw.text((52, y), line, fill="#09201b", font=body_font) - y += 31 - if y > 700: - break - image.save(target) - return target - - -def wrap_text(text: str, width: int) -> list[str]: - lines: list[str] = [] - current = "" - for char in text: - current += char - if len(current) >= width or char == "\n": - lines.append(current.strip()) - current = "" - if current.strip(): - lines.append(current.strip()) - return lines - - -def compact_text(text: str, limit: int) -> str: - compacted = re.sub(r"\s+", " ", text).strip() - if len(compacted) <= limit: - return compacted - return compacted[: limit - 3].rstrip() + "..." - - -def slugify(value: str) -> str: - slug = re.sub(r"[^a-zA-Z0-9\u4e00-\u9fff]+", "-", value).strip("-").lower() - return slug[:80] or "section" - - -def adapter_id(kind: TrainingAdapterKind, dataset_path: Path, model_id: str) -> str: - raw = f"{kind.value}:{model_id}:{dataset_path.resolve()}:{datetime.now(UTC).isoformat()}" - return f"{kind.value}_{sha256_bytes(raw.encode('utf-8'))[:12]}" - - -def relative_artifacts(root: Path) -> dict[str, str]: - artifacts: dict[str, str] = {} - for path in sorted(root.rglob("*")): - if path.is_file() and path.name != "adapter_manifest.json": - artifacts[path.stem] = path.relative_to(root).as_posix() - return artifacts - - -def sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - shutil.copyfileobj(handle, _DigestWriter(digest)) - return digest.hexdigest() - - -def sha256_tree(root: Path) -> str: - digest = hashlib.sha256() - for path in sorted(root.rglob("*")): - if path.is_file() and path.name != "adapter_manifest.json": - digest.update(path.relative_to(root).as_posix().encode("utf-8")) - with path.open("rb") as handle: - shutil.copyfileobj(handle, _DigestWriter(digest)) - return digest.hexdigest() - - -def sha256_bytes(data: bytes) -> str: - return hashlib.sha256(data).hexdigest() - - -class _DigestWriter: - def __init__(self, digest: hashlib._Hash) -> None: - self.digest = digest - - def write(self, data: bytes) -> int: - self.digest.update(data) - return len(data) diff --git a/backend/src/ely_eye/training_common.py b/backend/src/ely_eye/training_common.py new file mode 100644 index 0000000..03b70d4 --- /dev/null +++ b/backend/src/ely_eye/training_common.py @@ -0,0 +1,108 @@ +"""Shared constants, digests and trace writers: the vocabulary every training module speaks.""" + +from __future__ import annotations + +import hashlib +import json +import re +import shutil +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, TextIO + +from pydantic import BaseModel, Field + +from .schemas import TrainingAdapterKind + +SYSTEM_PROMPT = ( + "You are Ely-Eye. Answer from supplied evidence, return compact JSON, " + "and cite atom ids exactly." +) + +TRAINING_PROOF_VERSION = 2 + +TRAINING_METHOD = "BF16 LoRA" + +TRAINING_PRECISION = "bf16" + +TRAINING_FRAMEWORK = "unsloth_patched_transformers_peft" + +GRADIENT_CHECKPOINTING_MODE = "unsloth_patched_torch_non_reentrant" + +AUTOCAST_DTYPE = "bf16" + +class TrainingSample(BaseModel): + id: str + task: str + inputs: list[dict[str, Any]] + memory: dict[str, Any] = Field(default_factory=dict) + output: dict[str, Any] + +class _DigestWriter: + def __init__(self, digest: hashlib._Hash) -> None: + self.digest = digest + + def write(self, data: bytes) -> int: + self.digest.update(data) + return len(data) + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + shutil.copyfileobj(handle, _DigestWriter(digest)) + return digest.hexdigest() + +def sha256_tree(root: Path) -> str: + digest = hashlib.sha256() + for path in sorted(root.rglob("*")): + if path.is_file() and path.name != "adapter_manifest.json": + digest.update(path.relative_to(root).as_posix().encode("utf-8")) + with path.open("rb") as handle: + shutil.copyfileobj(handle, _DigestWriter(digest)) + return digest.hexdigest() + +def sha256_bytes(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + +def slugify(value: str) -> str: + slug = re.sub(r"[^a-zA-Z0-9\u4e00-\u9fff]+", "-", value).strip("-").lower() + return slug[:80] or "section" + +def adapter_id(kind: TrainingAdapterKind, dataset_path: Path, model_id: str) -> str: + raw = f"{kind.value}:{model_id}:{dataset_path.resolve()}:{datetime.now(UTC).isoformat()}" + return f"{kind.value}_{sha256_bytes(raw.encode('utf-8'))[:12]}" + +def relative_artifacts(root: Path) -> dict[str, str]: + artifacts: dict[str, str] = {} + for path in sorted(root.rglob("*")): + if path.is_file() and path.name != "adapter_manifest.json": + artifacts[path.stem] = path.relative_to(root).as_posix() + return artifacts + +def write_jsonl(path: Path, samples: list[TrainingSample]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8", newline="\n") as handle: + for sample in samples: + handle.write(sample.model_dump_json() + "\n") + +def write_trace(handle: TextIO, payload: dict[str, Any]) -> None: + handle.write(json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n") + handle.flush() + +def write_training_summary(path: Path, payload: dict[str, Any]) -> None: + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") + +def split_markdown_sections(text: str) -> list[dict[str, str]]: + matches = list(re.finditer(r"^(#{2,3})\s+(.+?)\s*$", text, flags=re.MULTILINE)) + sections: list[dict[str, str]] = [] + for index, match in enumerate(matches): + start = match.end() + end = matches[index + 1].start() if index + 1 < len(matches) else len(text) + sections.append({"title": match.group(2).strip(), "body": text[start:end].strip()}) + return sections + +def section_by_title(text: str, title_fragment: str) -> str: + for section in split_markdown_sections(text): + if title_fragment in section["title"]: + return section["body"] + return "" diff --git a/backend/src/ely_eye/training_datasets.py b/backend/src/ely_eye/training_datasets.py new file mode 100644 index 0000000..9b4df4a --- /dev/null +++ b/backend/src/ely_eye/training_datasets.py @@ -0,0 +1,387 @@ +"""Turn a store and a PRD into the SFT and retrieval datasets the adapters train on.""" + +from __future__ import annotations + +from pathlib import Path + +from .config import Settings +from .db import Database +from .schemas import EvidenceAtom, TrainingAdapterKind +from .training_common import ( + TrainingSample, + section_by_title, + sha256_bytes, + slugify, + split_markdown_sections, + write_jsonl, +) +from .training_encode import compact_text, prepare_training_image, render_prd_training_card + + +def validate_dataset(path: Path) -> list[TrainingSample]: + samples: list[TrainingSample] = [] + with path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + if not line.strip(): + continue + try: + samples.append(TrainingSample.model_validate_json(line)) + except Exception as exc: + raise ValueError(f"Invalid training sample at line {line_number}: {exc}") from exc + if not samples: + raise ValueError(f"Training dataset is empty: {path}") + return samples + +def atoms_for_path(database: Database, path: Path) -> list[EvidenceAtom]: + resolved = path.resolve() + atoms: list[EvidenceAtom] = [] + for atom in database.list_atoms(): + try: + if Path(atom.source).resolve() == resolved: + atoms.append(atom) + except OSError: + if atom.source == str(path): + atoms.append(atom) + return atoms + +def atom_id_for_path(database: Database, path: Path, text: str) -> str: + atoms = atoms_for_path(database, path) + return atoms[0].atom_id if atoms else f"{path.stem}:{sha256_bytes(text.encode('utf-8'))[:12]}" + +def latest_cartridge_id(database: Database) -> str | None: + rows = database.list_cartridges() + return str(rows[0]["cartridge_id"]) if rows else None + +def model_for_kind(settings: Settings, kind: TrainingAdapterKind) -> str: + if kind == TrainingAdapterKind.hme_router: + return settings.planner_model_id + return settings.model_id + +def build_prd_dataset( + settings: Settings, + database: Database, + prd_path: Path, + output_path: Path | None = None, + cartridge_id: str | None = None, + max_sections: int = 48, +) -> Path: + prd_path = prd_path.resolve() + text = prd_path.read_text(encoding="utf-8") + atoms = atoms_for_path(database, prd_path) + atom_id = atoms[0].atom_id if atoms else f"{prd_path.stem}:{sha256_bytes(text.encode('utf-8'))[:12]}" + cartridge_id = cartridge_id or latest_cartridge_id(database) + sections = split_markdown_sections(text) + if not sections: + raise ValueError(f"No markdown sections found in {prd_path}") + output_path = output_path or (settings.training_dir / "hme_prd_sft.jsonl") + output_path.parent.mkdir(parents=True, exist_ok=True) + + samples: list[TrainingSample] = [] + for index, section in enumerate(sections[:max_sections], start=1): + title = section["title"] + body = compact_text(section["body"], limit=5200) + if len(body) < 160: + continue + citation = f"{atom_id}#section={slugify(title)}" + answer = compact_text(body, limit=900) + samples.append( + TrainingSample( + id=f"hme_prd_{index:04d}", + task="hme_core_lora_evidence_json", + inputs=[ + { + "type": "text", + "text": ( + f"Section title: {title}\n" + f"Evidence atom: {atom_id}\n" + f"Evidence:\n{body}\n\n" + "Question: State the Ely-Eye requirement in this section. " + "Return JSON with answer, citations, contradiction_notes, confidence." + ), + } + ], + memory={ + "cartridge": cartridge_id, + "required_atoms": [atom_id], + "source": str(prd_path), + }, + output={ + "answer": answer, + "citations": [citation], + "contradiction_notes": [], + "confidence": 1.0, + "reasoning_policy": "evidence-first", + }, + ) + ) + + if not samples: + raise ValueError(f"No usable PRD training sections found in {prd_path}") + with output_path.open("w", encoding="utf-8", newline="\n") as handle: + for sample in samples: + handle.write(sample.model_dump_json() + "\n") + return output_path + +def build_visual_dataset( + settings: Settings, + database: Database, + prd_path: Path, + output_path: Path | None = None, + image_paths: list[Path] | None = None, + cartridge_id: str | None = None, +) -> Path: + prd_path = prd_path.resolve() + text = prd_path.read_text(encoding="utf-8") + atom_id = atom_id_for_path(database, prd_path, text) + cartridge_id = cartridge_id or latest_cartridge_id(database) + media_dir = settings.training_dir / "media" + media_dir.mkdir(parents=True, exist_ok=True) + + prepared_images: list[Path] = [] + for image_path in image_paths or []: + if image_path.exists(): + prepared_images.append(prepare_training_image(image_path, media_dir)) + prd_card = render_prd_training_card(text, atom_id, media_dir / "prd_training_card.png") + prepared_images.append(prd_card) + if not prepared_images: + raise ValueError("No visual training images are available") + + output_path = output_path or (settings.training_dir / "hme_visual_sft.jsonl") + samples: list[TrainingSample] = [] + visual_tasks = [ + ( + "hme_vision_lora_visual_grounding", + "12.1 训练目标", + ( + "Ely-Eye trains HME-Vision-LoRA on Qwen3.5-9B vision projector " + "and selected vision layers so UI, PDF, OCR, and chart evidence can be " + "grounded into cited multimodal memory." + ), + ), + ( + "hme_ttt_vl_visual_memory", + "11.2 X — TTT-VL：视觉语言测试时学习", + ( + "TTT-VL keeps Qwen3.5-9B frozen while mutable MLP and visual LoRA " + "adapters learn low-step visual memory updates that can be consolidated " + "into cartridge-local adapter deltas." + ), + ), + ( + "hme_visual_hashhop_grounding", + "11.3 X — Visual HashHop", + ( + "Visual HashHop converts random chain ids into non-text geometric image " + "blocks, then evaluates long-range visual addressing without OCR shortcuts." + ), + ), + ( + "hme_dvkv_visual_cache", + "11.4 S/H — DVKV：Diversity-Aware Visual KV Compression", + ( + "DVKV keeps visual KV tokens by combining attention importance, visual " + "diversity, and evidence value so OCR, layout, chart, and coordinate details " + "survive compression." + ), + ), + ( + "hme_visual_contradiction_lens", + "11.15 X — Visual Contradiction Lens", + ( + "The Visual Contradiction Lens compares visual versions and cites layout, " + "design-token, copy, temporal, and visual-code drift as evidence graph facts." + ), + ), + ] + for index, image_path in enumerate(prepared_images, start=1): + for task_index, (task, section_title, answer) in enumerate(visual_tasks, start=1): + requirement = section_by_title(text, section_title) + samples.append( + TrainingSample( + id=f"hme_visual_{index:02d}_{task_index:02d}", + task=task, + inputs=[ + { + "type": "ui_screenshot", + "path": str(image_path.resolve()), + "max_pixels": 262144, + }, + { + "type": "text", + "text": ( + f"Evidence atom: {atom_id}\n" + f"Visual source: {image_path.name}\n" + f"PRD visual requirement:\n" + f"{compact_text(requirement, 2200)}\n\n" + "Question: Map this visual evidence to the Ely-Eye PRD module. " + "Return JSON with answer, citations, contradiction_notes, confidence." + ), + }, + ], + memory={ + "cartridge": cartridge_id, + "required_atoms": [atom_id], + "source": str(prd_path), + }, + output={ + "answer": answer, + "citations": [f"{atom_id}#section={slugify(section_title)}"], + "contradiction_notes": [], + "confidence": 1.0, + "reasoning_policy": "evidence-first", + }, + ) + ) + write_jsonl(output_path, samples) + return output_path + +def build_router_dataset( + settings: Settings, + database: Database, + prd_path: Path, + output_path: Path | None = None, + cartridge_id: str | None = None, +) -> Path: + prd_path = prd_path.resolve() + text = prd_path.read_text(encoding="utf-8") + atom_id = atom_id_for_path(database, prd_path, text) + cartridge_id = cartridge_id or latest_cartridge_id(database) + output_path = output_path or (settings.training_dir / "hme_router_sft.jsonl") + scenarios = [ + ( + "Summarize a stable live-demo answer from local PRD evidence.", + "live_demo", + 262144, + ["mixed"], + "Prefix cache + sparse evidence packing", + ), + ( + "Plan a 1.01M extreme context run with YaRN and parked context.", + "extreme_context", + 1010000, + ["mixed"], + "YaRN + V-NSA sparse blocks + DVKV visual KV compression + CPU/NVMe parked context", + ), + ( + "Route a 100M library question through Context Cartridge memory.", + "library_100m", + 100000000, + ["mixed"], + "Context Cartridge retrieval + ColQwen-compatible dense vectors + BM25 + evidence graph pack", + ), + ( + "Audit UI screenshot and PRD differences for the research demo.", + "research_theater", + 262144, + ["image", "ui_screenshot"], + "Visual patch cache + DVKV evidence-biased region retention", + ), + ( + "Find visual token compression methods across a 100M library cartridge.", + "library_100m", + 100000000, + ["pdf_page", "image", "video_frame"], + "ColQwen-compatible retrieval + graph grouping + cited evidence pack", + ), + ( + "Compare UI versions and cite coordinates plus code evidence.", + "research_theater", + 262144, + ["ui_screenshot", "code", "text"], + "Visual Contradiction Lens + evidence graph writeback", + ), + ( + "Run Visual HashHop at 262K with a two-hop random image chain.", + "research_theater", + 262144, + ["image"], + "Visual HashHop arena + verifier contract + no OCR-dependent evidence", + ), + ( + "Answer from a TTT-VL cartridge after original images are absent.", + "library_100m", + 100000000, + ["mixed"], + "Load cartridge adapters + learned memory recall + cited adapter proof", + ), + ] + samples = [ + TrainingSample( + id=f"hme_router_{index:04d}", + task="hme_router_context_plan", + inputs=[ + { + "type": "text", + "text": ( + f"Evidence atom: {atom_id}\n" + f"PRD routing source:\n{compact_text(section_by_title(text, '1. 产品定义'), 1800)}\n\n" + f"User request: {request}\n" + "Return the Ely-Eye routing plan as compact JSON." + ), + } + ], + memory={ + "cartridge": cartridge_id, + "required_atoms": [atom_id], + "source": str(prd_path), + }, + output={ + "profile": profile, + "token_budget": token_budget, + "required_modalities": modalities, + "retrieval_query": request, + "compression_strategy": compression, + "citations": [f"{atom_id}#section=1-产品定义"], + "confidence": 1.0, + }, + ) + for index, (request, profile, token_budget, modalities, compression) in enumerate(scenarios, start=1) + ] + write_jsonl(output_path, samples) + return output_path + +def build_retrieval_dataset( + settings: Settings, + database: Database, + prd_path: Path, + output_path: Path | None = None, + cartridge_id: str | None = None, + max_sections: int = 32, +) -> Path: + prd_path = prd_path.resolve() + text = prd_path.read_text(encoding="utf-8") + atom_id = atom_id_for_path(database, prd_path, text) + cartridge_id = cartridge_id or latest_cartridge_id(database) + sections = [section for section in split_markdown_sections(text) if len(section["body"]) >= 160] + output_path = output_path or (settings.training_dir / "hme_retrieval_pairs.jsonl") + samples: list[TrainingSample] = [] + for index, section in enumerate(sections[:max_sections], start=1): + title = section["title"] + positive = compact_text(section["body"], 1600) + samples.append( + TrainingSample( + id=f"hme_retrieval_{index:04d}", + task="hme_retrieval_contrastive_pair", + inputs=[ + { + "type": "text", + "text": f"Find the PRD section that answers: {title}", + } + ], + memory={ + "cartridge": cartridge_id, + "required_atoms": [atom_id], + "source": str(prd_path), + }, + output={ + "positive": positive, + "positive_title": title, + "citations": [f"{atom_id}#section={slugify(title)}"], + "confidence": 1.0, + }, + ) + ) + if len(samples) < 2: + raise ValueError("Retrieval contrastive training needs at least two PRD sections") + write_jsonl(output_path, samples) + return output_path diff --git a/backend/src/ely_eye/training_encode.py b/backend/src/ely_eye/training_encode.py new file mode 100644 index 0000000..14d659d --- /dev/null +++ b/backend/src/ely_eye/training_encode.py @@ -0,0 +1,177 @@ +"""Turn a training sample into model inputs: chat text, images, pooled embeddings.""" + +from __future__ import annotations + +import json +import re +from pathlib import Path +from typing import Any + +from .training_common import SYSTEM_PROMPT, TrainingSample, section_by_title, sha256_file + + +def encode_training_sample( + sample: TrainingSample, + processor: Any, + device: Any, + max_length: int, +) -> dict[str, Any]: + messages = messages_for_sample(sample, include_answer=True) + prompt_messages = messages_for_sample(sample, include_answer=False) + image_inputs, video_inputs = vision_inputs(messages) + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + prompt_text = processor.apply_chat_template( + prompt_messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + inputs = processor( + text=[text], + images=image_inputs, + videos=video_inputs, + return_tensors="pt", + ) + prompt_inputs = processor( + text=[prompt_text], + images=image_inputs, + videos=video_inputs, + return_tensors="pt", + ) + labels = inputs["input_ids"].clone() + prompt_length = min(int(prompt_inputs["input_ids"].shape[-1]), int(labels.shape[-1])) + labels[:, :prompt_length] = -100 + pad_token_id = getattr(processor, "tokenizer", processor).pad_token_id + if pad_token_id is not None: + labels[labels == pad_token_id] = -100 + if int((labels != -100).sum().item()) == 0: + raise ValueError(f"Sample {sample.id} produced no trainable tokens") + inputs["labels"] = labels + if int(inputs["input_ids"].shape[-1]) > max_length: + for key in ("input_ids", "attention_mask", "labels"): + if key in inputs: + inputs[key] = inputs[key][:, -max_length:] + return {key: value.to(device) if hasattr(value, "to") else value for key, value in inputs.items()} + +def encode_retrieval_batch( + texts: list[str], + processor: Any, + device: Any, + max_length: int, +) -> dict[str, Any]: + tokenizer = getattr(processor, "tokenizer", processor) + inputs = tokenizer( + texts, + padding=True, + truncation=True, + max_length=max_length, + return_tensors="pt", + ) + return {key: value.to(device) if hasattr(value, "to") else value for key, value in inputs.items()} + +def pool_embeddings(outputs: Any, attention_mask: Any) -> Any: + import torch.nn.functional as F + + hidden = outputs.last_hidden_state + mask = attention_mask.unsqueeze(-1).to(hidden.dtype) + pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0) + return F.normalize(pooled, dim=-1) + +def messages_for_sample(sample: TrainingSample, include_answer: bool) -> list[dict[str, Any]]: + user_content: list[dict[str, Any]] = [] + for item in sample.inputs: + item_type = item.get("type") + if item_type in {"image", "pdf_page", "ui_screenshot"} and item.get("path"): + path = str(item["path"]) + if path.startswith("object://"): + user_content.append({"type": "text", "text": f"Image object reference: {path}"}) + else: + image_part: dict[str, Any] = {"type": "image", "image": path} + for key in ("min_pixels", "max_pixels"): + if key in item: + image_part[key] = item[key] + user_content.append(image_part) + elif item_type == "text" and item.get("text"): + user_content.append({"type": "text", "text": str(item["text"])}) + else: + user_content.append({"type": "text", "text": json.dumps(item, ensure_ascii=False)}) + messages: list[dict[str, Any]] = [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": user_content}, + ] + if include_answer: + messages.append( + { + "role": "assistant", + "content": json.dumps(sample.output, ensure_ascii=False, separators=(",", ":")), + } + ) + return messages + +def vision_inputs(messages: list[dict[str, Any]]) -> tuple[list[Any] | None, list[Any] | None]: + has_media = any( + isinstance(content, list) + and any(part.get("type") in {"image", "video"} for part in content if isinstance(part, dict)) + for content in (message.get("content") for message in messages) + ) + if not has_media: + return None, None + from qwen_vl_utils import process_vision_info + + return process_vision_info(messages) + +def prepare_training_image(source: Path, media_dir: Path, max_side: int = 896) -> Path: + from PIL import Image + + image = Image.open(source).convert("RGB") + image.thumbnail((max_side, max_side)) + target = media_dir / f"{source.stem}_{sha256_file(source)[:10]}.png" + image.save(target) + return target + +def render_prd_training_card(text: str, atom_id: str, target: Path) -> Path: + from PIL import Image, ImageDraw, ImageFont + + target.parent.mkdir(parents=True, exist_ok=True) + image = Image.new("RGB", (1024, 768), "#f7f5ef") + draw = ImageDraw.Draw(image) + font_path = Path("C:/Windows/Fonts/msyh.ttc") + title_font = ImageFont.truetype(str(font_path), 34) if font_path.exists() else ImageFont.load_default() + body_font = ImageFont.truetype(str(font_path), 21) if font_path.exists() else ImageFont.load_default() + small_font = ImageFont.truetype(str(font_path), 18) if font_path.exists() else ImageFont.load_default() + draw.rectangle((28, 28, 996, 740), outline="#1f5f4a", width=3) + draw.text((52, 48), "Ely-Eye PRD Training Card", fill="#09201b", font=title_font) + draw.text((52, 100), f"Evidence Atom: {atom_id}", fill="#1f5f4a", font=small_font) + source = ( + section_by_title(text, "12.1 训练目标") + + "\n" + + section_by_title(text, "12.2 训练策略") + + "\n" + + section_by_title(text, "12.3 训练数据结构") + ) + y = 148 + for line in wrap_text(compact_text(source, 1500), width=42): + draw.text((52, y), line, fill="#09201b", font=body_font) + y += 31 + if y > 700: + break + image.save(target) + return target + +def wrap_text(text: str, width: int) -> list[str]: + lines: list[str] = [] + current = "" + for char in text: + current += char + if len(current) >= width or char == "\n": + lines.append(current.strip()) + current = "" + if current.strip(): + lines.append(current.strip()) + return lines + +def compact_text(text: str, limit: int) -> str: + compacted = re.sub(r"\s+", " ", text).strip() + if len(compacted) <= limit: + return compacted + return compacted[: limit - 3].rstrip() + "..." diff --git a/backend/src/ely_eye/training_lora.py b/backend/src/ely_eye/training_lora.py new file mode 100644 index 0000000..97344ea --- /dev/null +++ b/backend/src/ely_eye/training_lora.py @@ -0,0 +1,408 @@ +"""Train a LoRA adapter on a frozen dataset and write the manifest that proves it.""" + +from __future__ import annotations + +import json +import time +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from .config import Settings +from .db import Database +from .schemas import AdapterManifest, TrainingAdapterKind +from .training_common import ( + AUTOCAST_DTYPE, + GRADIENT_CHECKPOINTING_MODE, + TRAINING_FRAMEWORK, + TRAINING_METHOD, + TRAINING_PRECISION, + TRAINING_PROOF_VERSION, + adapter_id, + relative_artifacts, + sha256_file, + sha256_tree, + write_trace, + write_training_summary, +) +from .training_datasets import model_for_kind, validate_dataset +from .training_encode import encode_training_sample +from .training_metrics import ( + activate_unsloth_training, + bitsandbytes_gaudi_probe_guard, + cuda_peak_memory_mb, + enable_gradient_checkpointing, + gradient_l2_norm, + infer_lora_targets, + make_optimizer, + model_type_for_proof, + optimizer_family, + package_version, + require_cuda_bf16, + safetensor_weight_stats, + trainable_l2_norm, + trainable_named_parameters, + trainable_state_cpu, + trainable_state_sha256, + trainable_update_l2, + training_contract, + training_environment, +) + + +def write_lora_config( + settings: Settings, + output_dir: Path, + kind: TrainingAdapterKind = TrainingAdapterKind.hme_core_lora, + rank: int = 16, + alpha: int = 32, + model_id: str | None = None, +) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + config = { + "base_model": model_id or settings.model_id, + "adapter_kind": kind.value, + "method": "BF16 LoRA", + "framework": TRAINING_FRAMEWORK, + "rank": rank, + "alpha": alpha, + "gradient_checkpointing": True, + "gradient_checkpointing_mode": GRADIENT_CHECKPOINTING_MODE, + "precision": "bf16", + "optimizer": "paged_adamw_8bit", + "ensure_weight_tying": kind == TrainingAdapterKind.hme_visual_mtp, + } + path = output_dir / "lora_config.json" + path.write_text(json.dumps(config, ensure_ascii=False, indent=2), encoding="utf-8") + return path + +def train_lora( + settings: Settings, + database: Database, + dataset_path: Path, + output_dir: Path, + max_steps: int, + kind: TrainingAdapterKind = TrainingAdapterKind.hme_core_lora, + model_id: str | None = None, + rank: int = 16, + alpha: int = 32, + max_length: int | None = None, + gradient_accumulation_steps: int | None = None, + learning_rate: float | None = None, + cartridge_id: str | None = None, +) -> AdapterManifest: + samples = validate_dataset(dataset_path) + output_dir.mkdir(parents=True, exist_ok=True) + model_id = model_id or model_for_kind(settings, kind) + max_length = max_length or settings.train_max_length + gradient_accumulation_steps = ( + gradient_accumulation_steps or settings.train_gradient_accumulation_steps + ) + learning_rate = learning_rate or settings.train_learning_rate + started_at = datetime.now(UTC) + wall_start = time.perf_counter() + trace_path = output_dir / "training_trace.jsonl" + + import torch + + activate_unsloth_training() + from peft import LoraConfig, get_peft_model + from transformers import AutoProcessor + + try: + from transformers import AutoModelForImageTextToText + + model_cls = AutoModelForImageTextToText + except ImportError: + from transformers import AutoModelForCausalLM + + model_cls = AutoModelForCausalLM + + require_cuda_bf16(torch) + processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) + dtype = torch.bfloat16 + device_map: dict[str, int] = {"": 0} + model = model_cls.from_pretrained( + model_id, + dtype=dtype, + device_map=device_map, + trust_remote_code=True, + ) + model.config.use_cache = False + input_require_grads = False + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + input_require_grads = True + + target_modules = infer_lora_targets(model, kind) + gradient_checkpointing = enable_gradient_checkpointing(model) + peft_config_options: dict[str, Any] = { + "r": rank, + "lora_alpha": alpha, + "target_modules": target_modules, + "lora_dropout": 0.0, + "bias": "none", + "task_type": "CAUSAL_LM", + } + if kind == TrainingAdapterKind.hme_visual_mtp: + peft_config_options["ensure_weight_tying"] = True + with bitsandbytes_gaudi_probe_guard(): + model = get_peft_model(model, LoraConfig(**peft_config_options)) + unsloth_model_type = model_type_for_proof(model) + trainable_params, total_params = model.get_nb_trainable_parameters() + optimizer, optimizer_name = make_optimizer(model, learning_rate) + model.train() + + losses: list[float] = [] + optimizer_updates: list[dict[str, Any]] = [] + trainable_tensors = trainable_named_parameters(model) + initial_trainable_sha = trainable_state_sha256(trainable_tensors) + initial_trainable_l2 = trainable_l2_norm(trainable_tensors) + optimizer.zero_grad(set_to_none=True) + environment = training_environment(torch) + contract = training_contract( + gradient_checkpointing=gradient_checkpointing, + input_require_grads=input_require_grads, + optimizer_name=optimizer_name, + unsloth_model_type=unsloth_model_type, + ) + with trace_path.open("w", encoding="utf-8", newline="\n") as trace: + write_trace( + trace, + { + "event": "run_start", + **contract, + "started_at": started_at.isoformat(), + "adapter_kind": kind.value, + "base_model": model_id, + "dataset_path": str(dataset_path.resolve()), + "dataset_sha256": sha256_file(dataset_path), + "sample_count": len(samples), + "max_steps": max_steps, + "max_length": max_length, + "rank": rank, + "alpha": alpha, + "gradient_accumulation_steps": gradient_accumulation_steps, + "learning_rate": learning_rate, + "target_modules": target_modules, + "trainable_params": int(trainable_params), + "total_params": int(total_params), + "optimizer": optimizer_name, + "environment": environment, + }, + ) + for step in range(max_steps): + sample = samples[step % len(samples)] + batch = encode_training_sample(sample, processor, model.device, max_length) + with torch.autocast( + device_type="cuda", + dtype=torch.bfloat16, + enabled=torch.cuda.is_available(), + ): + outputs = model(**batch) + loss = outputs.loss / gradient_accumulation_steps + loss.backward() + optimizer_step = (step + 1) % gradient_accumulation_steps == 0 or (step + 1) == max_steps + grad_l2 = gradient_l2_norm(trainable_tensors) + update_l2 = None + if optimizer_step: + before_update = trainable_state_cpu(trainable_tensors) + optimizer.step() + update_l2 = trainable_update_l2(before_update, trainable_tensors) + optimizer_updates.append( + { + "step": step + 1, + "gradient_l2": grad_l2, + "update_l2": update_l2, + } + ) + optimizer.zero_grad(set_to_none=True) + actual_loss = float(loss.detach().cpu()) * gradient_accumulation_steps + losses.append(actual_loss) + write_trace( + trace, + { + "event": "step", + "step": step + 1, + "sample_id": sample.id, + "task": sample.task, + "loss": actual_loss, + "optimizer_step": optimizer_step, + "gradient_l2": grad_l2, + "update_l2": update_l2, + }, + ) + + final_trainable_sha = trainable_state_sha256(trainable_tensors) + final_trainable_l2 = trainable_l2_norm(trainable_tensors) + max_update_l2 = max((item["update_l2"] for item in optimizer_updates), default=None) + loss_delta = losses[0] - losses[-1] if len(losses) >= 2 else None + model.save_pretrained(output_dir) + processor.save_pretrained(output_dir) + write_lora_config(settings, output_dir, kind, rank, alpha, model_id) + weight_stats = safetensor_weight_stats(output_dir / "adapter_model.safetensors") + finished_at = datetime.now(UTC) + wall_seconds = time.perf_counter() - wall_start + summary_path = output_dir / "training_summary.json" + proof_path = output_dir / "training_proof.json" + training_proof = { + **contract, + "adapter_kind": kind.value, + "base_model": model_id, + "dataset_path": str(dataset_path.resolve()), + "dataset_sha256": sha256_file(dataset_path), + "sample_count": len(samples), + "max_steps": max_steps, + "max_length": max_length, + "rank": rank, + "alpha": alpha, + "target_modules": target_modules, + "trainable_params": int(trainable_params), + "total_params": int(total_params), + "optimizer": optimizer_name, + "learning_rate": learning_rate, + "gradient_accumulation_steps": gradient_accumulation_steps, + "bf16": bool(torch.cuda.is_available()), + "environment": environment, + "initial_trainable_sha256": initial_trainable_sha, + "final_trainable_sha256": final_trainable_sha, + "initial_trainable_l2": initial_trainable_l2, + "final_trainable_l2": final_trainable_l2, + "optimizer_update_count": len(optimizer_updates), + "max_optimizer_update_l2": max_update_l2, + "loss_delta": loss_delta, + "loss_history": losses, + "optimizer_updates": optimizer_updates, + "adapter_weight_sha256": weight_stats["sha256"], + "adapter_weight_stats": weight_stats, + } + write_training_summary(proof_path, training_proof) + write_training_summary( + summary_path, + { + **contract, + "adapter_kind": kind.value, + "base_model": model_id, + "dataset_path": str(dataset_path.resolve()), + "dataset_sha256": sha256_file(dataset_path), + "sample_count": len(samples), + "max_steps": max_steps, + "max_length": max_length, + "rank": rank, + "alpha": alpha, + "target_modules": target_modules, + "trainable_params": int(trainable_params), + "total_params": int(total_params), + "optimizer": optimizer_name, + "learning_rate": learning_rate, + "gradient_accumulation_steps": gradient_accumulation_steps, + "bf16": bool(torch.cuda.is_available()), + "loss_history": losses, + "final_loss": losses[-1] if losses else None, + "initial_trainable_sha256": initial_trainable_sha, + "final_trainable_sha256": final_trainable_sha, + "initial_trainable_l2": initial_trainable_l2, + "final_trainable_l2": final_trainable_l2, + "optimizer_update_count": len(optimizer_updates), + "max_optimizer_update_l2": max_update_l2, + "loss_delta": loss_delta, + "started_at": started_at.isoformat(), + "finished_at": finished_at.isoformat(), + "wall_seconds": wall_seconds, + "environment": environment, + "peak_cuda_memory_mb": cuda_peak_memory_mb(torch), + "adapter_weight_sha256": weight_stats["sha256"], + "adapter_weight_stats": weight_stats, + }, + ) + with trace_path.open("a", encoding="utf-8", newline="\n") as trace: + write_trace( + trace, + { + "event": "run_finish", + **contract, + "finished_at": finished_at.isoformat(), + "wall_seconds": wall_seconds, + "final_loss": losses[-1] if losses else None, + "initial_trainable_sha256": initial_trainable_sha, + "final_trainable_sha256": final_trainable_sha, + "optimizer_update_count": len(optimizer_updates), + "max_optimizer_update_l2": max_update_l2, + "loss_delta": loss_delta, + "peak_cuda_memory_mb": cuda_peak_memory_mb(torch), + "adapter_weight_sha256": weight_stats["sha256"], + }, + ) + adapter_sha = sha256_tree(output_dir) + manifest = AdapterManifest( + adapter_id=adapter_id(kind, dataset_path, model_id), + kind=kind, + created_at=datetime.now(UTC), + proof_version=TRAINING_PROOF_VERSION, + training_method=TRAINING_METHOD, + precision=TRAINING_PRECISION, + framework=TRAINING_FRAMEWORK, + base_model=model_id, + dataset_path=str(dataset_path.resolve()), + dataset_sha256=sha256_file(dataset_path), + sample_count=len(samples), + max_steps=max_steps, + max_length=max_length, + rank=rank, + alpha=alpha, + target_modules=target_modules, + trainable_params=int(trainable_params), + total_params=int(total_params), + optimizer=optimizer_name, + optimizer_family=optimizer_family(optimizer_name), + learning_rate=learning_rate, + gradient_accumulation_steps=gradient_accumulation_steps, + gradient_checkpointing=gradient_checkpointing, + gradient_checkpointing_mode=GRADIENT_CHECKPOINTING_MODE if gradient_checkpointing else None, + autocast_dtype=AUTOCAST_DTYPE, + bf16=torch.cuda.is_available(), + unsloth_version=package_version("unsloth"), + triton_version=package_version("triton-windows") or package_version("triton"), + xformers_version=package_version("xformers"), + unsloth_model_type=unsloth_model_type, + final_loss=losses[-1] if losses else None, + loss_history=losses, + initial_trainable_sha256=initial_trainable_sha, + final_trainable_sha256=final_trainable_sha, + initial_trainable_l2=initial_trainable_l2, + final_trainable_l2=final_trainable_l2, + optimizer_update_count=len(optimizer_updates), + max_optimizer_update_l2=max_update_l2, + loss_delta=loss_delta, + training_proof_path=str(proof_path.resolve()), + training_proof_sha256=sha256_file(proof_path), + training_trace_path=str(trace_path.resolve()), + training_trace_sha256=sha256_file(trace_path), + training_summary_path=str(summary_path.resolve()), + training_summary_sha256=sha256_file(summary_path), + training_started_at=started_at, + training_finished_at=finished_at, + training_wall_seconds=wall_seconds, + training_device=str(environment.get("device_name") or environment.get("device") or ""), + torch_version=str(environment.get("torch_version") or ""), + cuda_version=( + str(environment["torch_cuda_version"]) + if environment.get("torch_cuda_version") is not None + else None + ), + adapter_weight_sha256=str(weight_stats["sha256"] or ""), + adapter_tensor_count=int(weight_stats["tensor_count"]), + adapter_lora_tensor_count=int(weight_stats["lora_tensor_count"]), + adapter_total_elements=int(weight_stats["total_elements"]), + adapter_nonzero_elements=int(weight_stats["nonzero_elements"]), + adapter_max_abs=weight_stats["max_abs"], + adapter_weights_finite=bool(weight_stats["all_finite"]), + artifacts=relative_artifacts(output_dir), + sha256=adapter_sha, + cartridge_id=cartridge_id, + ) + (output_dir / "adapter_manifest.json").write_text( + manifest.model_dump_json(indent=2), + encoding="utf-8", + ) + return manifest diff --git a/backend/src/ely_eye/training_metrics.py b/backend/src/ely_eye/training_metrics.py new file mode 100644 index 0000000..cc69c91 --- /dev/null +++ b/backend/src/ely_eye/training_metrics.py @@ -0,0 +1,307 @@ +"""Measure what a training run actually changed: trainable weights, gradients, optimizers, environment.""" + +from __future__ import annotations + +import hashlib +import math +import platform +import subprocess +from contextlib import contextmanager +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any + +from .schemas import TrainingAdapterKind +from .training_common import ( + AUTOCAST_DTYPE, + GRADIENT_CHECKPOINTING_MODE, + TRAINING_FRAMEWORK, + TRAINING_METHOD, + TRAINING_PRECISION, + TRAINING_PROOF_VERSION, + sha256_file, +) + + +def infer_lora_targets(model: Any, kind: TrainingAdapterKind) -> list[str]: + import torch + + core_leafs = {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"} + mlp_leafs = {"gate_proj", "up_proj", "down_proj"} + found_leafs: set[str] = set() + exact: set[str] = set() + for name, module in model.named_modules(): + if not isinstance(module, torch.nn.Linear): + continue + leaf = name.rsplit(".", 1)[-1] + lower = name.lower() + if kind == TrainingAdapterKind.hme_core_lora and leaf in core_leafs: + found_leafs.add(leaf) + elif kind == TrainingAdapterKind.hme_vision_lora and ("visual" in lower or "vision" in lower): + exact.add(name) + elif kind == TrainingAdapterKind.hme_ttt_vl and (leaf in mlp_leafs or "visual" in lower): + exact.add(name if "visual" in lower else leaf) + elif kind == TrainingAdapterKind.hme_visual_mtp and ( + "mtp" in lower or "draft" in lower or name == "lm_head" + ): + exact.add(name) + elif kind in {TrainingAdapterKind.hme_router, TrainingAdapterKind.hme_retrieval} and leaf in core_leafs: + found_leafs.add(leaf) + targets = sorted(exact or found_leafs) + if not targets: + for name, module in model.named_modules(): + if isinstance(module, torch.nn.Linear) and not name.endswith("lm_head"): + targets.append(name.rsplit(".", 1)[-1]) + targets = sorted(set(targets)) + if not targets: + raise ValueError(f"No LoRA target modules found for {kind.value}") + return targets + +def trainable_named_parameters(model: Any) -> list[tuple[str, Any]]: + return [(name, param) for name, param in model.named_parameters() if param.requires_grad] + +def trainable_state_cpu(named_parameters: list[tuple[str, Any]]) -> dict[str, Any]: + return {name: param.detach().float().cpu().clone() for name, param in named_parameters} + +def trainable_state_sha256(named_parameters: list[tuple[str, Any]]) -> str: + digest = hashlib.sha256() + for name, param in named_parameters: + tensor = param.detach().float().cpu().contiguous() + digest.update(name.encode("utf-8")) + digest.update(str(tuple(tensor.shape)).encode("utf-8")) + digest.update(tensor.numpy().tobytes()) + return digest.hexdigest() + +def trainable_l2_norm(named_parameters: list[tuple[str, Any]]) -> float: + total = 0.0 + for _, param in named_parameters: + tensor = param.detach().float() + total += float(tensor.square().sum().detach().cpu()) + return math.sqrt(total) + +def gradient_l2_norm(named_parameters: list[tuple[str, Any]]) -> float: + total = 0.0 + for _, param in named_parameters: + if param.grad is None: + continue + grad = param.grad.detach().float() + total += float(grad.square().sum().detach().cpu()) + return math.sqrt(total) + +def trainable_update_l2(before_update: dict[str, Any], named_parameters: list[tuple[str, Any]]) -> float: + total = 0.0 + for name, param in named_parameters: + delta = param.detach().float().cpu() - before_update[name] + total += float(delta.square().sum()) + return math.sqrt(total) + +def require_cuda_bf16(torch_module: Any) -> None: + if not torch_module.cuda.is_available(): + raise RuntimeError("training needs a CUDA device; none is available.") + bf16_supported = getattr(torch_module.cuda, "is_bf16_supported", lambda: False) + if not bool(bf16_supported()): + raise RuntimeError("PRD 12.2 training requires BF16-capable CUDA hardware.") + +def activate_unsloth_training() -> None: + required = { + "unsloth": package_version("unsloth"), + "unsloth_zoo": package_version("unsloth_zoo"), + "triton": package_version("triton-windows") or package_version("triton"), + "xformers": package_version("xformers"), + } + missing = [name for name, installed in required.items() if not installed] + if missing: + raise RuntimeError( + "PRD 12.2 training requires installed Unsloth dependencies: " + ", ".join(missing) + ) + with bitsandbytes_gaudi_probe_guard(): + __import__("unsloth") + +def enable_gradient_checkpointing(model: Any) -> bool: + if not hasattr(model, "gradient_checkpointing_enable"): + return False + try: + model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + except TypeError: + model.gradient_checkpointing_enable() + return True + +def model_type_for_proof(model: Any) -> str | None: + current = model + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + config = getattr(current, "config", None) + model_type = getattr(config, "model_type", None) + if isinstance(model_type, str) and model_type: + return model_type + current = getattr(current, "model", None) + return None + +def make_optimizer(model: Any, learning_rate: float) -> tuple[Any, str]: + params = [param for param in model.parameters() if param.requires_grad] + try: + with bitsandbytes_gaudi_probe_guard(): + import bitsandbytes as bnb + except Exception as exc: + raise RuntimeError("PRD 12.2 training requires bitsandbytes PagedAdamW8bit.") from exc + try: + return bnb.optim.PagedAdamW8bit(params, lr=learning_rate), "paged_adamw_8bit" + except Exception as exc: + raise RuntimeError("bitsandbytes PagedAdamW8bit failed to initialize.") from exc + +def optimizer_family(optimizer_name: str) -> str: + if optimizer_name == "paged_adamw_8bit": + return "8-bit AdamW / paged optimizer" + return optimizer_name + +@contextmanager +def bitsandbytes_gaudi_probe_guard() -> Any: + original_run = subprocess.run + + def guarded_run(*args: Any, **kwargs: Any) -> Any: + command = args[0] if args else kwargs.get("args") + if isinstance(command, str) and "habana-torch-plugin" in command: + return subprocess.CompletedProcess(command, 0, "", "") + return original_run(*args, **kwargs) + + subprocess.run = guarded_run + try: + yield + finally: + subprocess.run = original_run + +def training_contract( + *, + gradient_checkpointing: bool, + input_require_grads: bool, + optimizer_name: str, + unsloth_model_type: str | None, +) -> dict[str, Any]: + return { + "proof_version": TRAINING_PROOF_VERSION, + "training_method": TRAINING_METHOD, + "precision": TRAINING_PRECISION, + "framework": TRAINING_FRAMEWORK, + "optimizer_family": optimizer_family(optimizer_name), + "gradient_checkpointing": gradient_checkpointing, + "gradient_checkpointing_mode": GRADIENT_CHECKPOINTING_MODE if gradient_checkpointing else None, + "input_require_grads": input_require_grads, + "autocast_dtype": AUTOCAST_DTYPE, + "unsloth_version": package_version("unsloth"), + "triton_version": package_version("triton-windows") or package_version("triton"), + "xformers_version": package_version("xformers"), + "unsloth_model_type": unsloth_model_type, + } + +def training_environment(torch_module: Any) -> dict[str, Any]: + cuda_available = bool(torch_module.cuda.is_available()) + device_name = torch_module.cuda.get_device_name(0) if cuda_available else platform.processor() + capability = torch_module.cuda.get_device_capability(0) if cuda_available else None + bf16_supported = ( + bool(getattr(torch_module.cuda, "is_bf16_supported", lambda: False)()) + if cuda_available + else False + ) + return { + "python": platform.python_version(), + "platform": platform.platform(), + "device": "cuda:0" if cuda_available else "cpu", + "device_name": device_name, + "cuda_available": cuda_available, + "bf16_supported": bf16_supported, + "cuda_capability": list(capability) if capability else None, + "torch_version": getattr(torch_module, "__version__", None), + "torch_cuda_version": getattr(torch_module.version, "cuda", None), + "transformers_version": package_version("transformers"), + "peft_version": package_version("peft"), + "bitsandbytes_version": package_version("bitsandbytes"), + "unsloth_version": package_version("unsloth"), + "unsloth_zoo_version": package_version("unsloth_zoo"), + "triton_version": package_version("triton-windows") or package_version("triton"), + "xformers_version": package_version("xformers"), + } + +def package_version(package_name: str) -> str | None: + try: + return version(package_name) + except PackageNotFoundError: + return None + +def cuda_peak_memory_mb(torch_module: Any) -> float | None: + if not torch_module.cuda.is_available(): + return None + return float(torch_module.cuda.max_memory_allocated() / 1024 / 1024) + +def safetensor_weight_stats(path: Path) -> dict[str, Any]: + report: dict[str, Any] = { + "path": str(path), + "exists": path.exists(), + "sha256": sha256_file(path) if path.exists() else None, + "tensor_count": 0, + "lora_tensor_count": 0, + "total_elements": 0, + "nonzero_elements": 0, + "nonzero_tensor_count": 0, + "finite_tensor_count": 0, + "all_finite": False, + "max_abs": None, + "mean_abs": None, + "dtype_counts": {}, + "key_sample": [], + "valid": False, + } + if not path.exists(): + return report + + import torch + from safetensors import safe_open + + total_abs = 0.0 + max_abs = 0.0 + dtype_counts: dict[str, int] = {} + with safe_open(path, framework="pt", device="cpu") as handle: + keys = list(handle.keys()) + report["tensor_count"] = len(keys) + report["key_sample"] = keys[:12] + for key in keys: + tensor = handle.get_tensor(key) + numel = int(tensor.numel()) + dtype = str(tensor.dtype).replace("torch.", "") + dtype_counts[dtype] = dtype_counts.get(dtype, 0) + 1 + report["total_elements"] += numel + if "lora_" in key: + report["lora_tensor_count"] += 1 + if numel == 0: + continue + tensor_float = tensor.detach().float() + finite = bool(torch.isfinite(tensor_float).all().item()) + if finite: + report["finite_tensor_count"] += 1 + nonzero = int(torch.count_nonzero(tensor).item()) + report["nonzero_elements"] += nonzero + if nonzero > 0: + report["nonzero_tensor_count"] += 1 + abs_tensor = tensor_float.abs() + tensor_max_abs = float(abs_tensor.max().item()) + max_abs = max(max_abs, tensor_max_abs) + total_abs += float(abs_tensor.sum().item()) + + report["dtype_counts"] = dtype_counts + report["all_finite"] = ( + report["tensor_count"] > 0 and report["finite_tensor_count"] == report["tensor_count"] + ) + if int(report["total_elements"]) > 0: + report["max_abs"] = max_abs + report["mean_abs"] = total_abs / int(report["total_elements"]) + report["valid"] = ( + report["exists"] + and int(report["tensor_count"]) > 0 + and int(report["lora_tensor_count"]) > 0 + and int(report["total_elements"]) > 0 + and int(report["nonzero_elements"]) > 0 + and bool(report["all_finite"]) + and report["max_abs"] is not None + and float(report["max_abs"]) > 0.0 + ) + return report diff --git a/backend/src/ely_eye/training_retrieval.py b/backend/src/ely_eye/training_retrieval.py new file mode 100644 index 0000000..7aaba1b --- /dev/null +++ b/backend/src/ely_eye/training_retrieval.py @@ -0,0 +1,358 @@ +"""Train the embedding adapter: pairs in, cosine ranking out, weights measured after.""" + +from __future__ import annotations + +import time +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from .config import Settings +from .db import Database +from .schemas import AdapterManifest, TrainingAdapterKind +from .training_common import ( + AUTOCAST_DTYPE, + GRADIENT_CHECKPOINTING_MODE, + TRAINING_FRAMEWORK, + TRAINING_METHOD, + TRAINING_PRECISION, + TRAINING_PROOF_VERSION, + adapter_id, + relative_artifacts, + sha256_file, + sha256_tree, + write_trace, + write_training_summary, +) +from .training_datasets import validate_dataset +from .training_encode import encode_retrieval_batch, pool_embeddings +from .training_lora import write_lora_config +from .training_metrics import ( + activate_unsloth_training, + bitsandbytes_gaudi_probe_guard, + cuda_peak_memory_mb, + enable_gradient_checkpointing, + gradient_l2_norm, + infer_lora_targets, + make_optimizer, + model_type_for_proof, + optimizer_family, + package_version, + require_cuda_bf16, + safetensor_weight_stats, + trainable_l2_norm, + trainable_named_parameters, + trainable_state_cpu, + trainable_state_sha256, + trainable_update_l2, + training_contract, + training_environment, +) + + +def train_retrieval_lora( + settings: Settings, + database: Database, + dataset_path: Path, + output_dir: Path, + max_steps: int, + rank: int = 8, + alpha: int = 16, + max_length: int | None = None, + learning_rate: float | None = None, + cartridge_id: str | None = None, +) -> AdapterManifest: + samples = validate_dataset(dataset_path) + output_dir.mkdir(parents=True, exist_ok=True) + max_length = max_length or settings.train_max_length + learning_rate = learning_rate or settings.train_learning_rate + started_at = datetime.now(UTC) + wall_start = time.perf_counter() + trace_path = output_dir / "training_trace.jsonl" + + import torch + import torch.nn.functional as F + + activate_unsloth_training() + from peft import LoraConfig, TaskType, get_peft_model + from transformers import AutoModel, AutoProcessor + + model_id = settings.embedding_model_id + require_cuda_bf16(torch) + processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) + model = AutoModel.from_pretrained( + model_id, + dtype=torch.bfloat16, + device_map={"": 0}, + trust_remote_code=True, + ) + if hasattr(model.config, "use_cache"): + model.config.use_cache = False + target_modules = infer_lora_targets(model, TrainingAdapterKind.hme_retrieval) + gradient_checkpointing = enable_gradient_checkpointing(model) + with bitsandbytes_gaudi_probe_guard(): + model = get_peft_model( + model, + LoraConfig( + r=rank, + lora_alpha=alpha, + target_modules=target_modules, + lora_dropout=0.0, + bias="none", + task_type=TaskType.FEATURE_EXTRACTION, + ), + ) + unsloth_model_type = model_type_for_proof(model) + trainable_params, total_params = model.get_nb_trainable_parameters() + optimizer, optimizer_name = make_optimizer(model, learning_rate) + model.train() + + losses: list[float] = [] + optimizer_updates: list[dict[str, Any]] = [] + trainable_tensors = trainable_named_parameters(model) + initial_trainable_sha = trainable_state_sha256(trainable_tensors) + initial_trainable_l2 = trainable_l2_norm(trainable_tensors) + batch_size = min(4, len(samples)) + environment = training_environment(torch) + contract = training_contract( + gradient_checkpointing=gradient_checkpointing, + input_require_grads=False, + optimizer_name=optimizer_name, + unsloth_model_type=unsloth_model_type, + ) + with trace_path.open("w", encoding="utf-8", newline="\n") as trace: + write_trace( + trace, + { + "event": "run_start", + **contract, + "started_at": started_at.isoformat(), + "adapter_kind": TrainingAdapterKind.hme_retrieval.value, + "base_model": model_id, + "dataset_path": str(dataset_path.resolve()), + "dataset_sha256": sha256_file(dataset_path), + "sample_count": len(samples), + "max_steps": max_steps, + "max_length": max_length, + "rank": rank, + "alpha": alpha, + "target_modules": target_modules, + "trainable_params": int(trainable_params), + "total_params": int(total_params), + "optimizer": optimizer_name, + "learning_rate": learning_rate, + "environment": environment, + }, + ) + for step in range(max_steps): + window = [samples[(step + offset) % len(samples)] for offset in range(batch_size)] + queries = [str(sample.inputs[0]["text"]) for sample in window] + documents = [str(sample.output["positive"]) for sample in window] + query_batch = encode_retrieval_batch(queries, processor, model.device, max_length) + document_batch = encode_retrieval_batch(documents, processor, model.device, max_length) + query_embeddings = pool_embeddings(model(**query_batch), query_batch["attention_mask"]) + document_embeddings = pool_embeddings(model(**document_batch), document_batch["attention_mask"]) + logits = query_embeddings @ document_embeddings.T / 0.05 + labels = torch.arange(logits.shape[0], device=logits.device) + loss = F.cross_entropy(logits.float(), labels) + loss.backward() + grad_l2 = gradient_l2_norm(trainable_tensors) + before_update = trainable_state_cpu(trainable_tensors) + optimizer.step() + update_l2 = trainable_update_l2(before_update, trainable_tensors) + optimizer_updates.append( + { + "step": step + 1, + "gradient_l2": grad_l2, + "update_l2": update_l2, + } + ) + optimizer.zero_grad(set_to_none=True) + actual_loss = float(loss.detach().cpu()) + losses.append(actual_loss) + write_trace( + trace, + { + "event": "step", + "step": step + 1, + "sample_ids": [sample.id for sample in window], + "task": "hme_retrieval_contrastive_pair", + "loss": actual_loss, + "optimizer_step": True, + "gradient_l2": grad_l2, + "update_l2": update_l2, + }, + ) + + final_trainable_sha = trainable_state_sha256(trainable_tensors) + final_trainable_l2 = trainable_l2_norm(trainable_tensors) + max_update_l2 = max((item["update_l2"] for item in optimizer_updates), default=None) + loss_delta = losses[0] - losses[-1] if len(losses) >= 2 else None + model.save_pretrained(output_dir) + processor.save_pretrained(output_dir) + write_lora_config(settings, output_dir, TrainingAdapterKind.hme_retrieval, rank, alpha, model_id) + weight_stats = safetensor_weight_stats(output_dir / "adapter_model.safetensors") + finished_at = datetime.now(UTC) + wall_seconds = time.perf_counter() - wall_start + summary_path = output_dir / "training_summary.json" + proof_path = output_dir / "training_proof.json" + training_proof = { + **contract, + "adapter_kind": TrainingAdapterKind.hme_retrieval.value, + "base_model": model_id, + "dataset_path": str(dataset_path.resolve()), + "dataset_sha256": sha256_file(dataset_path), + "sample_count": len(samples), + "max_steps": max_steps, + "max_length": max_length, + "rank": rank, + "alpha": alpha, + "target_modules": target_modules, + "trainable_params": int(trainable_params), + "total_params": int(total_params), + "optimizer": optimizer_name, + "learning_rate": learning_rate, + "gradient_accumulation_steps": 1, + "bf16": bool(torch.cuda.is_available()), + "environment": environment, + "initial_trainable_sha256": initial_trainable_sha, + "final_trainable_sha256": final_trainable_sha, + "initial_trainable_l2": initial_trainable_l2, + "final_trainable_l2": final_trainable_l2, + "optimizer_update_count": len(optimizer_updates), + "max_optimizer_update_l2": max_update_l2, + "loss_delta": loss_delta, + "loss_history": losses, + "optimizer_updates": optimizer_updates, + "adapter_weight_sha256": weight_stats["sha256"], + "adapter_weight_stats": weight_stats, + } + write_training_summary(proof_path, training_proof) + write_training_summary( + summary_path, + { + **contract, + "adapter_kind": TrainingAdapterKind.hme_retrieval.value, + "base_model": model_id, + "dataset_path": str(dataset_path.resolve()), + "dataset_sha256": sha256_file(dataset_path), + "sample_count": len(samples), + "max_steps": max_steps, + "max_length": max_length, + "rank": rank, + "alpha": alpha, + "target_modules": target_modules, + "trainable_params": int(trainable_params), + "total_params": int(total_params), + "optimizer": optimizer_name, + "learning_rate": learning_rate, + "gradient_accumulation_steps": 1, + "bf16": bool(torch.cuda.is_available()), + "loss_history": losses, + "final_loss": losses[-1] if losses else None, + "initial_trainable_sha256": initial_trainable_sha, + "final_trainable_sha256": final_trainable_sha, + "initial_trainable_l2": initial_trainable_l2, + "final_trainable_l2": final_trainable_l2, + "optimizer_update_count": len(optimizer_updates), + "max_optimizer_update_l2": max_update_l2, + "loss_delta": loss_delta, + "started_at": started_at.isoformat(), + "finished_at": finished_at.isoformat(), + "wall_seconds": wall_seconds, + "environment": environment, + "peak_cuda_memory_mb": cuda_peak_memory_mb(torch), + "adapter_weight_sha256": weight_stats["sha256"], + "adapter_weight_stats": weight_stats, + }, + ) + with trace_path.open("a", encoding="utf-8", newline="\n") as trace: + write_trace( + trace, + { + "event": "run_finish", + **contract, + "finished_at": finished_at.isoformat(), + "wall_seconds": wall_seconds, + "final_loss": losses[-1] if losses else None, + "initial_trainable_sha256": initial_trainable_sha, + "final_trainable_sha256": final_trainable_sha, + "optimizer_update_count": len(optimizer_updates), + "max_optimizer_update_l2": max_update_l2, + "loss_delta": loss_delta, + "peak_cuda_memory_mb": cuda_peak_memory_mb(torch), + "adapter_weight_sha256": weight_stats["sha256"], + }, + ) + manifest = AdapterManifest( + adapter_id=adapter_id(TrainingAdapterKind.hme_retrieval, dataset_path, model_id), + kind=TrainingAdapterKind.hme_retrieval, + created_at=datetime.now(UTC), + proof_version=TRAINING_PROOF_VERSION, + training_method=TRAINING_METHOD, + precision=TRAINING_PRECISION, + framework=TRAINING_FRAMEWORK, + base_model=model_id, + dataset_path=str(dataset_path.resolve()), + dataset_sha256=sha256_file(dataset_path), + sample_count=len(samples), + max_steps=max_steps, + max_length=max_length, + rank=rank, + alpha=alpha, + target_modules=target_modules, + trainable_params=int(trainable_params), + total_params=int(total_params), + optimizer=optimizer_name, + optimizer_family=optimizer_family(optimizer_name), + learning_rate=learning_rate, + gradient_accumulation_steps=1, + gradient_checkpointing=gradient_checkpointing, + gradient_checkpointing_mode=GRADIENT_CHECKPOINTING_MODE if gradient_checkpointing else None, + autocast_dtype=AUTOCAST_DTYPE, + bf16=torch.cuda.is_available(), + unsloth_version=package_version("unsloth"), + triton_version=package_version("triton-windows") or package_version("triton"), + xformers_version=package_version("xformers"), + unsloth_model_type=unsloth_model_type, + final_loss=losses[-1] if losses else None, + loss_history=losses, + initial_trainable_sha256=initial_trainable_sha, + final_trainable_sha256=final_trainable_sha, + initial_trainable_l2=initial_trainable_l2, + final_trainable_l2=final_trainable_l2, + optimizer_update_count=len(optimizer_updates), + max_optimizer_update_l2=max_update_l2, + loss_delta=loss_delta, + training_proof_path=str(proof_path.resolve()), + training_proof_sha256=sha256_file(proof_path), + training_trace_path=str(trace_path.resolve()), + training_trace_sha256=sha256_file(trace_path), + training_summary_path=str(summary_path.resolve()), + training_summary_sha256=sha256_file(summary_path), + training_started_at=started_at, + training_finished_at=finished_at, + training_wall_seconds=wall_seconds, + training_device=str(environment.get("device_name") or environment.get("device") or ""), + torch_version=str(environment.get("torch_version") or ""), + cuda_version=( + str(environment["torch_cuda_version"]) + if environment.get("torch_cuda_version") is not None + else None + ), + adapter_weight_sha256=str(weight_stats["sha256"] or ""), + adapter_tensor_count=int(weight_stats["tensor_count"]), + adapter_lora_tensor_count=int(weight_stats["lora_tensor_count"]), + adapter_total_elements=int(weight_stats["total_elements"]), + adapter_nonzero_elements=int(weight_stats["nonzero_elements"]), + adapter_max_abs=weight_stats["max_abs"], + adapter_weights_finite=bool(weight_stats["all_finite"]), + artifacts=relative_artifacts(output_dir), + sha256=sha256_tree(output_dir), + cartridge_id=cartridge_id, + ) + (output_dir / "adapter_manifest.json").write_text( + manifest.model_dump_json(indent=2), + encoding="utf-8", + ) + return manifest diff --git a/backend/src/ely_eye/transactive.py b/backend/src/ely_eye/transactive.py index 1115fb1..9190cec 100644 --- a/backend/src/ely_eye/transactive.py +++ b/backend/src/ely_eye/transactive.py @@ -33,8 +33,7 @@ LIBRARY = "library" WHO_KNOWS_LIMIT = 8 MAX_HOPS_LIMIT = 10 -# Confidence is ordinal: it orders candidates, it does not estimate a probability. The weights -# say what the evidence is worth relative to itself, and every hit carries the raw numbers too. +# Confidence is ordinal: it orders candidates, it does not estimate a probability, and every hit carries the raw numbers too. CONFIDENCE_SATURATE = 3 WEIGHT_STRENGTH = 0.60 WEIGHT_COVERAGE = 0.25 diff --git a/backend/src/ely_eye/visual_index.py b/backend/src/ely_eye/visual_index.py index 2708ef6..69695dc 100644 --- a/backend/src/ely_eye/visual_index.py +++ b/backend/src/ely_eye/visual_index.py @@ -27,8 +27,7 @@ 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. +# Loading 48k vectors costs about as much as the search itself and the API builds a fresh service per request, so the array is cached per process and re-read when the file changes. _LOADED: dict[tuple[str, float], tuple[list[str], np.ndarray]] = {} diff --git a/backend/tests/test_host_profiles.py b/backend/tests/test_host_profiles.py index ac90410..7ec5a7e 100644 --- a/backend/tests/test_host_profiles.py +++ b/backend/tests/test_host_profiles.py @@ -11,7 +11,7 @@ import pytest from ely_eye import profiles from ely_eye.config import Settings from ely_eye.proofs.training_trace import cuda_bf16_training -from ely_eye.training import require_cuda_bf16 +from ely_eye.training_metrics import require_cuda_bf16 def _stub_nvidia_smi(tmp_path: Path, monkeypatch, output: str, exit_code: int = 0) -> None: diff --git a/backend/tests/test_memory_schema.py b/backend/tests/test_memory_schema.py index 5e2aef1..b273741 100644 --- a/backend/tests/test_memory_schema.py +++ b/backend/tests/test_memory_schema.py @@ -20,6 +20,7 @@ from ely_eye.schemas import ( TrustScores, Visibility, ) +from ely_eye.store_search import search_fts NOW = datetime(2026, 5, 1, 12, 0, tzinfo=UTC) @@ -235,9 +236,9 @@ def test_a_second_version_database_rebuilds_its_index_on_upgrade(tmp_path: Path) "INSERT INTO atom_fts(atom_id, source, text) VALUES ('key', 'notes/key.txt', '钥匙放在玄关抽屉里')" ) connection.execute("PRAGMA user_version = 2") - assert database.search_fts("钥匙", 5) == [], "the old index really could not answer this query" + assert search_fts(database, "钥匙", 5) == [], "the old index really could not answer this query" upgraded = _database(tmp_path) - assert [atom_id for atom_id, _ in upgraded.search_fts("钥匙", 5)] == ["key"] + assert [atom_id for atom_id, _ in search_fts(upgraded, "钥匙", 5)] == ["key"] with upgraded.connect() as connection: assert int(connection.execute("PRAGMA user_version").fetchone()[0]) == SCHEMA_VERSION diff --git a/backend/tests/test_proofs.py b/backend/tests/test_proofs.py index 754a02c..8a507ff 100644 --- a/backend/tests/test_proofs.py +++ b/backend/tests/test_proofs.py @@ -26,8 +26,7 @@ from ely_eye.proofs.contradictions import ( from ely_eye.proofs.table import render_table from ely_eye.schemas import MemoryOrigin -# The replay question is fixed inside the suite, so the fixture memory has to answer it in words -# the ranker can match. +# The replay question is fixed inside the suite, so the fixture memory must answer it in words the ranker can match. PRD_MEMORY = "PRD requirements prove long context memory, HashHop proof, and adapter training." OTHER_KEYS = "PRD requirements also prove cache invalidation for HashHop proof and adapters." @@ -95,8 +94,7 @@ def test_the_repository_check_fails_when_the_repository_never_reached_the_store( def test_the_replay_check_answers_only_from_the_cartridge_it_mounts(harness: Harness) -> None: inside = harness.remember(PRD_MEMORY) - # A second memory that answers the same question but stays out of the cartridge: an unmounted - # replay against the live store would have found it and called that a cartridge replay. + # A second memory answering the same question but outside the cartridge: an unmounted replay would have found it and called that a cartridge replay. harness.remember(OTHER_KEYS) cartridge_id = harness.cartridge("keys", [inside]) manifest = next( @@ -275,8 +273,7 @@ def test_the_hashhop_arena_marks_no_candidate_as_the_answer(harness: Harness, tm render_visual_hash_arena(tiles["top-left"], tiles, output) - # Every candidate box carries the same outline: a ring, an arrow, or a highlighted slot would - # tell a model the answer without reading a pixel of the source tile. + # Every candidate box carries the same outline: a ring, arrow or highlighted slot would give the answer without reading the tile. image = Image.open(output).convert("RGB") border_colors = {image.getpixel((box[0] + 2, box[1] + 40)) for box in boxes.values()} assert border_colors == {(180, 190, 205)} diff --git a/backend/tests/test_text_index.py b/backend/tests/test_text_index.py index d96d3c4..044020e 100644 --- a/backend/tests/test_text_index.py +++ b/backend/tests/test_text_index.py @@ -15,6 +15,7 @@ from ely_eye.config import Settings from ely_eye.db import Database from ely_eye.memory import MemoryService from ely_eye.schemas import MemoryOrigin +from ely_eye.store_search import search_atoms from ely_eye.text_index import ( RAREST_KEEP, fts_expression, @@ -190,7 +191,7 @@ def test_the_remembered_match_counts_are_voided_by_a_write(tmp_path: Path) -> No memory = MemoryService(settings=settings, database=Database(settings)) for index in range(20): memory.remember(f"the converter parses format {index} natively", holder_id="wife", origin=MemoryOrigin.self_) - memory.db.search_atoms("the converter", limit=5) + search_atoms(memory.db, "the converter", limit=5) with memory.db.connect() as connection: assert int(connection.execute("SELECT count(*) FROM term_df").fetchone()[0]) > 0 @@ -200,7 +201,7 @@ def test_the_remembered_match_counts_are_voided_by_a_write(tmp_path: Path) -> No with memory.db.connect() as connection: assert int(connection.execute("SELECT count(*) FROM term_df").fetchone()[0]) == 0 - page = memory.db.search_atoms("hashhop chain", limit=5) + page = search_atoms(memory.db, "hashhop chain", limit=5) assert "the parser follows the hashhop chain" in [ atom.text for atom in memory.db.get_atoms([row[0] for row in page.rows]) ], "a cleared cache must cost time, not answers" diff --git a/backend/tests/test_training_split.py b/backend/tests/test_training_split.py new file mode 100644 index 0000000..35f1543 --- /dev/null +++ b/backend/tests/test_training_split.py @@ -0,0 +1,152 @@ +"""Behaviour tests for the training modules after the 1630-line file was split.""" + +from __future__ import annotations + +import inspect +import json +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +import pytest + +from ely_eye import training_datasets, training_lora, training_retrieval +from ely_eye.config import Settings +from ely_eye.db import Database +from ely_eye.schemas import ( + EvidenceAtom, + Modality, + SourceRecord, + TrainingAdapterKind, + TrustScores, +) +from ely_eye.training import TrainingService +from ely_eye.training_common import sha256_bytes +from ely_eye.training_datasets import ( + atom_id_for_path, + atoms_for_path, + latest_cartridge_id, + model_for_kind, + validate_dataset, +) +from ely_eye.training_lora import write_lora_config + + +def _settings(root: Path) -> Settings: + return Settings(home=root, model_id="base/model", planner_model_id="planner/model") + + +def _database(root: Path) -> Database: + database = Database(_settings(root)) + database.init() + return database + + +def _sample(sample_id: str) -> dict[str, Any]: + return {"id": sample_id, "task": "recall", "inputs": [{"text": "where are the keys"}], "output": {"text": "the drawer"}} + + +def test_validate_dataset_reads_jsonl_and_names_the_line_that_broke(tmp_path: Path) -> None: + good = tmp_path / "data.jsonl" + good.write_text(json.dumps(_sample("a")) + "\n" + json.dumps(_sample("b")) + "\n", encoding="utf-8") + samples = validate_dataset(good) + assert [sample.id for sample in samples] == ["a", "b"] + + broken = tmp_path / "broken.jsonl" + broken.write_text(json.dumps(_sample("a")) + "\nnot json\n", encoding="utf-8") + with pytest.raises(ValueError, match="Invalid training sample at line 2"): + validate_dataset(broken) + + empty = tmp_path / "empty.jsonl" + empty.write_text("\n", encoding="utf-8") + with pytest.raises(ValueError, match="Training dataset is empty"): + validate_dataset(empty) + + +def test_model_for_kind_uses_the_planner_only_for_the_router(tmp_path: Path) -> None: + settings = _settings(tmp_path) + assert model_for_kind(settings, TrainingAdapterKind.hme_router) == "planner/model" + assert model_for_kind(settings, TrainingAdapterKind.hme_core_lora) == "base/model" + + +def test_latest_cartridge_id_is_none_until_a_cartridge_is_written(tmp_path: Path) -> None: + database = _database(tmp_path) + assert latest_cartridge_id(database) is None + database.write_cartridge("old", "old", tmp_path / "old", "{}", None, 10, datetime(2026, 1, 1, tzinfo=UTC)) + database.write_cartridge("new", "new", tmp_path / "new", "{}", None, 20, datetime(2026, 3, 1, tzinfo=UTC)) + assert latest_cartridge_id(database) == "new" + + +def test_atoms_for_path_follows_the_source_and_falls_back_to_a_hash(tmp_path: Path) -> None: + database = _database(tmp_path) + paper = tmp_path / "paper.md" + paper.write_text("the keys are in the drawer", encoding="utf-8") + other = tmp_path / "other.md" + database.upsert_sources( + [ + SourceRecord(source_id=source_id, path=str(path), kind="text", mime="text/markdown", sha256="a", size_bytes=1) + for source_id, path in (("paper", paper), ("other", other)) + ] + ) + trust = TrustScores(parser=1.0) + database.upsert_atoms( + [ + EvidenceAtom( + atom_id="paper:text", + modality=Modality.text, + source_id="paper", + source=str(paper), + text="drawer", + trust=trust, + ), + EvidenceAtom( + atom_id="other:text", + modality=Modality.text, + source_id="other", + source=str(other), + text="shelf", + trust=trust, + ), + ] + ) + assert [atom.atom_id for atom in atoms_for_path(database, paper)] == ["paper:text"] + expected = f"notes:{sha256_bytes(b'cold')[:12]}" + assert atom_id_for_path(database, tmp_path / "notes.md", "cold") == expected + assert atom_id_for_path(database, paper, "cold") == "paper:text" + + +def test_write_lora_config_ties_weights_only_for_the_mtp_adapter(tmp_path: Path) -> None: + settings = _settings(tmp_path) + core = write_lora_config(settings, tmp_path / "core", TrainingAdapterKind.hme_core_lora, rank=8, alpha=16) + payload = json.loads(core.read_text(encoding="utf-8")) + assert payload["base_model"] == "base/model" + assert payload["adapter_kind"] == "hme_core_lora" + assert (payload["rank"], payload["alpha"]) == (8, 16) + assert payload["ensure_weight_tying"] is False + + mtp = write_lora_config(settings, tmp_path / "mtp", TrainingAdapterKind.hme_visual_mtp, model_id="other/model") + tied = json.loads(mtp.read_text(encoding="utf-8")) + assert tied["base_model"] == "other/model" + assert tied["ensure_weight_tying"] is True + + +FACADE = { + "build_prd_dataset": training_datasets, + "build_visual_dataset": training_datasets, + "build_router_dataset": training_datasets, + "build_retrieval_dataset": training_datasets, + "train_lora": training_lora, + "train_retrieval_lora": training_retrieval, +} + + +def test_the_service_facade_passes_every_argument_the_module_expects(tmp_path: Path) -> None: + data = tmp_path / "data.jsonl" + data.write_text(json.dumps(_sample("a")) + "\n", encoding="utf-8") + service = TrainingService(_settings(tmp_path), _database(tmp_path)) + assert [sample.id for sample in service.validate_dataset(data)] == ["a"] + for name, module in FACADE.items(): + # A bound method has no ``self`` parameter; the module function carries the two injected ones. + method = list(inspect.signature(getattr(service, name)).parameters) + function = list(inspect.signature(getattr(module, name)).parameters)[2:] + assert method == function, name diff --git a/bench/audit.py b/bench/audit.py index 3479f26..c9c8172 100644 --- a/bench/audit.py +++ b/bench/audit.py @@ -26,12 +26,10 @@ ROOT = Path(__file__).resolve().parent REPORTS = ROOT / "reports" CORPUS_MANIFEST = ROOT / "corpus" / "manifest.jsonl" -# Words that only appear where a score is being compared to a target. A while-loop that re-rolls -# until one of them clears is best-of-N whatever the function is called. +# Words that only appear where a score is compared to a target; a loop that re-rolls until one clears is best-of-N whatever it is called. SCORE_WORDS = ("pass_rate", "pass_target", "accuracy", "recall", "hit_rate", "score", "target") RETRY_MARKERS = ("best_of", "bestof", "reroll", "re_roll", "until_pass", "retry_until") -# The scorer must stay a pure function of (query, ranked ids, gold ids). Anything else and the -# thing being graded starts grading itself. +# The scorer must stay a pure function of (query, ranked ids, gold ids), or the thing being graded starts grading itself. SCORER_FORBIDDEN = ("ely_eye", "transformers", "sentence_transformers", "torch", "bench.systems", "bench.counterpart") FROZEN_REPORTS = ("baseline.json", "baseline_external.json") diff --git a/bench/build_code.py b/bench/build_code.py index d5db48c..8bfc60c 100644 --- a/bench/build_code.py +++ b/bench/build_code.py @@ -124,8 +124,7 @@ def build_functions() -> tuple[list, list[dict[str, Any]], dict]: if not stripped.strip(): skipped["empty_after_strip"] += 1 continue - # One dataset row is one function; its index is the only thing that makes the - # identity unique, because a file may define the same name twice. + # One dataset row is one function; its index is the only unique identity, because a file may define the same name twice. tag = hashlib.sha256(f"{func_path}#{func_name}#{index}".encode()).hexdigest()[:10] relative = function_path(repository, func_path, func_name, tag) target = DATA_ROOT / relative diff --git a/bench/build_visrag.py b/bench/build_visrag.py index f46ab20..87cff8f 100644 --- a/bench/build_visrag.py +++ b/bench/build_visrag.py @@ -74,8 +74,7 @@ def image_name(corpus_id: str, raw: bytes) -> str: stem = corpus_id[: -len(known)] break if len(stem.encode()) > 120 or safe_segment(stem, f"corpus-id {corpus_id!r}") != stem: - # Any character we would have to rewrite means the id is not filesystem-safe: - # fall back to a digest so two ids can never land on the same file. + # Any character we would have to rewrite means the id is not filesystem-safe, so fall back to a digest that cannot collide. stem = hashlib.sha256(corpus_id.encode()).hexdigest()[:20] return stem + suffix diff --git a/bench/ingest_baseline.py b/bench/ingest_baseline.py index 1ee8125..afc0014 100644 --- a/bench/ingest_baseline.py +++ b/bench/ingest_baseline.py @@ -35,9 +35,7 @@ REPO = Path(__file__).resolve().parent.parent RUNS = REPO / "bench" / "runs" / "baseline" SHARDS = RUNS / "shards" TABLES = ("sources", "atoms", "atom_fts", "cartridges", "cache_events", "evidence_edges") -# Virtual-address ceiling per worker. ONNX Runtime and torch reserve far more address space than -# they touch, so a 3 GiB ceiling killed OCR on large images ("Unable to allocate 582. MiB for an -# array with shape (5632, 9024, 3)") even though resident memory stayed near 1.5 GiB. +# ONNX Runtime and torch reserve far more address space than they touch, so a 3 GiB ceiling killed OCR on large images ("Unable to allocate 582. MiB for an array with shape (5632, 9024, 3)") while resident memory stayed near 1.5 GiB. WORKER_MEMORY_BYTES = 6 * 1024**3 MAX_POOL_ATTEMPTS = 5 DELETE_STATEMENTS = ( @@ -125,8 +123,7 @@ def _pin_inference_threads() -> None: options = ort.SessionOptions() options.intra_op_num_threads = 1 options.inter_op_num_threads = 1 - # The CPU arena keeps freed blocks mapped, so a worker that OCRs thousands of pages grows its - # resident set until it hits its ceiling on an 11 MiB array. Hand the memory back instead. + # The CPU arena keeps freed blocks mapped, so a worker OCRing thousands of pages grows until it hits its ceiling on an 11 MiB array. options.enable_cpu_mem_arena = False kwargs["sess_options"] = options return original(model_path, *args, **kwargs) @@ -197,8 +194,7 @@ def ingest_shard(job: tuple[int, list[Path]]) -> dict: with ledger.open("a", encoding="utf-8") as handle: handle.write(json.dumps(outcome) + "\n") if ingested % 25 == 0: - # ONNX Runtime keeps the address space it has touched, and one shard runs thousands of - # files, so the arena has to be handed back or the worker hits its ceiling on an 11 MiB array. + # ONNX Runtime keeps the address space it touched, so one shard of thousands of files has to hand the arena back. gc.collect() elapsed = time.monotonic() - started print(f"shard-{index:02d} {ingested}/{len(paths)} files in {elapsed / 60:.1f} min", flush=True) @@ -310,9 +306,7 @@ def merge() -> dict: copied_objects += 1 connection.commit() - # A file processed in two attempts lands in two shards. The atoms table - # dedupes on atom_id; fts5 has no key and kept every copy (288,266 rows for - # 165,973 atoms), which would inflate the product's text search statistics. + # A file processed twice lands in two shards: atoms dedupe on atom_id but fts5 has no key and kept every copy (288,266 rows for 165,973 atoms). connection.execute("DELETE FROM atom_fts") connection.execute( "INSERT INTO atom_fts (atom_id, source, text) SELECT atom_id, source, text FROM atoms" @@ -335,8 +329,7 @@ def merge() -> dict: raise SystemExit(f"{main_db}: {len(unmapped)} sources are not named by the corpus index, first {unmapped[:3]}") product_paths = {path for path, _ in rows} product_ids = {source_id for _, source_id in rows} - # The product names a source by content hash, so two byte-identical files - # collapse into one source row (500 duplicate LongMemEval sessions). + # The product names a source by content hash, so byte-identical files collapse into one source row (500 duplicate LongMemEval sessions). collapsed_input = sorted(path for path in done - product_paths if safe_source_id(Path(path), sha256_bytes(Path(path).read_bytes())) in product_ids) unaccounted = sorted(done - product_paths - set(collapsed_input)) if unaccounted: @@ -363,8 +356,7 @@ def merge() -> dict: ) counts["uncovered_sources"] = len(total - covered) if counts["files_not_ingested"] == 0 and covered != total: - # Two honest reasons: byte-identical inputs collapse into one content-hashed - # source, and files the product cannot ingest at all. Anything else is a bug. + # Two honest reasons: byte-identical inputs collapse into one content-hashed source, and files the product cannot ingest at all. explained = counts["duplicate_content_files"] + counts["files_failed"] if counts["uncovered_sources"] > explained: raise SystemExit( diff --git a/bench/query_set.py b/bench/query_set.py new file mode 100644 index 0000000..edaa19b --- /dev/null +++ b/bench/query_set.py @@ -0,0 +1,83 @@ +"""The frozen query set: what is asked, who asks it, and which probes to time.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +from bench.sets_common import MODALITIES, SETS, read_jsonl + + +@dataclass(frozen=True) +class Query: + query_id: str + text: str + modality: str + gold: frozenset[str] + holder: str | None + answerable: bool + as_of: str | None + + + +def load_queries(modalities: Sequence[str] = MODALITIES) -> list[Query]: + """Every query the suite froze, in file order.""" + queries: list[Query] = [] + for modality in modalities: + path = SETS / f"{modality}.jsonl" + if not path.is_file(): + raise SystemExit(f"{path}: missing set, run python -m bench.build_sets first") + for record in read_jsonl(path): + gold = frozenset(hit["source_id"] for hit in record["gold"]) + answerable = not bool(record["unanswerable"]) + if answerable != bool(gold): + raise SystemExit( + f"{path}: {record['query_id']} carries {len(gold)} gold hits but " + f"unanswerable={record['unanswerable']}" + ) + queries.append( + Query( + query_id=record["query_id"], + text=record["question"], + modality=modality, + gold=gold, + holder=record["scope"]["holder"], + answerable=answerable, + as_of=record.get("as_of"), + ) + ) + if not queries: + raise SystemExit("no queries loaded") + return queries + + + +def holder_index(modalities: Sequence[str] = MODALITIES) -> dict[str, str]: + """source_id -> holder, from the corpus side of the suite.""" + holders: dict[str, str] = {} + for modality in modalities: + path = SETS / "corpus" / f"{modality}.jsonl" + if not path.is_file(): + raise SystemExit(f"{path}: missing corpus index, run python -m bench.build_sets first") + for record in read_jsonl(path): + holders[record["source_id"]] = record["holder"] + return holders + + + +SCALE_PROBES = 40 + + +def scale_probes(queries: Sequence[Query], count: int = SCALE_PROBES) -> list[Query]: + """Queries to time at every corpus size. + + Latency does not need gold hits, and a truncated corpus keeps almost none of them, + so selecting probes by gold would silently drop the small sizes entirely. + """ + if count <= 0: + raise SystemExit("the scale curve needs at least one probe query") + answerable = [query for query in queries if query.answerable] + if not answerable: + raise SystemExit("the scale curve needs answerable queries") + stride = max(1, len(answerable) // count) + return answerable[::stride][:count] diff --git a/bench/report.py b/bench/report.py new file mode 100644 index 0000000..73f5332 --- /dev/null +++ b/bench/report.py @@ -0,0 +1,211 @@ +"""The report the ablation leaves behind: metadata, the chance floor, and the markdown table.""" + +from __future__ import annotations + +import json +import platform +import subprocess +from collections.abc import Sequence +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from bench.metrics import chance_recall_at_k +from bench.query_set import Query +from bench.sets_common import MODALITIES, ROOT, SETS, sha256_file +from bench.systems import DENSE_MODEL, RERANK_MODEL, machine_facts + + +def chance_row(queries: Sequence[Query], corpus_sources: int) -> dict[str, float]: + """What a random ranking would score on the same queries.""" + answerable = [query for query in queries if query.answerable] + if not answerable: + raise SystemExit("chance needs answerable queries") + return { + "corpus_sources": corpus_sources, + "recall_at_1": sum(chance_recall_at_k(len(query.gold), corpus_sources, 1) for query in answerable) / len(answerable), + "recall_at_5": sum(chance_recall_at_k(len(query.gold), corpus_sources, 5) for query in answerable) / len(answerable), + "recall_at_10": sum(chance_recall_at_k(len(query.gold), corpus_sources, 10) for query in answerable) / len(answerable), + } + + + +def git_commit() -> str: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], cwd=ROOT, capture_output=True, text=True, check=True, timeout=30 + ) + commit = result.stdout.strip() + if len(commit) != 40: + raise SystemExit(f"git rev-parse HEAD returned {commit!r}, expected a 40 character sha") + return commit + + + +def store_schema_version(store: Path) -> int: + """The schema the measured store actually carries, read from the file.""" + import sqlite3 + + if not store.exists(): + raise SystemExit(f"{store}: the measured store is missing") + connection = sqlite3.connect(f"file:{store}?mode=ro", uri=True) + try: + row = connection.execute("PRAGMA user_version").fetchone() + finally: + connection.close() + return int(row[0]) + + + +def meta( + corpus_sources: int, queries: Sequence[Query], limit: int | None = None, store: Path | None = None +) -> dict[str, Any]: + manifest_path = SETS / "manifest.json" + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise SystemExit(f"{manifest_path}: cannot read the frozen set manifest: {error}") from error + return { + "generated_at": datetime.now(UTC).isoformat(timespec="seconds"), + "git_commit": git_commit(), + "python": platform.python_version(), + "machine": machine_facts(), + "corpus_manifest_sha256": manifest.get("corpus_manifest_sha256"), + "sets": {name: sha256_file(SETS / f"{name}.jsonl") for name in MODALITIES}, + "queries": len(queries), + "answerable": sum(1 for query in queries if query.answerable), + "limit_per_modality": limit, + "corpus_sources": corpus_sources, + "store_schema_version_at_start": None if store is None else store_schema_version(store), + "models": {"dense": DENSE_MODEL, "rerank": RERANK_MODEL}, + } + + + +def scale_arm_names(rows: Sequence[dict[str, Any]]) -> list[str]: + """Arms a scale row reports: only keys whose value is a latency dict, in order. + + A scale row also carries strings (the note explaining what the row means), and + treating those as arms is how the report writer first crashed. + """ + names: list[str] = [] + for row in rows: + for name, value in row.items(): + if isinstance(value, dict) and name not in names: + names.append(name) + if not names: + raise SystemExit("scale rows report no measured arm") + return names + + + +def write_markdown(report: dict[str, Any], path: Path, name: str = "BASELINE") -> None: + """Human-readable twin of the JSON report.""" + schema = report["meta"].get("store_schema_version_at_start") + lines = [ + f"# {name}: what the product retrieves at this build", + "", + f"- generated {report['meta']['generated_at']} at commit {report['meta']['git_commit'][:12]}", + ("" if schema is None else f"- measured store schema: v{schema}"), + ( + f"- corpus {report['meta']['corpus_sources']} sources, {report['meta']['queries']} queries " + f"({report['meta']['answerable']} answerable)" + ), + f"- models: dense `{report['meta']['models']['dense']}`, rerank `{report['meta']['models']['rerank']}`", + ( + f"- machine: {report['meta']['machine']['cpus']} cores, " + f"{report['meta']['machine']['threads_per_query']} threads per query, load average " + f"{report['meta']['machine']['load_average_1_5_15'][0]:.1f} when the report was written" + ), + "", + "| arm | recall@1 | recall@5 | recall@10 | nDCG@10 | MRR | hit@10 (95% lower) | p50 s | p95 s |", + "|---|---|---|---|---|---|---|---|---|", + ] + chance = report["chance"] + lines.append( + f"| chance (random ranking) | {chance['recall_at_1']:.3f} | {chance['recall_at_5']:.3f} " + f"| {chance['recall_at_10']:.3f} | - | - | - | - | - |" + ) + if note := next((arm["latency_seconds"].get("note") for arm in report["arms"] if arm["latency_seconds"].get("note")), None): + lines += ["", f"_{note} (n={report['arms'][0]['latency_seconds']['queries']} queries)_"] + for arm in report["arms"]: + overall = arm["overall"] + latency = arm["latency_seconds"] + lines.append( + f"| {arm['arm']} | {overall['recall_at_1']:.3f} | {overall['recall_at_5']:.3f} " + f"| {overall['recall_at_10']:.3f} | {overall['ndcg_at_10']:.3f} | {overall['mrr']:.3f} " + f"| {overall['hit_at_10']:.3f} ({overall['hit_at_10_wilson_lower']:.3f}) " + f"| {latency['p50']:.3f} | {latency['p95']:.3f} |" + ) + shown = [modality for modality in MODALITIES if any(modality in arm["by_modality"] for arm in report["arms"])] + if not shown: + raise SystemExit("no per-modality scores to write") + lines += ["", "## recall@10 by modality", "", "| arm | " + " | ".join(shown) + " |", "|---|" + "---|" * len(shown)] + for arm in report["arms"]: + cells = [ + f"{arm['by_modality'][modality]['recall_at_10']:.3f}" if modality in arm["by_modality"] else "-" + for modality in shown + ] + lines.append(f"| {arm['arm']} | " + " | ".join(cells) + " |") + total = report["meta"]["queries"] + if total <= 0: + raise SystemExit("the report records no queries") + share = 1 - report["meta"]["answerable"] / total + lines += ["", f"## unanswerable queries ({share:.1%} of the suite)", ""] + for arm in report["arms"]: + overall = arm["overall"] + lines.append( + f"- {arm['arm']}: abstains on {overall['abstention_rate_unanswerable']:.0%} of unanswerable queries, " + f"and its top hit comes from the query's own holder for {overall['top_hit_from_holder_unanswerable']:.0%} of them; " + "the rest are confident answers pulled out of somebody else's memory" + ) + if "live_product" in report: + live = report["live_product"] + lines += [ + "", + "## the product path, asked as its holder", + "", + ( + f"- RetrievalService.search as the query's own holder: recall@10 {live['recall_at_10']:.3f}, " + f"MRR {live['mrr']:.3f}, hit@10 {live['hit_at_10']:.3f} (95% lower {live['hit_at_10_wilson_lower']:.3f}) " + f"over n={live['queries']} answerable queries" + ), + ( + f"- latency p50 {live['p50']:.3f}s, p95 {live['p95']:.3f}s, mean {live['mean']:.3f}s, " + f"max {live['max']:.3f}s" + ), + ( + f"- top-10 overlap with the fast BM25 path: {live['top10_overlap_with_bm25_only']:.1%}, " + f"identical order {live['identical_top10']:.1%}" + ), + ( + f"- its top-{live['top_k']} atoms cover " + f"{live['distinct_sources_in_top10']['mean']:.1f} distinct sources on average " + f"(min {live['distinct_sources_in_top10']['min']}): the product ranks atoms and never diversifies, " + "so ten slots can fill with fragments of one document" + ), + ] + if report.get("scale"): + measured = scale_arm_names(report["scale"]) + header = " | ".join(f"{arm} p50/p95 s" for arm in measured) + lines += ["", "## latency vs corpus size", "", f"| sources | atoms | {header} |", "|" + "---|" * (2 + len(measured))] + for row in report["scale"]: + cells = [f"{row[arm]['p50']:.3f}/{row[arm]['p95']:.3f}" for arm in measured] + lines.append(f"| {row['sources']} | {row['atoms']} | " + " | ".join(cells) + " |") + note = report["scale"][0].get("note") + if note: + lines += ["", f"_{note}_"] + load = report["meta"]["machine"]["load_average_1_5_15"][0] + lines.append( + f"\n_latency was measured at load average {load:.1f} on {report['meta']['machine']['cpus']} cores; " + "a busy machine inflates these figures, so read them as a ceiling_" + ) + lines += [ + "", + "## reading", + "", + "- recall counts a source as found only when it is the source a third party marked correct.", + "- abstention is the only honest number for unanswerable queries: a retriever that always", + " returns ten results never says it does not know.", + "- chance is the floor any of these arms must beat; a number near chance is not a capability.", + ] + path.write_text("\n".join(lines) + "\n", encoding="utf-8") diff --git a/bench/run.py b/bench/run.py index dab452f..3259315 100644 --- a/bench/run.py +++ b/bench/run.py @@ -1,132 +1,36 @@ -"""Run the retrieval ablation over the frozen corpus and write the baseline report. - -Five arms, one database: sparse only, dense only, graph only, the product's weighted -fusion, and fusion plus a cross-encoder reranker. Every number in the report is either -measured or explicitly marked as chance. - -Run: python -m bench.run [--arms bm25_only] [--limit 200] [--no-live] [--out bench/reports] -""" +"""Run the retrieval ablation over the frozen corpus and write the report: python -m bench.run [--arms bm25_only] [--limit 200] [--no-live].""" from __future__ import annotations import argparse import hashlib import json -import platform -import subprocess from collections.abc import Sequence -from dataclasses import dataclass -from datetime import UTC, datetime from pathlib import Path from typing import Any import numpy as np -from bench.metrics import aggregate, chance_recall_at_k, score_query -from bench.sets_common import MODALITIES, SETS, read_jsonl, sha256_file +from bench.metrics import aggregate, score_query +from bench.query_set import Query, holder_index, load_queries, scale_probes +from bench.report import chance_row, meta, write_markdown +from bench.sets_common import MODALITIES, ROOT from bench.systems import ( ARMS, - DENSE_MODEL, MODELS, - RERANK_MODEL, ArmRunner, Atom, JointRunner, arm_needs, live_product_rankings, load_atoms, - machine_facts, ) -ROOT = SETS.parent REPORTS = ROOT / "reports" DB_HOME = ROOT / "runs" / "baseline" LIVE_QUERIES = 300 # the acceptance bar names n>=300; the SQL path answers in well under a second PROGRESS_EVERY = 250 # an arm takes tens of minutes over 5,308 queries; silence looks like a hang SCALE_SIZES = (500, 2_000, 10_000, 50_000) -SCALE_PROBES = 40 - - -@dataclass(frozen=True) -class Query: - query_id: str - text: str - modality: str - gold: frozenset[str] - holder: str | None - answerable: bool - as_of: str | None - - -def load_queries(modalities: Sequence[str] = MODALITIES) -> list[Query]: - """Every query the suite froze, in file order.""" - queries: list[Query] = [] - for modality in modalities: - path = SETS / f"{modality}.jsonl" - if not path.is_file(): - raise SystemExit(f"{path}: missing set, run python -m bench.build_sets first") - for record in read_jsonl(path): - gold = frozenset(hit["source_id"] for hit in record["gold"]) - answerable = not bool(record["unanswerable"]) - if answerable != bool(gold): - raise SystemExit( - f"{path}: {record['query_id']} carries {len(gold)} gold hits but " - f"unanswerable={record['unanswerable']}" - ) - queries.append( - Query( - query_id=record["query_id"], - text=record["question"], - modality=modality, - gold=gold, - holder=record["scope"]["holder"], - answerable=answerable, - as_of=record.get("as_of"), - ) - ) - if not queries: - raise SystemExit("no queries loaded") - return queries - - -def holder_index(modalities: Sequence[str] = MODALITIES) -> dict[str, str]: - """source_id -> holder, from the corpus side of the suite.""" - holders: dict[str, str] = {} - for modality in modalities: - path = SETS / "corpus" / f"{modality}.jsonl" - if not path.is_file(): - raise SystemExit(f"{path}: missing corpus index, run python -m bench.build_sets first") - for record in read_jsonl(path): - holders[record["source_id"]] = record["holder"] - return holders - - -def chance_row(queries: Sequence[Query], corpus_sources: int) -> dict[str, float]: - """What a random ranking would score on the same queries.""" - answerable = [query for query in queries if query.answerable] - if not answerable: - raise SystemExit("chance needs answerable queries") - return { - "corpus_sources": corpus_sources, - "recall_at_1": sum(chance_recall_at_k(len(query.gold), corpus_sources, 1) for query in answerable) / len(answerable), - "recall_at_5": sum(chance_recall_at_k(len(query.gold), corpus_sources, 5) for query in answerable) / len(answerable), - "recall_at_10": sum(chance_recall_at_k(len(query.gold), corpus_sources, 10) for query in answerable) / len(answerable), - } - - -def scale_probes(queries: Sequence[Query], count: int = SCALE_PROBES) -> list[Query]: - """Queries to time at every corpus size. - - Latency does not need gold hits, and a truncated corpus keeps almost none of them, - so selecting probes by gold would silently drop the small sizes entirely. - """ - if count <= 0: - raise SystemExit("the scale curve needs at least one probe query") - answerable = [query for query in queries if query.answerable] - if not answerable: - raise SystemExit("the scale curve needs answerable queries") - stride = max(1, len(answerable) // count) - return answerable[::stride][:count] def scale_curve(queries: Sequence[Query], atoms: Sequence[Atom], arms: Sequence[str]) -> list[dict[str, Any]]: @@ -209,183 +113,6 @@ def live_check( return metrics -def git_commit() -> str: - result = subprocess.run( - ["git", "rev-parse", "HEAD"], cwd=ROOT, capture_output=True, text=True, check=True, timeout=30 - ) - commit = result.stdout.strip() - if len(commit) != 40: - raise SystemExit(f"git rev-parse HEAD returned {commit!r}, expected a 40 character sha") - return commit - - -def store_schema_version(store: Path) -> int: - """The schema the measured store actually carries, read from the file.""" - import sqlite3 - - if not store.exists(): - raise SystemExit(f"{store}: the measured store is missing") - connection = sqlite3.connect(f"file:{store}?mode=ro", uri=True) - try: - row = connection.execute("PRAGMA user_version").fetchone() - finally: - connection.close() - return int(row[0]) - - -def meta( - corpus_sources: int, queries: Sequence[Query], limit: int | None = None, store: Path | None = None -) -> dict[str, Any]: - manifest_path = SETS / "manifest.json" - try: - manifest = json.loads(manifest_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as error: - raise SystemExit(f"{manifest_path}: cannot read the frozen set manifest: {error}") from error - return { - "generated_at": datetime.now(UTC).isoformat(timespec="seconds"), - "git_commit": git_commit(), - "python": platform.python_version(), - "machine": machine_facts(), - "corpus_manifest_sha256": manifest.get("corpus_manifest_sha256"), - "sets": {name: sha256_file(SETS / f"{name}.jsonl") for name in MODALITIES}, - "queries": len(queries), - "answerable": sum(1 for query in queries if query.answerable), - "limit_per_modality": limit, - "corpus_sources": corpus_sources, - "store_schema_version_at_start": None if store is None else store_schema_version(store), - "models": {"dense": DENSE_MODEL, "rerank": RERANK_MODEL}, - } - - -def scale_arm_names(rows: Sequence[dict[str, Any]]) -> list[str]: - """Arms a scale row reports: only keys whose value is a latency dict, in order. - - A scale row also carries strings (the note explaining what the row means), and - treating those as arms is how the report writer first crashed. - """ - names: list[str] = [] - for row in rows: - for name, value in row.items(): - if isinstance(value, dict) and name not in names: - names.append(name) - if not names: - raise SystemExit("scale rows report no measured arm") - return names - - -def write_markdown(report: dict[str, Any], path: Path, name: str = "BASELINE") -> None: - """Human-readable twin of the JSON report.""" - schema = report["meta"].get("store_schema_version_at_start") - lines = [ - f"# {name}: what the product retrieves at this build", - "", - f"- generated {report['meta']['generated_at']} at commit {report['meta']['git_commit'][:12]}", - ("" if schema is None else f"- measured store schema: v{schema}"), - ( - f"- corpus {report['meta']['corpus_sources']} sources, {report['meta']['queries']} queries " - f"({report['meta']['answerable']} answerable)" - ), - f"- models: dense `{report['meta']['models']['dense']}`, rerank `{report['meta']['models']['rerank']}`", - ( - f"- machine: {report['meta']['machine']['cpus']} cores, " - f"{report['meta']['machine']['threads_per_query']} threads per query, load average " - f"{report['meta']['machine']['load_average_1_5_15'][0]:.1f} when the report was written" - ), - "", - "| arm | recall@1 | recall@5 | recall@10 | nDCG@10 | MRR | hit@10 (95% lower) | p50 s | p95 s |", - "|---|---|---|---|---|---|---|---|---|", - ] - chance = report["chance"] - lines.append( - f"| chance (random ranking) | {chance['recall_at_1']:.3f} | {chance['recall_at_5']:.3f} " - f"| {chance['recall_at_10']:.3f} | - | - | - | - | - |" - ) - if note := next((arm["latency_seconds"].get("note") for arm in report["arms"] if arm["latency_seconds"].get("note")), None): - lines += ["", f"_{note} (n={report['arms'][0]['latency_seconds']['queries']} queries)_"] - for arm in report["arms"]: - overall = arm["overall"] - latency = arm["latency_seconds"] - lines.append( - f"| {arm['arm']} | {overall['recall_at_1']:.3f} | {overall['recall_at_5']:.3f} " - f"| {overall['recall_at_10']:.3f} | {overall['ndcg_at_10']:.3f} | {overall['mrr']:.3f} " - f"| {overall['hit_at_10']:.3f} ({overall['hit_at_10_wilson_lower']:.3f}) " - f"| {latency['p50']:.3f} | {latency['p95']:.3f} |" - ) - shown = [modality for modality in MODALITIES if any(modality in arm["by_modality"] for arm in report["arms"])] - if not shown: - raise SystemExit("no per-modality scores to write") - lines += ["", "## recall@10 by modality", "", "| arm | " + " | ".join(shown) + " |", "|---|" + "---|" * len(shown)] - for arm in report["arms"]: - cells = [ - f"{arm['by_modality'][modality]['recall_at_10']:.3f}" if modality in arm["by_modality"] else "-" - for modality in shown - ] - lines.append(f"| {arm['arm']} | " + " | ".join(cells) + " |") - total = report["meta"]["queries"] - if total <= 0: - raise SystemExit("the report records no queries") - share = 1 - report["meta"]["answerable"] / total - lines += ["", f"## unanswerable queries ({share:.1%} of the suite)", ""] - for arm in report["arms"]: - overall = arm["overall"] - lines.append( - f"- {arm['arm']}: abstains on {overall['abstention_rate_unanswerable']:.0%} of unanswerable queries, " - f"and its top hit comes from the query's own holder for {overall['top_hit_from_holder_unanswerable']:.0%} of them; " - "the rest are confident answers pulled out of somebody else's memory" - ) - if "live_product" in report: - live = report["live_product"] - lines += [ - "", - "## the product path, asked as its holder", - "", - ( - f"- RetrievalService.search as the query's own holder: recall@10 {live['recall_at_10']:.3f}, " - f"MRR {live['mrr']:.3f}, hit@10 {live['hit_at_10']:.3f} (95% lower {live['hit_at_10_wilson_lower']:.3f}) " - f"over n={live['queries']} answerable queries" - ), - ( - f"- latency p50 {live['p50']:.3f}s, p95 {live['p95']:.3f}s, mean {live['mean']:.3f}s, " - f"max {live['max']:.3f}s" - ), - ( - f"- top-10 overlap with the fast BM25 path: {live['top10_overlap_with_bm25_only']:.1%}, " - f"identical order {live['identical_top10']:.1%}" - ), - ( - f"- its top-{live['top_k']} atoms cover " - f"{live['distinct_sources_in_top10']['mean']:.1f} distinct sources on average " - f"(min {live['distinct_sources_in_top10']['min']}): the product ranks atoms and never diversifies, " - "so ten slots can fill with fragments of one document" - ), - ] - if report.get("scale"): - measured = scale_arm_names(report["scale"]) - header = " | ".join(f"{arm} p50/p95 s" for arm in measured) - lines += ["", "## latency vs corpus size", "", f"| sources | atoms | {header} |", "|" + "---|" * (2 + len(measured))] - for row in report["scale"]: - cells = [f"{row[arm]['p50']:.3f}/{row[arm]['p95']:.3f}" for arm in measured] - lines.append(f"| {row['sources']} | {row['atoms']} | " + " | ".join(cells) + " |") - note = report["scale"][0].get("note") - if note: - lines += ["", f"_{note}_"] - load = report["meta"]["machine"]["load_average_1_5_15"][0] - lines.append( - f"\n_latency was measured at load average {load:.1f} on {report['meta']['machine']['cpus']} cores; " - "a busy machine inflates these figures, so read them as a ceiling_" - ) - lines += [ - "", - "## reading", - "", - "- recall counts a source as found only when it is the source a third party marked correct.", - "- abstention is the only honest number for unanswerable queries: a retriever that always", - " returns ten results never says it does not know.", - "- chance is the floor any of these arms must beat; a number near chance is not a capability.", - ] - path.write_text("\n".join(lines) + "\n", encoding="utf-8") - - def load_config(parser: argparse.ArgumentParser, args: argparse.Namespace, path: Path) -> None: """Fill in whatever the command line left at its default; an explicit flag still wins.""" try: @@ -514,8 +241,7 @@ def main() -> int: load_config(parser, args, args.config) modalities = args.modality or list(MODALITIES) - # action="extend" appends to the default, so an explicit --arms would otherwise - # run every arm plus the requested ones. + # action="extend" appends to the default, so an explicit --arms would otherwise run every arm plus the requested ones. arms_requested: list[str] = list(args.arms or ARMS) queries = load_queries(modalities) if args.limit: @@ -533,8 +259,7 @@ def main() -> int: "meta": meta(corpus_sources, queries, args.limit, store=DB_HOME / "data" / "ely_eye.sqlite"), "chance": chance_row(queries, corpus_sources), } - # One pass over the queries: each signal is computed once and every arm is derived - # from it, and each query is checkpointed so a kill costs one query, not an arm. + # One pass over the queries: each signal is computed once, every arm derives from it, and each query is checkpointed so a kill costs one query. mark = fingerprint(arms_requested, queries, atoms) rankings_path = DB_HOME / "rankings.jsonl" done = load_rankings(rankings_path, mark) diff --git a/bench/run_external.py b/bench/run_external.py index ccdcab4..fa30c2f 100644 --- a/bench/run_external.py +++ b/bench/run_external.py @@ -27,7 +27,8 @@ from bench.counterpart import ( ) from bench.ingest_baseline import RUNS, database_path from bench.metrics import aggregate, chance_recall_at_k, score_query -from bench.run import Query, git_commit, holder_index, load_queries +from bench.query_set import Query, holder_index, load_queries +from bench.report import git_commit from bench.sets_common import ROOT from bench.systems import machine_facts diff --git a/bench/scale_check.py b/bench/scale_check.py index e76b211..d81e971 100644 --- a/bench/scale_check.py +++ b/bench/scale_check.py @@ -14,7 +14,8 @@ import time from pathlib import Path from typing import Any -from bench.run import latency, load_queries, scale_probes +from bench.query_set import load_queries, scale_probes +from bench.run import latency from bench.systems import machine_facts DEFAULT_HOME = Path("bench/runs/baseline") diff --git a/bench/sets_common.py b/bench/sets_common.py index 6d49041..27d3cdf 100644 --- a/bench/sets_common.py +++ b/bench/sets_common.py @@ -13,8 +13,7 @@ from typing import Any ROOT = Path(__file__).resolve().parent CORPUS = ROOT / "corpus" SETS = ROOT / "sets" -# Derived media (images pulled out of parquet, video frames, crops) lives outside -# the repository: tens of thousands of new files inside it swamps repo tooling. +# Derived media (parquet images, video frames, crops) lives outside the repository: tens of thousands of new files inside it swamps repo tooling. DATA_ROOT = Path(os.environ.get("ELY_EYE_BENCH_DATA", str(Path.home() / "ely_eye_bench_data"))) CORPUS_ROOTS = {"corpus": CORPUS, "derived": DATA_ROOT} MODALITIES = ("pdf", "image", "ui_screenshot", "video", "code", "conversation") diff --git a/bench/systems.py b/bench/systems.py index 7f0aa09..037a99f 100644 --- a/bench/systems.py +++ b/bench/systems.py @@ -324,8 +324,7 @@ class ArmRunner: return float(np.percentile(np.asarray(self.timings), fraction * 100)) -# Which measured signals each arm consumes; JointRunner times each signal once and -# attributes it to every arm that uses it. +# Which measured signals each arm consumes; JointRunner times each signal once and attributes it to every arm that uses it. ARM_COMPONENTS: dict[str, tuple[str, ...]] = { "bm25_only": ("sparse",), "dense_only": ("dense",), diff --git a/bench/tests/test_joint.py b/bench/tests/test_joint.py index 9cd574b..3dbb3ea 100644 --- a/bench/tests/test_joint.py +++ b/bench/tests/test_joint.py @@ -16,16 +16,9 @@ import numpy as np import pytest from bench import systems -from bench.run import ( - Query, - fingerprint, - latency, - load_rankings, - remeasure_latency, - scale_arm_names, - scale_probes, - write_markdown, -) +from bench.query_set import Query, scale_probes +from bench.report import scale_arm_names, write_markdown +from bench.run import fingerprint, latency, load_rankings, remeasure_latency from bench.systems import ARMS, ArmRunner, Atom, JointRunner from bench.tests.test_systems import ATOMS diff --git a/bench/tests/test_systems.py b/bench/tests/test_systems.py index 9ef5a70..6c31ab4 100644 --- a/bench/tests/test_systems.py +++ b/bench/tests/test_systems.py @@ -27,8 +27,7 @@ ATOMS = [ Atom(atom_id="a3", source_id="src_red", text="the key hangs on the red door"), ] -# The product names a source by content hash; the benchmark names the same file by -# its corpus id (the source_id above). load_atoms joins the two through sources.path. +# The product names a source by content hash, the benchmark by corpus id; load_atoms joins the two through sources.path. PRODUCT_SOURCES = {"prod_red_9f3a": "/fixture/red.png", "prod_blue_1c77": "/fixture/blue.png"} ATOM_PRODUCT_IDS = ["prod_red_9f3a", "prod_blue_1c77", "prod_red_9f3a"] @@ -123,8 +122,7 @@ def test_load_atoms_returns_every_row_in_a_stable_order(tmp_path: Path, monkeypa make_database(db) loaded = load_atoms(db) assert [atom.atom_id for atom in loaded] == ["a1", "a2", "a3"] - # The database stores the product's content-hash id; the caller gets the corpus id - # the gold annotations use. + # The database stores the product's content-hash id; the caller gets the corpus id the gold annotations use. assert [atom.source_id for atom in loaded] == ["src_red", "src_blue", "src_red"] assert loaded[2].text == "the key hangs on the red door"