| |
| """Evaluate one versioned single-task checkpoint on the released clean test.""" |
|
|
| import argparse |
| import contextlib |
| import gc |
| import json |
| import os |
| import sys |
| import time |
|
|
| SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) |
| REPO_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, "..")) |
| SRC_ROOT = os.path.join(REPO_ROOT, "src") |
| if SRC_ROOT not in sys.path: |
| sys.path.insert(0, SRC_ROOT) |
| if SCRIPT_DIR not in sys.path: |
| sys.path.insert(0, SCRIPT_DIR) |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from clean_eval_utils import ( |
| TASKS, |
| confusion_metrics, |
| discard_unused_full_code_features, |
| load_compressed_graph, |
| prepare_clean_records, |
| ) |
| from single_task_model import SingleTaskGAT |
|
|
|
|
| def graph_feature_dims(graph): |
| code_data = graph.nodes["code"].data |
| code_features = code_data.get("featmean", code_data.get("feat")) |
| if code_features is None: |
| raise RuntimeError("Clean graph is missing code featmean/feat.") |
|
|
| data_feat_dim = 1 |
| if "data" in graph.ntypes and graph.num_nodes("data") > 0: |
| data_features = graph.nodes["data"].data.get("feat") |
| if data_features is not None: |
| data_feat_dim = int(data_features.shape[1]) |
| return int(code_features.shape[1]), data_feat_dim |
|
|
|
|
| def load_checkpoint(path, task, probe_graph, device): |
| artifact = torch.load(path, map_location="cpu") |
| if not isinstance(artifact, dict) or artifact.get("artifact_format") != "icflownet_model_v1": |
| raise RuntimeError("Expected an icflownet_model_v1 single-task checkpoint.") |
| model_config = dict(artifact.get("model_config") or {}) |
| model_config.pop("architecture_version", None) |
| stored_task = model_config.get("task_type") |
| if stored_task != task: |
| raise RuntimeError( |
| "Checkpoint task {!r} does not match requested task {!r}.".format( |
| stored_task, task |
| ) |
| ) |
|
|
| code_feat_dim, data_feat_dim = graph_feature_dims(probe_graph) |
| model = SingleTaskGAT( |
| code_feat_dim=code_feat_dim, |
| data_feat_dim=data_feat_dim, |
| **model_config, |
| ) |
| model.load_state_dict(artifact["model_state_dict"], strict=True) |
| model.decision_threshold = 0.5 |
| model.to(device) |
| model.eval() |
| return model, model_config, artifact.get("artifact_format") |
|
|
|
|
| def evaluate_record( |
| model, |
| graph, |
| rows, |
| candidate_meta, |
| device, |
| amp_enabled, |
| ): |
| """Evaluate one clean-test graph and return aggregate edge metrics.""" |
| graph = graph.to(device) |
| pairs = torch.tensor( |
| [(row["src"], row["dst"]) for row in rows], |
| dtype=torch.long, |
| device=device, |
| ) |
| labels = torch.tensor( |
| [row["label"] for row in rows], dtype=torch.float32, device=device |
| ) |
|
|
| autocast_context = ( |
| torch.autocast(device_type=device.type, dtype=torch.float16) |
| if amp_enabled |
| else contextlib.nullcontext() |
| ) |
| with torch.no_grad(), autocast_context: |
| logits = model(graph, pairs, candidate_meta=candidate_meta) |
| logits = logits.reshape(-1) |
|
|
| scores = torch.sigmoid(logits) |
| predictions = scores >= 0.5 |
| positives = labels == 1 |
| negatives = ~positives |
| local_tn = int(((~predictions) & negatives).sum().item()) |
| local_fp = int((predictions & negatives).sum().item()) |
| local_fn = int(((~predictions) & positives).sum().item()) |
| local_tp = int((predictions & positives).sum().item()) |
| local_loss = float( |
| F.binary_cross_entropy_with_logits(logits, labels, reduction="sum").item() |
| ) |
| return local_tn, local_fp, local_fn, local_tp, local_loss |
|
|
|
|
| def evaluate(args): |
| if args.task not in TASKS: |
| raise ValueError("Unsupported task: {}".format(args.task)) |
| if args.gpu >= 0 and not torch.cuda.is_available(): |
| raise RuntimeError("CUDA was requested but is unavailable.") |
| device = torch.device("cuda:{}".format(args.gpu) if args.gpu >= 0 else "cpu") |
| if device.type == "cuda": |
| torch.cuda.set_device(device) |
|
|
| clean_root = os.path.abspath(args.clean_root) |
| resource_root = os.path.join(clean_root, "resources") |
| records, preflight = prepare_clean_records(clean_root, resource_root, args.scope) |
| all_eligible = [record for record in records if record["task_rows"].get(args.task)] |
| indexed_eligible = list(enumerate(all_eligible, start=1)) |
| end_record = args.end_record or len(indexed_eligible) |
| indexed_eligible = indexed_eligible[args.start_record - 1 : end_record] |
| if args.max_records is not None: |
| indexed_eligible = indexed_eligible[: args.max_records] |
| if not indexed_eligible: |
| raise RuntimeError("No clean-test records contain task {}.".format(args.task)) |
|
|
| probe_graph = discard_unused_full_code_features( |
| load_compressed_graph(indexed_eligible[0][1]["graph_path"]) |
| ) |
| model, model_config, artifact_format = load_checkpoint( |
| os.path.abspath(args.checkpoint), args.task, probe_graph, device |
| ) |
| del probe_graph |
|
|
| tn = fp = fn = tp = 0 |
| loss_sum = 0.0 |
| pair_count = 0 |
| details = [] |
| cpu_model = None |
| cpu_fallback_records = [] |
| started = time.time() |
| amp_enabled = device.type == "cuda" and args.enable_amp |
|
|
| for selected_index, (record_index, record) in enumerate(indexed_eligible, start=1): |
| rows = record["task_rows"][args.task] |
| graph = discard_unused_full_code_features( |
| load_compressed_graph(record["graph_path"]) |
| ) |
| if ( |
| graph.num_nodes("code") != record["expected_code_nodes"] |
| or graph.num_nodes("data") != record["expected_data_nodes"] |
| ): |
| raise RuntimeError( |
| "Clean graph node count mismatch for {}.".format(record["binary_key"]) |
| ) |
| try: |
| ( |
| local_tn, |
| local_fp, |
| local_fn, |
| local_tp, |
| local_loss, |
| ) = evaluate_record( |
| model, |
| graph, |
| rows, |
| record["candidate_meta"], |
| device, |
| amp_enabled, |
| ) |
| except RuntimeError as error: |
| if device.type != "cuda" or not args.cpu_fallback: |
| raise |
| print( |
| "[{}] CUDA evaluation failed for {}; retrying this record on CPU: {}".format( |
| record_index, record["binary_key"], error |
| ), |
| flush=True, |
| ) |
| gc.collect() |
| try: |
| torch.cuda.empty_cache() |
| except RuntimeError: |
| pass |
| if cpu_model is None: |
| cpu_device = torch.device("cpu") |
| cpu_model, _, _ = load_checkpoint( |
| os.path.abspath(args.checkpoint), args.task, graph, cpu_device |
| ) |
| ( |
| local_tn, |
| local_fp, |
| local_fn, |
| local_tp, |
| local_loss, |
| ) = evaluate_record( |
| cpu_model, |
| graph, |
| rows, |
| record["candidate_meta"], |
| torch.device("cpu"), |
| False, |
| ) |
| cpu_fallback_records.append( |
| { |
| "record_index": record_index, |
| "binary": record["binary_key"], |
| "graph_id": record["graph_id"], |
| "cuda_error": str(error), |
| } |
| ) |
| tn += local_tn |
| fp += local_fp |
| fn += local_fn |
| tp += local_tp |
| loss_sum += local_loss |
| pair_count += len(rows) |
|
|
| if args.save_details: |
| details.append( |
| { |
| "binary": record["binary_key"], |
| "graph_id": record["graph_id"], |
| "pairs": len(rows), |
| "confusion_matrix": [[local_tn, local_fp], [local_fn, local_tp]], |
| } |
| ) |
| if args.log_every and ( |
| selected_index % args.log_every == 0 |
| or selected_index == len(indexed_eligible) |
| ): |
| print( |
| "[{}/{} selected; source record {}/{}] {} records evaluated ({:.1f}s)".format( |
| selected_index, |
| len(indexed_eligible), |
| record_index, |
| len(all_eligible), |
| args.task, |
| time.time() - started, |
| ), |
| flush=True, |
| ) |
|
|
| del graph |
| if device.type == "cuda": |
| torch.cuda.empty_cache() |
|
|
| metrics = confusion_metrics(tn, fp, fn, tp) |
| metrics["loss"] = loss_sum / max(pair_count, 1) |
| result = { |
| "format": "icflownet_single_cleantest_eval_v1", |
| "scope": args.scope, |
| "task": args.task, |
| "checkpoint": os.path.abspath(args.checkpoint), |
| "checkpoint_artifact_format": artifact_format, |
| "model_config": model_config, |
| "threshold": 0.5, |
| "threshold_source": "fixed_evaluator", |
| "precision": "fp16_autocast" if amp_enabled else "fp32", |
| "amp_enabled": amp_enabled, |
| "clean_root": clean_root, |
| "preflight": preflight, |
| "available_eligible_records": len(all_eligible), |
| "eligible_records": len(indexed_eligible), |
| "selected_record_indices": [index for index, _ in indexed_eligible], |
| "cpu_fallback_records": cpu_fallback_records, |
| "metrics": metrics, |
| "elapsed_seconds": time.time() - started, |
| } |
| if args.save_details: |
| result["record_details"] = details |
| return result |
|
|
|
|
| def evaluate_both_scopes(args): |
| """Evaluate overall and long-range pairs with one graph encoding per record.""" |
| if args.task not in TASKS: |
| raise ValueError("Unsupported task: {}".format(args.task)) |
| if args.gpu >= 0 and not torch.cuda.is_available(): |
| raise RuntimeError("CUDA was requested but is unavailable.") |
| device = torch.device("cuda:{}".format(args.gpu) if args.gpu >= 0 else "cpu") |
| if device.type == "cuda": |
| torch.cuda.set_device(device) |
|
|
| clean_root = os.path.abspath(args.clean_root) |
| resource_root = os.path.join(clean_root, "resources") |
| scopes = ("overall", "long_range") |
| records_by_scope = {} |
| preflight_by_scope = {} |
| eligible_by_scope = {} |
| for scope in scopes: |
| records, preflight = prepare_clean_records(clean_root, resource_root, scope) |
| records_by_scope[scope] = { |
| record["binary_key"]: record for record in records |
| } |
| preflight_by_scope[scope] = preflight |
| eligible_by_scope[scope] = [ |
| record["binary_key"] |
| for record in records |
| if record["task_rows"].get(args.task) |
| ] |
|
|
| ordered_keys = [] |
| seen_keys = set() |
| for record in records_by_scope["overall"].values(): |
| key = record["binary_key"] |
| if any( |
| records_by_scope[scope][key]["task_rows"].get(args.task) |
| for scope in scopes |
| ): |
| ordered_keys.append(key) |
| seen_keys.add(key) |
| for key in records_by_scope["long_range"]: |
| if key not in seen_keys and records_by_scope["long_range"][key][ |
| "task_rows" |
| ].get(args.task): |
| ordered_keys.append(key) |
|
|
| end_record = args.end_record or len(ordered_keys) |
| selected_keys = ordered_keys[args.start_record - 1 : end_record] |
| if args.max_records is not None: |
| selected_keys = selected_keys[: args.max_records] |
| if not selected_keys: |
| raise RuntimeError("No clean-test records contain task {}.".format(args.task)) |
|
|
| probe_record = records_by_scope["overall"].get( |
| selected_keys[0], records_by_scope["long_range"][selected_keys[0]] |
| ) |
| probe_graph = discard_unused_full_code_features( |
| load_compressed_graph(probe_record["graph_path"]) |
| ) |
| model, model_config, artifact_format = load_checkpoint( |
| os.path.abspath(args.checkpoint), args.task, probe_graph, device |
| ) |
| del probe_graph |
|
|
| accumulators = { |
| scope: { |
| "tn": 0, |
| "fp": 0, |
| "fn": 0, |
| "tp": 0, |
| "loss_sum": 0.0, |
| "pair_count": 0, |
| "record_count": 0, |
| "details": [], |
| } |
| for scope in scopes |
| } |
| started = time.time() |
| amp_enabled = device.type == "cuda" and args.enable_amp |
|
|
| for selected_index, key in enumerate(selected_keys, start=1): |
| base_record = records_by_scope["overall"].get( |
| key, records_by_scope["long_range"][key] |
| ) |
| combined_rows = [] |
| scope_slices = {} |
| for scope in scopes: |
| record = records_by_scope[scope].get(key) |
| rows = record["task_rows"].get(args.task, []) if record else [] |
| if not rows: |
| continue |
| start = len(combined_rows) |
| combined_rows.extend(rows) |
| scope_slices[scope] = (start, len(combined_rows), rows) |
|
|
| graph = discard_unused_full_code_features( |
| load_compressed_graph(base_record["graph_path"]) |
| ) |
| if ( |
| graph.num_nodes("code") != base_record["expected_code_nodes"] |
| or graph.num_nodes("data") != base_record["expected_data_nodes"] |
| ): |
| raise RuntimeError( |
| "Clean graph node count mismatch for {}.".format(key) |
| ) |
|
|
| graph = graph.to(device) |
| pairs = torch.tensor( |
| [(row["src"], row["dst"]) for row in combined_rows], |
| dtype=torch.long, |
| device=device, |
| ) |
| labels = torch.tensor( |
| [row["label"] for row in combined_rows], |
| dtype=torch.float32, |
| device=device, |
| ) |
| autocast_context = ( |
| torch.autocast(device_type=device.type, dtype=torch.float16) |
| if amp_enabled |
| else contextlib.nullcontext() |
| ) |
| with torch.no_grad(), autocast_context: |
| logits = model( |
| graph, pairs, candidate_meta=base_record["candidate_meta"] |
| ).reshape(-1) |
|
|
| for scope, (start, end, rows) in scope_slices.items(): |
| local_logits = logits[start:end] |
| local_labels = labels[start:end] |
| predictions = torch.sigmoid(local_logits) >= 0.5 |
| positives = local_labels == 1 |
| negatives = ~positives |
| local_tn = int(((~predictions) & negatives).sum().item()) |
| local_fp = int((predictions & negatives).sum().item()) |
| local_fn = int(((~predictions) & positives).sum().item()) |
| local_tp = int((predictions & positives).sum().item()) |
| local_loss = float( |
| F.binary_cross_entropy_with_logits( |
| local_logits, local_labels, reduction="sum" |
| ).item() |
| ) |
| accumulator = accumulators[scope] |
| accumulator["tn"] += local_tn |
| accumulator["fp"] += local_fp |
| accumulator["fn"] += local_fn |
| accumulator["tp"] += local_tp |
| accumulator["loss_sum"] += local_loss |
| accumulator["pair_count"] += len(rows) |
| accumulator["record_count"] += 1 |
| if args.save_details: |
| accumulator["details"].append( |
| { |
| "binary": key, |
| "graph_id": base_record["graph_id"], |
| "pairs": len(rows), |
| "confusion_matrix": [ |
| [local_tn, local_fp], |
| [local_fn, local_tp], |
| ], |
| } |
| ) |
|
|
| if args.log_every and ( |
| selected_index % args.log_every == 0 |
| or selected_index == len(selected_keys) |
| ): |
| print( |
| "[{}/{} selected] {} records evaluated for both scopes ({:.1f}s)".format( |
| selected_index, |
| len(selected_keys), |
| args.task, |
| time.time() - started, |
| ), |
| flush=True, |
| ) |
|
|
| del graph |
| if device.type == "cuda": |
| torch.cuda.empty_cache() |
|
|
| elapsed_seconds = time.time() - started |
| results = {} |
| for scope in scopes: |
| accumulator = accumulators[scope] |
| metrics = confusion_metrics( |
| accumulator["tn"], |
| accumulator["fp"], |
| accumulator["fn"], |
| accumulator["tp"], |
| ) |
| metrics["loss"] = accumulator["loss_sum"] / max( |
| accumulator["pair_count"], 1 |
| ) |
| result = { |
| "format": "icflownet_single_cleantest_eval_v1", |
| "scope": scope, |
| "task": args.task, |
| "checkpoint": os.path.abspath(args.checkpoint), |
| "checkpoint_artifact_format": artifact_format, |
| "model_config": model_config, |
| "threshold": 0.5, |
| "threshold_source": "fixed_evaluator", |
| "precision": "fp16_autocast" if amp_enabled else "fp32", |
| "amp_enabled": amp_enabled, |
| "clean_root": clean_root, |
| "preflight": preflight_by_scope[scope], |
| "available_eligible_records": len(eligible_by_scope[scope]), |
| "eligible_records": accumulator["record_count"], |
| "selected_record_indices": list( |
| range(1, accumulator["record_count"] + 1) |
| ), |
| "cpu_fallback_records": [], |
| "metrics": metrics, |
| "elapsed_seconds": elapsed_seconds, |
| "joint_scope_evaluation": True, |
| } |
| if args.save_details: |
| result["record_details"] = accumulator["details"] |
| results[scope] = result |
| return results |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--checkpoint", required=True) |
| parser.add_argument("--clean_root", required=True) |
| parser.add_argument("--task", choices=TASKS, required=True) |
| parser.add_argument( |
| "--scope", choices=("overall", "long_range", "both"), default="overall" |
| ) |
| parser.add_argument("--output", required=True) |
| parser.add_argument("--gpu", type=int, default=0) |
| parser.add_argument("--max_records", type=int, default=None) |
| parser.add_argument("--start_record", type=int, default=1) |
| parser.add_argument("--end_record", type=int, default=None) |
| parser.add_argument("--log_every", type=int, default=25) |
| precision_group = parser.add_mutually_exclusive_group() |
| precision_group.add_argument( |
| "--enable_amp", |
| action="store_true", |
| help="Explicitly opt into FP16 autocast; FP32 is the default.", |
| ) |
| precision_group.add_argument( |
| "--disable_amp", |
| action="store_true", |
| help="Deprecated compatibility flag; FP32 is already the default.", |
| ) |
| parser.add_argument( |
| "--cpu_fallback", |
| action="store_true", |
| help="Retry an individual record on CPU if its CUDA evaluation fails.", |
| ) |
| parser.add_argument("--save_details", action="store_true") |
| args = parser.parse_args() |
| if args.enable_amp: |
| print( |
| "[WARNING] FP16 autocast was explicitly enabled; artifact results " |
| "must use the default FP32 mode.", |
| flush=True, |
| ) |
| if args.start_record < 1: |
| parser.error("--start_record must be at least 1") |
| if args.end_record is not None and args.end_record < args.start_record: |
| parser.error("--end_record must be greater than or equal to --start_record") |
|
|
| output = os.path.abspath(args.output) |
| if args.scope == "both": |
| results = evaluate_both_scopes(args) |
| os.makedirs(output, exist_ok=True) |
| summaries = [] |
| for scope, result in results.items(): |
| scope_output = os.path.join(output, scope + ".json") |
| with open(scope_output, "w") as handle: |
| json.dump(result, handle, indent=2) |
| handle.write("\n") |
| summaries.append( |
| { |
| "output": scope_output, |
| "task": result["task"], |
| "scope": result["scope"], |
| "eligible_records": result["eligible_records"], |
| "metrics": result["metrics"], |
| "elapsed_seconds": result["elapsed_seconds"], |
| } |
| ) |
| print(json.dumps(summaries, indent=2)) |
| else: |
| result = evaluate(args) |
| os.makedirs(os.path.dirname(output), exist_ok=True) |
| with open(output, "w") as handle: |
| json.dump(result, handle, indent=2) |
| handle.write("\n") |
| print(json.dumps({ |
| "output": output, |
| "task": result["task"], |
| "scope": result["scope"], |
| "eligible_records": result["eligible_records"], |
| "metrics": result["metrics"], |
| "elapsed_seconds": result["elapsed_seconds"], |
| }, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
