File size: 4,301 Bytes
c87881a | 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 | """Pinned ESM-2 protein embedding generation for external BGCs."""
from __future__ import annotations
import json
import time
from pathlib import Path
from typing import Any
import h5py
import numpy as np
import torch
from Bio import SeqIO
from .artifacts import sha256_file, write_json_immutable
MODEL_NAME = "esm2_t33_650M_UR50D"
MODEL_LAYER = 33
EMBEDDING_DIMENSION = 1280
MAX_SEQUENCE_LENGTH = 1022
def read_fasta(path: str | Path) -> list[tuple[str, str]]:
with Path(path).open("r", encoding="utf-8") as handle:
records = [(record.id, str(record.seq)) for record in SeqIO.parse(handle, "fasta")]
identifiers = [identifier for identifier, _ in records]
if len(identifiers) != len(set(identifiers)):
raise ValueError("External FASTA identifiers must be unique")
if not records:
raise ValueError("External FASTA is empty")
return records
def generate_esm2_embeddings(
fasta_path: str | Path,
output_h5: str | Path,
metadata_path: str | Path,
batch_size: int = 16,
resume: bool = False,
allow_cpu: bool = False,
) -> dict[str, Any]:
try:
import esm
except ImportError as error:
raise RuntimeError("Install the pinned fair-esm==2.0.0 dependency") from error
if not torch.cuda.is_available() and not allow_cpu:
raise RuntimeError("CUDA is required unless --allow-cpu is explicitly supplied")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
output = Path(output_h5)
partial = output.with_suffix(output.suffix + ".partial")
metadata_output = Path(metadata_path)
if output.exists() or metadata_output.exists():
raise FileExistsError("Refusing to overwrite completed external embeddings")
if partial.exists() and not resume:
raise FileExistsError(f"Partial output exists; pass --resume: {partial}")
output.parent.mkdir(parents=True, exist_ok=True)
sequences = read_fasta(fasta_path)
already_done: set[str] = set()
if resume and partial.exists():
with h5py.File(partial, "r") as handle:
already_done = set(handle.keys())
remaining = [(key, value) for key, value in sequences if key not in already_done]
remaining.sort(key=lambda item: (len(item[1]), item[0]))
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
model = model.eval().to(device)
if device.type == "cuda":
model = model.half()
converter = alphabet.get_batch_converter()
started = time.time()
truncated = sum(len(sequence) > MAX_SEQUENCE_LENGTH for _, sequence in sequences)
with h5py.File(partial, "a") as handle:
for start in range(0, len(remaining), batch_size):
batch = remaining[start : start + batch_size]
prepared = []
for identifier, sequence in batch:
prepared.append((identifier, sequence[:MAX_SEQUENCE_LENGTH]))
labels, _, tokens = converter(prepared)
tokens = tokens.to(device)
with torch.no_grad(), torch.autocast(
device_type=device.type,
dtype=torch.float16 if device.type == "cuda" else torch.bfloat16,
enabled=device.type == "cuda",
):
representations = model(
tokens, repr_layers=[MODEL_LAYER], return_contacts=False
)["representations"][MODEL_LAYER]
for row, (identifier, sequence) in enumerate(prepared):
embedding = representations[row, 1 : len(sequence) + 1].mean(dim=0)
handle.create_dataset(identifier, data=embedding.float().cpu().numpy().astype(np.float16))
handle.flush()
partial.replace(output)
metadata = {
"schema_version": 1,
"model": MODEL_NAME,
"fair_esm_version": "2.0.0",
"representation_layer": MODEL_LAYER,
"embedding_dimension": EMBEDDING_DIMENSION,
"maximum_sequence_length": MAX_SEQUENCE_LENGTH,
"proteins": len(sequences),
"truncated_proteins": truncated,
"device": str(device),
"elapsed_seconds": time.time() - started,
"fasta_sha256": sha256_file(fasta_path),
"h5_sha256": sha256_file(output),
}
write_json_immutable(metadata_output, metadata)
return metadata
|