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