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:
1 parent
56b94cd58e
commit
9cb46c18d3
47 files changed
+2678
-2272
No files matched your search
@@ -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
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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'"),
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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"] = (
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
@@ -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+")
|
||||
|
||||
+19
-1555
File diff suppressed because it is too large.
Load diff
@@ -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 ""
|
||||
@@ -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
|
||||
@@ -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() + "..."
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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]] = {}
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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)}
|
||||
@@ -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"
|
||||
@@ -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
|
||||
Reference in new issue
Block a user