File size: 6,281 Bytes
6dd9839 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | #!/usr/bin/env python3
"""Feature 2: protein + RNA → binding probability (RnaRealismClassifierV3)."""
from __future__ import annotations
import argparse
import json
import os
import sys
from pathlib import Path
import h5py
import torch
from transformers import AutoTokenizer
from common import PKG_ROOT, device, load_env, resolve
from encode_protein import encode_protein_to_h5
sys.path.insert(0, str(PKG_ROOT / "generator"))
sys.path.insert(0, str(PKG_ROOT / "classifier"))
from labeled_dataset import make_collate # noqa: E402
from model import RnaRealismClassifierV3 # noqa: E402
from src.model import EsmConfig, EsmForMaskedLM # noqa: E402
from src.utils import _remap_legacy_checkpoint_state_dict, base_config # noqa: E402
def _load_mlm(ckpt_path: Path, tokenizer, torch_device: torch.device, model_cfg: dict) -> EsmForMaskedLM:
mcfg = model_cfg["model"]
cfg = EsmConfig(**base_config)
cfg.vocab_size = len(tokenizer)
cfg.pad_token_id = tokenizer.pad_token_id
cfg.mask_token_id = tokenizer.mask_token_id
cfg.use_FiLM = mcfg["use_FiLM"]
cfg.use_AdaLN = mcfg["use_AdaLN"]
cfg.use_gated_bias = mcfg["use_gated_bias"]
cfg.use_protein_conditioning_attention = mcfg["use_protein_conditioning_attention"]
cfg.protein_dim = mcfg["protein_dim"]
mlm = EsmForMaskedLM(cfg)
ckpt = torch.load(ckpt_path, map_location="cpu")
sd = _remap_legacy_checkpoint_state_dict(ckpt["model_state_dict"])
mlm.load_state_dict(sd, strict=True)
mlm.to(torch_device)
return mlm
@torch.no_grad()
def score_binding(
protein_seq: str,
rna_seq: str,
*,
p_id: str = "QUERY",
protein_h5: Path | None = None,
) -> dict:
load_env()
torch_device = torch.device(device() if torch.cuda.is_available() else "cpu")
if protein_h5 is None:
work = resolve(os.environ.get("WORKSPACE", "workspace"))
work.mkdir(parents=True, exist_ok=True)
protein_h5 = encode_protein_to_h5(
protein_seq, p_id=p_id, output_h5=work / f"{p_id}_vesm3b.h5"
)
with h5py.File(protein_h5, "r") as f:
p_ids = [x.decode() if isinstance(x, bytes) else str(x) for x in f["p_ids"][:]]
if p_id not in p_ids:
p_id = p_ids[0]
idx = p_ids.index(p_id)
start = int(f["starts"][idx])
length = int(f["lengths"][idx])
emb = torch.from_numpy(f["embeddings"][start : start + length]).float()
rna = rna_seq.upper().replace("T", "U")
rna = "".join(ch for ch in rna if ch in "ACGU")
if not rna:
raise ValueError("RNA sequence is empty or invalid")
cls_ckpt = resolve(os.environ.get("CLASSIFIER_CKPT", "checkpoints/classifier.pt"))
backbone = resolve(
os.environ.get("BACKBONE_CKPT")
or os.environ.get("D3LM_BACKBONE_CKPT") # legacy env name
or "checkpoints/backbone.pt"
)
tok_dir = resolve(os.environ.get("TOKENIZER_DIR", "generator/tokenizer"))
bundle = torch.load(cls_ckpt, map_location="cpu")
# Checkpoints may still store config under the legacy key "d3lm_config".
model_cfg = bundle.get("model_config") or bundle["d3lm_config"]
train_args = bundle.get("train_args") or {}
data_cfg = model_cfg["data"]
tokenizer = AutoTokenizer.from_pretrained(str(tok_dir), trust_remote_code=True)
mlm = _load_mlm(backbone, tokenizer, torch_device, model_cfg)
model = RnaRealismClassifierV3(
mlm,
mlm.config.hidden_size,
int(model_cfg["model"]["protein_dim"]),
freeze_mode=train_args.get("freeze_mode", "conditioning_only"),
head_dropout=float(train_args.get("head_dropout", 0.1)),
cross_attn_heads=int(train_args.get("cross_attn_heads", 8)),
).to(torch_device)
state_key = "classifier_state_dict" if "classifier_state_dict" in bundle else "model_state_dict"
model.load_state_dict(bundle[state_key], strict=True)
model.eval()
sample = {
"r_id": "query",
"rna": rna,
"p_id": p_id,
"protein_emb": emb,
"rna_len": len(rna),
"protein_len": int(length),
"label": 0,
}
append_eos = bool(data_cfg.get("append_eos_token", False))
max_rna_nt = data_cfg.get("max_generated_rna_bp")
max_rna_nt = int(max_rna_nt) if max_rna_nt is not None else None
collate = make_collate(str(tok_dir), append_eos, max_rna_nt)
batch = collate([sample])
use_bf16 = torch_device.type == "cuda"
with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16, enabled=use_bf16):
logits = model(
protein_cond=batch["protein_cond"].to(torch_device),
protein_attention_mask=batch["protein_attention_mask"].to(torch_device),
input_ids=batch["rna_input_ids_clean"].to(torch_device),
attention_mask=batch["rna_attention_mask"].to(torch_device),
)
if isinstance(logits, (tuple, list)):
logits = logits[0]
logit_t = logits.reshape(-1)[0].float()
prob_t = torch.sigmoid(logit_t)
return {
"p_id": p_id,
"rna": rna,
"rna_len": len(rna),
"logit": float(logit_t.cpu()),
"binding_prob": float(prob_t.cpu()),
"protein_h5": str(protein_h5),
}
def main() -> None:
ap = argparse.ArgumentParser(description="Protein + RNA → binding probability")
ap.add_argument("--protein", required=True, help="AA sequence, or @path/to.fasta")
ap.add_argument("--rna", required=True, help="RNA sequence (ACGU/T)")
ap.add_argument("--p-id", default="QUERY")
ap.add_argument("--protein-h5", type=Path, default=None, help="Skip encoding if H5 exists")
args = ap.parse_args()
load_env()
protein = args.protein
if protein.startswith("@"):
text = Path(protein[1:]).read_text(encoding="utf-8")
lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
if lines and lines[0].startswith(">"):
protein = "".join(lines[1:])
if args.p_id == "QUERY":
args.p_id = lines[0][1:].split()[0]
else:
protein = "".join(lines)
out = score_binding(protein, args.rna, p_id=args.p_id, protein_h5=args.protein_h5)
print(json.dumps(out, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()
|