clean up legacy _smiles folders, stray diagnostic files, and half_life non-best models

training_classifiers/.ipynb_checkpoints/binding_affinity_iptm-checkpoint.py

@@ -1,132 +0,0 @@
-#!/usr/bin/env python3
-"""
-extract_iptm_affinity_csv_all.py
-
-Writes:
-  - out_dir/wt_iptm_affinity_all.csv
-  - out_dir/smiles_iptm_affinity_all.csv
-
-Also prints:
-  - N
-  - Spearman rho (affinity vs iptm)
-  - Pearson r (affinity vs iptm)
-"""
-
-from pathlib import Path
-import numpy as np
-import pandas as pd
-
-
-def corr_stats(df: pd.DataFrame, x: str, y: str):
-    # pandas handles NaNs if we already dropped them; still be safe
-    xx = pd.to_numeric(df[x], errors="coerce")
-    yy = pd.to_numeric(df[y], errors="coerce")
-    m = xx.notna() & yy.notna()
-    xx = xx[m]
-    yy = yy[m]
-    n = int(m.sum())
-
-    # Pearson r
-    pearson_r = float(xx.corr(yy, method="pearson")) if n > 1 else float("nan")
-    # Spearman rho
-    spearman_rho = float(xx.corr(yy, method="spearman")) if n > 1 else float("nan")
-
-    return {"n": n, "pearson_r": pearson_r, "spearman_rho": spearman_rho}
-
-
-def clean_one(
-    in_csv: Path,
-    out_csv: Path,
-    iptm_col: str,
-    affinity_col: str = "affinity",
-    keep_cols=(),
-):
-    df = pd.read_csv(in_csv)
-
-    # affinity + iptm must exist
-    need = [affinity_col, iptm_col]
-    missing = [c for c in need if c not in df.columns]
-    if missing:
-        raise ValueError(f"{in_csv} missing columns: {missing}. Found: {list(df.columns)}")
-
-    # coerce numeric
-    df[affinity_col] = pd.to_numeric(df[affinity_col], errors="coerce")
-    df[iptm_col] = pd.to_numeric(df[iptm_col], errors="coerce")
-
-    # drop NaNs in either
-    df = df.dropna(subset=[affinity_col, iptm_col]).reset_index(drop=True)
-
-    # output cols (standardize names)
-    out = pd.DataFrame({
-        "affinity": df[affinity_col].astype(float),
-        "iptm": df[iptm_col].astype(float),
-    })
-
-    # keep split if present (handy for coloring later, but not used for corr)
-    if "split" in df.columns:
-        out.insert(0, "split", df["split"].astype(str))
-
-    # optional extras for labeling/debug
-    for c in keep_cols:
-        if c in df.columns:
-            out[c] = df[c]
-
-    out_csv.parent.mkdir(parents=True, exist_ok=True)
-    out.to_csv(out_csv, index=False)
-
-    stats = corr_stats(out, "iptm", "affinity")
-    print(f"[write] {out_csv}")
-    print(f"  N={stats['n']} | Pearson r={stats['pearson_r']:.4f} | Spearman rho={stats['spearman_rho']:.4f}")
-
-    # also save stats json next to csv
-    stats_path = out_csv.with_suffix(".stats.json")
-    with open(stats_path, "w") as f:
-        import json
-        json.dump(
-            {
-                "input_csv": str(in_csv),
-                "output_csv": str(out_csv),
-                "iptm_col": iptm_col,
-                "affinity_col": affinity_col,
-                **stats,
-            },
-            f,
-            indent=2,
-        )
-
-
-def main():
-    import argparse
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--wt_meta_csv", type=str, required=True)
-    ap.add_argument("--smiles_meta_csv", type=str, required=True)
-    ap.add_argument("--out_dir", type=str, required=True)
-
-    ap.add_argument("--wt_iptm_col", type=str, default="wt_iptm_score")
-    ap.add_argument("--smiles_iptm_col", type=str, default="smiles_iptm_score")
-    ap.add_argument("--affinity_col", type=str, default="affinity")
-    args = ap.parse_args()
-
-    out_dir = Path(args.out_dir)
-
-    clean_one(
-        Path(args.wt_meta_csv),
-        out_dir / "wt_iptm_affinity_all.csv",
-        iptm_col=args.wt_iptm_col,
-        affinity_col=args.affinity_col,
-        keep_cols=("seq1", "seq2", "Fasta2SMILES", "REACT_SMILES"),
-    )
-
-    clean_one(
-        Path(args.smiles_meta_csv),
-        out_dir / "smiles_iptm_affinity_all.csv",
-        iptm_col=args.smiles_iptm_col,
-        affinity_col=args.affinity_col,
-        keep_cols=("seq1", "seq2", "Fasta2SMILES", "REACT_SMILES", "smiles_sequence"),
-    )
-
-    print(f"\n[DONE] CSVs + stats JSONs in: {out_dir}")
-
-
-if __name__ == "__main__":
-    main()

training_classifiers/.ipynb_checkpoints/binding_affinity_split-checkpoint.py

@@ -1,847 +0,0 @@
-#!/usr/bin/env python3
-import os
-import math
-from pathlib import Path
-import sys
-from contextlib import contextmanager
-
-import numpy as np
-import pandas as pd
-import torch
-
-# tqdm is optional; we’ll disable it by default in notebooks
-from tqdm import tqdm
-
-sys.path.append("/vast/projects/pranam/lab/yz927/projects/Classifier_Weight")
-from tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
-
-from datasets import Dataset, DatasetDict, Features, Value, Sequence as HFSequence
-from transformers import AutoTokenizer, EsmModel, AutoModelForMaskedLM
-
-# -------------------------
-# Config
-# -------------------------
-CSV_PATH = Path("/vast/projects/pranam/lab/yz927/projects/Classifier_Weight/c-binding_with_openfold_scores.csv")
-
-OUT_ROOT = Path(
-    "/vast/projects/pranam/lab/yz927/projects/Classifier_Weight/training_data_cleaned/binding_affinity"
-)
-
-# WT (seq) embedding model
-WT_MODEL_NAME = "facebook/esm2_t33_650M_UR50D"
-WT_MAX_LEN = 1022
-WT_BATCH = 32
-
-# SMILES embedding model + tokenizer
-SMI_MODEL_NAME = "aaronfeller/PeptideCLM-23M-all"
-TOKENIZER_VOCAB = "/vast/projects/pranam/lab/yz927/projects/Classifier_Weight/tokenizer/new_vocab.txt"
-TOKENIZER_SPLITS = "/vast/projects/pranam/lab/yz927/projects/Classifier_Weight/tokenizer/new_splits.txt"
-SMI_MAX_LEN = 768
-SMI_BATCH = 128
-
-# Split config
-TRAIN_FRAC = 0.80
-RANDOM_SEED = 1986
-AFFINITY_Q_BINS = 30
-
-# Columns expected in CSV
-COL_SEQ1 = "seq1"
-COL_SEQ2 = "seq2"
-COL_AFF = "affinity"
-COL_F2S = "Fasta2SMILES"
-COL_REACT = "REACT_SMILES"
-COL_WT_IPTM = "wt_iptm_score"
-COL_SMI_IPTM = "smiles_iptm_score"
-
-# Device
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-# -------------------------
-# Quiet / notebook-safe output controls
-# -------------------------
-QUIET = True       # suppress most prints
-USE_TQDM = False   # disable tqdm bars (recommended in Jupyter to avoid crashing)
-LOG_FILE = None    # optionally: OUT_ROOT / "build.log"
-
-def log(msg: str):
-    if LOG_FILE is not None:
-        Path(LOG_FILE).parent.mkdir(parents=True, exist_ok=True)
-        with open(LOG_FILE, "a") as f:
-            f.write(msg.rstrip() + "\n")
-    if not QUIET:
-        print(msg)
-
-def pbar(it, **kwargs):
-    return tqdm(it, **kwargs) if USE_TQDM else it
-
-@contextmanager
-def section(title: str):
-    log(f"\n=== {title} ===")
-    yield
-    log(f"=== done: {title} ===")
-
-
-# -------------------------
-# Helpers
-# -------------------------
-def has_uaa(seq: str) -> bool:
-    return "X" in str(seq).upper()
-
-def affinity_to_class(a: float) -> str:
-    # High: >= 9 ; Moderate: [7, 9) ; Low: < 7
-    if a >= 9.0:
-        return "High"
-    elif a >= 7.0:
-        return "Moderate"
-    else:
-        return "Low"
-
-def make_distribution_matched_split(df: pd.DataFrame) -> pd.DataFrame:
-    df = df.copy()
-
-    df[COL_AFF] = pd.to_numeric(df[COL_AFF], errors="coerce")
-    df = df.dropna(subset=[COL_AFF]).reset_index(drop=True)
-
-    df["affinity_class"] = df[COL_AFF].apply(affinity_to_class)
-
-    try:
-        df["aff_bin"] = pd.qcut(df[COL_AFF], q=AFFINITY_Q_BINS, duplicates="drop")
-        strat_col = "aff_bin"
-    except Exception:
-        df["aff_bin"] = df["affinity_class"]
-        strat_col = "aff_bin"
-
-    rng = np.random.RandomState(RANDOM_SEED)
-
-    df["split"] = None
-    for _, g in df.groupby(strat_col, observed=True):
-        idx = g.index.to_numpy()
-        rng.shuffle(idx)
-        n_train = int(math.floor(len(idx) * TRAIN_FRAC))
-        df.loc[idx[:n_train], "split"] = "train"
-        df.loc[idx[n_train:], "split"] = "val"
-
-    df["split"] = df["split"].fillna("train")
-    return df
-
-def _summ(x):
-    x = np.asarray(x, dtype=float)
-    x = x[~np.isnan(x)]
-    if len(x) == 0:
-        return {"n": 0, "mean": np.nan, "std": np.nan, "p50": np.nan, "p95": np.nan}
-    return {
-        "n": int(len(x)),
-        "mean": float(np.mean(x)),
-        "std": float(np.std(x)),
-        "p50": float(np.quantile(x, 0.50)),
-        "p95": float(np.quantile(x, 0.95)),
-    }
-
-def _len_stats(seqs):
-    lens = np.asarray([len(str(s)) for s in seqs], dtype=float)
-    if len(lens) == 0:
-        return {"n": 0, "mean": np.nan, "std": np.nan, "p50": np.nan, "p95": np.nan}
-    return {
-        "n": int(len(lens)),
-        "mean": float(lens.mean()),
-        "std": float(lens.std()),
-        "p50": float(np.quantile(lens, 0.50)),
-        "p95": float(np.quantile(lens, 0.95)),
-    }
-
-def verify_split_before_embedding(
-    df2: pd.DataFrame,
-    affinity_col: str,
-    split_col: str,
-    seq_col: str,
-    iptm_col: str,
-    aff_class_col: str = "affinity_class",
-    aff_bins: int = 30,
-    save_report_prefix: str | None = None,
-    verbose: bool = False,
-):
-    """
-    Notebook-safe: by default prints only ONE line via `log()`.
-    Optionally writes CSV reports (stats + class proportions).
-    """
-    df2 = df2.copy()
-    df2[affinity_col] = pd.to_numeric(df2[affinity_col], errors="coerce")
-    df2[iptm_col] = pd.to_numeric(df2[iptm_col], errors="coerce")
-
-    assert split_col in df2.columns, f"Missing split col: {split_col}"
-    assert set(df2[split_col].dropna().unique()).issubset({"train", "val"}), f"Unexpected split values: {df2[split_col].unique()}"
-    assert df2[affinity_col].notna().any(), "No valid affinity values after coercion."
-
-    try:
-        df2["_aff_bin_dbg"] = pd.qcut(df2[affinity_col], q=aff_bins, duplicates="drop")
-    except Exception:
-        df2["_aff_bin_dbg"] = df2[aff_class_col].astype(str)
-
-    tr = df2[df2[split_col] == "train"].reset_index(drop=True)
-    va = df2[df2[split_col] == "val"].reset_index(drop=True)
-
-    tr_aff = _summ(tr[affinity_col].to_numpy())
-    va_aff = _summ(va[affinity_col].to_numpy())
-    tr_len = _len_stats(tr[seq_col].tolist())
-    va_len = _len_stats(va[seq_col].tolist())
-
-    # bin drift
-    bin_ct = (
-        df2.groupby([split_col, "_aff_bin_dbg"])
-           .size()
-           .groupby(level=0)
-           .apply(lambda s: s / s.sum())
-    )
-    tr_bins = bin_ct.loc["train"]
-    va_bins = bin_ct.loc["val"]
-    all_bins = tr_bins.index.union(va_bins.index)
-    tr_bins = tr_bins.reindex(all_bins, fill_value=0.0)
-    va_bins = va_bins.reindex(all_bins, fill_value=0.0)
-    max_bin_diff = float(np.max(np.abs(tr_bins.values - va_bins.values)))
-
-    msg = (
-        f"[split-check] rows={len(df2)} train={len(tr)} val={len(va)} | "
-        f"aff(mean±std) train={tr_aff['mean']:.3f}±{tr_aff['std']:.3f} val={va_aff['mean']:.3f}±{va_aff['std']:.3f} | "
-        f"len(p50/p95) train={tr_len['p50']:.1f}/{tr_len['p95']:.1f} val={va_len['p50']:.1f}/{va_len['p95']:.1f} | "
-        f"max_bin_diff={max_bin_diff:.4f}"
-    )
-    log(msg)
-
-    if verbose and (not QUIET):
-        class_ct = df2.groupby([split_col, aff_class_col]).size().unstack(fill_value=0)
-        class_prop = class_ct.div(class_ct.sum(axis=1), axis=0)
-        print("\n[verbose] affinity_class counts:\n", class_ct)
-        print("\n[verbose] affinity_class proportions:\n", class_prop.round(4))
-
-    if save_report_prefix is not None:
-        out = Path(save_report_prefix)
-        out.parent.mkdir(parents=True, exist_ok=True)
-
-        stats_df = pd.DataFrame([
-            {"split": "train", **{f"aff_{k}": v for k, v in tr_aff.items()}, **{f"len_{k}": v for k, v in tr_len.items()}},
-            {"split": "val",   **{f"aff_{k}": v for k, v in va_aff.items()}, **{f"len_{k}": v for k, v in va_len.items()}},
-        ])
-        class_ct = df2.groupby([split_col, aff_class_col]).size().unstack(fill_value=0)
-        class_prop = class_ct.div(class_ct.sum(axis=1), axis=0).reset_index()
-
-        stats_df.to_csv(out.with_suffix(".stats.csv"), index=False)
-        class_prop.to_csv(out.with_suffix(".class_prop.csv"), index=False)
-
-
-# -------------------------
-# WT pooled (ESM2)
-# -------------------------
-@torch.no_grad()
-def wt_pooled_embeddings(seqs, tokenizer, model, batch_size=32, max_length=1022):
-    embs = []
-    for i in pbar(range(0, len(seqs), batch_size)):
-        batch = seqs[i:i + batch_size]
-        inputs = tokenizer(
-            batch,
-            padding=True,
-            truncation=True,
-            max_length=max_length,
-            return_tensors="pt",
-        )
-        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-        out = model(**inputs)
-        h = out.last_hidden_state  # (B, L, H)
-
-        attn = inputs["attention_mask"].unsqueeze(-1)  # (B, L, 1)
-        summed = (h * attn).sum(dim=1)                 # (B, H)
-        denom = attn.sum(dim=1).clamp(min=1e-9)        # (B, 1)
-        pooled = (summed / denom).detach().cpu().numpy()
-        embs.append(pooled)
-
-    return np.vstack(embs)
-
-
-# -------------------------
-# WT unpooled (ESM2)
-# -------------------------
-@torch.no_grad()
-def wt_unpooled_one(seq, tokenizer, model, cls_id, eos_id, max_length=1022):
-    tok = tokenizer(seq, padding=False, truncation=True, max_length=max_length, return_tensors="pt")
-    tok = {k: v.to(DEVICE) for k, v in tok.items()}
-    out = model(**tok)
-    h = out.last_hidden_state[0]           # (L, H)
-    attn = tok["attention_mask"][0].bool() # (L,)
-    ids = tok["input_ids"][0]
-
-    keep = attn.clone()
-    if cls_id is not None:
-        keep &= (ids != cls_id)
-    if eos_id is not None:
-        keep &= (ids != eos_id)
-
-    return h[keep].detach().cpu().to(torch.float16).numpy()
-
-def build_wt_unpooled_dataset(df_split: pd.DataFrame, out_dir: Path, tokenizer, model):
-    """
-    Expects df_split to have:
-      - target_sequence  (seq1)
-      - sequence         (binder seq2; WT binder)
-      - label, affinity_class, COL_AFF, COL_WT_IPTM
-    Saves a dataset where each row contains BOTH:
-      - target_embedding (Lt,H), target_attention_mask, target_length
-      - binder_embedding (Lb,H), binder_attention_mask, binder_length
-    """
-    cls_id = tokenizer.cls_token_id
-    eos_id = tokenizer.eos_token_id
-    H = model.config.hidden_size
-
-    features = Features({
-        "target_sequence": Value("string"),
-        "sequence": Value("string"),
-        "label": Value("float32"),
-        "affinity": Value("float32"),
-        "affinity_class": Value("string"),
-
-        "target_embedding": HFSequence(HFSequence(Value("float16"), length=H)),
-        "target_attention_mask": HFSequence(Value("int8")),
-        "target_length": Value("int64"),
-
-        "binder_embedding": HFSequence(HFSequence(Value("float16"), length=H)),
-        "binder_attention_mask": HFSequence(Value("int8")),
-        "binder_length": Value("int64"),
-
-        COL_WT_IPTM: Value("float32"),
-        COL_AFF: Value("float32"),
-    })
-
-    def gen_rows(df: pd.DataFrame):
-        for r in pbar(df.itertuples(index=False), total=len(df)):
-            tgt = str(getattr(r, "target_sequence")).strip()
-            bnd = str(getattr(r, "sequence")).strip()
-
-            y = float(getattr(r, "label"))
-            aff = float(getattr(r, COL_AFF))
-            acls = str(getattr(r, "affinity_class"))
-
-            iptm = getattr(r, COL_WT_IPTM)
-            iptm = float(iptm) if pd.notna(iptm) else np.nan
-
-            # token embeddings for target + binder (both ESM)
-            t_emb = wt_unpooled_one(tgt, tokenizer, model, cls_id, eos_id, max_length=WT_MAX_LEN)  # (Lt,H)
-            b_emb = wt_unpooled_one(bnd, tokenizer, model, cls_id, eos_id, max_length=WT_MAX_LEN)  # (Lb,H)
-
-            t_list = t_emb.tolist()
-            b_list = b_emb.tolist()
-            Lt = len(t_list)
-            Lb = len(b_list)
-
-            yield {
-                "target_sequence": tgt,
-                "sequence": bnd,
-                "label": np.float32(y),
-                "affinity": np.float32(aff),
-                "affinity_class": acls,
-
-                "target_embedding": t_list,
-                "target_attention_mask": [1] * Lt,
-                "target_length": int(Lt),
-
-                "binder_embedding": b_list,
-                "binder_attention_mask": [1] * Lb,
-                "binder_length": int(Lb),
-
-                COL_WT_IPTM: np.float32(iptm) if not np.isnan(iptm) else np.float32(np.nan),
-                COL_AFF: np.float32(aff),
-            }
-
-    out_dir.mkdir(parents=True, exist_ok=True)
-    ds = Dataset.from_generator(lambda: gen_rows(df_split), features=features)
-    ds.save_to_disk(str(out_dir), max_shard_size="1GB")
-    return ds
-
-def build_smiles_unpooled_paired_dataset(df_split: pd.DataFrame, out_dir: Path, wt_tokenizer, wt_model_unpooled,
-                                        smi_tok, smi_roformer):
-    """
-    df_split must have:
-      - target_sequence (seq1)
-      - sequence        (binder smiles string)
-      - label, affinity_class, COL_AFF, COL_SMI_IPTM
-    Saves rows with:
-      target_embedding (Lt,Ht) from ESM
-      binder_embedding (Lb,Hb) from PeptideCLM
-    """
-    cls_id = wt_tokenizer.cls_token_id
-    eos_id = wt_tokenizer.eos_token_id
-    Ht = wt_model_unpooled.config.hidden_size
-
-    # Infer Hb from one forward pass? easiest: run one mini batch outside in main if you want.
-    # Here: we’ll infer from model config if available.
-    Hb = getattr(smi_roformer.config, "hidden_size", None)
-    if Hb is None:
-        Hb = getattr(smi_roformer.config, "dim", None)
-    if Hb is None:
-        raise ValueError("Cannot infer Hb from smi_roformer config; print(smi_roformer.config) and set Hb manually.")
-
-    features = Features({
-        "target_sequence": Value("string"),
-        "sequence": Value("string"),
-        "label": Value("float32"),
-        "affinity": Value("float32"),
-        "affinity_class": Value("string"),
-
-        "target_embedding": HFSequence(HFSequence(Value("float16"), length=Ht)),
-        "target_attention_mask": HFSequence(Value("int8")),
-        "target_length": Value("int64"),
-
-        "binder_embedding": HFSequence(HFSequence(Value("float16"), length=Hb)),
-        "binder_attention_mask": HFSequence(Value("int8")),
-        "binder_length": Value("int64"),
-
-        COL_SMI_IPTM: Value("float32"),
-        COL_AFF: Value("float32"),
-    })
-
-    def gen_rows(df: pd.DataFrame):
-        for r in pbar(df.itertuples(index=False), total=len(df)):
-            tgt = str(getattr(r, "target_sequence")).strip()
-            bnd = str(getattr(r, "sequence")).strip()
-
-            y = float(getattr(r, "label"))
-            aff = float(getattr(r, COL_AFF))
-            acls = str(getattr(r, "affinity_class"))
-
-            iptm = getattr(r, COL_SMI_IPTM)
-            iptm = float(iptm) if pd.notna(iptm) else np.nan
-
-            # target token embeddings (ESM)
-            t_emb = wt_unpooled_one(tgt, wt_tokenizer, wt_model_unpooled, cls_id, eos_id, max_length=WT_MAX_LEN)
-            t_list = t_emb.tolist()
-            Lt = len(t_list)
-
-            # binder token embeddings (PeptideCLM) — single-item batch
-            _, tok_list, mask_list, lengths = smiles_embed_batch_return_both(
-                [bnd], smi_tok, smi_roformer, max_length=SMI_MAX_LEN
-            )
-            b_emb = tok_list[0]  # np.float16 (Lb, Hb)
-            b_list = b_emb.tolist()
-            Lb = int(lengths[0])
-            b_mask = mask_list[0].astype(np.int8).tolist()
-
-            yield {
-                "target_sequence": tgt,
-                "sequence": bnd,
-                "label": np.float32(y),
-                "affinity": np.float32(aff),
-                "affinity_class": acls,
-
-                "target_embedding": t_list,
-                "target_attention_mask": [1] * Lt,
-                "target_length": int(Lt),
-
-                "binder_embedding": b_list,
-                "binder_attention_mask": [int(x) for x in b_mask],
-                "binder_length": int(Lb),
-
-                COL_SMI_IPTM: np.float32(iptm) if not np.isnan(iptm) else np.float32(np.nan),
-                COL_AFF: np.float32(aff),
-            }
-
-    out_dir.mkdir(parents=True, exist_ok=True)
-    ds = Dataset.from_generator(lambda: gen_rows(df_split), features=features)
-    ds.save_to_disk(str(out_dir), max_shard_size="1GB")
-    return ds
-
-
-# -------------------------
-# SMILES pooled + unpooled (PeptideCLM)
-# -------------------------
-def get_special_ids(tokenizer_obj):
-    cand = [
-        getattr(tokenizer_obj, "pad_token_id", None),
-        getattr(tokenizer_obj, "cls_token_id", None),
-        getattr(tokenizer_obj, "sep_token_id", None),
-        getattr(tokenizer_obj, "bos_token_id", None),
-        getattr(tokenizer_obj, "eos_token_id", None),
-        getattr(tokenizer_obj, "mask_token_id", None),
-    ]
-    return sorted({x for x in cand if x is not None})
-
-@torch.no_grad()
-def smiles_embed_batch_return_both(batch_sequences, tokenizer_obj, model_roformer, max_length):
-    tok = tokenizer_obj(
-        batch_sequences,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=max_length,
-    )
-    input_ids = tok["input_ids"].to(DEVICE)
-    attention_mask = tok["attention_mask"].to(DEVICE)
-
-    outputs = model_roformer(input_ids=input_ids, attention_mask=attention_mask)
-    last_hidden = outputs.last_hidden_state  # (B, L, H)
-
-    special_ids = get_special_ids(tokenizer_obj)
-    valid = attention_mask.bool()
-    if len(special_ids) > 0:
-        sid = torch.tensor(special_ids, device=DEVICE, dtype=torch.long)
-        if hasattr(torch, "isin"):
-            valid = valid & (~torch.isin(input_ids, sid))
-        else:
-            m = torch.zeros_like(valid)
-            for s in special_ids:
-                m |= (input_ids == s)
-            valid = valid & (~m)
-
-    valid_f = valid.unsqueeze(-1).float()
-    summed = torch.sum(last_hidden * valid_f, dim=1)
-    denom = torch.clamp(valid_f.sum(dim=1), min=1e-9)
-    pooled = (summed / denom).detach().cpu().numpy()
-
-    token_emb_list, mask_list, lengths = [], [], []
-    for b in range(last_hidden.shape[0]):
-        emb = last_hidden[b, valid[b]]  # (Li, H)
-        token_emb_list.append(emb.detach().cpu().to(torch.float16).numpy())
-        li = emb.shape[0]
-        lengths.append(int(li))
-        mask_list.append(np.ones((li,), dtype=np.int8))
-
-    return pooled, token_emb_list, mask_list, lengths
-
-def smiles_generate_embeddings_batched_both(seqs, tokenizer_obj, model_roformer, batch_size, max_length):
-    pooled_all = []
-    token_emb_all = []
-    mask_all = []
-    lengths_all = []
-
-    for i in pbar(range(0, len(seqs), batch_size)):
-        batch = seqs[i:i + batch_size]
-        pooled, tok_list, m_list, lens = smiles_embed_batch_return_both(
-            batch, tokenizer_obj, model_roformer, max_length
-        )
-        pooled_all.append(pooled)
-        token_emb_all.extend(tok_list)
-        mask_all.extend(m_list)
-        lengths_all.extend(lens)
-
-    return np.vstack(pooled_all), token_emb_all, mask_all, lengths_all
-
-# -------------------------
-# Target embedding cache (NO extra ESM runs)
-# We will compute target pooled embeddings ONCE from WT view, then reuse for SMILES.
-# -------------------------
-def build_target_cache_from_wt_view(wt_view_train: pd.DataFrame, wt_view_val: pd.DataFrame):
-    wt_tok = AutoTokenizer.from_pretrained(WT_MODEL_NAME)
-    wt_model = EsmModel.from_pretrained(WT_MODEL_NAME).to(DEVICE).eval()
-
-    # compute target pooled embeddings once
-    tgt_wt_train = wt_view_train["target_sequence"].astype(str).tolist()
-    tgt_wt_val   = wt_view_val["target_sequence"].astype(str).tolist()
-
-    wt_train_tgt_emb = wt_pooled_embeddings(
-        tgt_wt_train, wt_tok, wt_model, batch_size=WT_BATCH, max_length=WT_MAX_LEN
-    )
-    wt_val_tgt_emb = wt_pooled_embeddings(
-        tgt_wt_val, wt_tok, wt_model, batch_size=WT_BATCH, max_length=WT_MAX_LEN
-    )
-
-    # build dict: target_sequence -> embedding (float32 array)
-    # if duplicates exist, last wins; you can add checks if needed
-    train_map = {s: e for s, e in zip(tgt_wt_train, wt_train_tgt_emb)}
-    val_map   = {s: e for s, e in zip(tgt_wt_val,   wt_val_tgt_emb)}
-    return wt_tok, wt_model, wt_train_tgt_emb, wt_val_tgt_emb, train_map, val_map
-# -------------------------
-# Main
-# -------------------------
-def main():
-    log(f"[INFO] DEVICE: {DEVICE}")
-    OUT_ROOT.mkdir(parents=True, exist_ok=True)
-
-    # 1) Load
-    with section("load csv + dedup"):
-        df = pd.read_csv(CSV_PATH)
-        for c in [COL_SEQ1, COL_SEQ2, COL_F2S, COL_REACT]:
-            if c in df.columns:
-                df[c] = df[c].apply(lambda x: x.strip() if isinstance(x, str) else x)
-        
-        # Dedup on the full identity tuple you want
-        DEDUP_COLS = [COL_SEQ1, COL_SEQ2, COL_F2S, COL_REACT]
-        df = df.drop_duplicates(subset=DEDUP_COLS).reset_index(drop=True)
-        
-        print("Rows after dedup on", DEDUP_COLS, ":", len(df))
-
-        need = [COL_SEQ1, COL_SEQ2, COL_AFF, COL_F2S, COL_REACT, COL_WT_IPTM, COL_SMI_IPTM]
-        missing = [c for c in need if c not in df.columns]
-        if missing:
-            raise ValueError(f"Missing required columns: {missing}")
-
-        # numeric affinity for both branches
-        df[COL_AFF] = pd.to_numeric(df[COL_AFF], errors="coerce")
-
-    # 2) Build WT subset + SMILES subset separately (NO global dropping)
-    with section("prepare wt/smiles subsets"):
-        # WT: requires a canonical peptide sequence (no X) + affinity
-        df_wt = df.copy()
-        df_wt["wt_sequence"] = df_wt[COL_SEQ2].astype(str).str.strip()
-        df_wt = df_wt.dropna(subset=[COL_AFF]).reset_index(drop=True)
-        df_wt = df_wt[df_wt["wt_sequence"].notna() & (df_wt["wt_sequence"] != "")]
-        df_wt = df_wt[~df_wt["wt_sequence"].str.contains("X", case=False, na=False)].reset_index(drop=True)
-
-        # SMILES: requires affinity + a usable picked SMILES (UAA->REACT, else->Fasta2SMILES)
-        df_smi = df.copy()
-        df_smi = df_smi.dropna(subset=[COL_AFF]).reset_index(drop=True)
-        df_smi = df_smi[
-            pd.to_numeric(df_smi[COL_SMI_IPTM], errors="coerce").notna()
-        ].reset_index(drop=True) # empty iptm means sth wrong with their smiles sequenc
-
-        is_uaa = df_smi[COL_SEQ2].astype(str).str.contains("X", case=False, na=False)
-        df_smi["smiles_sequence"] = np.where(is_uaa, df_smi[COL_REACT], df_smi[COL_F2S])
-        df_smi["smiles_sequence"] = df_smi["smiles_sequence"].astype(str).str.strip()
-        df_smi = df_smi[df_smi["smiles_sequence"].notna() & (df_smi["smiles_sequence"] != "")]
-        df_smi = df_smi[~df_smi["smiles_sequence"].isin(["nan", "None"])].reset_index(drop=True)
-
-        log(f"[counts] WT rows={len(df_wt)} | SMILES rows={len(df_smi)} (after per-branch filtering)")
-
-    # 3) Split separately (different sizes and memberships are expected)
-    with section("split wt and smiles separately"):
-        df_wt2 = make_distribution_matched_split(df_wt)
-        df_smi2 = make_distribution_matched_split(df_smi)
-
-        # save split tables
-        wt_split_csv = OUT_ROOT / "binding_affinity_wt_meta_with_split.csv"
-        smi_split_csv = OUT_ROOT / "binding_affinity_smiles_meta_with_split.csv"
-        df_wt2.to_csv(wt_split_csv, index=False)
-        df_smi2.to_csv(smi_split_csv, index=False)
-        log(f"Saved WT split meta: {wt_split_csv}")
-        log(f"Saved SMILES split meta: {smi_split_csv}")
-
-        # lightweight double-check (one-line)
-        verify_split_before_embedding(
-            df2=df_wt2,
-            affinity_col=COL_AFF,
-            split_col="split",
-            seq_col="wt_sequence",
-            iptm_col=COL_WT_IPTM,
-            aff_class_col="affinity_class",
-            aff_bins=AFFINITY_Q_BINS,
-            save_report_prefix=str(OUT_ROOT / "wt_split_doublecheck_report"),
-            verbose=False,
-        )
-        verify_split_before_embedding(
-            df2=df_smi2,
-            affinity_col=COL_AFF,
-            split_col="split",
-            seq_col="smiles_sequence",
-            iptm_col=COL_SMI_IPTM,
-            aff_class_col="affinity_class",
-            aff_bins=AFFINITY_Q_BINS,
-            save_report_prefix=str(OUT_ROOT / "smiles_split_doublecheck_report"),
-            verbose=False,
-        )
-
-    # Prepare split views
-    def prep_view(df_in: pd.DataFrame, binder_seq_col: str, iptm_col: str) -> pd.DataFrame:
-        out = df_in.copy()
-        out["target_sequence"] = out[COL_SEQ1].astype(str).str.strip()   # <-- NEW
-        out["sequence"] = out[binder_seq_col].astype(str).str.strip()   # binder
-        out["label"] = pd.to_numeric(out[COL_AFF], errors="coerce")
-        out[iptm_col] = pd.to_numeric(out[iptm_col], errors="coerce")
-        out[COL_AFF] = pd.to_numeric(out[COL_AFF], errors="coerce")
-        out = out.dropna(subset=["target_sequence", "sequence", "label"]).reset_index(drop=True)
-        return out[["target_sequence", "sequence", "label", "split", iptm_col, COL_AFF, "affinity_class"]]
-
-    wt_view = prep_view(df_wt2, "wt_sequence", COL_WT_IPTM)
-    smi_view = prep_view(df_smi2, "smiles_sequence", COL_SMI_IPTM)
-
-    # -------------------------
-    # Split views
-    # -------------------------
-    wt_train = wt_view[wt_view["split"] == "train"].reset_index(drop=True)
-    wt_val   = wt_view[wt_view["split"] == "val"].reset_index(drop=True)
-    smi_train = smi_view[smi_view["split"] == "train"].reset_index(drop=True)
-    smi_val   = smi_view[smi_view["split"] == "val"].reset_index(drop=True)
-    
-    
-    # =========================
-    # TARGET pooled embeddings (ESM) — SEPARATE per branch
-    # =========================
-    with section("TARGET pooled embeddings (ESM) — WT + SMILES separately"):
-        wt_tok = AutoTokenizer.from_pretrained(WT_MODEL_NAME)
-        wt_esm = EsmModel.from_pretrained(WT_MODEL_NAME).to(DEVICE).eval()
-    
-        # ---- WT targets ----
-        wt_train_tgt_emb = wt_pooled_embeddings(
-            wt_train["target_sequence"].astype(str).str.strip().tolist(),
-            wt_tok, wt_esm,
-            batch_size=WT_BATCH,
-            max_length=WT_MAX_LEN,
-        ).astype(np.float32)
-    
-        wt_val_tgt_emb = wt_pooled_embeddings(
-            wt_val["target_sequence"].astype(str).str.strip().tolist(),
-            wt_tok, wt_esm,
-            batch_size=WT_BATCH,
-            max_length=WT_MAX_LEN,
-        ).astype(np.float32)
-    
-        # ---- SMILES targets (independent; may include UAA-only targets) ----
-        smi_train_tgt_emb = wt_pooled_embeddings(
-            smi_train["target_sequence"].astype(str).str.strip().tolist(),
-            wt_tok, wt_esm,
-            batch_size=WT_BATCH,
-            max_length=WT_MAX_LEN,
-        ).astype(np.float32)
-    
-        smi_val_tgt_emb = wt_pooled_embeddings(
-            smi_val["target_sequence"].astype(str).str.strip().tolist(),
-            wt_tok, wt_esm,
-            batch_size=WT_BATCH,
-            max_length=WT_MAX_LEN,
-        ).astype(np.float32)
-    
-    
-    # =========================
-    # WT pooled binder embeddings (binder = WT peptide)
-    # =========================
-    with section("WT pooled binder embeddings + save"):
-        wt_train_emb = wt_pooled_embeddings(
-            wt_train["sequence"].astype(str).str.strip().tolist(),
-            wt_tok, wt_esm,
-            batch_size=WT_BATCH,
-            max_length=WT_MAX_LEN,
-        ).astype(np.float32)
-    
-        wt_val_emb = wt_pooled_embeddings(
-            wt_val["sequence"].astype(str).str.strip().tolist(),
-            wt_tok, wt_esm,
-            batch_size=WT_BATCH,
-            max_length=WT_MAX_LEN,
-        ).astype(np.float32)
-    
-        wt_train_ds = Dataset.from_dict({
-            "target_sequence": wt_train["target_sequence"].tolist(),
-            "sequence": wt_train["sequence"].tolist(),
-            "label": wt_train["label"].astype(float).tolist(),
-            "target_embedding": wt_train_tgt_emb,
-            "embedding": wt_train_emb,
-            COL_WT_IPTM: wt_train[COL_WT_IPTM].astype(float).tolist(),
-            COL_AFF: wt_train[COL_AFF].astype(float).tolist(),
-            "affinity_class": wt_train["affinity_class"].tolist(),
-        })
-    
-        wt_val_ds = Dataset.from_dict({
-            "target_sequence": wt_val["target_sequence"].tolist(),
-            "sequence": wt_val["sequence"].tolist(),
-            "label": wt_val["label"].astype(float).tolist(),
-            "target_embedding": wt_val_tgt_emb,
-            "embedding": wt_val_emb,
-            COL_WT_IPTM: wt_val[COL_WT_IPTM].astype(float).tolist(),
-            COL_AFF: wt_val[COL_AFF].astype(float).tolist(),
-            "affinity_class": wt_val["affinity_class"].tolist(),
-        })
-    
-        wt_pooled_dd = DatasetDict({"train": wt_train_ds, "val": wt_val_ds})
-        wt_pooled_out = OUT_ROOT / "pair_wt_wt_pooled"
-        wt_pooled_dd.save_to_disk(str(wt_pooled_out))
-        log(f"Saved WT pooled -> {wt_pooled_out}")
-    
-    
-    # =========================
-    # SMILES pooled binder embeddings (binder = SMILES via PeptideCLM)
-    # =========================
-    with section("SMILES pooled binder embeddings + save"):
-        smi_tok = SMILES_SPE_Tokenizer(TOKENIZER_VOCAB, TOKENIZER_SPLITS)
-        smi_roformer = (
-            AutoModelForMaskedLM
-            .from_pretrained(SMI_MODEL_NAME)
-            .roformer
-            .to(DEVICE)
-            .eval()
-        )
-    
-        smi_train_pooled, _, _, _ = smiles_generate_embeddings_batched_both(
-            smi_train["sequence"].astype(str).str.strip().tolist(),
-            smi_tok, smi_roformer,
-            batch_size=SMI_BATCH,
-            max_length=SMI_MAX_LEN,
-        )
-    
-        smi_val_pooled, _, _, _ = smiles_generate_embeddings_batched_both(
-            smi_val["sequence"].astype(str).str.strip().tolist(),
-            smi_tok, smi_roformer,
-            batch_size=SMI_BATCH,
-            max_length=SMI_MAX_LEN,
-        )
-    
-        smi_train_ds = Dataset.from_dict({
-            "target_sequence": smi_train["target_sequence"].tolist(),
-            "sequence": smi_train["sequence"].tolist(),
-            "label": smi_train["label"].astype(float).tolist(),
-            "target_embedding": smi_train_tgt_emb,
-            "embedding": smi_train_pooled.astype(np.float32),
-            COL_SMI_IPTM: smi_train[COL_SMI_IPTM].astype(float).tolist(),
-            COL_AFF: smi_train[COL_AFF].astype(float).tolist(),
-            "affinity_class": smi_train["affinity_class"].tolist(),
-        })
-    
-        smi_val_ds = Dataset.from_dict({
-            "target_sequence": smi_val["target_sequence"].tolist(),
-            "sequence": smi_val["sequence"].tolist(),
-            "label": smi_val["label"].astype(float).tolist(),
-            "target_embedding": smi_val_tgt_emb,
-            "embedding": smi_val_pooled.astype(np.float32),
-            COL_SMI_IPTM: smi_val[COL_SMI_IPTM].astype(float).tolist(),
-            COL_AFF: smi_val[COL_AFF].astype(float).tolist(),
-            "affinity_class": smi_val["affinity_class"].tolist(),
-        })
-    
-        smi_pooled_dd = DatasetDict({"train": smi_train_ds, "val": smi_val_ds})
-        smi_pooled_out = OUT_ROOT / "pair_wt_smiles_pooled"
-        smi_pooled_dd.save_to_disk(str(smi_pooled_out))
-        log(f"Saved SMILES pooled -> {smi_pooled_out}")
-
-
-        # =========================
-    # WT unpooled paired (ESM target + ESM binder) + save
-    # =========================
-    with section("WT unpooled paired embeddings + save"):
-        wt_tok_unpooled = wt_tok                       # reuse tokenizer
-        wt_esm_unpooled = wt_esm                       # reuse model
-
-        wt_unpooled_out = OUT_ROOT / "pair_wt_wt_unpooled"
-        wt_unpooled_dd = DatasetDict({
-            "train": build_wt_unpooled_dataset(wt_train, wt_unpooled_out / "train",
-                                               wt_tok_unpooled, wt_esm_unpooled),
-            "val":   build_wt_unpooled_dataset(wt_val,   wt_unpooled_out / "val",
-                                               wt_tok_unpooled, wt_esm_unpooled),
-        })
-        # (Optional) also save as DatasetDict root if you want a single load_from_disk path:
-        wt_unpooled_dd.save_to_disk(str(wt_unpooled_out))
-        log(f"Saved WT unpooled -> {wt_unpooled_out}")
-
-
-    # =========================
-    # SMILES unpooled paired (ESM target + PeptideCLM binder) + save
-    # =========================
-    with section("SMILES unpooled paired embeddings + save"):
-        # reuse already-loaded smi_tok/smi_roformer from pooled section if still in scope;
-        # otherwise re-init here:
-        # smi_tok = SMILES_SPE_Tokenizer(TOKENIZER_VOCAB, TOKENIZER_SPLITS)
-        # smi_roformer = AutoModelForMaskedLM.from_pretrained(SMI_MODEL_NAME).roformer.to(DEVICE).eval()
-
-        smi_unpooled_out = OUT_ROOT / "pair_wt_smiles_unpooled"
-        smi_unpooled_dd = DatasetDict({
-            "train": build_smiles_unpooled_paired_dataset(
-                smi_train, smi_unpooled_out / "train",
-                wt_tok, wt_esm,
-                smi_tok, smi_roformer
-            ),
-            "val": build_smiles_unpooled_paired_dataset(
-                smi_val, smi_unpooled_out / "val",
-                wt_tok, wt_esm,
-                smi_tok, smi_roformer
-            ),
-        })
-        smi_unpooled_dd.save_to_disk(str(smi_unpooled_out))
-        log(f"Saved SMILES unpooled -> {smi_unpooled_out}")
-
-    log(f"\n[DONE] All datasets saved under: {OUT_ROOT}")
-
-
-if __name__ == "__main__":
-    main()

