Files
aituner/runs/frontier-code-trace-v0/prepare_code_window.py

117 lines
4.2 KiB
Python

#!/usr/bin/env python3
"""Materialize the stable code window selected by audit_code_trace.py."""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
from typing import Any
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--audit", type=Path, required=True)
parser.add_argument("--output-root", type=Path, required=True)
parser.add_argument("--sample-seed", type=int, default=20260723)
return parser.parse_args()
def session_uniform(seed: int, window_id: str, session_root: Any) -> float:
payload = json.dumps(
{"seed": seed, "window_id": window_id, "session_root": session_root},
sort_keys=True,
separators=(",", ":"),
).encode()
return int.from_bytes(hashlib.blake2b(payload, digest_size=8).digest(), "big") / (
1 << 64
)
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as stream:
for chunk in iter(lambda: stream.read(1 << 20), b""):
digest.update(chunk)
return digest.hexdigest()
def main() -> None:
args = parse_args()
audit = json.loads(args.audit.read_text())
if audit["data_gate"] != "PASS":
raise ValueError(f"trace data gate is not PASS: {audit['data_gate']}")
selected = audit["selected"]
source = Path(selected["source"])
window = selected["stable_window"]
start = float(window["start_timestamp"])
end = float(window["end_timestamp"])
if args.output_root.exists():
raise ValueError(f"refusing to overwrite {args.output_root}")
args.output_root.mkdir(parents=True)
destination = args.output_root / "code-raw-window.jsonl"
root_of: dict[Any, Any] = {}
request_count = 0
with source.open() as input_stream, destination.open("w") as output_stream:
for source_index, line in enumerate(input_stream):
if not line.strip():
continue
row = json.loads(line)
timestamp = float(row["timestamp"])
if timestamp < start:
continue
if timestamp >= end:
break
if int(row["input_length"]) <= 0 or int(row["output_length"]) <= 0:
continue
chat = row.get("chat_id", source_index)
parent = row.get("parent_chat_id")
has_parent = parent not in (None, "", -1, "-1")
session_root = root_of.get(parent, parent) if has_parent else chat
root_of[chat] = session_root
materialized = {
**row,
"source_index": source_index,
"session_root": session_root,
"sampling_u": session_uniform(
args.sample_seed,
f"code-{start:.6f}-{end:.6f}",
session_root,
),
}
output_stream.write(
json.dumps(materialized, ensure_ascii=False, separators=(",", ":"))
+ "\n"
)
request_count += 1
expected = int(selected["selected_window_stats"]["requests"])
if request_count != expected:
raise ValueError(f"window request mismatch: materialized={request_count}, audit={expected}")
manifest = {
"schema": "frontier-code-window-v1",
"audit": str(args.audit.resolve()),
"audit_sha256": sha256(args.audit),
"source": str(source.resolve()),
"source_block_size": selected["hash_contract"]["exact_source_block_size"],
"target_block_size": 16,
"start_timestamp": start,
"end_timestamp": end,
"duration_s": end - start,
"requests": request_count,
"sample_seed": args.sample_seed,
"sampling_rule": "session-coherent deterministic sampling_u",
"max_model_len": audit["max_model_len_recommendation"],
"window_stats": selected["selected_window_stats"],
"raw_window": str(destination.resolve()),
"raw_window_sha256": sha256(destination),
}
(args.output_root / "window-manifest.json").write_text(
json.dumps(manifest, indent=2, sort_keys=True) + "\n"
)
print(json.dumps({"requests": request_count, "output_root": str(args.output_root)}))
if __name__ == "__main__":
main()