evo2_case_job.py

#!/usr/bin/env python
"""Evo 2 frozen-detection case study (CASE_STUDY_PLAN.md).
NGC 25.03 image (torch 2.7, flash-attn 2.7.3). evo2_7b_base.
Reads /workspace/pools/evo2_windows.csv; writes /workspace/out35/evo2_case_metrics.json
+ evo2_windows_scored.csv (per-window NLL, kNN, flags).
"""
import os, json, time
import numpy as np
import pandas as pd
import torch

POOLS, OUT = "/workspace/pools", "/workspace/out35"
os.makedirs(OUT, exist_ok=True)
MODEL_NAME = "evo2_7b_base"
BATCH = 8

try:
    from transformer_engine.common.recipe import _OverrideLinearPrecision
    torch.serialization.add_safe_globals([_OverrideLinearPrecision])
except Exception:
    pass

# ---------------- inference ----------------
def run_inference(df):
    cache = os.path.join(OUT, "evo2_windows_features.npz")
    if os.path.exists(cache):
        d = np.load(cache)
        return d["emb"], d["nll"]
    from evo2 import Evo2
    m = Evo2(MODEL_NAME)
    # hook the module feeding the logits to capture hidden states
    hidden = {}
    core = m.model
    cand = None
    for name, mod in core.named_modules():
        ln = name.lower()
        if "unembed" in ln or "output" in ln or "embed_out" in ln or "logits" in ln:
            cand = (name, mod)
    if cand is None:  # fall back to last *norm module
        for name, mod in core.named_modules():
            if "norm" in name.lower():
                cand = (name, mod)
    print("hooking:", cand[0], type(cand[1]).__name__, flush=True)
    def pre_hook(module, args):
        hidden["h"] = args[0].detach()
    cand[1].register_forward_pre_hook(pre_hook)

    n = len(df)
    d_model = None
    embs, nlls = [], []
    t0 = time.time()
    with torch.no_grad():
        for s in range(0, n, BATCH):
            seqs = df.seq.iloc[s:s + BATCH].tolist()
            ids = torch.tensor([m.tokenizer.tokenize(x) for x in seqs],
                               dtype=torch.long).cuda()
            logits, _ = core(ids)
            h = hidden["h"]  # (B, T, d)
            lp = torch.log_softmax(logits[:, :-1].float(), dim=-1)
            tgt = ids[:, 1:]
            nll = -lp.gather(-1, tgt.unsqueeze(-1)).squeeze(-1).mean(1)  # per window
            emb = h[:, 1:-1].float().mean(1)  # mean-pool, drop BOS/EOS
            embs.append(emb.cpu().numpy()); nlls.append(nll.cpu().numpy())
            d_model = h.shape[-1]
            if (s // BATCH) % 100 == 0:
                print(f"{s}/{n} ({(s+1)/(time.time()-t0):.1f} w/s)", flush=True)
    emb = np.concatenate(embs); nll = np.concatenate(nlls)
    print("emb shape:", emb.shape, "nll shape:", nll.shape, flush=True)
    np.savez_compressed(cache, emb=emb, nll=nll)
    return emb, nll

# ---------------- analysis ----------------
def main():
    from sklearn.preprocessing import StandardScaler
    from sklearn.decomposition import PCA
    from sklearn.covariance import LedoitWolf
    from sklearn.metrics import roc_auc_score
    from scipy.stats import spearmanr

    df = pd.read_csv(os.path.join(POOLS, "evo2_windows.csv"))
    emb, nll = run_inference(df)
    df["nll"] = nll
    res = {"model": MODEL_NAME, "n_windows": int(len(df))}

    # ---- FROZEN detector: fit on ref_bacteria only; freeze on dev_bacteria ----
    ref = df.cls == "ref_bacteria"
    dev = df.cls == "dev_bacteria"
    sc = StandardScaler().fit(emb[ref])
    Z = sc.transform(emb)
    Zref = Z[ref]
    Zref_n = Zref / np.linalg.norm(Zref, axis=1, keepdims=True)
    Z_n = Z / np.linalg.norm(Z, axis=1, keepdims=True)
    knn = 1 - np.sort(Z_n @ Zref_n.T, axis=1)[:, -10:].mean(1)
    pca = PCA(n_components=64, random_state=0).fit(Zref)
    mahal = LedoitWolf().fit(pca.transform(Zref)).mahalanobis(pca.transform(Z)) ** 0.5
    df["knn"], df["mahal"] = knn, mahal
    thr = np.quantile(knn[ref], 0.99)
    res["frozen_threshold_knn"] = float(thr)
    res["flag_rate_ref_check"] = float((knn[ref] > thr).mean())
    res["flag_rate_dev_bacteria"] = float((knn[dev] > thr).mean())

    # ---- apply (nothing viral touched above this line) ----
    for cls in ["phage", "excluded_virus"]:
        res[f"flag_rate_{cls}"] = float((knn[df.cls == cls] > thr).mean())
    res["flag_ratio_excluded_vs_dev"] = res["flag_rate_excluded_virus"] / max(res["flag_rate_dev_bacteria"], 1e-9)
    res["flag_ratio_excluded_vs_phage"] = res["flag_rate_excluded_virus"] / max(res["flag_rate_phage"], 1e-9)

    # ---- detection AUROCs (raw + GC-matched) ----
    def aurocs(sub, name):
        out = {}
        y = (sub.cls == "excluded_virus").to_numpy()
        for sig in ["knn", "mahal", "nll"]:
            for ctrl in ["phage", "ref_bacteria"]:
                m_ = sub.cls.isin(["excluded_virus", ctrl]).to_numpy()
                out[f"auroc_{sig}_excl_vs_{ctrl}_{name}"] = float(
                    roc_auc_score(y[m_], sub[sig].to_numpy()[m_]))
        return out
    res.update(aurocs(df, "raw"))
    gcm = df[(df.gc >= 0.35) & (df.gc <= 0.60)]
    res.update(aurocs(gcm, "gcmatch"))
    res["gcmatch_counts"] = gcm.cls.value_counts().to_dict()

    # ---- magnitude ranking: per-genome mean kNN vs per-genome mean NLL ----
    vir = df[df.cls.isin(["phage", "excluded_virus"])]
    g = vir.groupby("genome").agg(knn=("knn", "mean"), nll=("nll", "mean"),
                                  cls=("cls", "first"), n=("seq", "size"))
    res["per_genome"] = g.reset_index().to_dict("records")
    res["spearman_genome_knn_vs_nll"] = float(spearmanr(g.knn, g.nll).statistic)

    # ---- NLL by class (ED Fig 2a direction check) ----
    res["nll_mean_by_class"] = df.groupby("cls").nll.mean().to_dict()
    res["nll_median_by_class"] = df.groupby("cls").nll.median().to_dict()

    df.to_csv(os.path.join(OUT, "evo2_windows_scored.csv"), index=False)
    with open(os.path.join(OUT, "evo2_case_metrics.json"), "w") as f:
        json.dump(res, f, indent=2)
    print(json.dumps(res, indent=2), flush=True)
    print("JOB DONE", flush=True)

if __name__ == "__main__":
    main()