training_classifiers/.ipynb_checkpoints/binding_training-checkpoint.py

@@ -1,414 +0,0 @@
-import os, json
-from pathlib import Path
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-import optuna
-from datasets import load_from_disk, DatasetDict
-from scipy.stats import spearmanr
-from lightning.pytorch import seed_everything
-seed_everything(1986)
-
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-
-def safe_spearmanr(y_true: np.ndarray, y_pred: np.ndarray) -> float:
-    rho = spearmanr(y_true, y_pred).correlation
-    if rho is None or np.isnan(rho):
-        return 0.0
-    return float(rho)
-
-
-# -----------------------------
-# Affinity class thresholds (final spec)
-# High >= 9 ; Moderate 7-9 ; Low < 7
-# 0=High, 1=Moderate, 2=Low
-# -----------------------------
-def affinity_to_class_tensor(y: torch.Tensor) -> torch.Tensor:
-    high = y >= 9.0
-    low  = y < 7.0
-    mid  = ~(high | low)
-    cls = torch.zeros_like(y, dtype=torch.long)
-    cls[mid] = 1
-    cls[low] = 2
-    return cls
-
-
-# -----------------------------
-# Load paired DatasetDict
-# -----------------------------
-def load_split_paired(path: str):
-    dd = load_from_disk(path)
-    if not isinstance(dd, DatasetDict):
-        raise ValueError(f"Expected DatasetDict at {path}")
-    if "train" not in dd or "val" not in dd:
-        raise ValueError(f"DatasetDict missing train/val at {path}")
-    return dd["train"], dd["val"]
-
-
-# -----------------------------
-# Collate: pooled paired
-# -----------------------------
-def collate_pair_pooled(batch):
-    Pt = torch.tensor([x["target_embedding"] for x in batch], dtype=torch.float32)  # (B,Ht)
-    Pb = torch.tensor([x["binder_embedding"] for x in batch], dtype=torch.float32)  # (B,Hb)
-    y  = torch.tensor([float(x["label"]) for x in batch], dtype=torch.float32)
-    return Pt, Pb, y
-
-
-# -----------------------------
-# Collate: unpooled paired
-# -----------------------------
-def collate_pair_unpooled(batch):
-    B = len(batch)
-    Ht = len(batch[0]["target_embedding"][0])
-    Hb = len(batch[0]["binder_embedding"][0])
-    Lt_max = max(int(x["target_length"]) for x in batch)
-    Lb_max = max(int(x["binder_length"]) for x in batch)
-
-    Pt = torch.zeros(B, Lt_max, Ht, dtype=torch.float32)
-    Pb = torch.zeros(B, Lb_max, Hb, dtype=torch.float32)
-    Mt = torch.zeros(B, Lt_max, dtype=torch.bool)
-    Mb = torch.zeros(B, Lb_max, dtype=torch.bool)
-    y  = torch.tensor([float(x["label"]) for x in batch], dtype=torch.float32)
-
-    for i, x in enumerate(batch):
-        t = torch.tensor(x["target_embedding"], dtype=torch.float32)
-        b = torch.tensor(x["binder_embedding"], dtype=torch.float32)
-        lt, lb = t.shape[0], b.shape[0]
-        Pt[i, :lt] = t
-        Pb[i, :lb] = b
-        Mt[i, :lt] = torch.tensor(x["target_attention_mask"][:lt], dtype=torch.bool)
-        Mb[i, :lb] = torch.tensor(x["binder_attention_mask"][:lb], dtype=torch.bool)
-
-    return Pt, Mt, Pb, Mb, y
-
-
-# -----------------------------
-# Cross-attention models
-# -----------------------------
-class CrossAttnPooled(nn.Module):
-    """
-    pooled vectors -> treat as single-token sequences for cross attention
-    """
-    def __init__(self, Ht, Hb, hidden=512, n_heads=8, n_layers=3, dropout=0.1):
-        super().__init__()
-        self.t_proj = nn.Sequential(nn.Linear(Ht, hidden), nn.LayerNorm(hidden))
-        self.b_proj = nn.Sequential(nn.Linear(Hb, hidden), nn.LayerNorm(hidden))
-
-        self.layers = nn.ModuleList([])
-        for _ in range(n_layers):
-            self.layers.append(nn.ModuleDict({
-                "attn_tb": nn.MultiheadAttention(hidden, n_heads, dropout=dropout, batch_first=False),
-                "attn_bt": nn.MultiheadAttention(hidden, n_heads, dropout=dropout, batch_first=False),
-                "n1t": nn.LayerNorm(hidden),
-                "n2t": nn.LayerNorm(hidden),
-                "n1b": nn.LayerNorm(hidden),
-                "n2b": nn.LayerNorm(hidden),
-                "fft": nn.Sequential(nn.Linear(hidden, 4*hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(4*hidden, hidden)),
-                "ffb": nn.Sequential(nn.Linear(hidden, 4*hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(4*hidden, hidden)),
-            }))
-
-        self.shared = nn.Sequential(nn.Linear(2*hidden, hidden), nn.GELU(), nn.Dropout(dropout))
-        self.reg = nn.Linear(hidden, 1)
-        self.cls = nn.Linear(hidden, 3)
-
-    def forward(self, t_vec, b_vec):
-        # (B,Ht),(B,Hb)
-        t = self.t_proj(t_vec).unsqueeze(0)  # (1,B,H)
-        b = self.b_proj(b_vec).unsqueeze(0)  # (1,B,H)
-
-        for L in self.layers:
-            t_attn, _ = L["attn_tb"](t, b, b)
-            t = L["n1t"]((t + t_attn).transpose(0,1)).transpose(0,1)
-            t = L["n2t"]((t + L["fft"](t)).transpose(0,1)).transpose(0,1)
-
-            b_attn, _ = L["attn_bt"](b, t, t)
-            b = L["n1b"]((b + b_attn).transpose(0,1)).transpose(0,1)
-            b = L["n2b"]((b + L["ffb"](b)).transpose(0,1)).transpose(0,1)
-
-        t0 = t[0]
-        b0 = b[0]
-        z = torch.cat([t0, b0], dim=-1)
-        h = self.shared(z)
-        return self.reg(h).squeeze(-1), self.cls(h)
-
-
-class CrossAttnUnpooled(nn.Module):
-    """
-    token sequences with masks; alternating cross attention.
-    """
-    def __init__(self, Ht, Hb, hidden=512, n_heads=8, n_layers=3, dropout=0.1):
-        super().__init__()
-        self.t_proj = nn.Sequential(nn.Linear(Ht, hidden), nn.LayerNorm(hidden))
-        self.b_proj = nn.Sequential(nn.Linear(Hb, hidden), nn.LayerNorm(hidden))
-
-        self.layers = nn.ModuleList([])
-        for _ in range(n_layers):
-            self.layers.append(nn.ModuleDict({
-                "attn_tb": nn.MultiheadAttention(hidden, n_heads, dropout=dropout, batch_first=True),
-                "attn_bt": nn.MultiheadAttention(hidden, n_heads, dropout=dropout, batch_first=True),
-                "n1t": nn.LayerNorm(hidden),
-                "n2t": nn.LayerNorm(hidden),
-                "n1b": nn.LayerNorm(hidden),
-                "n2b": nn.LayerNorm(hidden),
-                "fft": nn.Sequential(nn.Linear(hidden, 4*hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(4*hidden, hidden)),
-                "ffb": nn.Sequential(nn.Linear(hidden, 4*hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(4*hidden, hidden)),
-            }))
-
-        self.shared = nn.Sequential(nn.Linear(2*hidden, hidden), nn.GELU(), nn.Dropout(dropout))
-        self.reg = nn.Linear(hidden, 1)
-        self.cls = nn.Linear(hidden, 3)
-
-    def masked_mean(self, X, M):
-        Mf = M.unsqueeze(-1).float()
-        denom = Mf.sum(dim=1).clamp(min=1.0)
-        return (X * Mf).sum(dim=1) / denom
-
-    def forward(self, T, Mt, B, Mb):
-        # T:(B,Lt,Ht), Mt:(B,Lt) ; B:(B,Lb,Hb), Mb:(B,Lb)
-        T = self.t_proj(T)
-        Bx = self.b_proj(B)
-
-        kp_t = ~Mt  # key_padding_mask True = pad
-        kp_b = ~Mb
-
-        for L in self.layers:
-            # T attends to B
-            T_attn, _ = L["attn_tb"](T, Bx, Bx, key_padding_mask=kp_b)
-            T = L["n1t"](T + T_attn)
-            T = L["n2t"](T + L["fft"](T))
-
-            # B attends to T
-            B_attn, _ = L["attn_bt"](Bx, T, T, key_padding_mask=kp_t)
-            Bx = L["n1b"](Bx + B_attn)
-            Bx = L["n2b"](Bx + L["ffb"](Bx))
-
-        t_pool = self.masked_mean(T, Mt)
-        b_pool = self.masked_mean(Bx, Mb)
-        z = torch.cat([t_pool, b_pool], dim=-1)
-        h = self.shared(z)
-        return self.reg(h).squeeze(-1), self.cls(h)
-
-
-# -----------------------------
-# Train/eval
-# -----------------------------
-@torch.no_grad()
-def eval_spearman_pooled(model, loader):
-    model.eval()
-    ys, ps = [], []
-    for t, b, y in loader:
-        t = t.to(DEVICE, non_blocking=True)
-        b = b.to(DEVICE, non_blocking=True)
-        pred, _ = model(t, b)
-        ys.append(y.numpy())
-        ps.append(pred.detach().cpu().numpy())
-    return safe_spearmanr(np.concatenate(ys), np.concatenate(ps))
-
-@torch.no_grad()
-def eval_spearman_unpooled(model, loader):
-    model.eval()
-    ys, ps = [], []
-    for T, Mt, B, Mb, y in loader:
-        T = T.to(DEVICE, non_blocking=True)
-        Mt = Mt.to(DEVICE, non_blocking=True)
-        B = B.to(DEVICE, non_blocking=True)
-        Mb = Mb.to(DEVICE, non_blocking=True)
-        pred, _ = model(T, Mt, B, Mb)
-        ys.append(y.numpy())
-        ps.append(pred.detach().cpu().numpy())
-    return safe_spearmanr(np.concatenate(ys), np.concatenate(ps))
-
-def train_one_epoch_pooled(model, loader, opt, loss_reg, loss_cls, cls_w=1.0, clip=1.0):
-    model.train()
-    for t, b, y in loader:
-        t = t.to(DEVICE, non_blocking=True)
-        b = b.to(DEVICE, non_blocking=True)
-        y = y.to(DEVICE, non_blocking=True)
-        y_cls = affinity_to_class_tensor(y)
-
-        opt.zero_grad(set_to_none=True)
-        pred, logits = model(t, b)
-        L = loss_reg(pred, y) + cls_w * loss_cls(logits, y_cls)
-        L.backward()
-        if clip is not None:
-            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
-        opt.step()
-
-def train_one_epoch_unpooled(model, loader, opt, loss_reg, loss_cls, cls_w=1.0, clip=1.0):
-    model.train()
-    for T, Mt, B, Mb, y in loader:
-        T = T.to(DEVICE, non_blocking=True)
-        Mt = Mt.to(DEVICE, non_blocking=True)
-        B = B.to(DEVICE, non_blocking=True)
-        Mb = Mb.to(DEVICE, non_blocking=True)
-        y = y.to(DEVICE, non_blocking=True)
-        y_cls = affinity_to_class_tensor(y)
-
-        opt.zero_grad(set_to_none=True)
-        pred, logits = model(T, Mt, B, Mb)
-        L = loss_reg(pred, y) + cls_w * loss_cls(logits, y_cls)
-        L.backward()
-        if clip is not None:
-            torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
-        opt.step()
-
-
-# -----------------------------
-# Optuna objective
-# -----------------------------
-def objective_crossattn(trial: optuna.Trial, mode: str, train_ds, val_ds) -> float:
-    lr = trial.suggest_float("lr", 1e-5, 3e-3, log=True)
-    wd = trial.suggest_float("weight_decay", 1e-10, 1e-2, log=True)
-    dropout = trial.suggest_float("dropout", 0.0, 0.4)
-    hidden = trial.suggest_categorical("hidden_dim", [256, 384, 512, 768])
-    n_heads = trial.suggest_categorical("n_heads", [4, 8])
-    n_layers = trial.suggest_int("n_layers", 1, 4)
-    cls_w = trial.suggest_float("cls_weight", 0.1, 2.0, log=True)
-    batch = trial.suggest_categorical("batch_size", [16, 32, 64, 128])
-
-    # infer dims from first row
-    if mode == "pooled":
-        Ht = len(train_ds[0]["target_embedding"])
-        Hb = len(train_ds[0]["binder_embedding"])
-        collate = collate_pair_pooled
-        model = CrossAttnPooled(Ht, Hb, hidden=hidden, n_heads=n_heads, n_layers=n_layers, dropout=dropout).to(DEVICE)
-        train_loader = DataLoader(train_ds, batch_size=batch, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate)
-        val_loader   = DataLoader(val_ds,   batch_size=batch, shuffle=False, num_workers=4, pin_memory=True, collate_fn=collate)
-        eval_fn = eval_spearman_pooled
-        train_fn = train_one_epoch_pooled
-
-    else:
-        Ht = len(train_ds[0]["target_embedding"][0])
-        Hb = len(train_ds[0]["binder_embedding"][0])
-        collate = collate_pair_unpooled
-        model = CrossAttnUnpooled(Ht, Hb, hidden=hidden, n_heads=n_heads, n_layers=n_layers, dropout=dropout).to(DEVICE)
-        train_loader = DataLoader(train_ds, batch_size=batch, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate)
-        val_loader   = DataLoader(val_ds,   batch_size=batch, shuffle=False, num_workers=4, pin_memory=True, collate_fn=collate)
-        eval_fn = eval_spearman_unpooled
-        train_fn = train_one_epoch_unpooled
-
-    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
-    loss_reg = nn.MSELoss()
-    loss_cls = nn.CrossEntropyLoss()
-
-    best = -1e9
-    bad = 0
-    patience = 10
-
-    for ep in range(1, 61):
-        train_fn(model, train_loader, opt, loss_reg, loss_cls, cls_w=cls_w)
-        rho = eval_fn(model, val_loader)
-
-        trial.report(rho, ep)
-        if trial.should_prune():
-            raise optuna.TrialPruned()
-
-        if rho > best + 1e-6:
-            best = rho
-            bad = 0
-        else:
-            bad += 1
-            if bad >= patience:
-                break
-
-    return float(best)
-
-
-# -----------------------------
-# Run: optuna + refit best
-# -----------------------------
-def run(dataset_path: str, out_dir: str, mode: str, n_trials: int = 50):
-    out_dir = Path(out_dir)
-    out_dir.mkdir(parents=True, exist_ok=True)
-
-    train_ds, val_ds = load_split_paired(dataset_path)
-    print(f"[Data] Train={len(train_ds)} Val={len(val_ds)} | mode={mode}")
-
-    study = optuna.create_study(direction="maximize", pruner=optuna.pruners.MedianPruner())
-    study.optimize(lambda t: objective_crossattn(t, mode, train_ds, val_ds), n_trials=n_trials)
-
-    study.trials_dataframe().to_csv(out_dir / "optuna_trials.csv", index=False)
-    best = study.best_trial
-    best_params = dict(best.params)
-
-    # refit longer
-    lr = float(best_params["lr"])
-    wd = float(best_params["weight_decay"])
-    dropout = float(best_params["dropout"])
-    hidden = int(best_params["hidden_dim"])
-    n_heads = int(best_params["n_heads"])
-    n_layers = int(best_params["n_layers"])
-    cls_w = float(best_params["cls_weight"])
-    batch = int(best_params["batch_size"])
-
-    loss_reg = nn.MSELoss()
-    loss_cls = nn.CrossEntropyLoss()
-
-    if mode == "pooled":
-        Ht = len(train_ds[0]["target_embedding"])
-        Hb = len(train_ds[0]["binder_embedding"])
-        model = CrossAttnPooled(Ht, Hb, hidden=hidden, n_heads=n_heads, n_layers=n_layers, dropout=dropout).to(DEVICE)
-        collate = collate_pair_pooled
-        train_loader = DataLoader(train_ds, batch_size=batch, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate)
-        val_loader   = DataLoader(val_ds,   batch_size=batch, shuffle=False, num_workers=4, pin_memory=True, collate_fn=collate)
-        eval_fn = eval_spearman_pooled
-        train_fn = train_one_epoch_pooled
-    else:
-        Ht = len(train_ds[0]["target_embedding"][0])
-        Hb = len(train_ds[0]["binder_embedding"][0])
-        model = CrossAttnUnpooled(Ht, Hb, hidden=hidden, n_heads=n_heads, n_layers=n_layers, dropout=dropout).to(DEVICE)
-        collate = collate_pair_unpooled
-        train_loader = DataLoader(train_ds, batch_size=batch, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate)
-        val_loader   = DataLoader(val_ds,   batch_size=batch, shuffle=False, num_workers=4, pin_memory=True, collate_fn=collate)
-        eval_fn = eval_spearman_unpooled
-        train_fn = train_one_epoch_unpooled
-
-    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
-
-    best_rho = -1e9
-    bad = 0
-    patience = 20
-    best_state = None
-
-    for ep in range(1, 201):
-        train_fn(model, train_loader, opt, loss_reg, loss_cls, cls_w=cls_w)
-        rho = eval_fn(model, val_loader)
-
-        if rho > best_rho + 1e-6:
-            best_rho = rho
-            bad = 0
-            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
-        else:
-            bad += 1
-            if bad >= patience:
-                break
-
-    if best_state is not None:
-        model.load_state_dict(best_state)
-
-    # save
-    torch.save({"mode": mode, "best_params": best_params, "state_dict": model.state_dict()}, out_dir / "best_model.pt")
-    with open(out_dir / "best_params.json", "w") as f:
-        json.dump(best_params, f, indent=2)
-
-    print(f"[DONE] {out_dir} | best_optuna_rho={study.best_value:.4f} | refit_best_rho={best_rho:.4f}")
-
-
-if __name__ == "__main__":
-    import argparse
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--dataset_path", type=str, required=True, help="Paired DatasetDict path (pair_*)")
-    ap.add_argument("--mode", type=str, choices=["pooled", "unpooled"], required=True)
-    ap.add_argument("--out_dir", type=str, required=True)
-    ap.add_argument("--n_trials", type=int, default=50)
-    args = ap.parse_args()
-
-    run(
-        dataset_path=args.dataset_path,
-        out_dir=args.out_dir,
-        mode=args.mode,
-        n_trials=args.n_trials,
-    )

training_classifiers/.ipynb_checkpoints/binding_wt-checkpoint.bash

@@ -1,31 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=b-data
-#SBATCH --partition=dgx-b200
-#SBATCH --gpus=1
-#SBATCH --cpus-per-task=10
-#SBATCH --mem=100G
-#SBATCH --time=48:00:00
-#SBATCH --output=%x_%j.out
-
-HOME_LOC=/vast/projects/pranam/lab/yz927
-SCRIPT_LOC=$HOME_LOC/projects/Classifier_Weight/training_classifiers
-DATA_LOC=$HOME_LOC/projects/Classifier_Weight/training_data_cleaned
-OBJECTIVE='binding_affinity'
-WT='smiles' #wt/smiles
-STATUS='pooled' #pooled/unpooled
-DATA_FILE="pair_wt_${WT}_${STATUS}"
-LOG_LOC=$SCRIPT_LOC
-DATE=$(date +%m_%d)
-SPECIAL_PREFIX="binding_affinity_data_generation"
-
-# Create log directory if it doesn't exist
-mkdir -p $LOG_LOC
-
-cd $SCRIPT_LOC
-source /vast/projects/pranam/lab/shared/miniconda3/etc/profile.d/conda.sh
-conda activate /vast/projects/pranam/lab/shared/miniconda3/envs/metal
-
-python -u binding_affinity_split.py > "${LOG_LOC}/${DATE}_${SPECIAL_PREFIX}.log" 2>&1
-
-echo "Script completed at $(date)"
-conda deactivate

training_classifiers/.ipynb_checkpoints/finetune_boost-checkpoint.py

@@ -1,508 +0,0 @@
-#!/usr/bin/env python3
-# finetune_xgb_halflife_cv_optuna.py
-
-import os
-import json
-import math
-import hashlib
-from dataclasses import dataclass
-from typing import Dict, Any, Optional, Tuple, List
-
-import numpy as np
-import pandas as pd
-import optuna
-
-from sklearn.model_selection import KFold
-from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
-from scipy.stats import spearmanr
-
-import torch
-from transformers import AutoTokenizer, AutoModel
-
-import xgboost as xgb
-
-
-# -----------------------------
-# Repro
-# -----------------------------
-SEED = 1986
-np.random.seed(SEED)
-torch.manual_seed(SEED)
-
-
-# -----------------------------
-# Metrics (mirrors your stability script style)
-# -----------------------------
-def safe_spearmanr(y_true: np.ndarray, y_pred: np.ndarray) -> float:
-    rho = spearmanr(y_true, y_pred).correlation
-    if rho is None or np.isnan(rho):
-        return 0.0
-    return float(rho)
-
-def eval_regression(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
-    rmse = float(np.sqrt(mean_squared_error(y_true, y_pred)))
-    mae = float(mean_absolute_error(y_true, y_pred))
-    r2 = float(r2_score(y_true, y_pred))
-    rho = float(safe_spearmanr(y_true, y_pred))
-    return {"rmse": rmse, "mae": mae, "r2": r2, "spearman_rho": rho}
-
-
-# -----------------------------
-# ESM-2 embeddings (cached)
-# -----------------------------
-@dataclass
-class ESMEmbedderConfig:
-    model_name: str = "facebook/esm2_t33_650M_UR50D"
-    batch_size: int = 8
-    max_length: int = 1024   # truncate very long proteins
-    fp16: bool = True
-
-class ESM2Embedder:
-    """
-    Mean-pooled last hidden state (excluding special tokens) -> (H,) per sequence.
-    """
-    def __init__(self, cfg: ESMEmbedderConfig, device: str = "cuda"):
-        self.cfg = cfg
-        self.device = device if (device == "cuda" and torch.cuda.is_available()) else "cpu"
-        self.tokenizer = AutoTokenizer.from_pretrained(cfg.model_name, do_lower_case=False)
-        self.model = AutoModel.from_pretrained(cfg.model_name)
-        self.model.eval()
-        self.model.to(self.device)
-
-        # Turn off gradients
-        for p in self.model.parameters():
-            p.requires_grad = False
-
-    @torch.inference_mode()
-    def embed(self, seqs: List[str]) -> np.ndarray:
-        out = []
-        bs = self.cfg.batch_size
-
-        use_amp = (self.cfg.fp16 and self.device == "cuda")
-        autocast = torch.cuda.amp.autocast if use_amp else torch.cpu.amp.autocast  # safe fallback
-
-        for i in range(0, len(seqs), bs):
-            batch = [s.strip().upper() for s in seqs[i:i+bs]]
-            toks = self.tokenizer(
-                batch,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=self.cfg.max_length,
-                add_special_tokens=True,
-            )
-            toks = {k: v.to(self.device) for k, v in toks.items()}
-            attn = toks["attention_mask"]  # (B, L)
-
-            with autocast(enabled=use_amp):
-                h = self.model(**toks).last_hidden_state  # (B, L, H)
-
-            # mask out special tokens: first token is <cls>; last non-pad token is usually <eos>
-            mask = attn.clone()
-            mask[:, 0] = 0
-            lengths = attn.sum(dim=1)  # includes special tokens
-            # zero out last real token position per sequence
-            eos_pos = (lengths - 1).clamp(min=0)
-            mask[torch.arange(mask.size(0), device=mask.device), eos_pos] = 0
-
-            denom = mask.sum(dim=1).clamp(min=1).unsqueeze(-1)  # (B,1)
-            pooled = (h * mask.unsqueeze(-1)).sum(dim=1) / denom  # (B,H)
-            out.append(pooled.float().detach().cpu().numpy())
-
-        return np.concatenate(out, axis=0).astype(np.float32)
-
-
-def dataset_fingerprint(seqs: List[str], y: np.ndarray, extra: str = "") -> str:
-    h = hashlib.sha256()
-    for s in seqs:
-        h.update(s.encode("utf-8"))
-        h.update(b"\n")
-    h.update(np.asarray(y, dtype=np.float32).tobytes())
-    h.update(extra.encode("utf-8"))
-    return h.hexdigest()[:16]
-
-
-def load_or_compute_embeddings(
-    df: pd.DataFrame,
-    out_dir: str,
-    embed_cfg: ESMEmbedderConfig,
-    device: str,
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    os.makedirs(out_dir, exist_ok=True)
-
-    seqs = df["sequence"].astype(str).tolist()
-    y = df["half_life_hours"].astype(float).to_numpy(dtype=np.float32)
-
-    fp = dataset_fingerprint(seqs, y, extra=f"{embed_cfg.model_name}|{embed_cfg.max_length}")
-    emb_path = os.path.join(out_dir, f"esm2_embeddings_{fp}.npy")
-    meta_path = os.path.join(out_dir, f"esm2_embeddings_{fp}.json")
-
-    if os.path.exists(emb_path) and os.path.exists(meta_path):
-        X = np.load(emb_path).astype(np.float32)
-        return X, y, np.asarray(seqs)
-
-    embedder = ESM2Embedder(embed_cfg, device=device)
-    X = embedder.embed(seqs)  # (N,H)
-
-    np.save(emb_path, X)
-    with open(meta_path, "w") as f:
-        json.dump(
-            {
-                "fingerprint": fp,
-                "model_name": embed_cfg.model_name,
-                "max_length": embed_cfg.max_length,
-                "n": len(seqs),
-                "dim": int(X.shape[1]),
-            },
-            f,
-            indent=2,
-        )
-    return X, y, np.asarray(seqs)
-
-
-# -----------------------------
-# XGBoost training (supports "finetune" via xgb_model)
-# -----------------------------
-def train_xgb_reg(
-    X_train: np.ndarray,
-    y_train: np.ndarray,
-    X_val: np.ndarray,
-    y_val: np.ndarray,
-    params: Dict[str, Any],
-    base_model_json: Optional[str] = None,
-) -> Tuple[xgb.Booster, np.ndarray, np.ndarray, int]:
-    dtrain = xgb.DMatrix(X_train, label=y_train)
-    dval = xgb.DMatrix(X_val, label=y_val)
-
-    num_boost_round = int(params.pop("num_boost_round"))
-    early_stopping_rounds = int(params.pop("early_stopping_rounds"))
-
-    # Important: load a fresh base model each fold (avoid leakage)
-    xgb_model = None
-    if base_model_json is not None:
-        booster0 = xgb.Booster()
-        booster0.load_model(base_model_json)
-        xgb_model = booster0
-
-    booster = xgb.train(
-        params=params,
-        dtrain=dtrain,
-        num_boost_round=num_boost_round,
-        evals=[(dval, "val")],
-        early_stopping_rounds=early_stopping_rounds,
-        verbose_eval=False,
-        xgb_model=xgb_model,  # <-- "finetune": continue boosting from base model
-    )
-
-    p_train = booster.predict(dtrain)
-    p_val = booster.predict(dval)
-    best_iter = int(getattr(booster, "best_iteration", num_boost_round - 1))
-    return booster, p_train, p_val, best_iter
-
-
-# -----------------------------
-# Optuna objective: 5-fold mean Spearman rho
-# -----------------------------
-def make_cv_objective(
-    X: np.ndarray,
-    y: np.ndarray,
-    n_splits: int,
-    device: str,
-    base_model_json: Optional[str],
-    target_transform: str,
-):
-    kf = KFold(n_splits=n_splits, shuffle=True, random_state=SEED)
-
-    # Optional target transform (sometimes helps with heavy-tailed half-life)
-    if target_transform == "log1p":
-        y_used = np.log1p(np.clip(y, a_min=0.0, a_max=None)).astype(np.float32)
-    elif target_transform == "none":
-        y_used = y.astype(np.float32)
-    else:
-        raise ValueError(f"Unknown target_transform: {target_transform}")
-
-    def objective(trial: optuna.Trial) -> float:
-        # Hyperparam ranges patterned after your stability script :contentReference[oaicite:1]{index=1}
-        params = {
-            "objective": "reg:squarederror",
-            "eval_metric": "rmse",
-
-            "lambda": trial.suggest_float("lambda", 1e-10, 100.0, log=True),
-            "alpha":  trial.suggest_float("alpha",  1e-10, 100.0, log=True),
-            "gamma":  trial.suggest_float("gamma",  0.0, 10.0),
-
-            "max_depth": trial.suggest_int("max_depth", 2, 12),
-            "min_child_weight": trial.suggest_float("min_child_weight", 1e-3, 200.0, log=True),
-            "subsample": trial.suggest_float("subsample", 0.5, 1.0),
-            "colsample_bytree": trial.suggest_float("colsample_bytree", 0.3, 1.0),
-
-            "learning_rate": trial.suggest_float("learning_rate", 1e-3, 0.2, log=True),
-
-            "tree_method": "hist",
-            "device": "cuda" if (device == "cuda" and torch.cuda.is_available()) else "cpu",
-        }
-        params["num_boost_round"] = trial.suggest_int("num_boost_round", 30, 1500)
-        params["early_stopping_rounds"] = trial.suggest_int("early_stopping_rounds", 10, 150)
-
-        fold_metrics = []
-        fold_best_iters = []
-
-        for fold, (tr_idx, va_idx) in enumerate(kf.split(X), start=1):
-            Xtr, ytr = X[tr_idx], y_used[tr_idx]
-            Xva, yva = X[va_idx], y_used[va_idx]
-
-            _, _, p_va, best_iter = train_xgb_reg(
-                Xtr, ytr, Xva, yva, params.copy(),
-                base_model_json=base_model_json,
-            )
-
-            m = eval_regression(yva, p_va)
-            fold_metrics.append(m)
-            fold_best_iters.append(best_iter)
-
-        mean_rho = float(np.mean([m["spearman_rho"] for m in fold_metrics]))
-        mean_rmse = float(np.mean([m["rmse"] for m in fold_metrics]))
-        mean_mae = float(np.mean([m["mae"] for m in fold_metrics]))
-        mean_r2 = float(np.mean([m["r2"] for m in fold_metrics]))
-        mean_best_iter = float(np.mean(fold_best_iters))
-
-        trial.set_user_attr("cv_spearman_rho", mean_rho)
-        trial.set_user_attr("cv_rmse", mean_rmse)
-        trial.set_user_attr("cv_mae", mean_mae)
-        trial.set_user_attr("cv_r2", mean_r2)
-        trial.set_user_attr("cv_mean_best_iter", mean_best_iter)
-
-        # maximize Spearman rho (same as your stability workflow :contentReference[oaicite:2]{index=2})
-        return mean_rho
-
-    return objective
-
-
-def refit_and_save(
-    X: np.ndarray,
-    y: np.ndarray,
-    seqs: np.ndarray,
-    out_dir: str,
-    best_params: Dict[str, Any],
-    n_splits: int,
-    device: str,
-    base_model_json: Optional[str],
-    target_transform: str,
-):
-    os.makedirs(out_dir, exist_ok=True)
-
-    # Transform target consistently
-    if target_transform == "log1p":
-        y_used = np.log1p(np.clip(y, a_min=0.0, a_max=None)).astype(np.float32)
-    else:
-        y_used = y.astype(np.float32)
-
-    kf = KFold(n_splits=n_splits, shuffle=True, random_state=SEED)
-
-    # 1) get OOF preds + average best_iteration
-    oof_pred = np.zeros_like(y_used, dtype=np.float32)
-    best_iters = []
-    fold_rows = []
-
-    for fold, (tr_idx, va_idx) in enumerate(kf.split(X), start=1):
-        Xtr, ytr = X[tr_idx], y_used[tr_idx]
-        Xva, yva = X[va_idx], y_used[va_idx]
-
-        _, _, p_va, best_iter = train_xgb_reg(
-            Xtr, ytr, Xva, yva, best_params.copy(),
-            base_model_json=base_model_json,
-        )
-        oof_pred[va_idx] = p_va.astype(np.float32)
-        best_iters.append(best_iter)
-
-        m = eval_regression(yva, p_va)
-        fold_rows.append({"fold": fold, **m, "best_iter": int(best_iter)})
-
-    fold_df = pd.DataFrame(fold_rows)
-    fold_df.to_csv(os.path.join(out_dir, "cv_fold_metrics.csv"), index=False)
-
-    cv_metrics = eval_regression(y_used, oof_pred)
-    with open(os.path.join(out_dir, "cv_oof_summary.json"), "w") as f:
-        json.dump(cv_metrics, f, indent=2)
-
-    oof_df = pd.DataFrame({
-        "sequence": seqs,
-        "y_true_used": y_used.astype(float),
-        "y_pred_oof": oof_pred.astype(float),
-        "residual": (y_used - oof_pred).astype(float),
-    })
-    oof_df.to_csv(os.path.join(out_dir, "cv_oof_predictions.csv"), index=False)
-
-    mean_best_iter = int(round(float(np.mean(best_iters))))
-    final_rounds = max(mean_best_iter + 1, 10)
-
-    # 2) train final model on ALL data (no early stopping here; use final_rounds)
-    dtrain_all = xgb.DMatrix(X, label=y_used)
-
-    xgb_model = None
-    if base_model_json is not None:
-        booster0 = xgb.Booster()
-        booster0.load_model(base_model_json)
-        xgb_model = booster0
-
-    final_params = best_params.copy()
-    final_params.pop("early_stopping_rounds", None)
-    final_params["device"] = "cuda" if (device == "cuda" and torch.cuda.is_available()) else "cpu"
-
-    booster = xgb.train(
-        params=final_params,
-        dtrain=dtrain_all,
-        num_boost_round=int(final_params.pop("num_boost_round", final_rounds)),
-        evals=[],
-        verbose_eval=False,
-        xgb_model=xgb_model,
-    )
-
-    model_path = os.path.join(out_dir, "best_model_finetuned.json")
-    booster.save_model(model_path)
-
-    with open(os.path.join(out_dir, "final_training_notes.json"), "w") as f:
-        json.dump(
-            {
-                "target_transform": target_transform,
-                "final_rounds_used": int(final_rounds),
-                "cv_oof_metrics_on_used_target": cv_metrics,
-                "model_path": model_path,
-            },
-            f,
-            indent=2,
-        )
-
-    print("=" * 72)
-    print("[Final] CV OOF metrics (on transformed target if enabled):")
-    print(json.dumps(cv_metrics, indent=2))
-    print(f"[Final] Saved finetuned model -> {model_path}")
-    print("=" * 72)
-
-
-def main():
-    import argparse
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--csv_path", type=str, default="/scratch/pranamlab/tong/data/halflife/wt_halflife_merged_dedup.csv")
-    parser.add_argument("--out_dir", type=str, default="/scratch/pranamlab/tong/PeptiVerse/src/halflife/finetune_stability_xgb")
-
-    # If provided, we will "finetune" by continuing boosting from this model
-    parser.add_argument("--base_model_json", type=str, default='/scratch/pranamlab/tong/PeptiVerse/src/stability/xgboost/best_model.json', help="Path to an existing XGBoost .json model to continue training from")
-
-    # ESM embedding config
-    parser.add_argument("--esm_model", type=str, default="facebook/esm2_t33_650M_UR50D")
-    parser.add_argument("--esm_batch_size", type=int, default=8)
-    parser.add_argument("--esm_max_length", type=int, default=1024)
-    parser.add_argument("--no_fp16", action="store_true")
-
-    # Training config
-    parser.add_argument("--n_trials", type=int, default=200)
-    parser.add_argument("--n_splits", type=int, default=5)
-    parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"])
-    parser.add_argument("--target_transform", type=str, default="none", choices=["none", "log1p"])
-
-    args = parser.parse_args()
-    os.makedirs(args.out_dir, exist_ok=True)
-
-    # Load data
-    df = pd.read_csv(args.csv_path)
-    if "sequence" not in df.columns or "half_life_hours" not in df.columns:
-        raise ValueError("CSV must contain columns: sequence, half_life_hours")
-
-    df = df.dropna(subset=["sequence", "half_life_hours"]).copy()
-    df["sequence"] = df["sequence"].astype(str).str.strip()
-    df = df[df["sequence"].str.len() > 0]
-    df = df.drop_duplicates(subset=["sequence"], keep="first").reset_index(drop=True)
-
-    print(f"[Data] N={len(df)} from {args.csv_path}")
-
-    # Embeddings (cached)
-    embed_cfg = ESMEmbedderConfig(
-        model_name=args.esm_model,
-        batch_size=args.esm_batch_size,
-        max_length=args.esm_max_length,
-        fp16=(not args.no_fp16),
-    )
-    X, y, seqs = load_or_compute_embeddings(df, args.out_dir, embed_cfg, device=args.device)
-    print(f"[Embeddings] X={X.shape} (float32)")
-
-    # Optuna study
-    sampler = optuna.samplers.TPESampler(seed=SEED)
-    study = optuna.create_study(
-        direction="maximize",  # like your stability script :contentReference[oaicite:3]{index=3}
-        sampler=sampler,
-        pruner=optuna.pruners.MedianPruner(),
-    )
-
-    objective = make_cv_objective(
-        X=X,
-        y=y,
-        n_splits=args.n_splits,
-        device=args.device,
-        base_model_json=args.base_model_json,
-        target_transform=args.target_transform,
-    )
-    study.optimize(objective, n_trials=args.n_trials)
-
-    # Save trials
-    trials_df = study.trials_dataframe()
-    trials_df.to_csv(os.path.join(args.out_dir, "study_trials.csv"), index=False)
-
-    best = study.best_trial
-    best_params = dict(best.params)
-
-    # Build full param dict for refit
-    best_xgb_params = {
-        "objective": "reg:squarederror",
-        "eval_metric": "rmse",
-        "lambda": best_params["lambda"],
-        "alpha": best_params["alpha"],
-        "gamma": best_params["gamma"],
-        "max_depth": best_params["max_depth"],
-        "min_child_weight": best_params["min_child_weight"],
-        "subsample": best_params["subsample"],
-        "colsample_bytree": best_params["colsample_bytree"],
-        "learning_rate": best_params["learning_rate"],
-        "tree_method": "hist",
-        "device": "cuda" if (args.device == "cuda" and torch.cuda.is_available()) else "cpu",
-        "num_boost_round": best_params["num_boost_round"],
-        "early_stopping_rounds": best_params["early_stopping_rounds"],
-    }
-
-    # Summary
-    summary = {
-        "best_trial_number": int(best.number),
-        "best_value_cv_spearman_rho": float(best.value),
-        "best_user_attrs": best.user_attrs,
-        "best_params": best_params,
-        "best_xgb_params_full": best_xgb_params,
-        "base_model_json": args.base_model_json,
-        "target_transform": args.target_transform,
-        "esm_model": args.esm_model,
-        "esm_max_length": args.esm_max_length,
-    }
-    with open(os.path.join(args.out_dir, "optimization_summary.json"), "w") as f:
-        json.dump(summary, f, indent=2)
-
-    print("=" * 72)
-    print("[Optuna] Best CV Spearman rho:", float(best.value))
-    print("[Optuna] Best params:\n", json.dumps(best_params, indent=2))
-    print("=" * 72)
-
-    # Refit + save final finetuned model + OOF predictions
-    refit_and_save(
-        X=X,
-        y=y,
-        seqs=seqs,
-        out_dir=args.out_dir,
-        best_params=best_xgb_params,
-        n_splits=args.n_splits,
-        device=args.device,
-        base_model_json=args.base_model_json,
-        target_transform=args.target_transform,
-    )
-
-
-if __name__ == "__main__":
-    main()

training_classifiers/.ipynb_checkpoints/generate_binding_val-checkpoint.py

@@ -1,309 +0,0 @@
-#!/usr/bin/env python3
-# export_val_preds_csv.py
-
-import argparse
-from pathlib import Path
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-from datasets import load_from_disk, DatasetDict
-
-# -----------------------------
-# Repro / device
-# -----------------------------
-def seed_all(seed=1986):
-    import random
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-seed_all(1986)
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-
-# -----------------------------
-# Load paired DatasetDict
-# -----------------------------
-def load_split_paired(path: str):
-    dd = load_from_disk(path)
-    if not isinstance(dd, DatasetDict):
-        raise ValueError(f"Expected DatasetDict at {path}")
-    if "train" not in dd or "val" not in dd:
-        raise ValueError(f"DatasetDict missing train/val at {path}")
-    return dd["train"], dd["val"]
-
-
-# -----------------------------
-# Collate fns (same as yours)
-# -----------------------------
-def collate_pair_pooled(batch):
-    Pt = torch.tensor([x["target_embedding"] for x in batch], dtype=torch.float32)
-    Pb = torch.tensor([x["binder_embedding"] for x in batch], dtype=torch.float32)
-    y  = torch.tensor([float(x["label"]) for x in batch], dtype=torch.float32)
-    return Pt, Pb, y
-
-def collate_pair_unpooled(batch):
-    B = len(batch)
-    Ht = len(batch[0]["target_embedding"][0])
-    Hb = len(batch[0]["binder_embedding"][0])
-    Lt_max = max(int(x["target_length"]) for x in batch)
-    Lb_max = max(int(x["binder_length"]) for x in batch)
-
-    Pt = torch.zeros(B, Lt_max, Ht, dtype=torch.float32)
-    Pb = torch.zeros(B, Lb_max, Hb, dtype=torch.float32)
-    Mt = torch.zeros(B, Lt_max, dtype=torch.bool)
-    Mb = torch.zeros(B, Lb_max, dtype=torch.bool)
-    y  = torch.tensor([float(x["label"]) for x in batch], dtype=torch.float32)
-
-    for i, x in enumerate(batch):
-        t = torch.tensor(x["target_embedding"], dtype=torch.float32)
-        b = torch.tensor(x["binder_embedding"], dtype=torch.float32)
-        lt, lb = t.shape[0], b.shape[0]
-        Pt[i, :lt] = t
-        Pb[i, :lb] = b
-        Mt[i, :lt] = torch.tensor(x["target_attention_mask"][:lt], dtype=torch.bool)
-        Mb[i, :lb] = torch.tensor(x["binder_attention_mask"][:lb], dtype=torch.bool)
-
-    return Pt, Mt, Pb, Mb, y
-
-
-# -----------------------------
-# Models (same as yours)
-# -----------------------------
-class CrossAttnPooled(nn.Module):
-    def __init__(self, Ht, Hb, hidden=512, n_heads=8, n_layers=3, dropout=0.1):
-        super().__init__()
-        self.t_proj = nn.Sequential(nn.Linear(Ht, hidden), nn.LayerNorm(hidden))
-        self.b_proj = nn.Sequential(nn.Linear(Hb, hidden), nn.LayerNorm(hidden))
-
-        self.layers = nn.ModuleList([])
-        for _ in range(n_layers):
-            self.layers.append(nn.ModuleDict({
-                "attn_tb": nn.MultiheadAttention(hidden, n_heads, dropout=dropout, batch_first=False),
-                "attn_bt": nn.MultiheadAttention(hidden, n_heads, dropout=dropout, batch_first=False),
-                "n1t": nn.LayerNorm(hidden),
-                "n2t": nn.LayerNorm(hidden),
-                "n1b": nn.LayerNorm(hidden),
-                "n2b": nn.LayerNorm(hidden),
-                "fft": nn.Sequential(nn.Linear(hidden, 4*hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(4*hidden, hidden)),
-                "ffb": nn.Sequential(nn.Linear(hidden, 4*hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(4*hidden, hidden)),
-            }))
-
-        self.shared = nn.Sequential(nn.Linear(2*hidden, hidden), nn.GELU(), nn.Dropout(dropout))
-        self.reg = nn.Linear(hidden, 1)
-        self.cls = nn.Linear(hidden, 3)
-
-    def forward(self, t_vec, b_vec):
-        t = self.t_proj(t_vec).unsqueeze(0)  # (1,B,H)
-        b = self.b_proj(b_vec).unsqueeze(0)  # (1,B,H)
-
-        for L in self.layers:
-            t_attn, _ = L["attn_tb"](t, b, b)
-            t = L["n1t"]((t + t_attn).transpose(0,1)).transpose(0,1)
-            t = L["n2t"]((t + L["fft"](t)).transpose(0,1)).transpose(0,1)
-
-            b_attn, _ = L["attn_bt"](b, t, t)
-            b = L["n1b"]((b + b_attn).transpose(0,1)).transpose(0,1)
-            b = L["n2b"]((b + L["ffb"](b)).transpose(0,1)).transpose(0,1)
-
-        z = torch.cat([t[0], b[0]], dim=-1)
-        h = self.shared(z)
-        return self.reg(h).squeeze(-1), self.cls(h)
-
-
-class CrossAttnUnpooled(nn.Module):
-    def __init__(self, Ht, Hb, hidden=512, n_heads=8, n_layers=3, dropout=0.1):
-        super().__init__()
-        self.t_proj = nn.Sequential(nn.Linear(Ht, hidden), nn.LayerNorm(hidden))
-        self.b_proj = nn.Sequential(nn.Linear(Hb, hidden), nn.LayerNorm(hidden))
-
-        self.layers = nn.ModuleList([])
-        for _ in range(n_layers):
-            self.layers.append(nn.ModuleDict({
-                "attn_tb": nn.MultiheadAttention(hidden, n_heads, dropout=dropout, batch_first=True),
-                "attn_bt": nn.MultiheadAttention(hidden, n_heads, dropout=dropout, batch_first=True),
-                "n1t": nn.LayerNorm(hidden),
-                "n2t": nn.LayerNorm(hidden),
-                "n1b": nn.LayerNorm(hidden),
-                "n2b": nn.LayerNorm(hidden),
-                "fft": nn.Sequential(nn.Linear(hidden, 4*hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(4*hidden, hidden)),
-                "ffb": nn.Sequential(nn.Linear(hidden, 4*hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(4*hidden, hidden)),
-            }))
-
-        self.shared = nn.Sequential(nn.Linear(2*hidden, hidden), nn.GELU(), nn.Dropout(dropout))
-        self.reg = nn.Linear(hidden, 1)
-        self.cls = nn.Linear(hidden, 3)
-
-    def masked_mean(self, X, M):
-        Mf = M.unsqueeze(-1).float()
-        denom = Mf.sum(dim=1).clamp(min=1.0)
-        return (X * Mf).sum(dim=1) / denom
-
-    def forward(self, T, Mt, B, Mb):
-        T = self.t_proj(T)
-        Bx = self.b_proj(B)
-
-        kp_t = ~Mt
-        kp_b = ~Mb
-
-        for L in self.layers:
-            T_attn, _ = L["attn_tb"](T, Bx, Bx, key_padding_mask=kp_b)
-            T = L["n1t"](T + T_attn)
-            T = L["n2t"](T + L["fft"](T))
-
-            B_attn, _ = L["attn_bt"](Bx, T, T, key_padding_mask=kp_t)
-            Bx = L["n1b"](Bx + B_attn)
-            Bx = L["n2b"](Bx + L["ffb"](Bx))
-
-        t_pool = self.masked_mean(T, Mt)
-        b_pool = self.masked_mean(Bx, Mb)
-        z = torch.cat([t_pool, b_pool], dim=-1)
-        h = self.shared(z)
-        return self.reg(h).squeeze(-1), self.cls(h)
-
-
-# -----------------------------
-# Helpers
-# -----------------------------
-def softmax_np(logits: np.ndarray) -> np.ndarray:
-    x = logits - logits.max(axis=1, keepdims=True)
-    ex = np.exp(x)
-    return ex / ex.sum(axis=1, keepdims=True)
-
-def expected_score_from_probs(probs: np.ndarray, class_centers=(9.5, 8.0, 6.0)) -> np.ndarray:
-    centers = np.asarray(class_centers, dtype=np.float32)[None, :]  # (1,3)
-    return (probs * centers).sum(axis=1)
-
-def load_checkpoint(ckpt_path: str, mode: str, train_ds):
-    ckpt = torch.load(ckpt_path, map_location="cpu")
-    params = ckpt.get("best_params", {})
-
-    hidden = int(params.get("hidden_dim", 512))
-    n_heads = int(params.get("n_heads", 8))
-    n_layers = int(params.get("n_layers", 3))
-    dropout = float(params.get("dropout", 0.1))
-
-    if mode == "pooled":
-        Ht = len(train_ds[0]["target_embedding"])
-        Hb = len(train_ds[0]["binder_embedding"])
-        model = CrossAttnPooled(Ht, Hb, hidden=hidden, n_heads=n_heads, n_layers=n_layers, dropout=dropout)
-    else:
-        Ht = len(train_ds[0]["target_embedding"][0])
-        Hb = len(train_ds[0]["binder_embedding"][0])
-        model = CrossAttnUnpooled(Ht, Hb, hidden=hidden, n_heads=n_heads, n_layers=n_layers, dropout=dropout)
-
-    model.load_state_dict(ckpt["state_dict"], strict=True)
-    model.to(DEVICE).eval()
-    return model
-
-
-@torch.no_grad()
-def export_val_preds_csv(dataset_path: str, ckpt_path: str, mode: str,
-                         out_csv: str, batch_size: int, num_workers: int,
-                         class_centers=(9.5, 8.0, 6.0)):
-    train_ds, val_ds = load_split_paired(dataset_path)
-    model = load_checkpoint(ckpt_path, mode, train_ds)
-
-    if mode == "pooled":
-        loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,
-                            num_workers=num_workers, pin_memory=True,
-                            collate_fn=collate_pair_pooled)
-        y_all, pred_reg_all, logits_all = [], [], []
-        for t, b, y in loader:
-            t = t.to(DEVICE, non_blocking=True)
-            b = b.to(DEVICE, non_blocking=True)
-            pred_reg, logits = model(t, b)
-            y_all.append(y.numpy())
-            pred_reg_all.append(pred_reg.detach().cpu().numpy())
-            logits_all.append(logits.detach().cpu().numpy())
-
-    else:
-        loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,
-                            num_workers=num_workers, pin_memory=True,
-                            collate_fn=collate_pair_unpooled)
-        y_all, pred_reg_all, logits_all = [], [], []
-        for T, Mt, B, Mb, y in loader:
-            T = T.to(DEVICE, non_blocking=True)
-            Mt = Mt.to(DEVICE, non_blocking=True)
-            B = B.to(DEVICE, non_blocking=True)
-            Mb = Mb.to(DEVICE, non_blocking=True)
-            pred_reg, logits = model(T, Mt, B, Mb)
-            y_all.append(y.numpy())
-            pred_reg_all.append(pred_reg.detach().cpu().numpy())
-            logits_all.append(logits.detach().cpu().numpy())
-
-    y_true = np.concatenate(y_all)
-    y_pred_reg = np.concatenate(pred_reg_all)
-    logits = np.concatenate(logits_all)
-
-    probs = softmax_np(logits)  # (N,3)
-    y_pred_cls_score = expected_score_from_probs(probs, class_centers=class_centers)
-
-    # Build CSV rows
-    out = Path(out_csv)
-    out.parent.mkdir(parents=True, exist_ok=True)
-
-    header = [
-        "split", "mode",
-        "y_true",
-        "y_pred_reg",
-        "p_high", "p_moderate", "p_low",
-        "y_pred_cls_score",
-        "center_high", "center_moderate", "center_low",
-    ]
-
-    centers = list(class_centers)
-    rows = np.column_stack([
-        y_true,
-        y_pred_reg,
-        probs[:, 0], probs[:, 1], probs[:, 2],
-        y_pred_cls_score,
-        np.full_like(y_true, centers[0], dtype=np.float32),
-        np.full_like(y_true, centers[1], dtype=np.float32),
-        np.full_like(y_true, centers[2], dtype=np.float32),
-    ])
-
-    with out.open("w") as f:
-        f.write(",".join(header) + "\n")
-        for i in range(rows.shape[0]):
-            f.write(
-                "val," + mode + "," +
-                ",".join(f"{rows[i, j]:.8f}" for j in range(rows.shape[1])) +
-                "\n"
-            )
-
-    print(f"[Data] Val N={len(y_true)} | mode={mode}")
-    print(f"[Saved] {out}")
-
-
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--dataset_path", required=True, help="Paired DatasetDict path (pair_*)")
-    ap.add_argument("--ckpt", required=True, help="Path to best_model.pt")
-    ap.add_argument("--mode", choices=["pooled", "unpooled"], required=True)
-    ap.add_argument("--out_csv", required=True)
-    ap.add_argument("--batch_size", type=int, default=128)
-    ap.add_argument("--num_workers", type=int, default=4)
-
-    # Optional: choose class-centers for expected-score conversion
-    ap.add_argument("--center_high", type=float, default=9.5)
-    ap.add_argument("--center_moderate", type=float, default=8.0)
-    ap.add_argument("--center_low", type=float, default=6.0)
-
-    args = ap.parse_args()
-
-    export_val_preds_csv(
-        dataset_path=args.dataset_path,
-        ckpt_path=args.ckpt,
-        mode=args.mode,
-        out_csv=args.out_csv,
-        batch_size=args.batch_size,
-        num_workers=args.num_workers,
-        class_centers=(args.center_high, args.center_moderate, args.center_low),
-    )
-
-
-if __name__ == "__main__":
-    main()

training_classifiers/.ipynb_checkpoints/peptiverse_filelist-checkpoint.txt

@@ -1,234 +0,0 @@
-./hemolysis/cnn_smiles/optimization_summary.txt
-./hemolysis/cnn_smiles/pr_curve.png
-./hemolysis/cnn_smiles/roc_curve.png
-./hemolysis/cnn_smiles/study_trials.csv
-./hemolysis/cnn_smiles/train_predictions.csv
-./hemolysis/cnn_smiles/val_predictions.csv
-./hemolysis/cnn_wt/optimization_summary.txt
-./hemolysis/cnn_wt/pr_curve.png
-./hemolysis/cnn_wt/roc_curve.png
-./hemolysis/cnn_wt/study_trials.csv
-./hemolysis/cnn_wt/train_predictions.csv
-./hemolysis/cnn_wt/val_predictions.csv
-./hemolysis/enet_gpu/optimization_summary.txt
-./hemolysis/enet_gpu/pr_curve.png
-./hemolysis/enet_gpu/roc_curve.png
-./hemolysis/enet_gpu/study_trials.csv
-./hemolysis/enet_gpu/train_predictions.csv
-./hemolysis/enet_gpu/val_predictions.csv
-./hemolysis/enet_gpu_smiles/optimization_summary.txt
-./hemolysis/enet_gpu_smiles/pr_curve.png
-./hemolysis/enet_gpu_smiles/roc_curve.png
-./hemolysis/enet_gpu_smiles/study_trials.csv
-./hemolysis/enet_gpu_smiles/train_predictions.csv
-./hemolysis/enet_gpu_smiles/val_predictions.csv
-./hemolysis/enet_gpu_wt/optimization_summary.txt
-./hemolysis/enet_gpu_wt/pr_curve.png
-./hemolysis/enet_gpu_wt/roc_curve.png
-./hemolysis/enet_gpu_wt/study_trials.csv
-./hemolysis/enet_gpu_wt/train_predictions.csv
-./hemolysis/enet_gpu_wt/val_predictions.csv
-./hemolysis/mlp_smiles/optimization_summary.txt
-./hemolysis/mlp_smiles/pr_curve.png
-./hemolysis/mlp_smiles/roc_curve.png
-./hemolysis/mlp_smiles/study_trials.csv
-./hemolysis/mlp_smiles/train_predictions.csv
-./hemolysis/mlp_smiles/val_predictions.csv
-./hemolysis/mlp_wt/optimization_summary.txt
-./hemolysis/mlp_wt/pr_curve.png
-./hemolysis/mlp_wt/roc_curve.png
-./hemolysis/mlp_wt/study_trials.csv
-./hemolysis/mlp_wt/train_predictions.csv
-./hemolysis/mlp_wt/val_predictions.csv
-./hemolysis/svm_gpu_wt/optimization_summary.txt
-./hemolysis/svm_gpu_wt/pr_curve.png
-./hemolysis/svm_gpu_wt/roc_curve.png
-./hemolysis/svm_gpu_wt/study_trials.csv
-./hemolysis/svm_gpu_wt/train_predictions.csv
-./hemolysis/svm_gpu_wt/val_predictions.csv
-./hemolysis/transformer_smiles/optimization_summary.txt
-./hemolysis/transformer_smiles/pr_curve.png
-./hemolysis/transformer_smiles/roc_curve.png
-./hemolysis/transformer_smiles/study_trials.csv
-./hemolysis/transformer_smiles/train_predictions.csv
-./hemolysis/transformer_smiles/val_predictions.csv
-./hemolysis/transformer_wt/optimization_summary.txt
-./hemolysis/transformer_wt/pr_curve.png
-./hemolysis/transformer_wt/roc_curve.png
-./hemolysis/transformer_wt/study_trials.csv
-./hemolysis/transformer_wt/train_predictions.csv
-./hemolysis/transformer_wt/val_predictions.csv
-./hemolysis/xgb/optimization_summary.txt
-./hemolysis/xgb/pr_curve.png
-./hemolysis/xgb/roc_curve.png
-./hemolysis/xgb/study_trials.csv
-./hemolysis/xgb/train_predictions.csv
-./hemolysis/xgb/val_predictions.csv
-./hemolysis/xgb_smiles/optimization_summary.txt
-./hemolysis/xgb_smiles/pr_curve.png
-./hemolysis/xgb_smiles/roc_curve.png
-./hemolysis/xgb_smiles/study_trials.csv
-./hemolysis/xgb_smiles/train_predictions.csv
-./hemolysis/xgb_smiles/val_predictions.csv
-./hemolysis/xgb_wt/optimization_summary.txt
-./hemolysis/xgb_wt/pr_curve.png
-./hemolysis/xgb_wt/roc_curve.png
-./hemolysis/xgb_wt/study_trials.csv
-./hemolysis/xgb_wt/train_predictions.csv
-./hemolysis/xgb_wt/val_predictions.csv
-./nf/cnn/optimization_summary.txt
-./nf/cnn/pr_curve.png
-./nf/cnn/roc_curve.png
-./nf/cnn/study_trials.csv
-./nf/cnn/train_predictions.csv
-./nf/cnn/val_predictions.csv
-./nf/cnn_wt/optimization_summary.txt
-./nf/cnn_wt/pr_curve.png
-./nf/cnn_wt/roc_curve.png
-./nf/cnn_wt/study_trials.csv
-./nf/cnn_wt/train_predictions.csv
-./nf/cnn_wt/val_predictions.csv
-./nf/enet_gpu/optimization_summary.txt
-./nf/enet_gpu/pr_curve.png
-./nf/enet_gpu/roc_curve.png
-./nf/enet_gpu/study_trials.csv
-./nf/enet_gpu/train_predictions.csv
-./nf/enet_gpu/val_predictions.csv
-./nf/enet_gpu_smiles/optimization_summary.txt
-./nf/enet_gpu_smiles/pr_curve.png
-./nf/enet_gpu_smiles/roc_curve.png
-./nf/enet_gpu_smiles/study_trials.csv
-./nf/enet_gpu_smiles/train_predictions.csv
-./nf/enet_gpu_smiles/val_predictions.csv
-./nf/enet_gpu_wt/optimization_summary.txt
-./nf/enet_gpu_wt/pr_curve.png
-./nf/enet_gpu_wt/roc_curve.png
-./nf/enet_gpu_wt/study_trials.csv
-./nf/enet_gpu_wt/train_predictions.csv
-./nf/enet_gpu_wt/val_predictions.csv
-./nf/mlp/optimization_summary.txt
-./nf/mlp/pr_curve.png
-./nf/mlp/roc_curve.png
-./nf/mlp/study_trials.csv
-./nf/mlp/train_predictions.csv
-./nf/mlp/val_predictions.csv
-./nf/mlp_wt/optimization_summary.txt
-./nf/mlp_wt/pr_curve.png
-./nf/mlp_wt/roc_curve.png
-./nf/mlp_wt/study_trials.csv
-./nf/mlp_wt/train_predictions.csv
-./nf/mlp_wt/val_predictions.csv
-./nf/svm_gpu/optimization_summary.txt
-./nf/svm_gpu/pr_curve.png
-./nf/svm_gpu/roc_curve.png
-./nf/svm_gpu/study_trials.csv
-./nf/svm_gpu/train_predictions.csv
-./nf/svm_gpu/val_predictions.csv
-./nf/svm_gpu_wt/optimization_summary.txt
-./nf/svm_gpu_wt/pr_curve.png
-./nf/svm_gpu_wt/roc_curve.png
-./nf/svm_gpu_wt/study_trials.csv
-./nf/svm_gpu_wt/train_predictions.csv
-./nf/svm_gpu_wt/val_predictions.csv
-./nf/transformer/optimization_summary.txt
-./nf/transformer/pr_curve.png
-./nf/transformer/roc_curve.png
-./nf/transformer/study_trials.csv
-./nf/transformer/train_predictions.csv
-./nf/transformer/val_predictions.csv
-./nf/transformer_wt/optimization_summary.txt
-./nf/transformer_wt/pr_curve.png
-./nf/transformer_wt/roc_curve.png
-./nf/transformer_wt/study_trials.csv
-./nf/transformer_wt/train_predictions.csv
-./nf/transformer_wt/val_predictions.csv
-./nf/xgb_wt/optimization_summary.txt
-./nf/xgb_wt/pr_curve.png
-./nf/xgb_wt/roc_curve.png
-./nf/xgb_wt/study_trials.csv
-./nf/xgb_wt/train_predictions.csv
-./nf/xgb_wt/val_predictions.csv
-./permeability_caco2/cnn_smiles/optimization_summary.txt
-./permeability_caco2/cnn_smiles/study_trials.csv
-./permeability_caco2/cnn_smiles/train_predictions.csv
-./permeability_caco2/cnn_smiles/val_predictions.csv
-./permeability_caco2/enet_gpu_smiles/optimization_summary.txt
-./permeability_caco2/enet_gpu_smiles/study_trials.csv
-./permeability_caco2/enet_gpu_smiles/train_predictions.csv
-./permeability_caco2/enet_gpu_smiles/val_predictions.csv
-./permeability_caco2/mlp_smiles/optimization_summary.txt
-./permeability_caco2/mlp_smiles/study_trials.csv
-./permeability_caco2/mlp_smiles/train_predictions.csv
-./permeability_caco2/mlp_smiles/val_predictions.csv
-./permeability_caco2/svr_smiles/optimization_summary.txt
-./permeability_caco2/svr_smiles/study_trials.csv
-./permeability_caco2/svr_smiles/train_predictions.csv
-./permeability_caco2/svr_smiles/val_predictions.csv
-./permeability_caco2/transformer_smiles/optimization_summary.txt
-./permeability_caco2/transformer_smiles/study_trials.csv
-./permeability_caco2/transformer_smiles/train_predictions.csv
-./permeability_caco2/transformer_smiles/val_predictions.csv
-./permeability_caco2/xgb_reg_smiles/optimization_summary.txt
-./permeability_caco2/xgb_reg_smiles/study_trials.csv
-./permeability_caco2/xgb_reg_smiles/train_predictions.csv
-./permeability_caco2/xgb_reg_smiles/val_predictions.csv
-./permeability_pampa/cnn_smiles/optimization_summary.txt
-./permeability_pampa/cnn_smiles/study_trials.csv
-./permeability_pampa/cnn_smiles/train_predictions.csv
-./permeability_pampa/cnn_smiles/val_predictions.csv
-./permeability_pampa/enet_gpu_smiles/optimization_summary.txt
-./permeability_pampa/enet_gpu_smiles/study_trials.csv
-./permeability_pampa/enet_gpu_smiles/train_predictions.csv
-./permeability_pampa/enet_gpu_smiles/val_predictions.csv
-./permeability_pampa/mlp_smiles/optimization_summary.txt
-./permeability_pampa/mlp_smiles/study_trials.csv
-./permeability_pampa/mlp_smiles/train_predictions.csv
-./permeability_pampa/mlp_smiles/val_predictions.csv
-./permeability_pampa/transformer_smiles/optimization_summary.txt
-./permeability_pampa/transformer_smiles/study_trials.csv
-./permeability_pampa/transformer_smiles/train_predictions.csv
-./permeability_pampa/transformer_smiles/val_predictions.csv
-./permeability_pampa/xgb_reg_smiles/optimization_summary.txt
-./permeability_pampa/xgb_reg_smiles/study_trials.csv
-./permeability_pampa/xgb_reg_smiles/train_predictions.csv
-./permeability_pampa/xgb_reg_smiles/val_predictions.csv
-./solubility/cnn_wt/optimization_summary.txt
-./solubility/cnn_wt/pr_curve.png
-./solubility/cnn_wt/roc_curve.png
-./solubility/cnn_wt/study_trials.csv
-./solubility/cnn_wt/train_predictions.csv
-./solubility/cnn_wt/val_predictions.csv
-./solubility/enet_gpu/optimization_summary.txt
-./solubility/enet_gpu/pr_curve.png
-./solubility/enet_gpu/roc_curve.png
-./solubility/enet_gpu/study_trials.csv
-./solubility/enet_gpu/train_predictions.csv
-./solubility/enet_gpu/val_predictions.csv
-./solubility/mlp_wt/optimization_summary.txt
-./solubility/mlp_wt/pr_curve.png
-./solubility/mlp_wt/roc_curve.png
-./solubility/mlp_wt/study_trials.csv
-./solubility/mlp_wt/train_predictions.csv
-./solubility/mlp_wt/val_predictions.csv
-./solubility/svm_gpu/optimization_summary.txt
-./solubility/svm_gpu/pr_curve.png
-./solubility/svm_gpu/roc_curve.png
-./solubility/svm_gpu/study_trials.csv
-./solubility/svm_gpu/train_predictions.csv
-./solubility/svm_gpu/val_predictions.csv
-./solubility/transformer_wt/optimization_summary.txt
-./solubility/transformer_wt/pr_curve.png
-./solubility/transformer_wt/roc_curve.png
-./solubility/transformer_wt/study_trials.csv
-./solubility/transformer_wt/train_predictions.csv
-./solubility/transformer_wt/val_predictions.csv
-./solubility/xgb/optimization_summary.txt
-./solubility/xgb/pr_curve.png
-./solubility/xgb/roc_curve.png
-./solubility/xgb/study_trials.csv
-./solubility/xgb/train_predictions.csv
-./solubility/xgb/val_predictions.csv
-./binding_affinity/wt_wt_pooled/optuna_trials.csv
-./binding_affinity/wt_smiles_pooled/optuna_trials.csv
-./binding_affinity/wt_smiles_unpooled/optuna_trials.csv
-./binding_affinity/wt_wt_unpooled/optuna_trials.csv

training_classifiers/.ipynb_checkpoints/train_boost-checkpoint.py

@@ -1,417 +0,0 @@
-import os
-import json
-import joblib
-import optuna
-import numpy as np
-import pandas as pd
-import matplotlib.pyplot as plt
-
-from dataclasses import dataclass
-from typing import Dict, Any, Tuple, Optional
-
-from datasets import load_from_disk, DatasetDict
-from sklearn.metrics import (
-    f1_score, roc_auc_score, average_precision_score,
-    precision_recall_curve, roc_curve
-)
-from sklearn.linear_model import LogisticRegression
-from sklearn.ensemble import AdaBoostClassifier
-from sklearn.tree import DecisionTreeClassifier
-from linearboost import LinearBoostClassifier
-
-import xgboost as xgb
-from lightning.pytorch import seed_everything
-
-seed_everything(1986)
-
-# -----------------------------
-# Data loading
-# -----------------------------
-@dataclass
-class SplitData:
-    X_train: np.ndarray
-    y_train: np.ndarray
-    seq_train: Optional[np.ndarray]
-    X_val: np.ndarray
-    y_val: np.ndarray
-    seq_val: Optional[np.ndarray]
-
-
-def _stack_embeddings(col) -> np.ndarray:
-    # HF datasets often store embeddings as list-of-floats per row
-    arr = np.asarray(col, dtype=np.float32)
-    if arr.ndim != 2:
-        arr = np.stack(col).astype(np.float32)
-    return arr
-
-
-def load_split_data(dataset_path: str) -> SplitData:
-    ds = load_from_disk(dataset_path)
-
-    # Case A: DatasetDict with train/val
-    if isinstance(ds, DatasetDict) and "train" in ds and "val" in ds:
-        train_ds, val_ds = ds["train"], ds["val"]
-    else:
-        # Case B: Single dataset with "split" column
-        if "split" not in ds.column_names:
-            raise ValueError(
-                "Dataset must be a DatasetDict(train/val) or have a 'split' column."
-            )
-        train_ds = ds.filter(lambda x: x["split"] == "train")
-        val_ds   = ds.filter(lambda x: x["split"] == "val")
-
-    for required in ["embedding", "label"]:
-        if required not in train_ds.column_names:
-            raise ValueError(f"Missing column '{required}' in train split.")
-        if required not in val_ds.column_names:
-            raise ValueError(f"Missing column '{required}' in val split.")
-
-    X_train = _stack_embeddings(train_ds["embedding"])
-    y_train = np.asarray(train_ds["label"], dtype=np.int64)
-
-    X_val = _stack_embeddings(val_ds["embedding"])
-    y_val = np.asarray(val_ds["label"], dtype=np.int64)
-
-    seq_train = None
-    seq_val = None
-    if "sequence" in train_ds.column_names:
-        seq_train = np.asarray(train_ds["sequence"])
-    if "sequence" in val_ds.column_names:
-        seq_val = np.asarray(val_ds["sequence"])
-
-    return SplitData(X_train, y_train, seq_train, X_val, y_val, seq_val)
-
-
-# -----------------------------
-# Metrics + thresholding
-# -----------------------------
-def best_f1_threshold(y_true: np.ndarray, y_prob: np.ndarray) -> Tuple[float, float]:
-    """
-    Find threshold maximizing F1 on the given set.
-    Returns (best_threshold, best_f1).
-    """
-    precision, recall, thresholds = precision_recall_curve(y_true, y_prob)
-    # precision_recall_curve returns thresholds of length n-1
-    # compute F1 for those thresholds
-    f1s = (2 * precision[:-1] * recall[:-1]) / (precision[:-1] + recall[:-1] + 1e-12)
-    best_idx = int(np.nanargmax(f1s))
-    return float(thresholds[best_idx]), float(f1s[best_idx])
-
-
-def eval_binary(y_true: np.ndarray, y_prob: np.ndarray, threshold: float) -> Dict[str, float]:
-    y_pred = (y_prob >= threshold).astype(int)
-    return {
-        "f1": float(f1_score(y_true, y_pred)),
-        "auc": float(roc_auc_score(y_true, y_prob)),
-        "ap": float(average_precision_score(y_true, y_prob)),
-        "threshold": float(threshold),
-    }
-
-
-# -----------------------------
-# Model factories
-# -----------------------------
-def train_xgb(
-    X_train, y_train, X_val, y_val, params: Dict[str, Any]
-) -> Tuple[xgb.Booster, np.ndarray, np.ndarray]:
-    dtrain = xgb.DMatrix(X_train, label=y_train)
-    dval   = xgb.DMatrix(X_val, label=y_val)
-
-    num_boost_round = int(params.pop("num_boost_round"))
-    early_stopping_rounds = int(params.pop("early_stopping_rounds"))
-
-    booster = xgb.train(
-        params=params,
-        dtrain=dtrain,
-        num_boost_round=num_boost_round,
-        evals=[(dval, "val")],
-        early_stopping_rounds=early_stopping_rounds,
-        verbose_eval=False,
-    )
-
-    p_train = booster.predict(dtrain)
-    p_val   = booster.predict(dval)
-    return booster, p_train, p_val
-
-
-def train_adaboost(
-    X_train, y_train, X_val, y_val, params: Dict[str, Any]
-) -> Tuple[AdaBoostClassifier, np.ndarray, np.ndarray]:
-    base_depth = int(params.pop("base_depth"))
-    clf = AdaBoostClassifier(
-        estimator=DecisionTreeClassifier(max_depth=base_depth),
-        n_estimators=int(params["n_estimators"]),
-        learning_rate=float(params["learning_rate"]),
-        algorithm="SAMME",
-    )
-    clf.fit(X_train, y_train)
-    p_train = clf.predict_proba(X_train)[:, 1]
-    p_val   = clf.predict_proba(X_val)[:, 1]
-    return clf, p_train, p_val
-
-
-def train_linearboost(X_train, y_train, X_val, y_val, params):
-    clf = LinearBoostClassifier(**params)
-    clf.fit(X_train, y_train)
-    p_train = clf.predict_proba(X_train)[:, 1]
-    p_val   = clf.predict_proba(X_val)[:, 1]
-    return clf, p_train, p_val
-
-
-def suggest_linearboost_params(trial):
-    # Core boosting params
-    params = {
-        "n_estimators": trial.suggest_int("n_estimators", 50, 800),
-        "learning_rate": trial.suggest_float("learning_rate", 0.01, 1.0, log=True),
-        "algorithm": trial.suggest_categorical("algorithm", ["SAMME.R", "SAMME"]),
-        # Scaling choices from docs (you can expand this list if you want)
-        "scaler": trial.suggest_categorical(
-            "scaler",
-            ["minmax", "standard", "robust", "quantile-uniform", "quantile-normal", "power"]
-        ),
-        # useful for imbalanced splits
-        "class_weight": trial.suggest_categorical("class_weight", [None, "balanced"]),
-        # kernel trick
-        "kernel": trial.suggest_categorical("kernel", ["linear", "rbf", "poly", "sigmoid"]),
-    }
-
-    # Kernel-specific params (only when relevant)
-    if params["kernel"] in ["rbf", "poly"]:
-        params["gamma"] = trial.suggest_float("gamma", 1e-6, 10.0, log=True)
-    else:
-        params["gamma"] = None  # docs: default treated as 1/n_features for rbf/poly :contentReference[oaicite:5]{index=5}
-
-    if params["kernel"] == "poly":
-        params["degree"] = trial.suggest_int("degree", 2, 6)  # docs default=3 :contentReference[oaicite:6]{index=6}
-        params["coef0"]  = trial.suggest_float("coef0", 0.0, 5.0)  # docs default=1 :contentReference[oaicite:7]{index=7}
-    else:
-        # safe defaults
-        params["degree"] = 3
-        params["coef0"]  = 1.0
-
-    return params
-# -----------------------------
-# Saving artifacts
-# -----------------------------
-def save_predictions_csv(
-    out_dir: str,
-    split_name: str,
-    y_true: np.ndarray,
-    y_prob: np.ndarray,
-    threshold: float,
-    sequences: Optional[np.ndarray] = None,
-):
-    os.makedirs(out_dir, exist_ok=True)
-    df = pd.DataFrame({
-        "y_true": y_true.astype(int),
-        "y_prob": y_prob.astype(float),
-        "y_pred": (y_prob >= threshold).astype(int),
-    })
-    if sequences is not None:
-        df.insert(0, "sequence", sequences)
-    df.to_csv(os.path.join(out_dir, f"{split_name}_predictions.csv"), index=False)
-
-
-def plot_curves(out_dir: str, y_true: np.ndarray, y_prob: np.ndarray):
-    os.makedirs(out_dir, exist_ok=True)
-
-    # PR
-    precision, recall, _ = precision_recall_curve(y_true, y_prob)
-    plt.figure()
-    plt.plot(recall, precision)
-    plt.xlabel("Recall")
-    plt.ylabel("Precision")
-    plt.title("Precision-Recall Curve")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "pr_curve.png"))
-    plt.close()
-
-    # ROC
-    fpr, tpr, _ = roc_curve(y_true, y_prob)
-    plt.figure()
-    plt.plot(fpr, tpr)
-    plt.xlabel("False Positive Rate")
-    plt.ylabel("True Positive Rate")
-    plt.title("ROC Curve")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "roc_curve.png"))
-    plt.close()
-
-
-# -----------------------------
-# Optuna objectives
-# -----------------------------
-def make_objective(model_name: str, data: SplitData, out_dir: str):
-    Xtr, ytr, Xva, yva = data.X_train, data.y_train, data.X_val, data.y_val
-
-    def objective(trial: optuna.Trial) -> float:
-        if model_name == "xgb":
-            params = {
-                "objective": "binary:logistic",
-                "eval_metric": "logloss",
-                "lambda": trial.suggest_float("lambda", 1e-8, 50.0, log=True),
-                "alpha": trial.suggest_float("alpha", 1e-8, 50.0, log=True),
-                "colsample_bytree": trial.suggest_float("colsample_bytree", 0.3, 1.0),
-                "subsample": trial.suggest_float("subsample", 0.5, 1.0),
-                "learning_rate": trial.suggest_float("learning_rate", 1e-3, 0.3, log=True),
-                "max_depth": trial.suggest_int("max_depth", 2, 15),
-                "min_child_weight": trial.suggest_int("min_child_weight", 1, 500),
-                "gamma": trial.suggest_float("gamma", 0.0, 10.0),
-                "tree_method": "hist",
-                "device": "cuda",
-            }
-
-            # Optional GPU: set env CUDA_VISIBLE_DEVICES externally if you want.
-            # If you *know* you want GPU and your xgboost supports it:
-            # params["device"] = "cuda"
-
-            params["num_boost_round"] = trial.suggest_int("num_boost_round", 50, 1500)
-            params["early_stopping_rounds"] = trial.suggest_int("early_stopping_rounds", 20, 200)
-
-            model, p_tr, p_va = train_xgb(Xtr, ytr, Xva, yva, params.copy())
-
-        elif model_name == "adaboost":
-            params = {
-                "n_estimators": trial.suggest_int("n_estimators", 50, 800),
-                "learning_rate": trial.suggest_float("learning_rate", 1e-3, 2.0, log=True),
-                "base_depth": trial.suggest_int("base_depth", 1, 4),
-            }
-            model, p_tr, p_va = train_adaboost(Xtr, ytr, Xva, yva, params)
-
-        elif model_name == "linearboost":
-            params = suggest_linearboost_params(trial)
-            model, p_tr, p_va = train_linearboost(Xtr, ytr, Xva, yva, params)
-        else:
-            raise ValueError(f"Unknown model_name={model_name}")
-
-        # Threshold picked on val for fair comparison across models
-        thr, f1_at_thr = best_f1_threshold(yva, p_va)
-        metrics = eval_binary(yva, p_va, thr)
-
-        # Track best trial artifacts inside the study directory
-        trial.set_user_attr("threshold", thr)
-        trial.set_user_attr("auc", metrics["auc"])
-        trial.set_user_attr("ap", metrics["ap"])
-
-        return f1_at_thr
-
-    return objective
-
-# -----------------------------
-# Main runner
-# -----------------------------
-def run_optuna_and_refit(
-    dataset_path: str,
-    out_dir: str,
-    model_name: str,
-    n_trials: int = 200,
-):
-    os.makedirs(out_dir, exist_ok=True)
-
-    data = load_split_data(dataset_path)
-    print(f"[Data] Train: {data.X_train.shape}, Val: {data.X_val.shape}")
-
-    study = optuna.create_study(direction="maximize", pruner=optuna.pruners.MedianPruner())
-    study.optimize(make_objective(model_name, data, out_dir), n_trials=n_trials)
-
-    # Save trials table
-    trials_df = study.trials_dataframe()
-    trials_df.to_csv(os.path.join(out_dir, "study_trials.csv"), index=False)
-
-    best = study.best_trial
-    best_params = dict(best.params)
-    best_thr = float(best.user_attrs["threshold"])
-    best_auc = float(best.user_attrs["auc"])
-    best_ap  = float(best.user_attrs["ap"])
-    best_f1  = float(best.value)
-
-    # Refit best model on train (same protocol as objective)
-    if model_name == "xgb":
-        # Reconstruct full param dict
-        params = {
-            "objective": "binary:logistic",
-            "eval_metric": "logloss",
-            "lambda": best_params["lambda"],
-            "alpha": best_params["alpha"],
-            "colsample_bytree": best_params["colsample_bytree"],
-            "subsample": best_params["subsample"],
-            "learning_rate": best_params["learning_rate"],
-            "max_depth": best_params["max_depth"],
-            "min_child_weight": best_params["min_child_weight"],
-            "gamma": best_params["gamma"],
-            "tree_method": "hist",
-            "num_boost_round": best_params["num_boost_round"],
-            "early_stopping_rounds": best_params["early_stopping_rounds"],
-        }
-        model, p_tr, p_va = train_xgb(
-            data.X_train, data.y_train, data.X_val, data.y_val, params
-        )
-        model_path = os.path.join(out_dir, "best_model.json")
-        model.save_model(model_path)
-
-    elif model_name == "adaboost":
-        params = best_params
-        model, p_tr, p_va = train_adaboost(
-            data.X_train, data.y_train, data.X_val, data.y_val, params
-        )
-        model_path = os.path.join(out_dir, "best_model.joblib")
-        joblib.dump(model, model_path)
-
-    elif model_name == "linearboost":
-        params = best_params
-    
-        model, p_tr, p_va = train_linearboost(
-            data.X_train, data.y_train, data.X_val, data.y_val, params
-        )
-    
-        model_path = os.path.join(out_dir, "best_model.joblib")
-        joblib.dump(model, model_path)
-    else:
-        raise ValueError(model_name)
-
-    # Save predictions CSVs
-    save_predictions_csv(out_dir, "train", data.y_train, p_tr, best_thr, data.seq_train)
-    save_predictions_csv(out_dir, "val",   data.y_val,   p_va, best_thr, data.seq_val)
-
-    # Plots on val
-    plot_curves(out_dir, data.y_val, p_va)
-
-    # Summary
-    summary = [
-        "=" * 72,
-        f"MODEL: {model_name}",
-        f"Best trial: {best.number}",
-        f"Best F1 (val @ best-threshold): {best_f1:.4f}",
-        f"Val AUC: {best_auc:.4f}",
-        f"Val AP:  {best_ap:.4f}",
-        f"Best threshold (picked on val): {best_thr:.4f}",
-        f"Model saved to: {model_path}",
-        "Best params:",
-        json.dumps(best_params, indent=2),
-        "=" * 72,
-    ]
-    with open(os.path.join(out_dir, "optimization_summary.txt"), "w") as f:
-        f.write("\n".join(summary))
-    print("\n".join(summary))
-
-
-if __name__ == "__main__":
-    # Example usage:
-    # dataset_path = "/vast/projects/pranam/lab/yz927/projects/Classifier_Weight/training_classifiers/data/solubility"
-    # out_dir = "/vast/projects/pranam/lab/yz927/projects/Classifier_Weight/training_classifiers/src/solubility/xgb"
-    # run_optuna_and_refit(dataset_path, out_dir, model_name="xgb", n_trials=200)
-
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--dataset_path", type=str, required=True)
-    parser.add_argument("--out_dir", type=str, required=True)
-    parser.add_argument("--model", type=str, choices=["xgb", "adaboost", "linearboost"], required=True)
-    parser.add_argument("--n_trials", type=int, default=200)
-    args = parser.parse_args()
-
-    run_optuna_and_refit(
-        dataset_path=args.dataset_path,
-        out_dir=args.out_dir,
-        model_name=args.model,
-        n_trials=args.n_trials,
-    )

training_classifiers/.ipynb_checkpoints/train_ml-checkpoint.py

@@ -1,468 +0,0 @@
-import os
-import json
-import joblib
-import optuna
-import numpy as np
-import pandas as pd
-import matplotlib.pyplot as plt
-from dataclasses import dataclass
-from typing import Dict, Any, Tuple, Optional
-from datasets import load_from_disk, DatasetDict
-from sklearn.metrics import (
-    f1_score, roc_auc_score, average_precision_score,
-    precision_recall_curve, roc_curve
-)
-from sklearn.linear_model import LogisticRegression
-from sklearn.svm import SVC, LinearSVC
-from sklearn.calibration import CalibratedClassifierCV
-import torch
-import time
-import xgboost as xgb
-from lightning.pytorch import seed_everything
-import cupy as cp
-from cuml.svm import SVC as cuSVC
-from cuml.linear_model import LogisticRegression as cuLogReg
-seed_everything(1986)
-
-
-def to_gpu(X: np.ndarray):
-    if isinstance(X, cp.ndarray):
-        return X
-    return cp.asarray(X, dtype=cp.float32)
-
-def to_cpu(x):
-    if isinstance(x, cp.ndarray):
-        return cp.asnumpy(x)
-    return np.asarray(x)
-                     
-@dataclass
-class SplitData:
-    X_train: np.ndarray
-    y_train: np.ndarray
-    seq_train: Optional[np.ndarray]
-    X_val: np.ndarray
-    y_val: np.ndarray
-    seq_val: Optional[np.ndarray]
-
-
-def _stack_embeddings(col) -> np.ndarray:
-    arr = np.asarray(col, dtype=np.float32)
-    if arr.ndim != 2:
-        arr = np.stack(col).astype(np.float32)
-    return arr
-
-
-def load_split_data(dataset_path: str) -> SplitData:
-    ds = load_from_disk(dataset_path)
-
-    # Case A: DatasetDict with train/val
-    if isinstance(ds, DatasetDict) and "train" in ds and "val" in ds:
-        train_ds, val_ds = ds["train"], ds["val"]
-    else:
-        # Case B: Single dataset with "split" column
-        if "split" not in ds.column_names:
-            raise ValueError(
-                "Dataset must be a DatasetDict(train/val) or have a 'split' column."
-            )
-        train_ds = ds.filter(lambda x: x["split"] == "train")
-        val_ds   = ds.filter(lambda x: x["split"] == "val")
-
-    for required in ["embedding", "label"]:
-        if required not in train_ds.column_names:
-            raise ValueError(f"Missing column '{required}' in train split.")
-        if required not in val_ds.column_names:
-            raise ValueError(f"Missing column '{required}' in val split.")
-
-    X_train = _stack_embeddings(train_ds["embedding"])
-    y_train = np.asarray(train_ds["label"], dtype=np.int64)
-
-    X_val = _stack_embeddings(val_ds["embedding"])
-    y_val = np.asarray(val_ds["label"], dtype=np.int64)
-
-    seq_train = None
-    seq_val = None
-    if "sequence" in train_ds.column_names:
-        seq_train = np.asarray(train_ds["sequence"])
-    if "sequence" in val_ds.column_names:
-        seq_val = np.asarray(val_ds["sequence"])
-
-    return SplitData(X_train, y_train, seq_train, X_val, y_val, seq_val)
-
-
-def best_f1_threshold(y_true: np.ndarray, y_prob: np.ndarray) -> Tuple[float, float]:
-    """
-    Find threshold maximizing F1 on the given set.
-    Returns (best_threshold, best_f1).
-    """
-    precision, recall, thresholds = precision_recall_curve(y_true, y_prob)
-    f1s = (2 * precision[:-1] * recall[:-1]) / (precision[:-1] + recall[:-1] + 1e-12)
-    best_idx = int(np.nanargmax(f1s))
-    return float(thresholds[best_idx]), float(f1s[best_idx])
-
-
-def eval_binary(y_true: np.ndarray, y_prob: np.ndarray, threshold: float) -> Dict[str, float]:
-    y_pred = (y_prob >= threshold).astype(int)
-    return {
-        "f1": float(f1_score(y_true, y_pred)),
-        "auc": float(roc_auc_score(y_true, y_prob)),
-        "ap": float(average_precision_score(y_true, y_prob)),
-        "threshold": float(threshold),
-    }
-
-
-# -----------------------------
-# Model
-# -----------------------------
-def train_xgb(
-    X_train, y_train, X_val, y_val, params: Dict[str, Any]
-) -> Tuple[xgb.Booster, np.ndarray, np.ndarray]:
-    dtrain = xgb.DMatrix(X_train, label=y_train)
-    dval   = xgb.DMatrix(X_val, label=y_val)
-
-    num_boost_round = int(params.pop("num_boost_round"))
-    early_stopping_rounds = int(params.pop("early_stopping_rounds"))
-
-    booster = xgb.train(
-        params=params,
-        dtrain=dtrain,
-        num_boost_round=num_boost_round,
-        evals=[(dval, "val")],
-        early_stopping_rounds=early_stopping_rounds,
-        verbose_eval=False,
-    )
-
-    p_train = booster.predict(dtrain)
-    p_val   = booster.predict(dval)
-    return booster, p_train, p_val
-
-def train_cuml_svc(X_train, y_train, X_val, y_val, params):
-    Xtr = to_gpu(X_train)
-    Xva = to_gpu(X_val)
-    ytr = to_gpu(y_train).astype(cp.int32)
-
-    clf = cuSVC(
-        C=float(params["C"]),
-        kernel=params["kernel"],
-        gamma=params.get("gamma", "scale"),
-        class_weight=params.get("class_weight", None),
-        probability=bool(params.get("probability", True)),
-        random_state=1986,
-        max_iter=int(params.get("max_iter", 1000)),
-        tol=float(params.get("tol", 1e-4)),
-    )
-
-    clf.fit(Xtr, ytr)
-
-    p_train = to_cpu(clf.predict_proba(Xtr)[:, 1])
-    p_val   = to_cpu(clf.predict_proba(Xva)[:, 1])
-    return clf, p_train, p_val
-
-def train_cuml_elastic_net(X_train, y_train, X_val, y_val, params):
-    Xtr = to_gpu(X_train)
-    Xva = to_gpu(X_val)
-    ytr = to_gpu(y_train).astype(cp.int32)
-
-    clf = cuLogReg(
-        penalty="elasticnet",
-        C=float(params["C"]),
-        l1_ratio=float(params["l1_ratio"]),
-        class_weight=params.get("class_weight", None),
-        max_iter=int(params.get("max_iter", 1000)),
-        tol=float(params.get("tol", 1e-4)),
-        solver="qn",
-        fit_intercept=True,
-    )
-    clf.fit(Xtr, ytr)
-
-    p_train = to_cpu(clf.predict_proba(Xtr)[:, 1])
-    p_val   = to_cpu(clf.predict_proba(Xva)[:, 1])
-    return clf, p_train, p_val
-
-
-def train_svm(X_train, y_train, X_val, y_val, params):
-    """
-    Kernel SVM via SVC. CPU only in sklearn.
-    probability=True enables predict_proba but is slower.
-    """
-    clf = SVC(
-        C=float(params["C"]),
-        kernel=params["kernel"],
-        gamma=params.get("gamma", "scale"),
-        class_weight=params.get("class_weight", None),
-        probability=True,
-        random_state=1986,
-    )
-    clf.fit(X_train, y_train)
-    p_train = clf.predict_proba(X_train)[:, 1]
-    p_val   = clf.predict_proba(X_val)[:, 1]
-    return clf, p_train, p_val
-
-
-def train_linearsvm_calibrated(X_train, y_train, X_val, y_val, params):
-    """
-    Fast linear SVM (LinearSVC) + probability calibration.
-    Usually much faster than SVC on large datasets.
-    """
-    base = LinearSVC(
-        C=float(params["C"]),
-        class_weight=params.get("class_weight", None),
-        max_iter=int(params.get("max_iter", 5000)),
-        random_state=1986,
-    )
-    # calibration to get probabilities for PR/ROC + thresholding
-    clf = CalibratedClassifierCV(base, method="sigmoid", cv=3)
-    clf.fit(X_train, y_train)
-    p_train = clf.predict_proba(X_train)[:, 1]
-    p_val   = clf.predict_proba(X_val)[:, 1]
-    return clf, p_train, p_val
-
-# -----------------------------
-# Saving artifacts
-# -----------------------------
-def save_predictions_csv(
-    out_dir: str,
-    split_name: str,
-    y_true: np.ndarray,
-    y_prob: np.ndarray,
-    threshold: float,
-    sequences: Optional[np.ndarray] = None,
-):
-    os.makedirs(out_dir, exist_ok=True)
-    df = pd.DataFrame({
-        "y_true": y_true.astype(int),
-        "y_prob": y_prob.astype(float),
-        "y_pred": (y_prob >= threshold).astype(int),
-    })
-    if sequences is not None:
-        df.insert(0, "sequence", sequences)
-    df.to_csv(os.path.join(out_dir, f"{split_name}_predictions.csv"), index=False)
-
-
-def plot_curves(out_dir: str, y_true: np.ndarray, y_prob: np.ndarray):
-    os.makedirs(out_dir, exist_ok=True)
-
-    # PR
-    precision, recall, _ = precision_recall_curve(y_true, y_prob)
-    plt.figure()
-    plt.plot(recall, precision)
-    plt.xlabel("Recall")
-    plt.ylabel("Precision")
-    plt.title("Precision-Recall Curve")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "pr_curve.png"))
-    plt.close()
-
-    # ROC
-    fpr, tpr, _ = roc_curve(y_true, y_prob)
-    plt.figure()
-    plt.plot(fpr, tpr)
-    plt.xlabel("False Positive Rate")
-    plt.ylabel("True Positive Rate")
-    plt.title("ROC Curve")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "roc_curve.png"))
-    plt.close()
-
-
-# -----------------------------
-# Optuna objectives
-# -----------------------------
-def make_objective(model_name: str, data: SplitData, out_dir: str):
-    Xtr, ytr, Xva, yva = data.X_train, data.y_train, data.X_val, data.y_val
-
-    def objective(trial: optuna.Trial) -> float:
-        if model_name == "xgb":
-            params = {
-                "objective": "binary:logistic",
-                "eval_metric": "logloss",
-                "lambda": trial.suggest_float("lambda", 1e-8, 50.0, log=True),
-                "alpha": trial.suggest_float("alpha", 1e-8, 50.0, log=True),
-                "colsample_bytree": trial.suggest_float("colsample_bytree", 0.3, 1.0),
-                "subsample": trial.suggest_float("subsample", 0.5, 1.0),
-                "learning_rate": trial.suggest_float("learning_rate", 1e-3, 0.3, log=True),
-                "max_depth": trial.suggest_int("max_depth", 2, 15),
-                "min_child_weight": trial.suggest_int("min_child_weight", 1, 500),
-                "gamma": trial.suggest_float("gamma", 0.0, 10.0),
-                "tree_method": "hist",
-                "device": "cuda",
-            }
-            params["num_boost_round"] = trial.suggest_int("num_boost_round", 50, 1500)
-            params["early_stopping_rounds"] = trial.suggest_int("early_stopping_rounds", 20, 200)
-
-            model, p_tr, p_va = train_xgb(Xtr, ytr, Xva, yva, params.copy())
-
-        elif model_name == "svm":
-            svm_kind = trial.suggest_categorical("svm_kind", ["svc", "linear_calibrated"])
-
-            if svm_kind == "svc":
-                params = {
-                    "C": trial.suggest_float("C", 1e-3, 1e3, log=True),
-                    "kernel": trial.suggest_categorical("kernel", ["rbf", "linear", "poly", "sigmoid"]),
-                    "class_weight": trial.suggest_categorical("class_weight", [None, "balanced"]),
-                }
-                if params["kernel"] in ["rbf", "poly", "sigmoid"]:
-                    params["gamma"] = trial.suggest_float("gamma", 1e-6, 10.0, log=True)
-                else:
-                    params["gamma"] = "scale"
-
-                model, p_tr, p_va = train_svm(Xtr, ytr, Xva, yva, params)
-
-            else:
-                params = {
-                    "C": trial.suggest_float("C", 1e-3, 1e3, log=True),
-                    "class_weight": trial.suggest_categorical("class_weight", [None, "balanced"]),
-                    "max_iter": trial.suggest_int("max_iter", 2000, 20000),
-                }
-                model, p_tr, p_va = train_linearsvm_calibrated(Xtr, ytr, Xva, yva, params)
-        elif model_name == "svm_gpu":
-            params = {
-                "C": trial.suggest_float("C", 1e-3, 1e3, log=True),
-                "kernel": trial.suggest_categorical("kernel", ["rbf", "linear", "poly", "sigmoid"]),
-                "class_weight": trial.suggest_categorical("class_weight", [None, "balanced"]),
-                "probability": True,
-                "max_iter": trial.suggest_int("max_iter", 200, 5000),
-                "tol": trial.suggest_float("tol", 1e-6, 1e-2, log=True),
-            }
-            if params["kernel"] in ["rbf", "poly", "sigmoid"]:
-                params["gamma"] = trial.suggest_float("gamma", 1e-6, 10.0, log=True)
-            else:
-                params["gamma"] = "scale"
-        
-            model, p_tr, p_va = train_cuml_svc(Xtr, ytr, Xva, yva, params)
-        
-        elif model_name == "enet_gpu":
-            params = {
-                "C": trial.suggest_float("C", 1e-4, 1e3, log=True),
-                "l1_ratio": trial.suggest_float("l1_ratio", 0.0, 1.0),
-                "class_weight": trial.suggest_categorical("class_weight", [None, "balanced"]),
-                "max_iter": trial.suggest_int("max_iter", 200, 5000),
-                "tol": trial.suggest_float("tol", 1e-6, 1e-2, log=True),
-            }
-            model, p_tr, p_va = train_cuml_elastic_net(Xtr, ytr, Xva, yva, params)
-        else:
-            raise ValueError(f"Unknown model_name={model_name}")
-
-        thr, f1_at_thr = best_f1_threshold(yva, p_va)
-        metrics = eval_binary(yva, p_va, thr)
-        trial.set_user_attr("threshold", thr)
-        trial.set_user_attr("auc", metrics["auc"])
-        trial.set_user_attr("ap", metrics["ap"])
-        return f1_at_thr
-
-    return objective
-
-# -----------------------------
-# Main
-# -----------------------------
-def run_optuna_and_refit(
-    dataset_path: str,
-    out_dir: str,
-    model_name: str,
-    n_trials: int = 200,
-):
-    os.makedirs(out_dir, exist_ok=True)
-
-    data = load_split_data(dataset_path)
-    print(f"[Data] Train: {data.X_train.shape}, Val: {data.X_val.shape}")
-
-    study = optuna.create_study(direction="maximize", pruner=optuna.pruners.MedianPruner())
-    study.optimize(make_objective(model_name, data, out_dir), n_trials=n_trials)
-
-    trials_df = study.trials_dataframe()
-    trials_df.to_csv(os.path.join(out_dir, "study_trials.csv"), index=False)
-
-    best = study.best_trial
-    best_params = dict(best.params)
-    best_thr = float(best.user_attrs["threshold"])
-    best_auc = float(best.user_attrs["auc"])
-    best_ap  = float(best.user_attrs["ap"])
-    best_f1  = float(best.value)
-
-    # Refit best model on train
-    if model_name == "xgb":
-        params = {
-            "objective": "binary:logistic",
-            "eval_metric": "logloss",
-            "lambda": best_params["lambda"],
-            "alpha": best_params["alpha"],
-            "colsample_bytree": best_params["colsample_bytree"],
-            "subsample": best_params["subsample"],
-            "learning_rate": best_params["learning_rate"],
-            "max_depth": best_params["max_depth"],
-            "min_child_weight": best_params["min_child_weight"],
-            "gamma": best_params["gamma"],
-            "tree_method": "hist",
-            "num_boost_round": best_params["num_boost_round"],
-            "early_stopping_rounds": best_params["early_stopping_rounds"],
-        }
-        model, p_tr, p_va = train_xgb(
-            data.X_train, data.y_train, data.X_val, data.y_val, params
-        )
-        model_path = os.path.join(out_dir, "best_model.json")
-        model.save_model(model_path)
-
-    elif model_name == "svm":
-        svm_kind = best_params["svm_kind"]
-        if svm_kind == "svc":
-            model, p_tr, p_va = train_svm(data.X_train, data.y_train, data.X_val, data.y_val, best_params)
-        else:
-            model, p_tr, p_va = train_linearsvm_calibrated(data.X_train, data.y_train, data.X_val, data.y_val, best_params)
-
-        model_path = os.path.join(out_dir, "best_model.joblib")
-        joblib.dump(model, model_path)
-    elif model_name == "svm_gpu":
-        model, p_tr, p_va = train_cuml_svc(
-            data.X_train, data.y_train, data.X_val, data.y_val, best_params
-        )
-        model_path = os.path.join(out_dir, "best_model_cuml_svc.joblib")
-        joblib.dump(model, model_path)
-
-    elif model_name == "enet_gpu":
-        model, p_tr, p_va = train_cuml_elastic_net(
-            data.X_train, data.y_train, data.X_val, data.y_val, best_params
-        )
-        model_path = os.path.join(out_dir, "best_model_cuml_enet.joblib")
-        joblib.dump(model, model_path)
-    else:
-        raise ValueError(model_name)
-
-    # Save predictions CSVs
-    save_predictions_csv(out_dir, "train", data.y_train, p_tr, best_thr, data.seq_train)
-    save_predictions_csv(out_dir, "val",   data.y_val,   p_va, best_thr, data.seq_val)
-
-    # Plots on val
-    plot_curves(out_dir, data.y_val, p_va)
-
-    summary = [
-        "=" * 72,
-        f"MODEL: {model_name}",
-        f"Best trial: {best.number}",
-        f"Best F1 (val @ best-threshold): {best_f1:.4f}",
-        f"Val AUC: {best_auc:.4f}",
-        f"Val AP:  {best_ap:.4f}",
-        f"Best threshold (picked on val): {best_thr:.4f}",
-        f"Model saved to: {model_path}",
-        "Best params:",
-        json.dumps(best_params, indent=2),
-        "=" * 72,
-    ]
-    with open(os.path.join(out_dir, "optimization_summary.txt"), "w") as f:
-        f.write("\n".join(summary))
-    print("\n".join(summary))
-
-
-if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--dataset_path", type=str, required=True)
-    parser.add_argument("--out_dir", type=str, required=True)
-    parser.add_argument("--model", type=str, choices=["xgb", "svm_gpu", "enet_gpu"], required=True)
-    parser.add_argument("--n_trials", type=int, default=200)
-    args = parser.parse_args()
-
-    run_optuna_and_refit(
-        dataset_path=args.dataset_path,
-        out_dir=args.out_dir,
-        model_name=args.model,
-        n_trials=args.n_trials,
-    )

