auto: sync run_qwen_injection_layer_ablation.py

scripts/run_qwen_injection_layer_ablation.py

@@ -0,0 +1,357 @@
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+
+ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_INPUT = ROOT / "archive" / "qwen35_4b_injection_geometry_medium.json"
+DEFAULT_OUTPUT_JSON = ROOT / "archive" / "qwen_injection_layer_ablation.json"
+DEFAULT_OUTPUT_MD = ROOT / "docs" / "research" / "qwen_injection_layer_ablation.md"
+
+
+def compute_auc(negative_scores: list[float], positive_scores: list[float]) -> float | None:
+    if not negative_scores or not positive_scores:
+        return None
+    wins = 0.0
+    total = 0
+    for negative in negative_scores:
+        for positive in positive_scores:
+            total += 1
+            if positive > negative:
+                wins += 1.0
+            elif positive == negative:
+                wins += 0.5
+    return float(wins / total) if total else None
+
+
+def finite_float(value: Any) -> float | None:
+    try:
+        out = float(value)
+    except (TypeError, ValueError):
+        return None
+    return out if math.isfinite(out) else None
+
+
+def mean_vector(vectors: list[list[float]]) -> list[float]:
+    if not vectors:
+        return []
+    dims = len(vectors[0])
+    return [
+        float(sum(vector[dim] for vector in vectors) / len(vectors))
+        for dim in range(dims)
+    ]
+
+
+def compute_scale(vectors: list[list[float]]) -> list[float]:
+    if not vectors:
+        return []
+    dims = len(vectors[0])
+    scales: list[float] = []
+    for dim in range(dims):
+        column = [float(vector[dim]) for vector in vectors]
+        mean_value = sum(column) / len(column)
+        variance = sum((value - mean_value) ** 2 for value in column) / max(1, len(column))
+        scales.append(max(variance ** 0.5, 1e-6))
+    return scales
+
+
+def standardized_l1_distance(vector: list[float], prototype: list[float], scale: list[float]) -> float:
+    if not vector or not prototype or not scale:
+        return 0.0
+    total = 0.0
+    for value, proto, denom in zip(vector, prototype, scale):
+        total += abs(float(value) - float(proto)) / max(float(denom), 1e-6)
+    return float(total / len(vector))
+
+
+def vector_for_subset(row: dict[str, Any], features: list[tuple[int, str]]) -> list[float] | None:
+    layer_metrics = row.get("layer_metrics")
+    if not isinstance(layer_metrics, dict):
+        return None
+    values: list[float] = []
+    for layer, metric in features:
+        metric_map = layer_metrics.get(str(layer))
+        if not isinstance(metric_map, dict):
+            return None
+        value = finite_float(metric_map.get(metric))
+        if value is None:
+            return None
+        values.append(value)
+    return values
+
+
+def distance_auc_for_features(rows: list[dict[str, Any]], features: list[tuple[int, str]]) -> dict[str, Any]:
+    legit_rows = [row for row in rows if row.get("label") == "legit" and row.get("status") == "ok"]
+    injected_rows = [row for row in rows if row.get("label") == "injected" and row.get("status") == "ok"]
+
+    row_vectors: dict[int, list[float]] = {}
+    for row in legit_rows + injected_rows:
+        vector = vector_for_subset(row, features)
+        if vector is not None:
+            row_vectors[id(row)] = vector
+
+    legit_vectors = [row_vectors[id(row)] for row in legit_rows if id(row) in row_vectors]
+    scale = compute_scale(legit_vectors)
+    legit_by_group: dict[str, list[list[float]]] = {}
+    for row in legit_rows:
+        vector = row_vectors.get(id(row))
+        if vector is None:
+            continue
+        legit_by_group.setdefault(str(row.get("anchor_group")), []).append(vector)
+
+    legit_scores: list[float] = []
+    injected_scores: list[float] = []
+    for row in legit_rows:
+        vector = row_vectors.get(id(row))
+        if vector is None:
+            continue
+        group = str(row.get("anchor_group"))
+        candidates = [
+            other_vector
+            for other in legit_rows
+            if other is not row and str(other.get("anchor_group")) == group
+            for other_vector in [row_vectors.get(id(other))]
+            if other_vector is not None
+        ]
+        if not candidates:
+            continue
+        legit_scores.append(standardized_l1_distance(vector, mean_vector(candidates), scale))
+
+    for row in injected_rows:
+        vector = row_vectors.get(id(row))
+        if vector is None:
+            continue
+        prototype = mean_vector(legit_by_group.get(str(row.get("anchor_group")), []))
+        if not prototype:
+            continue
+        injected_scores.append(standardized_l1_distance(vector, prototype, scale))
+
+    auc = compute_auc(legit_scores, injected_scores)
+    return {
+        "auc": auc,
+        "n_legit": len(legit_scores),
+        "n_injected": len(injected_scores),
+        "legit_mean_distance": mean_or_none(legit_scores),
+        "injected_mean_distance": mean_or_none(injected_scores),
+    }
+
+
+def raw_metric_auc(rows: list[dict[str, Any]], layer: int, metric: str) -> dict[str, Any]:
+    legit_values: list[float] = []
+    injected_values: list[float] = []
+    for row in rows:
+        if row.get("status") != "ok":
+            continue
+        value = vector_for_subset(row, [(layer, metric)])
+        if value is None:
+            continue
+        if row.get("label") == "legit":
+            legit_values.append(value[0])
+        elif row.get("label") == "injected":
+            injected_values.append(value[0])
+    auc = compute_auc(legit_values, injected_values)
+    if auc is None:
+        return {
+            "layer": layer,
+            "metric": metric,
+            "auc": None,
+            "separation_auc": None,
+            "direction": "unknown",
+        }
+    if auc >= 0.5:
+        direction = "higher_for_injected"
+        separation_auc = auc
+    else:
+        direction = "lower_for_injected"
+        separation_auc = 1.0 - auc
+    return {
+        "layer": layer,
+        "metric": metric,
+        "auc": auc,
+        "separation_auc": separation_auc,
+        "direction": direction,
+        "legit_mean": mean_or_none(legit_values),
+        "injected_mean": mean_or_none(injected_values),
+    }
+
+
+def mean_or_none(values: list[float]) -> float | None:
+    if not values:
+        return None
+    return float(sum(values) / len(values))
+
+
+def infer_layers_and_metrics(rows: list[dict[str, Any]]) -> tuple[list[int], list[str]]:
+    for row in rows:
+        layer_metrics = row.get("layer_metrics")
+        if not isinstance(layer_metrics, dict):
+            continue
+        layers = sorted(int(layer) for layer in layer_metrics)
+        metrics: list[str] = []
+        for layer in layers:
+            metric_map = layer_metrics.get(str(layer))
+            if isinstance(metric_map, dict):
+                metrics = list(metric_map.keys())
+                break
+        if layers and metrics:
+            return layers, metrics
+    raise ValueError("no layer_metrics found in input samples")
+
+
+def build_layer_features(layers: list[int], metrics: list[str]) -> list[tuple[int, str]]:
+    return [(layer, metric) for layer in layers for metric in metrics]
+
+
+def analyze(payload: dict[str, Any]) -> dict[str, Any]:
+    rows = payload.get("samples")
+    if not isinstance(rows, list):
+        raise ValueError("input JSON must contain a samples list")
+    layers, metrics = infer_layers_and_metrics(rows)
+    all_features = build_layer_features(layers, metrics)
+
+    per_layer = [
+        {
+            "layer": layer,
+            **distance_auc_for_features(rows, build_layer_features([layer], metrics)),
+        }
+        for layer in layers
+    ]
+    per_metric = [
+        {
+            "metric": metric,
+            **distance_auc_for_features(rows, [(layer, metric) for layer in layers]),
+        }
+        for metric in metrics
+    ]
+    per_layer_metric = [
+        raw_metric_auc(rows, layer, metric)
+        for layer in layers
+        for metric in metrics
+    ]
+
+    crystal_layers = [layer for layer in layers if 4 <= layer <= 8]
+    handoff_layers = [layer for layer in layers if layer >= 24]
+    mid_layers = [layer for layer in layers if 9 <= layer < 24]
+
+    subsets = {
+        "all_probe_layers": all_features,
+        "crystallization_zone_4_8": build_layer_features(crystal_layers, metrics),
+        "mid_layers_9_23": build_layer_features(mid_layers, metrics),
+        "handoff_layers_24_plus": build_layer_features(handoff_layers, metrics),
+    }
+    subset_results = {
+        name: distance_auc_for_features(rows, features) if features else {"auc": None}
+        for name, features in subsets.items()
+    }
+
+    best_layer = max(per_layer, key=lambda item: item.get("auc") if item.get("auc") is not None else -1.0)
+    best_metric = max(per_metric, key=lambda item: item.get("auc") if item.get("auc") is not None else -1.0)
+    best_layer_metric = max(
+        per_layer_metric,
+        key=lambda item: item.get("separation_auc") if item.get("separation_auc") is not None else -1.0,
+    )
+
+    return {
+        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
+        "source_metadata": payload.get("metadata", {}),
+        "layers": layers,
+        "metrics": metrics,
+        "summary": {
+            "source_detection_auc": payload.get("summary", {}).get("detection_auc"),
+            "all_probe_layers_auc": subset_results["all_probe_layers"].get("auc"),
+            "crystallization_zone_auc": subset_results["crystallization_zone_4_8"].get("auc"),
+            "mid_layers_auc": subset_results["mid_layers_9_23"].get("auc"),
+            "handoff_layers_auc": subset_results["handoff_layers_24_plus"].get("auc"),
+            "best_single_layer": best_layer.get("layer"),
+            "best_single_layer_auc": best_layer.get("auc"),
+            "best_metric": best_metric.get("metric"),
+            "best_metric_auc": best_metric.get("auc"),
+            "best_layer_metric": {
+                "layer": best_layer_metric.get("layer"),
+                "metric": best_layer_metric.get("metric"),
+                "separation_auc": best_layer_metric.get("separation_auc"),
+                "direction": best_layer_metric.get("direction"),
+            },
+        },
+        "subsets": subset_results,
+        "per_layer": sorted(per_layer, key=lambda item: item.get("auc") if item.get("auc") is not None else -1.0, reverse=True),
+        "per_metric": sorted(per_metric, key=lambda item: item.get("auc") if item.get("auc") is not None else -1.0, reverse=True),
+        "per_layer_metric": sorted(
+            per_layer_metric,
+            key=lambda item: item.get("separation_auc") if item.get("separation_auc") is not None else -1.0,
+            reverse=True,
+        ),
+    }
+
+
+def write_markdown(result: dict[str, Any], path: Path) -> None:
+    summary = result["summary"]
+    lines = [
+        "# Qwen injection geometry layer ablation",
+        "",
+        f"Generated: `{result['generated_at_utc']}`",
+        "",
+        "## Summary",
+        "",
+        f"- source_detection_auc: `{summary.get('source_detection_auc')}`",
+        f"- all_probe_layers_auc: `{summary.get('all_probe_layers_auc')}`",
+        f"- crystallization_zone_auc: `{summary.get('crystallization_zone_auc')}`",
+        f"- mid_layers_auc: `{summary.get('mid_layers_auc')}`",
+        f"- handoff_layers_auc: `{summary.get('handoff_layers_auc')}`",
+        f"- best_single_layer: `L{summary.get('best_single_layer')}` auc=`{summary.get('best_single_layer_auc')}`",
+        f"- best_metric: `{summary.get('best_metric')}` auc=`{summary.get('best_metric_auc')}`",
+        f"- best_layer_metric: `{summary.get('best_layer_metric')}`",
+        "",
+        "## Per-layer distance AUC",
+        "",
+        "| Layer | AUC | Legit mean dist | Injected mean dist |",
+        "|---:|---:|---:|---:|",
+    ]
+    for row in result["per_layer"]:
+        lines.append(
+            f"| {row.get('layer')} | {row.get('auc')} | "
+            f"{row.get('legit_mean_distance')} | {row.get('injected_mean_distance')} |"
+        )
+    lines.extend([
+        "",
+        "## Top layer-metric raw separations",
+        "",
+        "| Layer | Metric | Separation AUC | Direction |",
+        "|---:|---|---:|---|",
+    ])
+    for row in result["per_layer_metric"][:12]:
+        lines.append(
+            f"| {row.get('layer')} | {row.get('metric')} | "
+            f"{row.get('separation_auc')} | {row.get('direction')} |"
+        )
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text("\n".join(lines) + "\n", encoding="utf-8")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Layer ablation for saved Qwen injection geometry samples.")
+    parser.add_argument("--input_json", "--input-json", dest="input_json", type=Path, default=DEFAULT_INPUT)
+    parser.add_argument("--output_json", "--output-json", dest="output_json", type=Path, default=DEFAULT_OUTPUT_JSON)
+    parser.add_argument("--output_md", "--output-md", dest="output_md", type=Path, default=DEFAULT_OUTPUT_MD)
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    payload = json.loads(args.input_json.read_text(encoding="utf-8"))
+    result = analyze(payload)
+    args.output_json.parent.mkdir(parents=True, exist_ok=True)
+    args.output_json.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
+    write_markdown(result, args.output_md)
+    print(f"saved_json={args.output_json}")
+    print(f"saved_md={args.output_md}")
+    print(f"best_single_layer_auc={result['summary']['best_single_layer_auc']}")
+
+
+if __name__ == "__main__":
+    main()