recovery/scripts/embed_sample.py

#!/usr/bin/env python3
"""Cache pinned frozen-encoder embeddings for a MitoInteract sample."""

from __future__ import annotations

import argparse
import json
import time
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer

DEFAULT_PROTEIN_MODEL = "facebook/esm2_t12_35M_UR50D"
DEFAULT_PROTEIN_REVISION = "6fbf070e65b0b7291e7bbcd451118c216cff79d8"
DEFAULT_LIGAND_MODEL = "DeepChem/ChemBERTa-77M-MLM"
DEFAULT_LIGAND_REVISION = "ed8a5374f2024ec8da53760af91a33fb8f6a15ff"


def read_rows(path: Path, limit: int | None) -> list[dict]:
    rows = []
    with path.open() as handle:
        for line in handle:
            if line.strip():
                rows.append(json.loads(line))
            if limit and len(rows) >= limit:
                break
    return rows


def masked_mean(
    last_hidden: torch.Tensor, attention: torch.Tensor, special: torch.Tensor
) -> torch.Tensor:
    mask = attention.bool() & ~special.bool()
    weights = mask.unsqueeze(-1).to(last_hidden.dtype)
    return (last_hidden * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1)


def encode_texts(
    model, tokenizer, texts: list[str], batch_size: int, device: torch.device
) -> np.ndarray:
    outputs = []
    for start in range(0, len(texts), batch_size):
        batch = texts[start : start + batch_size]
        encoded = tokenizer(
            batch,
            padding=True,
            truncation=True,
            max_length=min(getattr(tokenizer, "model_max_length", 512), 512),
            return_special_tokens_mask=True,
            return_tensors="pt",
        )
        special = encoded.pop("special_tokens_mask")
        encoded = {key: value.to(device) for key, value in encoded.items()}
        with torch.inference_mode():
            hidden = model(**encoded).last_hidden_state
        pooled = masked_mean(hidden, encoded["attention_mask"], special.to(device))
        outputs.append(pooled.float().cpu().numpy())
    return np.concatenate(outputs, axis=0)


def encode_proteins(
    model,
    tokenizer,
    entities: list[tuple[str, str]],
    batch_size: int,
    device: torch.device,
    chunk_residues: int,
) -> dict[str, np.ndarray]:
    chunks: list[str] = []
    owners: list[str] = []
    weights: list[int] = []
    for entity_id, sequence in entities:
        for start in range(0, len(sequence), chunk_residues):
            chunk = sequence[start : start + chunk_residues]
            chunks.append(chunk)
            owners.append(entity_id)
            weights.append(len(chunk))

    chunk_embeddings = []
    for start in range(0, len(chunks), batch_size):
        batch = chunks[start : start + batch_size]
        encoded = tokenizer(
            batch,
            padding=True,
            truncation=True,
            max_length=chunk_residues + 2,
            return_special_tokens_mask=True,
            return_tensors="pt",
        )
        special = encoded.pop("special_tokens_mask")
        encoded = {key: value.to(device) for key, value in encoded.items()}
        with torch.inference_mode():
            hidden = model(**encoded).last_hidden_state
        pooled = masked_mean(hidden, encoded["attention_mask"], special.to(device))
        chunk_embeddings.extend(pooled.float().cpu().numpy())

    accum: dict[str, list[tuple[np.ndarray, int]]] = defaultdict(list)
    for owner, embedding, weight in zip(owners, chunk_embeddings, weights, strict=True):
        accum[owner].append((embedding, weight))
    result = {}
    for owner, values in accum.items():
        matrix = np.stack([value for value, _ in values])
        entity_weights = np.asarray([weight for _, weight in values], dtype=np.float32)
        result[owner] = np.average(matrix, axis=0, weights=entity_weights)
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--sample", type=Path, default=Path("artifacts/dev-10k/sample.jsonl")
    )
    parser.add_argument("--target-key", default="paffinity")
    parser.add_argument("--target-name", default="pAffinity")
    parser.add_argument("--limit", type=int)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--protein-chunk-residues", type=int, default=1022)
    parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto")
    parser.add_argument("--protein-model", default=DEFAULT_PROTEIN_MODEL)
    parser.add_argument("--protein-revision", default=DEFAULT_PROTEIN_REVISION)
    parser.add_argument("--ligand-model", default=DEFAULT_LIGAND_MODEL)
    parser.add_argument("--ligand-revision", default=DEFAULT_LIGAND_REVISION)
    parser.add_argument("--output", type=Path, default=Path("artifacts/embeddings.npz"))
    args = parser.parse_args()

    device_name = (
        "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device
    )
    if device_name == "auto":
        device_name = "cpu"
    device = torch.device(device_name)
    rows = read_rows(args.sample, args.limit)
    proteins = sorted({row["protein_id"]: row["sequence"] for row in rows}.items())
    ligands = sorted({row["ligand_id"]: row["smiles"] for row in rows}.items())

    started = time.monotonic()
    protein_tokenizer = AutoTokenizer.from_pretrained(
        args.protein_model, revision=args.protein_revision
    )
    protein_model = (
        AutoModel.from_pretrained(args.protein_model, revision=args.protein_revision)
        .eval()
        .to(device)
    )
    protein_embeddings = encode_proteins(
        protein_model,
        protein_tokenizer,
        proteins,
        args.batch_size,
        device,
        args.protein_chunk_residues,
    )
    del protein_model

    ligand_tokenizer = AutoTokenizer.from_pretrained(
        args.ligand_model, revision=args.ligand_revision
    )
    ligand_model = (
        AutoModel.from_pretrained(args.ligand_model, revision=args.ligand_revision)
        .eval()
        .to(device)
    )
    ligand_matrix = encode_texts(
        ligand_model,
        ligand_tokenizer,
        [smiles for _, smiles in ligands],
        args.batch_size,
        device,
    )
    ligand_embeddings = {
        entity_id: embedding
        for (entity_id, _), embedding in zip(ligands, ligand_matrix, strict=True)
    }

    pair_protein = np.stack(
        [protein_embeddings[row["protein_id"]] for row in rows]
    ).astype(np.float32)
    pair_ligand = np.stack(
        [ligand_embeddings[row["ligand_id"]] for row in rows]
    ).astype(np.float32)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        args.output,
        pair_ids=np.asarray([row["pair_id"] for row in rows]),
        target=np.asarray([row[args.target_key] for row in rows], dtype=np.float32),
        protein=pair_protein,
        ligand=pair_ligand,
    )
    metadata = {
        "rows": len(rows),
        "unique_proteins": len(proteins),
        "unique_ligands": len(ligands),
        "protein_dim": int(pair_protein.shape[1]),
        "ligand_dim": int(pair_ligand.shape[1]),
        "target_key": args.target_key,
        "target_name": args.target_name,
        "protein_model": args.protein_model,
        "protein_revision": args.protein_revision,
        "ligand_model": args.ligand_model,
        "ligand_revision": args.ligand_revision,
        "protein_chunk_residues": args.protein_chunk_residues,
        "device": str(device),
        "seconds": time.monotonic() - started,
    }
    args.output.with_suffix(".json").write_text(json.dumps(metadata, indent=2) + "\n")
    print(json.dumps(metadata, indent=2))


if __name__ == "__main__":
    main()