training_classifiers/.ipynb_checkpoints/train_ml_regression-checkpoint.py

@@ -1,410 +0,0 @@
-import os
-import json
-import joblib
-import optuna
-import numpy as np
-import pandas as pd
-import matplotlib.pyplot as plt
-from dataclasses import dataclass
-from typing import Dict, Any, Tuple, Optional
-from datasets import load_from_disk, DatasetDict
-from sklearn.preprocessing import StandardScaler
-from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
-from sklearn.svm import SVR
-import xgboost as xgb
-from lightning.pytorch import seed_everything
-import cupy as cp
-from cuml.linear_model import ElasticNet as cuElasticNet
-from scipy.stats import spearmanr
-seed_everything(1986)
-
-
-# -----------------------------
-# GPU/CPU helpers
-# -----------------------------
-def to_gpu(X: np.ndarray):
-    if isinstance(X, cp.ndarray):
-        return X
-    return cp.asarray(X, dtype=cp.float32)
-
-def to_cpu(x):
-    if isinstance(x, cp.ndarray):
-        return cp.asnumpy(x)
-    return np.asarray(x)
-
-
-# -----------------------------
-# Data loading
-# -----------------------------
-@dataclass
-class SplitData:
-    X_train: np.ndarray
-    y_train: np.ndarray
-    seq_train: Optional[np.ndarray]
-    X_val: np.ndarray
-    y_val: np.ndarray
-    seq_val: Optional[np.ndarray]
-
-def _stack_embeddings(col) -> np.ndarray:
-    arr = np.asarray(col, dtype=np.float32)
-    if arr.ndim != 2:
-        arr = np.stack(col).astype(np.float32)
-    return arr
-
-def load_split_data(dataset_path: str) -> SplitData:
-    ds = load_from_disk(dataset_path)
-
-    if isinstance(ds, DatasetDict) and "train" in ds and "val" in ds:
-        train_ds, val_ds = ds["train"], ds["val"]
-    else:
-        if "split" not in ds.column_names:
-            raise ValueError("Dataset must be a DatasetDict(train/val) or have a 'split' column.")
-        train_ds = ds.filter(lambda x: x["split"] == "train")
-        val_ds   = ds.filter(lambda x: x["split"] == "val")
-
-    for required in ["embedding", "label"]:
-        if required not in train_ds.column_names:
-            raise ValueError(f"Missing column '{required}' in train split.")
-        if required not in val_ds.column_names:
-            raise ValueError(f"Missing column '{required}' in val split.")
-
-    X_train = _stack_embeddings(train_ds["embedding"]).astype(np.float32)
-    X_val   = _stack_embeddings(val_ds["embedding"]).astype(np.float32)
-
-    y_train = np.asarray(train_ds["label"], dtype=np.float32)
-    y_val   = np.asarray(val_ds["label"], dtype=np.float32)
-
-    seq_train = None
-    seq_val = None
-    if "sequence" in train_ds.column_names:
-        seq_train = np.asarray(train_ds["sequence"])
-    if "sequence" in val_ds.column_names:
-        seq_val = np.asarray(val_ds["sequence"])
-
-    return SplitData(X_train, y_train, seq_train, X_val, y_val, seq_val)
-
-
-# -----------------------------
-# Metrics
-# -----------------------------
-def safe_spearmanr(y_true: np.ndarray, y_pred: np.ndarray) -> float:
-    rho = spearmanr(y_true, y_pred).correlation
-    if rho is None or np.isnan(rho):
-        return 0.0
-    return float(rho)
-
-def eval_regression(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
-    # RMSE
-    try:
-        from sklearn.metrics import root_mean_squared_error
-        rmse = root_mean_squared_error(y_true, y_pred)
-    except Exception:
-        rmse = float(np.sqrt(mean_squared_error(y_true, y_pred)))
-
-    mae = float(mean_absolute_error(y_true, y_pred))
-    r2  = float(r2_score(y_true, y_pred))
-    rho = float(safe_spearmanr(y_true, y_pred))
-    return {"rmse": rmse, "mae": mae, "r2": r2, "spearman_rho": rho}
-
-
-# -----------------------------
-# Model
-# -----------------------------
-def train_xgb_reg(
-    X_train, y_train, X_val, y_val, params: Dict[str, Any]
-) -> Tuple[xgb.Booster, np.ndarray, np.ndarray]:
-    dtrain = xgb.DMatrix(X_train, label=y_train)
-    dval   = xgb.DMatrix(X_val, label=y_val)
-
-    num_boost_round = int(params.pop("num_boost_round"))
-    early_stopping_rounds = int(params.pop("early_stopping_rounds"))
-
-    booster = xgb.train(
-        params=params,
-        dtrain=dtrain,
-        num_boost_round=num_boost_round,
-        evals=[(dval, "val")],
-        early_stopping_rounds=early_stopping_rounds,
-        verbose_eval=False,
-    )
-
-    p_train = booster.predict(dtrain)
-    p_val   = booster.predict(dval)
-    return booster, p_train, p_val
-
-
-def train_cuml_elasticnet_reg(
-    X_train, y_train, X_val, y_val, params: Dict[str, Any]
-):
-    Xtr = to_gpu(X_train)
-    Xva = to_gpu(X_val)
-    ytr = to_gpu(y_train).astype(cp.float32)
-
-    model = cuElasticNet(
-        alpha=float(params["alpha"]),
-        l1_ratio=float(params["l1_ratio"]),
-        fit_intercept=True,
-        max_iter=int(params.get("max_iter", 5000)),
-        tol=float(params.get("tol", 1e-4)),
-        selection=params.get("selection", "cyclic"),
-    )
-    model.fit(Xtr, ytr)
-
-    p_train = to_cpu(model.predict(Xtr))
-    p_val   = to_cpu(model.predict(Xva))
-    return model, p_train, p_val
-
-
-def train_svr_reg(
-    X_train, y_train, X_val, y_val, params: Dict[str, Any]
-):
-    model = SVR(
-        C=float(params["C"]),
-        epsilon=float(params["epsilon"]),
-        kernel=params["kernel"],
-        gamma=params.get("gamma", "scale"),
-    )
-    model.fit(X_train, y_train)
-    p_train = model.predict(X_train)
-    p_val   = model.predict(X_val)
-    return model, p_train, p_val
-
-
-# -----------------------------
-# Saving + plots
-# -----------------------------
-def save_predictions_csv(
-    out_dir: str,
-    split_name: str,
-    y_true: np.ndarray,
-    y_pred: np.ndarray,
-    sequences: Optional[np.ndarray] = None,
-):
-    os.makedirs(out_dir, exist_ok=True)
-    df = pd.DataFrame({
-        "y_true": y_true.astype(float),
-        "y_pred": y_pred.astype(float),
-        "residual": (y_true - y_pred).astype(float),
-    })
-    if sequences is not None:
-        df.insert(0, "sequence", sequences)
-    df.to_csv(os.path.join(out_dir, f"{split_name}_predictions.csv"), index=False)
-
-def plot_regression_diagnostics(out_dir: str, y_true: np.ndarray, y_pred: np.ndarray):
-    os.makedirs(out_dir, exist_ok=True)
-
-    plt.figure()
-    plt.scatter(y_true, y_pred, s=8, alpha=0.5)
-    plt.xlabel("y_true")
-    plt.ylabel("y_pred")
-    plt.title("Predicted vs True")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "pred_vs_true.png"))
-    plt.close()
-
-    resid = y_true - y_pred
-    plt.figure()
-    plt.hist(resid, bins=50)
-    plt.xlabel("residual (y_true - y_pred)")
-    plt.ylabel("count")
-    plt.title("Residual Histogram")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "residual_hist.png"))
-    plt.close()
-
-    plt.figure()
-    plt.scatter(y_pred, resid, s=8, alpha=0.5)
-    plt.xlabel("y_pred")
-    plt.ylabel("residual")
-    plt.title("Residuals vs Prediction")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "residual_vs_pred.png"))
-    plt.close()
-
-
-# -----------------------------
-# Optuna objective (OPTIMIZE SPEARMAN RHO)
-# -----------------------------
-def make_objective(model_name: str, data: SplitData):
-    Xtr, ytr, Xva, yva = data.X_train, data.y_train, data.X_val, data.y_val
-
-    def objective(trial: optuna.Trial) -> float:
-        if model_name == "xgb_reg":
-            params = {
-                "objective": "reg:squarederror",
-                "eval_metric": "rmse",
-                "lambda": trial.suggest_float("lambda", 1e-10, 100.0, log=True),
-                "alpha":  trial.suggest_float("alpha",  1e-10, 100.0, log=True),
-                "gamma":  trial.suggest_float("gamma",  0.0, 10.0),
-                "max_depth": trial.suggest_int("max_depth", 2, 16),
-                "min_child_weight": trial.suggest_float("min_child_weight", 1e-3, 500.0, log=True),
-                "subsample": trial.suggest_float("subsample", 0.5, 1.0),
-                "colsample_bytree": trial.suggest_float("colsample_bytree", 0.3, 1.0),
-                "learning_rate": trial.suggest_float("learning_rate", 1e-3, 0.3, log=True),
-                "tree_method": "hist",
-                "device": "cuda",
-            }
-            params["num_boost_round"] = trial.suggest_int("num_boost_round", 50, 2000)
-            params["early_stopping_rounds"] = trial.suggest_int("early_stopping_rounds", 20, 200)
-
-            model, p_tr, p_va = train_xgb_reg(Xtr, ytr, Xva, yva, params.copy())
-
-        elif model_name == "enet_gpu":
-            params = {
-                "alpha": trial.suggest_float("alpha", 1e-8, 10.0, log=True),
-                "l1_ratio": trial.suggest_float("l1_ratio", 0.0, 1.0),
-                "max_iter": trial.suggest_int("max_iter", 1000, 20000),
-                "tol": trial.suggest_float("tol", 1e-6, 1e-2, log=True),
-                "selection": trial.suggest_categorical("selection", ["cyclic", "random"]),
-            }
-            model, p_tr, p_va = train_cuml_elasticnet_reg(Xtr, ytr, Xva, yva, params)
-
-        elif model_name == "svr":
-            params = {
-                "kernel": trial.suggest_categorical("kernel", ["rbf", "linear", "poly", "sigmoid"]),
-                "C": trial.suggest_float("C", 1e-3, 1e3, log=True),
-                "epsilon": trial.suggest_float("epsilon", 1e-4, 1.0, log=True),
-            }
-            if params["kernel"] in ["rbf", "poly", "sigmoid"]:
-                params["gamma"] = trial.suggest_float("gamma", 1e-6, 10.0, log=True)
-            else:
-                params["gamma"] = "scale"
-
-            model, p_tr, p_va = train_svr_reg(Xtr, ytr, Xva, yva, params)
-
-        else:
-            raise ValueError(f"Unknown model_name={model_name}")
-
-        metrics = eval_regression(yva, p_va)
-        trial.set_user_attr("spearman_rho", metrics["spearman_rho"])
-        trial.set_user_attr("rmse", metrics["rmse"])
-        trial.set_user_attr("mae", metrics["mae"])
-        trial.set_user_attr("r2", metrics["r2"])
-
-        # OPTUNA OBJECTIVE = maximize Spearman rho
-        return metrics["spearman_rho"]
-
-    return objective
-
-
-# -----------------------------
-# Main
-# -----------------------------
-def run_optuna_and_refit(
-    dataset_path: str,
-    out_dir: str,
-    model_name: str,
-    n_trials: int = 200,
-    standardize_X: bool = True,
-):
-    os.makedirs(out_dir, exist_ok=True)
-
-    data = load_split_data(dataset_path)
-    print(f"[Data] Train: {data.X_train.shape}, Val: {data.X_val.shape}")
-
-    # Standardize features (SVR + ElasticNet)
-    if standardize_X:
-        scaler = StandardScaler()
-        data.X_train = scaler.fit_transform(data.X_train).astype(np.float32)
-        data.X_val   = scaler.transform(data.X_val).astype(np.float32)
-        joblib.dump(scaler, os.path.join(out_dir, "scaler.joblib"))
-        print("[Preprocess] Saved StandardScaler -> scaler.joblib")
-
-    study = optuna.create_study(
-        direction="maximize",
-        pruner=optuna.pruners.MedianPruner()
-    )
-    study.optimize(make_objective(model_name, data), n_trials=n_trials)
-
-    trials_df = study.trials_dataframe()
-    trials_df.to_csv(os.path.join(out_dir, "study_trials.csv"), index=False)
-
-    best = study.best_trial
-    best_params = dict(best.params)
-
-    best_rho  = float(best.user_attrs.get("spearman_rho", best.value))
-    best_rmse = float(best.user_attrs.get("rmse", np.nan))
-    best_mae  = float(best.user_attrs.get("mae", np.nan))
-    best_r2   = float(best.user_attrs.get("r2", np.nan))
-
-    # Refit best model on train
-    if model_name == "xgb_reg":
-        params = {
-            "objective": "reg:squarederror",
-            "eval_metric": "rmse",
-            "lambda": best_params["lambda"],
-            "alpha": best_params["alpha"],
-            "gamma": best_params["gamma"],
-            "max_depth": best_params["max_depth"],
-            "min_child_weight": best_params["min_child_weight"],
-            "subsample": best_params["subsample"],
-            "colsample_bytree": best_params["colsample_bytree"],
-            "learning_rate": best_params["learning_rate"],
-            "tree_method": "hist",
-            "device": "cuda",
-            "num_boost_round": best_params["num_boost_round"],
-            "early_stopping_rounds": best_params["early_stopping_rounds"],
-        }
-        model, p_tr, p_va = train_xgb_reg(
-            data.X_train, data.y_train, data.X_val, data.y_val, params
-        )
-        model_path = os.path.join(out_dir, "best_model.json")
-        model.save_model(model_path)
-
-    elif model_name == "enet_gpu":
-        model, p_tr, p_va = train_cuml_elasticnet_reg(
-            data.X_train, data.y_train, data.X_val, data.y_val, best_params
-        )
-        model_path = os.path.join(out_dir, "best_model_cuml_enet.joblib")
-        joblib.dump(model, model_path)
-
-    elif model_name == "svr":
-        model, p_tr, p_va = train_svr_reg(
-            data.X_train, data.y_train, data.X_val, data.y_val, best_params
-        )
-        model_path = os.path.join(out_dir, "best_model_svr.joblib")
-        joblib.dump(model, model_path)
-
-    else:
-        raise ValueError(model_name)
-
-    save_predictions_csv(out_dir, "train", data.y_train, p_tr, data.seq_train)
-    save_predictions_csv(out_dir, "val",   data.y_val,   p_va, data.seq_val)
-
-    plot_regression_diagnostics(out_dir, data.y_val, p_va)
-
-    summary = [
-        "=" * 72,
-        f"MODEL: {model_name}",
-        f"Best trial: {best.number}",
-        f"Val Spearman rho (objective): {best_rho:.6f}",
-        f"Val RMSE:                      {best_rmse:.6f}",
-        f"Val MAE:                       {best_mae:.6f}",
-        f"Val R2:                        {best_r2:.6f}",
-        f"Model saved to:                {model_path}",
-        "Best params:",
-        json.dumps(best_params, indent=2),
-        "=" * 72,
-    ]
-    with open(os.path.join(out_dir, "optimization_summary.txt"), "w") as f:
-        f.write("\n".join(summary))
-    print("\n".join(summary))
-
-
-if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--dataset_path", type=str, required=True)
-    parser.add_argument("--out_dir", type=str, required=True)
-    parser.add_argument("--model", type=str, choices=["xgb_reg", "enet_gpu", "svr"], required=True)
-    parser.add_argument("--n_trials", type=int, default=200)
-    parser.add_argument("--no_standardize", action="store_true", help="Disable StandardScaler on X")
-    args = parser.parse_args()
-
-    run_optuna_and_refit(
-        dataset_path=args.dataset_path,
-        out_dir=args.out_dir,
-        model_name=args.model,
-        n_trials=args.n_trials,
-        standardize_X=(not args.no_standardize),
-    )

