plugin_experiment.py

# =============================================================================
# DEPRECATED — JEPA-era research artifact. Kept for provenance, not for use.
#
# This file is part of the original experiments (2026-06-30 to 2026-07-01) that
# the Aporia Engine grew out of: bridging a frozen LLM with a JEPA energy model
# to detect confident-wrong claims. The JEPA fact-grammar approach was SHELVED
# on 2026-07-01 (per-fact density beat it under a pre-registered bar); its
# surviving legacy is the typicality signal in the lacunae ledger's triage.
# Nothing here is wired into the running system. Where it started: this file.
# Where it went: docs/design/ and the verification/ladder/orchestrator code.
# =============================================================================
"""Sub-step 2: first REAL cartridge. Train the energy-JEPA on RiftSuite plugin code
(embedded via Ollama), then test whether its energy separates RiftSuite from:
  - Python  (mid OOD: code, very different language)  -> expect easy separation
  - Lyra    (hard OOD: same UE C++ idiom, different codebase) -> the meaningful test

Metric: AUROC = P(energy(OOD) > energy(in-dist)). 1.0 = perfect, 0.5 = no separation.
"""

from __future__ import annotations

import numpy as np
import torch

from corpus import (CPP_EXTS, LYRA_ROOTS, PY_EXTS, PYTHON_ROOTS, RIFTSUITE_ROOTS,
                    file_chunk_lists)
from embed import Embedder
from jepa_core import EnergyJEPA

T = 5
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def make_samples(fcl: list[list[str]], emb: Embedder) -> np.ndarray:
    """Embed chunks, then form within-file sliding windows of T chunks -> [n, T, D]."""
    flat = [c for f in fcl for c in f]
    vecs = emb.embed(flat)                       # [total, D]
    samples, idx = [], 0
    for f in fcl:
        fv = vecs[idx : idx + len(f)]
        idx += len(f)
        for s in range(0, len(f) - T + 1):
            samples.append(fv[s : s + T])
    return np.stack(samples) if samples else np.empty((0, T, vecs.shape[1]), np.float32)


def auroc(higher: np.ndarray, lower: np.ndarray) -> float:
    """P(higher > lower) via rank statistic. higher = OOD energies, lower = in-dist."""
    allv = np.concatenate([higher, lower])
    ranks = np.empty(len(allv))
    ranks[allv.argsort()] = np.arange(1, len(allv) + 1)
    r = ranks[: len(higher)].sum()
    return (r - len(higher) * (len(higher) + 1) / 2) / (len(higher) * len(lower))


def main() -> int:
    emb = Embedder()
    print("embedding corpora (cached after first run)...")
    rift = make_samples(file_chunk_lists(RIFTSUITE_ROOTS, CPP_EXTS), emb)
    lyra = make_samples(file_chunk_lists(LYRA_ROOTS, CPP_EXTS, max_files=250), emb)
    pyth = make_samples(file_chunk_lists(PYTHON_ROOTS, PY_EXTS, max_files=60), emb)
    print(f"samples  rift={len(rift)}  lyra={len(lyra)}  python={len(pyth)}")

    # split RiftSuite train/test
    rng = np.random.default_rng(0)
    perm = rng.permutation(len(rift))
    n_test = len(rift) // 5
    test_idx, train_idx = perm[:n_test], perm[n_test:]
    train, test_in = rift[train_idx], rift[test_idx]

    # standardize per-feature using TRAIN stats only
    mu = train.reshape(-1, train.shape[-1]).mean(0)
    sd = train.reshape(-1, train.shape[-1]).std(0) + 1e-6
    norm = lambda x: (x - mu) / sd
    train_t = torch.tensor(norm(train), dtype=torch.float32)

    print(f"device={DEVICE}  train={tuple(train_t.shape)}")
    model = EnergyJEPA(in_dim=train.shape[-1], n_pos=T, learn_encoder=False).fit(
        train_t, steps=1500, batch=256, device=DEVICE, log_every=250
    )

    def energy(x):
        return model.energy(torch.tensor(norm(x), dtype=torch.float32, device=DEVICE)).cpu().numpy()

    e_in, e_lyra, e_py = energy(test_in), energy(lyra), energy(pyth)

    print("\n=== energy (lower = fits RiftSuite world) ===")
    for name, e in [("rift_test", e_in), ("lyra", e_lyra), ("python", e_py)]:
        print(f"  {name:10s} mean={e.mean():.4f}  std={e.std():.4f}")
    print("\n=== separation AUROC (1.0 = perfect, 0.5 = none) ===")
    auc_py = auroc(e_py, e_in)
    auc_lyra = auroc(e_lyra, e_in)
    print(f"  rift vs python (mid OOD) : {auc_py:.3f}")
    print(f"  rift vs lyra   (hard OOD): {auc_lyra:.3f}   <-- the meaningful test")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())