from __future__ import annotations import argparse import hashlib import json import re import sys from pathlib import Path from typing import Any SCRIPT_DIR = Path(__file__).resolve().parent V3P5_SCRIPT_DIR = SCRIPT_DIR.parents[1] / "v3p5_static" / "scripts" if str(SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(SCRIPT_DIR)) if str(V3P5_SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(V3P5_SCRIPT_DIR)) from smoke_v3p4_architecture import load_json from train_v4p4_cloud import DEFAULT_STAGE0_DIR, DEFAULT_TENSOR_DIR, build_color_map AUTHORITATIVE_SCRIPTS = [ "model_v4p4.py", "loss_v4p4.py", "config_v4p4_train.py", "train_v4p4_cloud.py", "hf_submit_v4p4_world_model_train_job.py", "hf_submit_v4p4_eval_export_job.py", "export_v4p4_visit_outputs.py", "smoke_v4p4_architecture.py", "rollout_v4p4_closed_loop.py", "evaluate_v4p4_independent.py", "evaluate_v4p4_closed_loop_rollout.py", "evaluate_v4p4_final_metric_tables.py", "evaluate_v4p4_world_model.py", "audit_v4p4_metric_coverage.py", "audit_v4p4_finalization.py", ] def sha256_file(path: Path) -> str: h = hashlib.sha256() with path.open("rb") as f: for chunk in iter(lambda: f.read(1024 * 1024), b""): h.update(chunk) return h.hexdigest() def json_ready(value: Any) -> Any: if isinstance(value, Path): return str(value) if isinstance(value, dict): return {str(k): json_ready(v) for k, v in value.items()} if isinstance(value, (list, tuple)): return [json_ready(v) for v in value] return value def read_text(name: str) -> str: return (SCRIPT_DIR / name).read_text(encoding="utf-8") def read_optional_text(name: str) -> str: path = SCRIPT_DIR / name return path.read_text(encoding="utf-8") if path.exists() else "" def read_compat_script_text(name: str, legacy_dir: str) -> str: candidates = [ SCRIPT_DIR / name, SCRIPT_DIR.parents[1] / legacy_dir / "scripts" / name, ] for path in candidates: if path.exists(): return path.read_text(encoding="utf-8") raise FileNotFoundError( f"Could not find {name}; checked: " + ", ".join(str(path) for path in candidates) ) def audit_code_version() -> dict[str, Any]: rows = [] for name in AUTHORITATIVE_SCRIPTS: path = SCRIPT_DIR / name rows.append( { "script": name, "exists": path.exists(), "sha256": sha256_file(path) if path.exists() else "", "bytes": path.stat().st_size if path.exists() else 0, } ) py_files = sorted(p.name for p in SCRIPT_DIR.glob("*.py")) return { "status": "pass" if all(row["exists"] for row in rows) else "fail", "authoritative_scripts": rows, "python_files_in_script_dir": py_files, "legacy_policy": "Use v4p4 scripts for v4p4 training/evaluation; v4p3/v3p5 imports are compatibility helpers only.", } def audit_export_sources() -> dict[str, Any]: src = read_text("export_v4p4_visit_outputs.py") checks = { "pre_visit_uses_history_deployable": 'out["z_history"], torch.sigmoid(out["deployable_risk_logits_pre"]), out["pwe_log_lambda_pre"]' in src, "contact_uses_contact_deployable": 'out["z_contact"], torch.sigmoid(out["deployable_risk_logits_contact"]), out["pwe_log_lambda_contact"]' in src, "field_end_uses_post_deployable": 'out["z_post"], torch.sigmoid(out["deployable_risk_logits_post"]), out["pwe_log_lambda_post"]' in src, "export_disables_reconstruction": "compute_reconstruction=False" in src, "export_disables_teacher_forcing": "compute_teacher_forced=False" in src, "export_does_not_read_event_logits": 'out["event_logits"]' not in src and "out['event_logits']" not in src, } return { "status": "pass" if all(checks.values()) else "fail", "checks": checks, "expected_mode_sources": { "pre_visit": {"embedding": "z_history", "risk": "deployable_risk_logits_pre", "pwe": "pwe_log_lambda_pre"}, "contact": {"embedding": "z_contact", "risk": "deployable_risk_logits_contact", "pwe": "pwe_log_lambda_contact"}, "field_end": {"embedding": "z_post", "risk": "deployable_risk_logits_post", "pwe": "pwe_log_lambda_post"}, }, } def audit_model_factorization() -> dict[str, Any]: src = read_text("model_v4p4.py") checks = { "active_state_head_uses_time_context": 'out["active_state_logits"] = self.active_state_head(time_ctx)' in src, "downstream_heads_condition_on_active_context": "obs_ctx = self._condition_on_next_active(time_ctx, active_ids)" in src, "rollout_downstream_conditions_on_sampled_active": "active_ids = torch.where(next_contact, active_state" in src and "obs_ctx = self._condition_on_next_active(time_ctx, active_ids)" in src, "event_generation_head_dedicated": "self.event_generation_head" in src and "generation_event_logits = self.event_generation_head(obs_ctx)" in src, "year_uses_absolute_elapsed_time": "start_year + torch.floor(next_time / 365.25).long()" in src, "event_logits_removed": '"event_logits"' not in re.sub(r"event_logits_absent", "", src), } return { "status": "pass" if all(checks.values()) else "fail", "checks": checks, "factorization": "p(K,dt|z_post) -> p(S_next|z_post,dt) -> p(M_next,X_next,Y_next|z_post,dt,S_next)", } def audit_pwe_math() -> dict[str, Any]: src = read_compat_script_text("loss_v4p3.py", "v4p3_l1l2") v4p4_src = read_text("loss_v4p4.py") checks = { "exact_time_density_log_survival_plus_log_lambda": "event_loglik = log_survival_total + torch.log(lambda_k_star.clamp(min=1.0e-8))" in src, "admin_censor_survival_only": "torch.where(is_censored, log_survival_total, event_loglik)" in src, "v4p4_scores_pre_contact_post_pwe": all(token in v4p4_src for token in ["pwe_log_lambda_pre", "pwe_log_lambda_contact", "pwe_log_lambda_post"]), "no_interval_bin_probability_term_in_nll": "log1mexp_neg(" not in src.split("def pwe_target_from_terminal", 1)[0].split("def log1mexp_neg", 1)[-1], } return { "status": "pass" if all(checks.values()) else "fail", "checks": checks, "math": "Exact-time piecewise exponential competing-risk density: log S(t) + log lambda_k(t); admin censoring uses log S(t).", } def audit_color_order(tensor_dir: Path) -> dict[str, Any]: meta = load_json(tensor_dir / "tensor_metadata.json") vocab = load_json(tensor_dir / "cat_value_vocab.json") color_field_index, pairs = build_color_map(meta, vocab) id_to_label = {int(v): k for k, v in vocab.items()} rows = [] for value_id, cls in pairs: token = id_to_label.get(int(value_id), "") rows.append({"value_id": int(value_id), "token": token, "ordinal_class": int(cls)}) rows_sorted = sorted(rows, key=lambda x: (x["ordinal_class"], x["value_id"])) observed_order = [row["token"].split("=", 1)[-1] for row in rows_sorted] checks = { "color_field_present": color_field_index >= 0, "five_color_values_found": len(rows_sorted) == 5, "green_blue_yellow_orange_red": observed_order == ["绿色", "蓝色", "黄色", "橙色", "红色"] or observed_order == ["绿", "蓝", "黄", "橙", "红"], } return { "status": "pass" if all(checks.values()) else "fail", "checks": checks, "color_field_index": int(color_field_index), "official_order_implemented": "green < blue < yellow < orange < red", "value_to_class": rows_sorted, "evidence": "train_v4p4_cloud.build_color_map plus tensor cat_value_vocab.json", } def audit_training_defaults() -> dict[str, Any]: train_src = read_text("train_v4p4_cloud.py") submit_src = read_optional_text("hf_submit_v4p4_world_model_train_job.py") checks = { "train_default_stage2a": 'parser.add_argument("--stage", choices=["stage1", "stage2a"], default="stage2a")' in train_src, "submit_script_present": bool(submit_src), "submit_default_stage2a": 'parser.add_argument("--stage", choices=["stage1", "stage2a"], default="stage2a")' in submit_src, "metrics_are_sparse": "return_metrics=need_step_metrics" in train_src, "loss_can_skip_metrics": "return_metrics: bool = False" in read_text("loss_v4p4.py"), "permutation_sampler_present": "class TrainIndexSampler" in train_src and "self.rng.permutation" in train_src, "formal_eval_frequency": '--eval-every", type=int, default=500' in train_src and '--eval-every", type=int, default=500' in submit_src, "no_replacement_submit_default": "parser.set_defaults(sample_with_replacement=False)" in submit_src, } return { "status": "pass" if all(checks.values()) else "fail", "checks": checks, "recommended_cloud_run": "stage2a/joint_one_step, bf16, fused AdamW, no replacement sampler, log_every=25, eval_every=500, save_every=1000.", } def main() -> None: parser = argparse.ArgumentParser(description="Write SCTM-v4.4 finalization audit files.") parser.add_argument("--tensor-dir", type=Path, default=DEFAULT_TENSOR_DIR) parser.add_argument("--stage0-dir", type=Path, default=DEFAULT_STAGE0_DIR) parser.add_argument("--out-dir", type=Path, required=True) args = parser.parse_args() args.out_dir.mkdir(parents=True, exist_ok=True) audits = { "code_version_audit": audit_code_version(), "export_source_audit": audit_export_sources(), "model_factorization_audit": audit_model_factorization(), "pwe_math_audit": audit_pwe_math(), "color_signal_order_audit": audit_color_order(args.tensor_dir), "training_engineering_audit": audit_training_defaults(), } for name, payload in audits.items(): (args.out_dir / f"{name}.json").write_text(json.dumps(json_ready(payload), ensure_ascii=False, indent=2), encoding="utf-8") summary = { "status": "pass" if all(payload.get("status") == "pass" for payload in audits.values()) else "fail", "tensor_dir": str(args.tensor_dir), "stage0_dir": str(args.stage0_dir), "audit_files": {name: str(args.out_dir / f"{name}.json") for name in audits}, } (args.out_dir / "finalization_audit_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(summary, ensure_ascii=False), flush=True) if __name__ == "__main__": main()