training_classifiers/.ipynb_checkpoints/train_nn-checkpoint.py

@@ -1,426 +0,0 @@
-import numpy as np
-import torch
-from torch.utils.data import DataLoader
-from datasets import load_from_disk, DatasetDict
-from sklearn.metrics import roc_auc_score, precision_recall_curve, f1_score
-import torch.nn as nn
-import optuna
-import os
-from typing import Dict, Any, Tuple, Optional
-import matplotlib.pyplot as plt
-from sklearn.metrics import (
-    f1_score, roc_auc_score, average_precision_score,
-    precision_recall_curve, roc_curve
-)
-import json
-import joblib
-import pandas as pd
-import time
-
-def infer_in_dim_from_unpooled_ds(ds) -> int:
-    ex = ds[0]
-    # ex["embedding"] is (L, H) list/array
-    return int(len(ex["embedding"][0]))
-    
-def load_split(dataset_path):
-    ds = load_from_disk(dataset_path)
-
-    if isinstance(ds, DatasetDict):
-        return ds["train"], ds["val"]
-
-    raise ValueError("Expected DatasetDict with 'train' and 'val' splits")
-
-def collate_unpooled(batch):
-    # batch: list of dicts
-    lengths = [int(x["length"]) for x in batch]
-    Lmax = max(lengths)
-    H = len(batch[0]["embedding"][0])  # 1280
-
-    X = torch.zeros(len(batch), Lmax, H, dtype=torch.float32)
-    M = torch.zeros(len(batch), Lmax, dtype=torch.bool)
-    y = torch.tensor([x["label"] for x in batch], dtype=torch.float32)
-
-    for i, x in enumerate(batch):
-        emb = torch.tensor(x["embedding"], dtype=torch.float32)  # (L, H)
-        L = emb.shape[0]
-        X[i, :L] = emb
-        if "attention_mask" in x:
-            m = torch.tensor(x["attention_mask"], dtype=torch.bool)
-            M[i, :L] = m[:L]
-        else:
-            M[i, :L] = True
-
-    return X, M, y
-
-# ======================== Helper functions =========================================
-def save_predictions_csv(
-    out_dir: str,
-    split_name: str,
-    y_true: np.ndarray,
-    y_prob: np.ndarray,
-    threshold: float,
-    sequences: Optional[np.ndarray] = None,
-):
-    os.makedirs(out_dir, exist_ok=True)
-    df = pd.DataFrame({
-        "y_true": y_true.astype(int),
-        "y_prob": y_prob.astype(float),
-        "y_pred": (y_prob >= threshold).astype(int),
-    })
-    if sequences is not None:
-        df.insert(0, "sequence", sequences)
-    df.to_csv(os.path.join(out_dir, f"{split_name}_predictions.csv"), index=False)
-
-
-def plot_curves(out_dir: str, y_true: np.ndarray, y_prob: np.ndarray):
-    os.makedirs(out_dir, exist_ok=True)
-
-    # PR
-    precision, recall, _ = precision_recall_curve(y_true, y_prob)
-    plt.figure()
-    plt.plot(recall, precision)
-    plt.xlabel("Recall")
-    plt.ylabel("Precision")
-    plt.title("Precision-Recall Curve")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "pr_curve.png"))
-    plt.close()
-
-    # ROC
-    fpr, tpr, _ = roc_curve(y_true, y_prob)
-    plt.figure()
-    plt.plot(fpr, tpr)
-    plt.xlabel("False Positive Rate")
-    plt.ylabel("True Positive Rate")
-    plt.title("ROC Curve")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "roc_curve.png"))
-    plt.close()
-
-# ======================== Shared OPTUNA training scheme =========================================
-def best_f1_threshold(y_true, y_prob):
-    p, r, thr = precision_recall_curve(y_true, y_prob)
-    f1s = (2*p[:-1]*r[:-1])/(p[:-1]+r[:-1]+1e-12)
-    i = int(np.nanargmax(f1s))
-    return float(thr[i]), float(f1s[i])
-
-@torch.no_grad()
-def eval_probs(model, loader, device):
-    model.eval()
-    ys, ps = [], []
-    for X, M, y in loader:
-        X, M = X.to(device), M.to(device)
-        logits = model(X, M)
-        prob = torch.sigmoid(logits).detach().cpu().numpy()
-        ys.append(y.numpy())
-        ps.append(prob)
-    return np.concatenate(ys), np.concatenate(ps)
-
-def train_one_epoch(model, loader, optim, criterion, device):
-    model.train()
-    for X, M, y in loader:
-        X, M, y = X.to(device), M.to(device), y.to(device)
-        optim.zero_grad(set_to_none=True)
-        logits = model(X, M)
-        loss = criterion(logits, y)
-        loss.backward()
-        optim.step()
-
-# ======================== MLP =========================================
-# Still need mean pooling along lengths
-class MaskedMeanPool(nn.Module):
-    def forward(self, X, M):  # X: (B,L,H), M: (B,L)
-        Mf = M.unsqueeze(-1).float()
-        denom = Mf.sum(dim=1).clamp(min=1.0)
-        return (X * Mf).sum(dim=1) / denom  # (B,H)
-
-class MLPClassifier(nn.Module):
-    def __init__(self, in_dim, hidden=512, dropout=0.1):
-        super().__init__()
-        self.pool = MaskedMeanPool()
-        self.net = nn.Sequential(
-            nn.Linear(in_dim, hidden),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden, 1),
-        )
-    def forward(self, X, M):
-        z = self.pool(X, M)
-        return self.net(z).squeeze(-1)  # logits
-
-# ======================== CNN =========================================
-# Treat 1280 dimensions as channels
-class CNNClassifier(nn.Module):
-    def __init__(self, in_ch, c=256, k=5, layers=2, dropout=0.1):
-        super().__init__()
-        blocks = []
-        ch = in_ch
-        for _ in range(layers):
-            blocks += [
-                nn.Conv1d(ch, c, kernel_size=k, padding=k//2),
-                nn.GELU(),
-                nn.Dropout(dropout),
-            ]
-            ch = c
-        self.conv = nn.Sequential(*blocks)
-        self.head = nn.Linear(c, 1)
-
-    def forward(self, X, M):
-        # X: (B,L,H) -> (B,H,L)
-        Xc = X.transpose(1, 2)
-        Y = self.conv(Xc).transpose(1, 2)  # (B,L,C)
-
-        # masked mean pool over L
-        Mf = M.unsqueeze(-1).float()
-        denom = Mf.sum(dim=1).clamp(min=1.0)
-        pooled = (Y * Mf).sum(dim=1) / denom  # (B,C)
-        return self.head(pooled).squeeze(-1)
-
-# ========================== Transformer ====================================
-class TransformerClassifier(nn.Module):
-    def __init__(self, in_dim, d_model=256, nhead=8, layers=2, ff=512, dropout=0.1):
-        super().__init__()
-        self.proj = nn.Linear(in_dim, d_model)
-        enc_layer = nn.TransformerEncoderLayer(
-            d_model=d_model, nhead=nhead, dim_feedforward=ff,
-            dropout=dropout, batch_first=True, activation="gelu"
-        )
-        self.enc = nn.TransformerEncoder(enc_layer, num_layers=layers)
-        self.head = nn.Linear(d_model, 1)
-
-    def forward(self, X, M):
-        # src_key_padding_mask: True = pad positions
-        pad_mask = ~M
-        Z = self.proj(X)             # (B,L,d)
-        Z = self.enc(Z, src_key_padding_mask=pad_mask)  # (B,L,d)
-
-        Mf = M.unsqueeze(-1).float()
-        denom = Mf.sum(dim=1).clamp(min=1.0)
-        pooled = (Z * Mf).sum(dim=1) / denom
-        return self.head(pooled).squeeze(-1)
-
-# ========================== OPTUNA ====================================
-
-def objective_nn(trial, model_name, train_ds, val_ds, device="cuda:0"):
-    # hyperparams shared
-    lr = trial.suggest_float("lr", 1e-5, 3e-3, log=True)
-    wd = trial.suggest_float("weight_decay", 1e-8, 1e-2, log=True)
-    dropout = trial.suggest_float("dropout", 0.0, 0.5)
-    batch_size = trial.suggest_categorical("batch_size", [16, 32, 64])
-
-    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
-                              collate_fn=collate_unpooled, num_workers=4, pin_memory=True)
-    val_loader = DataLoader(val_ds, batch_size=64, shuffle=False,
-                            collate_fn=collate_unpooled, num_workers=4, pin_memory=True)
-    
-    in_dim = infer_in_dim_from_unpooled_ds(train_ds)
-
-    if model_name == "mlp":
-        hidden = trial.suggest_categorical("hidden", [256, 512, 1024, 2048])
-        model = MLPClassifier(in_dim=in_dim, hidden=hidden, dropout=dropout)
-    elif model_name == "cnn":
-        c = trial.suggest_categorical("channels", [128, 256, 512])
-        k = trial.suggest_categorical("kernel", [3, 5, 7])
-        layers = trial.suggest_int("layers", 1, 4)
-        model = CNNClassifier(in_ch=in_dim, c=c, k=k, layers=layers, dropout=dropout)
-    elif model_name == "transformer":
-        d = trial.suggest_categorical("d_model", [128, 256, 384])
-        nhead = trial.suggest_categorical("nhead", [4, 8])
-        layers = trial.suggest_int("layers", 1, 4)
-        ff = trial.suggest_categorical("ff", [256, 512, 1024, 1536])
-        model = TransformerClassifier(in_dim=in_dim, d_model=d, nhead=nhead, layers=layers, ff=ff, dropout=dropout)
-    else:
-        raise ValueError(model_name)
-
-    model = model.to(device)
-
-    # class imbalance handling
-    ytr = np.asarray(train_ds["label"], dtype=np.int64)
-    pos = ytr.sum()
-    neg = len(ytr) - pos
-    pos_weight = torch.tensor([neg / max(pos, 1)], device=device, dtype=torch.float32)
-    criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
-
-    optim = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
-
-    best_f1 = -1.0
-    patience = 8
-    bad = 0
-
-    for epoch in range(1, 51):
-        train_one_epoch(model, train_loader, optim, criterion, device)
-    
-        y_true, y_prob = eval_probs(model, val_loader, device)
-        auc = roc_auc_score(y_true, y_prob)
-    
-        thr, f1 = best_f1_threshold(y_true, y_prob)
-    
-        trial.set_user_attr("val_auc", float(auc))
-        trial.set_user_attr("val_f1", float(f1))
-        trial.set_user_attr("val_thr", float(thr))
-    
-        # prune
-        trial.report(f1, epoch)
-        if trial.should_prune():
-            raise optuna.TrialPruned()
-    
-        if f1 > best_f1 + 1e-4:
-            best_f1 = f1
-            bad = 0
-        else:
-            bad += 1
-            if bad >= patience:
-                break
-    
-    return best_f1
-
-def run_optuna_and_refit_nn(dataset_path: str, out_dir: str, model_name: str, n_trials: int = 50, device="cuda:0"):
-    os.makedirs(out_dir, exist_ok=True)
-
-    train_ds, val_ds = load_split(dataset_path)
-    print(f"[Data] Train: {len(train_ds)}, Val: {len(val_ds)}")
-
-    study = optuna.create_study(direction="maximize", pruner=optuna.pruners.MedianPruner())
-    study.optimize(lambda trial: objective_nn(trial, model_name, train_ds, val_ds, device=device), n_trials=n_trials)
-
-    trials_df = study.trials_dataframe()
-    trials_df.to_csv(os.path.join(out_dir, "study_trials.csv"), index=False)
-
-    best = study.best_trial
-    best_params = dict(best.params)
-    best_f1_optuna = float(best.value)
-    best_auc_optuna = float(best.user_attrs.get("val_auc", np.nan))
-    best_thr = float(best.user_attrs.get("val_thr", 0.5))
-
-    in_dim = infer_in_dim_from_unpooled_ds(train_ds)
-    
-    # --- Refit best model  ---
-    batch_size = int(best_params.get("batch_size", 32))
-    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
-                              collate_fn=collate_unpooled, num_workers=4, pin_memory=True)
-    val_loader = DataLoader(val_ds, batch_size=64, shuffle=False,
-                            collate_fn=collate_unpooled, num_workers=4, pin_memory=True)
-
-    # Rebuild
-    dropout = float(best_params.get("dropout", 0.1))
-    if model_name == "mlp":
-        model = MLPClassifier(
-            in_dim=in_dim,                       
-            hidden=int(best_params["hidden"]),
-            dropout=dropout,
-        )
-
-    elif model_name == "cnn":
-        model = CNNClassifier(
-            in_ch=in_dim,                         
-            c=int(best_params["channels"]),
-            k=int(best_params["kernel"]),
-            layers=int(best_params["layers"]),
-            dropout=dropout,
-        )
-
-    elif model_name == "transformer":
-        model = TransformerClassifier(
-            in_dim=in_dim,   
-            d_model=int(best_params["d_model"]),
-            nhead=int(best_params["nhead"]),
-            layers=int(best_params["layers"]),
-            ff=int(best_params["ff"]),
-            dropout=dropout,
-        )
-    else:
-        raise ValueError(model_name)
-
-    model = model.to(device)
-
-    # loss + optimizer
-    ytr = np.asarray(train_ds["label"], dtype=np.int64)
-    pos = ytr.sum()
-    neg = len(ytr) - pos
-    pos_weight = torch.tensor([neg / max(pos, 1)], device=device, dtype=torch.float32)
-    criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
-
-    lr = float(best_params["lr"])
-    wd = float(best_params["weight_decay"])
-    optim = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
-
-    # train longer with early stopping on AUC
-    best_f1_seen, bad, patience = -1.0, 0, 12
-    best_state = None
-    best_thr_seen = 0.5
-    best_auc_seen = -1.0
-    
-    for epoch in range(1, 151):
-        train_one_epoch(model, train_loader, optim, criterion, device)
-    
-        y_true, y_prob = eval_probs(model, val_loader, device)
-        auc = roc_auc_score(y_true, y_prob)
-        thr, f1 = best_f1_threshold(y_true, y_prob)
-    
-        if f1 > best_f1_seen + 1e-4:
-            best_f1_seen = f1
-            best_thr_seen = thr
-            best_auc_seen = auc
-            bad = 0
-            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
-        else:
-            bad += 1
-            if bad >= patience:
-                break
-    
-    if best_state is not None:
-        model.load_state_dict(best_state)
-    
-    # final preds + threshold picked on val
-    y_true_val, y_prob_val = eval_probs(model, val_loader, device)
-    best_thr_final, best_f1_final = best_f1_threshold(y_true_val, y_prob_val)
-
-    # save model
-    model_path = os.path.join(out_dir, "best_model.pt")
-    torch.save({"state_dict": model.state_dict(), "best_params": best_params}, model_path)
-
-    # train preds
-    y_true_tr, y_prob_tr = eval_probs(model, DataLoader(train_ds, batch_size=64, shuffle=False,
-                                                       collate_fn=collate_unpooled, num_workers=4, pin_memory=True), device)
-
-    save_predictions_csv(out_dir, "train", y_true_tr, y_prob_tr, best_thr_final,
-                         sequences=np.asarray(train_ds["sequence"]) if "sequence" in train_ds.column_names else None)
-    save_predictions_csv(out_dir, "val", y_true_val, y_prob_val, best_thr_final,
-                         sequences=np.asarray(val_ds["sequence"]) if "sequence" in val_ds.column_names else None)
-
-    plot_curves(out_dir, y_true_val, y_prob_val)
-
-    summary = [
-        "=" * 72,
-        f"MODEL: {model_name}",
-    
-        # Optuna results (objective = F1)
-        f"Best Optuna F1 (objective): {best_f1_optuna:.4f}",
-        f"Best Optuna AUC (val, recorded): {best_auc_optuna:.4f}",
-        f"Best Optuna threshold (val): {best_thr:.4f}",
-    
-        # Refit results
-        f"Refit best AUC (val): {best_auc_seen:.4f}",
-        f"Refit best F1@thr (val): {best_f1_final:.4f} at thr={best_thr_final:.4f}",
-    
-        "Best params:",
-        json.dumps(best_params, indent=2),
-        f"Saved model: {model_path}",
-        "=" * 72,
-    ]
-
-    with open(os.path.join(out_dir, "optimization_summary.txt"), "w") as f:
-        f.write("\n".join(summary))
-    print("\n".join(summary))
-
-if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--dataset_path", type=str, required=True)
-    parser.add_argument("--out_dir", type=str, required=True)
-    parser.add_argument("--model", type=str, choices=["mlp", "cnn", "transformer"], required=True)
-    parser.add_argument("--n_trials", type=int, default=50)
-    args = parser.parse_args()
-
-    if args.model in ["mlp", "cnn", "transformer"]:
-        run_optuna_and_refit_nn(args.dataset_path, args.out_dir, args.model, args.n_trials, device="cuda:0")

