File size: 4,266 Bytes
6c4ade3 4b63894 6c4ade3 e882a80 6c4ade3 3972381 4b63894 6c4ade3 4b63894 6c4ade3 4b63894 6c4ade3 66a69dd 6c4ade3 66a69dd 6c4ade3 | 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 | """PHLAM genome/protein embeddings from a genome's ESM2 protein embeddings.
phlam_genome: h_global -- one vector per genome.
phlam_protein: h_protein - mean(h_protein over the genome) -- one residual
vector per protein, i.e. each protein's deviation from its
own genome's average.
"""
import os
import shutil
import torch
import torch.nn.functional as F
from phlam.model import PHLAM
from utils.retry import retry
# matches training (phagevault/pipelines/embed_phlam.py: raw +1 -> 2, raw -1 -> 1)
STRAND_TO_ID = {"+": 2, "-": 1}
MODELS_DIR = os.environ.get("MODELS_DATA_DIR", "/data/models")
CKPT_FILENAME = os.environ.get("PHLAM_CKPT_FILE", "phlam_model.pt")
MODEL_CKPT_PATH = os.environ.get("PHLAM_CKPT_PATH", os.path.join(MODELS_DIR, CKPT_FILENAME))
LOCAL_CACHE_DIR = os.environ.get("LOCAL_CACHE_DIR", "/tmp/phlam_local_cache")
NUM_LAYERS = 10
NUM_HEADS = 12
def _ensure_local(mount_path):
"""Copies a small file off the (possibly flaky) network mount once, so
every load after the first reads a plain local file instead of racing
the mount again. Checkpoint-sized files only -- the multi-GB FAISS
indexes stay mmap'd straight off the mount, that's a different tradeoff."""
os.makedirs(LOCAL_CACHE_DIR, exist_ok=True)
local_path = os.path.join(LOCAL_CACHE_DIR, os.path.basename(mount_path))
if not os.path.exists(local_path):
tmp_path = local_path + ".tmp"
def _copy():
shutil.copyfile(mount_path, tmp_path)
os.replace(tmp_path, local_path)
retry(_copy, label=f"copy {mount_path} to local cache")
return local_path
class PHLAMEmbedder:
def __init__(self, ckpt_path=MODEL_CKPT_PATH, device="cpu"):
local_path = _ensure_local(ckpt_path)
ckpt = torch.load(local_path, map_location="cpu", weights_only=False)
state_dict = {k.replace("_orig_mod.", ""): v for k, v in ckpt["model_state_dict"].items()}
model = PHLAM(input_dim=480, d_model=480, output_dim=480,
num_layers=NUM_LAYERS, num_heads=NUM_HEADS, num_global=1)
model.load_state_dict(state_dict, strict=True)
model.eval()
self.model = model.to(device)
self.device = device
def embed(self, esm_embeds, strands, is_circular=True):
"""esm_embeds: [L, 480] tensor/array of raw (pre-normalization) ESM
embeddings for one genome, in gene order. strands: list[str] of "+"/"-"
of length L. Returns {"phlam_genome": [480], "phlam_protein": [L, 480]}
as numpy arrays."""
strand_ids = [STRAND_TO_ID[s] for s in strands]
esm_t = torch.as_tensor(esm_embeds, dtype=torch.float32)
esm_in = F.normalize(esm_t, p=2, dim=-1).unsqueeze(0).to(self.device)
strand_t = torch.tensor([strand_ids], dtype=torch.long, device=self.device)
with torch.no_grad():
out = self.model(
strand_ids=strand_t, esm_embeds=esm_in,
is_circular=1.0 if is_circular else 0.0,
output_attentions=False,
)
h_protein = out["h_protein"][0] # [L, 480]
h_global = out["h_global"][0] # [480]
phlam_protein = h_protein - h_protein.mean(dim=0, keepdim=True)
# Reconstruction quality per protein: cosine distance between
# PHLAM's denoised reconstruction and the (normalized) ESM input
# it was asked to recover -- low = familiar/well-modeled, high =
# novel. strand_flip_prob: no corruption is applied at
# inference, so a high value flags a protein whose orientation
# looks inconsistent with the local synteny PHLAM has learned.
cos_sim = F.cosine_similarity(out["recon_esm"][0], esm_in[0], dim=-1)
cosine_loss = (1.0 - cos_sim).float().cpu().numpy()
strand_probs = F.softmax(out["strand_logits"][0], dim=-1)
strand_flip_prob = strand_probs[:, 1].float().cpu().numpy()
return {
"phlam_genome": h_global.float().cpu().numpy(),
"phlam_protein": phlam_protein.float().cpu().numpy(),
"cosine_loss": cosine_loss,
"strand_flip_prob": strand_flip_prob,
}
|