utils/phlam_embedder.py

"""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,
        }