training_classifiers/.ipynb_checkpoints/train_nn_regression-checkpoint.py

@@ -1,420 +0,0 @@
-import os, json, time
-import numpy as np
-import pandas as pd
-import matplotlib.pyplot as plt
-
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-from datasets import load_from_disk, DatasetDict
-import optuna
-from dataclasses import dataclass
-from typing import Dict, Any, Tuple, Optional
-from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
-from scipy.stats import spearmanr
-from torch.cuda.amp import autocast
-from torch.cuda.amp import autocast, GradScaler
-scaler = GradScaler(enabled=torch.cuda.is_available())
-from lightning.pytorch import seed_everything
-seed_everything(1986)
-
-
-def load_split(dataset_path):
-    ds = load_from_disk(dataset_path)
-    if isinstance(ds, DatasetDict):
-        return ds["train"], ds["val"]
-    raise ValueError("Expected DatasetDict with 'train' and 'val' splits")
-
-def collate_unpooled_reg(batch):
-    lengths = [int(x["length"]) for x in batch]
-    Lmax = max(lengths)
-    H = len(batch[0]["embedding"][0])
-
-    X = torch.zeros(len(batch), Lmax, H, dtype=torch.float32)
-    M = torch.zeros(len(batch), Lmax, dtype=torch.bool)
-    y = torch.tensor([float(x["label"]) for x in batch], dtype=torch.float32)
-
-    for i, x in enumerate(batch):
-        emb = torch.tensor(x["embedding"], dtype=torch.float32)  # (L,H)
-        L = emb.shape[0]
-        X[i, :L] = emb
-        if "attention_mask" in x:
-            m = torch.tensor(x["attention_mask"], dtype=torch.bool)
-            M[i, :L] = m[:L]
-        else:
-            M[i, :L] = True
-    return X, M, y
-
-def infer_in_dim(ds) -> int:
-    ex = ds[0]
-    return int(len(ex["embedding"][0]))
-
-# ============================
-# Metrics
-# ============================
-def safe_spearmanr(y_true: np.ndarray, y_pred: np.ndarray) -> float:
-    rho = spearmanr(y_true, y_pred).correlation
-    if rho is None or np.isnan(rho):
-        return 0.0
-    return float(rho)
-
-def eval_regression(y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, float]:
-    # ---- RMSE ----
-    try:
-        from sklearn.metrics import root_mean_squared_error
-        rmse = root_mean_squared_error(y_true, y_pred)
-    except Exception:
-        mse = mean_squared_error(y_true, y_pred)
-        rmse = float(np.sqrt(mse))
-
-    mae  = float(mean_absolute_error(y_true, y_pred))
-    r2   = float(r2_score(y_true, y_pred))
-    rho  = float(safe_spearmanr(y_true, y_pred))
-    return {"rmse": float(rmse), "mae": mae, "r2": r2, "spearman_rho": rho}
-
-
-# ============================
-# Models
-# ============================
-class MaskedMeanPool(nn.Module):
-    def forward(self, X, M):
-        Mf = M.unsqueeze(-1).float()
-        denom = Mf.sum(dim=1).clamp(min=1.0)
-        return (X * Mf).sum(dim=1) / denom
-
-class MLPRegressor(nn.Module):
-    def __init__(self, in_dim, hidden=512, dropout=0.1):
-        super().__init__()
-        self.pool = MaskedMeanPool()
-        self.net = nn.Sequential(
-            nn.Linear(in_dim, hidden),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden, 1),
-        )
-    def forward(self, X, M):
-        z = self.pool(X, M)
-        return self.net(z).squeeze(-1)  # y_pred
-
-class CNNRegressor(nn.Module):
-    def __init__(self, in_ch, c=256, k=5, layers=2, dropout=0.1):
-        super().__init__()
-        blocks = []
-        ch = in_ch
-        for _ in range(layers):
-            blocks += [
-                nn.Conv1d(ch, c, kernel_size=k, padding=k//2),
-                nn.GELU(),
-                nn.Dropout(dropout),
-            ]
-            ch = c
-        self.conv = nn.Sequential(*blocks)
-        self.head = nn.Linear(c, 1)
-
-    def forward(self, X, M):
-        Xc = X.transpose(1, 2)                # (B,H,L)
-        Y = self.conv(Xc).transpose(1, 2)     # (B,L,C)
-        Mf = M.unsqueeze(-1).float()
-        denom = Mf.sum(dim=1).clamp(min=1.0)
-        pooled = (Y * Mf).sum(dim=1) / denom  # (B,C)
-        return self.head(pooled).squeeze(-1)
-
-class TransformerRegressor(nn.Module):
-    def __init__(self, in_dim, d_model=256, nhead=8, layers=2, ff=512, dropout=0.1):
-        super().__init__()
-        self.proj = nn.Linear(in_dim, d_model)
-        enc_layer = nn.TransformerEncoderLayer(
-            d_model=d_model, nhead=nhead, dim_feedforward=ff,
-            dropout=dropout, batch_first=True, activation="gelu"
-        )
-        self.enc = nn.TransformerEncoder(enc_layer, num_layers=layers)
-        self.head = nn.Linear(d_model, 1)
-
-    def forward(self, X, M):
-        pad_mask = ~M
-        Z = self.proj(X)
-        Z = self.enc(Z, src_key_padding_mask=pad_mask)
-        Mf = M.unsqueeze(-1).float()
-        denom = Mf.sum(dim=1).clamp(min=1.0)
-        pooled = (Z * Mf).sum(dim=1) / denom
-        return self.head(pooled).squeeze(-1)
-
-# ============================
-# Train / eval
-# ============================
-@torch.no_grad()
-def eval_preds(model, loader, device):
-    model.eval()
-    ys, ps = [], []
-    for X, M, y in loader:
-        X, M = X.to(device), M.to(device)
-        pred = model(X, M).detach().cpu().numpy()
-        ys.append(y.numpy())
-        ps.append(pred)
-    return np.concatenate(ys), np.concatenate(ps)
-
-def train_one_epoch_reg(model, loader, optim, criterion, device):
-    model.train()
-    for X, M, y in loader:
-        X, M, y = X.to(device), M.to(device), y.to(device)
-        optim.zero_grad(set_to_none=True)
-        with autocast(enabled=torch.cuda.is_available()):
-            pred = model(X, M)
-            loss = criterion(pred, y)
-        scaler.scale(loss).backward()
-        scaler.step(optim)
-        scaler.update()
-        
-# ============================
-# Saving + plots
-# ============================
-def save_predictions_csv(out_dir, split_name, y_true, y_pred, sequences=None):
-    os.makedirs(out_dir, exist_ok=True)
-    df = pd.DataFrame({
-        "y_true": y_true.astype(float),
-        "y_pred": y_pred.astype(float),
-        "residual": (y_true - y_pred).astype(float),
-    })
-    if sequences is not None:
-        df.insert(0, "sequence", sequences)
-    df.to_csv(os.path.join(out_dir, f"{split_name}_predictions.csv"), index=False)
-
-def plot_regression_diagnostics(out_dir, y_true, y_pred):
-    os.makedirs(out_dir, exist_ok=True)
-
-    plt.figure()
-    plt.scatter(y_true, y_pred, s=8, alpha=0.5)
-    plt.xlabel("y_true"); plt.ylabel("y_pred")
-    plt.title("Predicted vs True")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "pred_vs_true.png"))
-    plt.close()
-
-    resid = y_true - y_pred
-    plt.figure()
-    plt.hist(resid, bins=50)
-    plt.xlabel("residual (y_true - y_pred)"); plt.ylabel("count")
-    plt.title("Residual Histogram")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "residual_hist.png"))
-    plt.close()
-
-    plt.figure()
-    plt.scatter(y_pred, resid, s=8, alpha=0.5)
-    plt.xlabel("y_pred"); plt.ylabel("residual")
-    plt.title("Residuals vs Prediction")
-    plt.tight_layout()
-    plt.savefig(os.path.join(out_dir, "residual_vs_pred.png"))
-    plt.close()
-
-# ============================
-# Optuna objective
-# ============================
-def score_from_metrics(metrics: Dict[str, float], objective: str) -> float:
-    if objective == "spearman":
-        return metrics["spearman_rho"]
-    if objective == "r2":
-        return metrics["r2"]
-    if objective == "neg_rmse":
-        return -metrics["rmse"]
-    raise ValueError(f"Unknown objective={objective}")
-
-def objective_nn_reg(trial, model_name, train_ds, val_ds, device="cuda:0", objective="spearman"):
-    lr = trial.suggest_float("lr", 1e-5, 3e-3, log=True)
-    wd = trial.suggest_float("weight_decay", 1e-10, 1e-2, log=True)
-    dropout = trial.suggest_float("dropout", 0.0, 0.5)
-    batch_size = trial.suggest_categorical("batch_size", [16, 32, 64])
-
-    in_dim = infer_in_dim(train_ds)
-
-    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
-                              collate_fn=collate_unpooled_reg, num_workers=4, pin_memory=True)
-    val_loader = DataLoader(val_ds, batch_size=64, shuffle=False,
-                            collate_fn=collate_unpooled_reg, num_workers=4, pin_memory=True)
-
-    if model_name == "mlp":
-        hidden = trial.suggest_categorical("hidden", [256, 512, 1024, 2048])
-        model = MLPRegressor(in_dim=in_dim, hidden=hidden, dropout=dropout)
-    elif model_name == "cnn":
-        c = trial.suggest_categorical("channels", [128, 256, 512])
-        k = trial.suggest_categorical("kernel", [3, 5, 7])
-        layers = trial.suggest_int("layers", 1, 4)
-        model = CNNRegressor(in_ch=in_dim, c=c, k=k, layers=layers, dropout=dropout)
-    elif model_name == "transformer":
-        d = trial.suggest_categorical("d_model", [128, 256, 384])
-        nhead = trial.suggest_categorical("nhead", [4, 8])
-        layers = trial.suggest_int("layers", 1, 4)
-        ff = trial.suggest_categorical("ff", [256, 512, 1024, 1536])
-        model = TransformerRegressor(in_dim=in_dim, d_model=d, nhead=nhead, layers=layers, ff=ff, dropout=dropout)
-    else:
-        raise ValueError(model_name)
-
-    model = model.to(device)
-
-    loss_name = trial.suggest_categorical("loss", ["mse", "huber"])
-    if loss_name == "mse":
-        criterion = nn.MSELoss()
-    else:
-        delta = trial.suggest_float("huber_delta", 0.5, 5.0, log=True)
-        criterion = nn.HuberLoss(delta=delta)
-
-    optim = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
-
-    best_score = -1e18
-    patience = 10
-    bad = 0
-
-    for epoch in range(1, 61):
-        train_one_epoch_reg(model, train_loader, optim, criterion, device)
-
-        y_true, y_pred = eval_preds(model, val_loader, device)
-        metrics = eval_regression(y_true, y_pred)
-        score = score_from_metrics(metrics, objective)
-
-        # log attrs
-        for k, v in metrics.items():
-            trial.set_user_attr(f"val_{k}", float(v))
-
-        trial.report(score, epoch)
-        if trial.should_prune():
-            raise optuna.TrialPruned()
-
-        if score > best_score + 1e-6:
-            best_score = score
-            bad = 0
-        else:
-            bad += 1
-            if bad >= patience:
-                break
-
-    return float(best_score)
-
-# ============================
-# Main runner
-# ============================
-def run_optuna_and_refit_nn_reg(dataset_path, out_dir, model_name, n_trials=80, device="cuda:0",
-                                objective="spearman"):
-    os.makedirs(out_dir, exist_ok=True)
-
-    train_ds, val_ds = load_split(dataset_path)
-    print(f"[Data] Train: {len(train_ds)}, Val: {len(val_ds)}")
-
-    study = optuna.create_study(direction="maximize", pruner=optuna.pruners.MedianPruner())
-    study.optimize(lambda t: objective_nn_reg(t, model_name, train_ds, val_ds, device=device, objective=objective),
-                   n_trials=n_trials)
-
-    trials_df = study.trials_dataframe()
-    trials_df.to_csv(os.path.join(out_dir, "study_trials.csv"), index=False)
-
-    best = study.best_trial
-    best_params = dict(best.params)
-
-    # rebuild model from best params
-    in_dim = infer_in_dim(train_ds)
-    dropout = float(best_params.get("dropout", 0.1))
-    if model_name == "mlp":
-        model = MLPRegressor(in_dim=in_dim, hidden=int(best_params["hidden"]), dropout=dropout)
-    elif model_name == "cnn":
-        model = CNNRegressor(in_ch=in_dim, c=int(best_params["channels"]),
-                             k=int(best_params["kernel"]), layers=int(best_params["layers"]),
-                             dropout=dropout)
-    elif model_name == "transformer":
-        model = TransformerRegressor(in_dim=in_dim, d_model=int(best_params["d_model"]),
-                                     nhead=int(best_params["nhead"]), layers=int(best_params["layers"]),
-                                     ff=int(best_params["ff"]), dropout=dropout)
-    else:
-        raise ValueError(model_name)
-
-    model = model.to(device)
-
-    batch_size = int(best_params.get("batch_size", 32))
-    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
-                              collate_fn=collate_unpooled_reg, num_workers=4, pin_memory=True)
-    val_loader = DataLoader(val_ds, batch_size=64, shuffle=False,
-                            collate_fn=collate_unpooled_reg, num_workers=4, pin_memory=True)
-
-    # loss
-    if best_params.get("loss", "mse") == "mse":
-        criterion = nn.MSELoss()
-    else:
-        criterion = nn.HuberLoss(delta=float(best_params["huber_delta"]))
-
-    optim = torch.optim.AdamW(model.parameters(), lr=float(best_params["lr"]),
-                              weight_decay=float(best_params["weight_decay"]))
-
-    # refit longer with early stopping on the SAME objective
-    best_score, bad, patience = -1e18, 0, 15
-    best_state = None
-
-    for epoch in range(1, 201):
-        train_one_epoch_reg(model, train_loader, optim, criterion, device)
-
-        y_true, y_pred = eval_preds(model, val_loader, device)
-        metrics = eval_regression(y_true, y_pred)
-        score = score_from_metrics(metrics, objective)
-
-        if score > best_score + 1e-6:
-            best_score = score
-            bad = 0
-            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
-            best_metrics = metrics
-        else:
-            bad += 1
-            if bad >= patience:
-                break
-
-    if best_state is not None:
-        model.load_state_dict(best_state)
-
-    # preds
-    y_true_tr, y_pred_tr = eval_preds(model, DataLoader(train_ds, batch_size=64, shuffle=False,
-                                                       collate_fn=collate_unpooled_reg, num_workers=4, pin_memory=True), device)
-    y_true_va, y_pred_va = eval_preds(model, val_loader, device)
-
-    seq_train = np.asarray(train_ds["sequence"]) if "sequence" in train_ds.column_names else None
-    seq_val   = np.asarray(val_ds["sequence"])   if "sequence" in val_ds.column_names else None
-    save_predictions_csv(out_dir, "train", y_true_tr, y_pred_tr, seq_train)
-    save_predictions_csv(out_dir, "val",   y_true_va, y_pred_va, seq_val)
-    plot_regression_diagnostics(out_dir, y_true_va, y_pred_va)
-
-    # save model
-    model_path = os.path.join(out_dir, "best_model.pt")
-    torch.save({"state_dict": model.state_dict(), "best_params": best_params, "in_dim": in_dim}, model_path)
-
-    summary = [
-        "=" * 72,
-        f"MODEL: {model_name}",
-        f"OPTUNA objective: {objective} (direction=maximize)",
-        f"Best trial: {best.number}",
-        "Best val metrics:",
-        json.dumps({k: float(v) for k, v in best_metrics.items()}, indent=2),
-        f"Saved model: {model_path}",
-        "Best params:",
-        json.dumps(best_params, indent=2),
-        "=" * 72,
-    ]
-    with open(os.path.join(out_dir, "optimization_summary.txt"), "w") as f:
-        f.write("\n".join(summary))
-    print("\n".join(summary))
-
-
-if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--dataset_path", type=str, required=True)
-    parser.add_argument("--out_dir", type=str, required=True)
-    parser.add_argument("--model", type=str, choices=["mlp","cnn","transformer"], required=True)
-    parser.add_argument("--n_trials", type=int, default=80)
-    parser.add_argument("--objective", type=str, default="spearman",
-                        choices=["spearman","neg_rmse","r2"])
-    parser.add_argument("--device", type=str, default="cuda:0")
-    args = parser.parse_args()
-
-    run_optuna_and_refit_nn_reg(
-        dataset_path=args.dataset_path,
-        out_dir=args.out_dir,
-        model_name=args.model,
-        n_trials=args.n_trials,
-        device=args.device,
-        objective=args.objective,
-    )

