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"