Skills › Scientific Agent Skills › esm › esm_embeddings.py
esm_embeddings.py
106 lines · skills/esm/scripts/esm_embeddings.py at 5f2bc03
1"""Residue-only ESMC embeddings, without loading weights or contacting a service.
3Use with esm 3.4.1.post1. Callers supply a loaded model/tokenizer or SDK tensors.
4Only nonempty, unaligned, single-chain amino-acid sequences are accepted here.
5"""
7from collections.abc import Sequence
9import torch
11AMINO_ACIDS = frozenset("ACDEFGHIKLMNPQRSTVWYXBZUO")
14def validate_sequence(sequence: str) -> None:
15 if not isinstance(sequence, str) or not sequence:
16 raise ValueError("Expected a nonempty amino-acid sequence")
17 if len(sequence) > 2046:
18 raise ValueError("ESMC's 2048-token context allows at most 2046 plain residues")
19 invalid = set(sequence) - AMINO_ACIDS
20 if invalid:
21 raise ValueError(f"Unsupported sequence characters: {sorted(invalid)!r}")
24def mean_pool_residues(
25 hidden: torch.Tensor,
26 attention_mask: torch.Tensor,
27 special_tokens_mask: torch.Tensor,
28) -> torch.Tensor:
29 """Pool B,T,D states over actual residues; exclude CLS/EOS/padding.
31 Masks must use 0/1 entries. Invalid/padded states are masked before summing,
32 so a NaN padding value cannot contaminate a valid sequence embedding.
33 """
34 if hidden.ndim != 3 or hidden.shape[-1] == 0:
35 raise ValueError("Expected nonempty hidden states with shape (B, T, D)")
36 if attention_mask.shape != hidden.shape[:2] or special_tokens_mask.shape != hidden.shape[:2]:
37 raise ValueError("Masks must match the hidden-state batch and token axes")
38 for mask in (attention_mask, special_tokens_mask):
39 if not torch.all((mask == 0) | (mask == 1)):
40 raise ValueError("Masks must contain only 0 and 1")
41 keep = attention_mask.to(hidden.device).bool() & ~special_tokens_mask.to(hidden.device).bool()
42 counts = keep.sum(dim=1)
43 if torch.any(counts == 0):
44 raise ValueError("Each sequence must have at least one residue")
45 if not torch.isfinite(hidden[keep]).all():
46 raise ValueError("Nonfinite residue embeddings")
47 clean = hidden.float().masked_fill(~keep.unsqueeze(-1), 0)
48 return clean.sum(dim=1) / counts.unsqueeze(-1)
51def embed_sequences(model, tokenizer, sequences: Sequence[str]) -> torch.Tensor:
52 """Return CPU (B,D) residue means using the native EsmcForMaskedLM API.
54 This is one padded batch. The caller chooses a batch size fitting memory.
55 No truncation, implicit sequence cleaning, model loading, or network calls.
56 """
57 if isinstance(sequences, str) or not sequences:
58 raise ValueError("Provide a nonempty collection of sequences")
59 for sequence in sequences:
60 validate_sequence(sequence)
61 if model.training:
62 raise ValueError("Put the model in evaluation mode before embedding")
63 inputs = tokenizer(
64 list(sequences), padding=True, return_tensors="pt",
65 return_special_tokens_mask=True, truncation=False,
66 )
67 special = inputs.pop("special_tokens_mask")
68 residue_counts = (inputs["attention_mask"].bool() & ~special.bool()).sum(1)
69 if residue_counts.tolist() != [len(sequence) for sequence in sequences]:
70 raise ValueError("Tokenizer changed the number of residues")
71 device = next(model.parameters()).device
72 inputs = {key: value.to(device) for key, value in inputs.items()}
73 with torch.inference_mode():
74 output = model(**inputs)
75 pooled = mean_pool_residues(
76 output.last_hidden_state, inputs["attention_mask"], special
77 )
78 return pooled.cpu()
81def sdk_residue_embeddings(
82 tokens: torch.Tensor,
83 embeddings: torch.Tensor,
84 tokenizer,
85 sequence_length: int,
86) -> torch.Tensor:
87 """Extract L,D residues from a single-chain SDK (1,L+2,D) response.
89 Reject unexpected layouts instead of silently assigning BOS to residue 1.
90 This helper intentionally excludes padded/multichain SDK inputs.
91 """
92 if sequence_length < 1 or tokens.ndim != 1:
93 raise ValueError("Expected a positive sequence length and 1D tokens")
94 if tokens.numel() != sequence_length + 2:
95 raise ValueError("Expected exactly CLS, L residues, EOS")
96 if embeddings.ndim != 3 or embeddings.shape[:2] != (1, tokens.numel()):
97 raise ValueError("Expected SDK embeddings with shape (1, L+2, D)")
98 if tokens[0].item() != tokenizer.cls_token_id or tokens[-1].item() != tokenizer.eos_token_id:
99 raise ValueError("Missing expected CLS/EOS boundary tokens")
100 special_ids = set(tokenizer.all_special_ids)
101 if any(token in special_ids for token in tokens[1:-1].tolist()):
102 raise ValueError("Unexpected special token within sequence")
103 residues = embeddings[0, 1:-1].float()
104 if not torch.isfinite(residues).all():
105 raise ValueError("Nonfinite residue embeddings")
106 return residues