training_classifiers/binding_affinity/val_smiles_pooled.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5410a45a7b65def6cfb94c167b07537abd33b5aac4ecdffe162b7ce4e9bc3d19
-size 36525

training_classifiers/binding_affinity/val_smiles_unpooled.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:cdf71fbb3e7b3b8e8dbfe4ed45b32a2da0049df851f09ee32564825f626cb86c
-size 37187

training_classifiers/binding_affinity/val_wt_pooled.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b194e7b2b97258320323021b3ffe6143133070212a0215ade22fa91b87a3a861
-size 33224

training_classifiers/binding_affinity/val_wt_unpooled.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:051325790047e749fbf1daf7bf25a08178297b0c37acaf9439816d09f2b6c1e3
-size 33826

training_classifiers/binding_affinity/wt_smiles_pooled/best_model.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:12f956a7bf04ed602c11fd275377afa73f3f0af1982dbe06c607d8ada304b01c
-size 21617397

training_classifiers/binding_affinity/wt_smiles_unpooled/best_model.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3d7ae3d2190b034352a65bda1bce86aa5a96ce3daf74cf10a166f8d9e9af51f0
-size 181183221

training_classifiers/binding_affinity/wt_wt_pooled/.ipynb_checkpoints/optuna_trials-checkpoint.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b685b92714882d618b42b582000574d83c3be2fbecbec5e0de6b5476948b96c5
-size 40700

