Keep every file inside the rules the repo already wrote down

The 500-line rule and the one-line-comment rule were enforced only from
memory, so they drifted. bench/filesize_check.py now checks the budget on
every tracked source file, and it passes.

Split, with behaviour proven unchanged by comparing each moved function body
against HEAD byte for byte (53 training bodies, 21 benchmark bodies, 0
differences):

  db.py        503 -> 396   store_search.py     105
  bench/run.py 609 -> 351   query_set.py         80, report.py 211
  benchmarks.py 619 -> 424  benchmark_scoring.py 262
  training.py 1630 ->  94   training_{common,encode,metrics,datasets,lora,retrieval}

TrainingService stays as a thin facade for the CLI; the facade's parameter
lists are pinned against the module functions by inspect, which is what
catches a leaked self.

Deleted the eight top-level functions no file called, after fixing the scan
that had reported 201 false positives by skipping the defining file. The
v1/v2/v3 sweep found nothing to delete and says so: every remaining version
tag is an on-disk format marker, a migration step, a fixture name, or a
third-party model name.

Collapsed 39 multi-line comment blocks (86 lines) to one line each, keeping
the measured facts and error strings that justify the code.

269 tests pass; ruff clean; mypy 24 errors in 9 files against a 24-in-6
baseline, net zero; pyright 4, the same four HEAD measures.
This commit is contained in:
ZacharyZhang-NY committed 2026-10-08 08:20:21 -04:00
1 parent 56b94cd58e
commit 9cb46c18d3
47 files changed
+2678 -2272

No files matched your search

+16 -6
View File
@@ -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 干净
---
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -3
View File
@@ -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
+262
View File
@@ -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":"<recognized text>"}}
"""
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":"<prediction>"}}
"""
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"
+16 -211
View File
@@ -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":"<recognized text>"}}
"""
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":"<prediction>"}}
"""
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"
+1 -2
View File
@@ -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")
+1 -2
View File
@@ -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
+1 -2
View File
@@ -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"}
)
+1 -108
View File
@@ -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()
+2 -4
View File
@@ -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,
+2 -6
View File
@@ -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'"),
+1 -2
View File
@@ -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")
+4 -10
View File
@@ -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)
@@ -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),
+1 -3
View File
@@ -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"] = (
+2 -2
View File
@@ -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:
+4 -4
View File
@@ -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)
)
+106
View File
@@ -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)
+2 -4
View File
@@ -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+")
File diff suppressed because it is too large. Load diff
+108
View File
@@ -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 ""
+387
View File
@@ -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
+177
View File
@@ -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() + "..."
+408
View File
@@ -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
+307
View File
@@ -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
+358
View File
@@ -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
+1 -2
View File
@@ -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
+1 -2
View File
@@ -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]] = {}
+1 -1
View File
@@ -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:
+3 -2
View File
@@ -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
+3 -6
View File
@@ -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)}
+3 -2
View File
@@ -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"
+152
View File
@@ -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
+2 -4
View File
@@ -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")
+1 -2
View File
@@ -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
+1 -2
View File
@@ -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
+6 -14
View File
@@ -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(
+83
View File
@@ -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]
+211
View File
@@ -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")
+7 -282
View File
@@ -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)
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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")
+1 -2
View File
@@ -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")
+1 -2
View File
@@ -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",),
+3 -10
View File
@@ -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
+2 -4
View File
@@ -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"