training_classifiers/half_life/cnn_smiles/cv_oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0a8a57d44cac3fcd701b550a4eaf9e29910540bfb7580a9b8ee997a7227375d2
-size 13748

training_classifiers/half_life/cnn_unpooled_peptideclm/best_model.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:9eaaafffe02663f7cfe67fde25cdebd7d4315af69b393b433c4291b700bc5063
-size 16525563

training_classifiers/half_life/cnn_unpooled_smiles/cv_oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:879d8c6f47b02c1ddd86fbe3982d8b0134167521f9f71d2450957dc3bbbb6bd1
-size 13705

training_classifiers/half_life/enet_gpu_smiles/cv_oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:fd5c84e9788e2db949c6be785f8539178e72fc6fa6bc703daf9574ad0622e0f1
-size 13649

training_classifiers/half_life/enet_peptideclm/smiles_halflife_best_enet.joblib

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0eb93bcb27436e80bce2a6433cbd7502b90de4962731250972eef08a5d96ce69
-size 22698

training_classifiers/half_life/mlp_smiles/cv_oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:14b1010a2c0b5d065fa9b82636085806ec7f6f6091c7c2355c6c4717d07fa79b
-size 13724

training_classifiers/half_life/mlp_unpooled_peptideclm/best_model.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ef78a5e5c555768f91dc646652a39e367287e851a14e2cf85e4006c9355a8368
-size 2368455

training_classifiers/half_life/mlp_unpooled_smiles/cv_oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2a037cf36528fc8e04a375c0443577830733636fca83fa9ce44e457e28e4f771
-size 13745

training_classifiers/half_life/svr_gpu_smiles/cv_oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a0dd42537c1a5589b78451de8645bfc089b8f7f5839808222bb1e9e033d78c66
-size 13746

training_classifiers/half_life/svr_peptideclm/smiles_halflife_best_svr.joblib

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5579f1407fc8dfd1e42b4ea2a6b619dea8b0eff4ce9a4c0869890cbd1b321851
-size 1530479

training_classifiers/half_life/transformer_smiles/cv_oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:6c6dcaff0d542c3a6bbaf499aba56e5f440c50aa18b55271ee85feb43851fe92
-size 13694

training_classifiers/half_life/transformer_unpooled_peptideclm/best_model.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:9a19240f8067e68a6e2eaff139f90b6d2f37ab5431197c5496894937c01918f7
-size 931353

training_classifiers/half_life/transformer_unpooled_smiles/cv_oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7b1111c0b57288092ab97b940598b2d3b44c2ff5299fe55a50a2312d8c2e45af
-size 13683

training_classifiers/half_life/transformer_wt_log/oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8ec7b8dee908ef43ba7633a887a988834e24f952711b906472e1b41b833de714
-size 14100

training_classifiers/half_life/transformer_wt_log/optimization_summary.txt

@@ -1,33 +0,0 @@
-========================================================================
-MODEL: transformer
-Dataset: /scratch/pranamlab/tong/data/halflife/halflife_embedding_unpooled
-Target column: log_label
-CV folds: 5
-Optuna objective: spearman (direction=maximize)
-Best trial: 45
-OOF metrics:
-{
-  "rmse": 1.0389505624771118,
-  "mae": 0.722099244594574,
-  "r2": 0.30950748920440674,
-  "spearman_rho": 0.5818272477094295
-}
-OOF score (spearman): 0.581827
-Best params:
-{
-  "lr": 0.0003603824115240561,
-  "weight_decay": 2.9442493502916885e-09,
-  "dropout": 0.3851371373367485,
-  "batch_size": 16
-}
-Final refit epochs (all data): 15
-Saved final model: /scratch/pranamlab/tong/PeptiVerse/src/halflife/finetune_stability_transformer/final_model.pt
-Benchmark (final model on full data):
-{
-  "n_samples": 130,
-  "wall_time_s": 1.9577592574059963,
-  "throughput_samples_per_s": 66.40244427818372,
-  "gpu_ms_per_sample": 0.28296443315652703,
-  "gpu_peak_mem_MB": 77.5693359375
-}
-========================================================================

training_classifiers/half_life/transformer_wt_log/study_trials.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5726f6b8541c7ca85eda9f0e526db0cb10156eadb2c440fd7a66f7a7d1209175
-size 10154

training_classifiers/half_life/transformer_wt_raw/oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c3df70b094757f34fa380a28727877694fcb1ec367bbbef28c63b257ecec74e6
-size 13516

training_classifiers/half_life/transformer_wt_raw/optimization_summary.txt

@@ -1,33 +0,0 @@
-========================================================================
-MODEL: transformer
-Dataset: /scratch/pranamlab/tong/data/halflife/halflife_embedding_unpooled
-Target column: label
-CV folds: 5
-Optuna objective: spearman (direction=maximize)
-Best trial: 22
-OOF metrics:
-{
-  "rmse": 45.00321578979492,
-  "mae": 11.352466583251953,
-  "r2": 0.02070075273513794,
-  "spearman_rho": 0.3759734508605516
-}
-OOF score (spearman): 0.375973
-Best params:
-{
-  "lr": 0.00019977882554167927,
-  "weight_decay": 1.102955470301081e-07,
-  "dropout": 1.2707176359392082e-05,
-  "batch_size": 16
-}
-Final refit epochs (all data): 14
-Saved final model: /scratch/pranamlab/tong/PeptiVerse/src/halflife/finetune_stability_transformer_raw/final_model.pt
-Benchmark (final model on full data):
-{
-  "n_samples": 130,
-  "wall_time_s": 1.6299039730802178,
-  "throughput_samples_per_s": 79.7593000244818,
-  "gpu_ms_per_sample": 0.23774326214423547,
-  "gpu_peak_mem_MB": 77.5693359375
-}
-========================================================================

training_classifiers/half_life/transformer_wt_raw/study_trials.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7c38352854ad4142c02a4bcb33caee9fe8fa22fca86dcb8c17c05c24f3fa5bca
-size 10895

training_classifiers/half_life/xgb_smiles/cv_oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e2df43f71aad2cf791b49daa0b3353f524d5a3f3e132fecf1251e96242639ca5
-size 13675

training_classifiers/half_life/xgb_wt_log/oof_predictions.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2293b5752ef6bdc7b1ec8ae2f56e11ccbf32aee024d777c86c1e63f390fa89cf
-size 14805

training_classifiers/half_life/xgb_wt_log/optimization_summary.txt

@@ -1,27 +0,0 @@
-{
-  "model": "xgb_reg",
-  "dataset_path": "/scratch/pranamlab/tong/data/halflife/halflife_embedding",
-  "target_col": "log_label",
-  "n_folds": 5,
-  "best_trial_number": 20,
-  "best_objective_cv_spearman": 0.5879000126060311,
-  "oof_metrics": {
-    "rmse": 1.0810768604278564,
-    "mae": 0.7866008281707764,
-    "r2": 0.2524225115776062,
-    "spearman_rho": 0.557870619380726
-  },
-  "model_path": "/scratch/pranamlab/tong/PeptiVerse/src/halflife/finetune_stability_xgb_log/best_model.json",
-  "best_params": {
-    "lambda": 0.0006291983667746282,
-    "alpha": 0.0820082035401697,
-    "gamma": 1.2243543209914751,
-    "max_depth": 3,
-    "min_child_weight": 1.7773959178614585,
-    "subsample": 0.568291807635477,
-    "colsample_bytree": 0.8597778117881122,
-    "learning_rate": 0.0512590763008084,
-    "num_boost_round": 1728,
-    "early_stopping_rounds": 121
-  }
-}