Synthyra/FastESMFold
065
1import os
2import sqlite3
3import networkx as nx
4import numpy as np
5import torch
6from tqdm.auto import tqdm
7from typing import Callable, Dict, List, Optional, Set
8from torch.utils.data import DataLoader
9from torch.utils.data import Dataset as TorchDataset
10from transformers import PreTrainedTokenizerBase
11
12
13class Pooler:
14 def __init__(self, pooling_types: List[str]) -> None:
15 self.pooling_types = pooling_types
16 self.pooling_options: Dict[str, Callable] = {
17 'mean': self.mean_pooling,
18 'max': self.max_pooling,
19 'norm': self.norm_pooling,
20 'median': self.median_pooling,
21 'std': self.std_pooling,
22 'var': self.var_pooling,
23 'cls': self.cls_pooling,
24 'parti': self._pool_parti,
25 }
26
27 def _create_pooled_matrices_across_layers(self, attentions: torch.Tensor) -> torch.Tensor:
28 assert isinstance(attentions, torch.Tensor)
29 maxed_attentions = torch.max(attentions, dim=1)[0]
30 return maxed_attentions
31
32 def _page_rank(self, attention_matrix: np.ndarray, personalization: Optional[dict] = None, nstart: Optional[dict] = None, prune_type: str = "top_k_outdegree") -> Dict[int, float]:
33 # Run PageRank on the attention matrix converted to a graph.
34 # Raises exceptions if the graph doesn't match the token sequence or has no edges.
35 # Returns the PageRank scores for each token node.
36 G = self._convert_to_graph(attention_matrix)
37 if G.number_of_nodes() != attention_matrix.shape[0]:
38 raise Exception(
39 f"The number of nodes in the graph should be equal to the number of tokens in sequence! You have {G.number_of_nodes()} nodes for {attention_matrix.shape[0]} tokens.")
40 if G.number_of_edges() == 0:
41 raise Exception(f"You don't seem to have any attention edges left in the graph.")
42
43 return nx.pagerank(G, alpha=0.85, tol=1e-06, weight='weight', personalization=personalization, nstart=nstart, max_iter=100)
44
45 def _convert_to_graph(self, matrix: np.ndarray) -> nx.DiGraph:
46 # Convert a matrix (e.g., attention scores) to a directed graph using networkx.
47 # Each element in the matrix represents a directed edge with a weight.
48 G = nx.from_numpy_array(matrix, create_using=nx.DiGraph)
49 return G
50
51 def _calculate_importance_weights(self, dict_importance: Dict[int, float], attention_mask: Optional[torch.Tensor] = None) -> np.ndarray:
52 # Remove keys where attention_mask is 0
53 if attention_mask is not None:
54 for k in list(dict_importance.keys()):
55 if attention_mask[k] == 0:
56 del dict_importance[k]
57
58 #dict_importance[0] # remove cls
59 #dict_importance[-1] # remove eos
60 total = sum(dict_importance.values())
61 return np.array([v / total for _, v in dict_importance.items()])
62
63 def _pool_parti(self, emb: torch.Tensor, attentions: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: # (b, L, d) -> (b, d)
64 maxed_attentions = self._create_pooled_matrices_across_layers(attentions).numpy()
65 # emb is (b, L, d), maxed_attentions is (b, L, L)
66 emb_pooled = []
67 for e, a, mask in zip(emb, maxed_attentions, attention_mask):
68 dict_importance = self._page_rank(a)
69 importance_weights = self._calculate_importance_weights(dict_importance, mask)
70 num_tokens = int(mask.sum().item())
71 emb_pooled.append(np.average(e[:num_tokens], weights=importance_weights, axis=0))
72 pooled = torch.tensor(np.array(emb_pooled))
73 return pooled
74
75 def mean_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: # (b, L, d) -> (b, d)
76 if attention_mask is None:
77 return emb.mean(dim=1)
78 else:
79 attention_mask = attention_mask.unsqueeze(-1)
80 return (emb * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
81
82 def max_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: # (b, L, d) -> (b, d)
83 if attention_mask is None:
84 return emb.max(dim=1).values
85 else:
86 attention_mask = attention_mask.unsqueeze(-1)
87 return (emb * attention_mask).max(dim=1).values
88
89 def norm_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: # (b, L, d) -> (b, d)
90 if attention_mask is None:
91 return emb.norm(dim=1, p=2)
92 else:
93 attention_mask = attention_mask.unsqueeze(-1)
94 return (emb * attention_mask).norm(dim=1, p=2)
95
96 def median_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: # (b, L, d) -> (b, d)
97 if attention_mask is None:
98 return emb.median(dim=1).values
99 else:
100 attention_mask = attention_mask.unsqueeze(-1)
101 return (emb * attention_mask).median(dim=1).values
102
103 def std_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: # (b, L, d) -> (b, d)
104 if attention_mask is None:
105 return emb.std(dim=1)
106 else:
107 # Compute variance correctly over non-masked positions, then take sqrt
108 var = self.var_pooling(emb, attention_mask, **kwargs)
109 return torch.sqrt(var)
110
111 def var_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: # (b, L, d) -> (b, d)
112 if attention_mask is None:
113 return emb.var(dim=1)
114 else:
115 # Correctly compute variance over only non-masked positions
116 attention_mask = attention_mask.unsqueeze(-1) # (b, L, 1)
117 # Compute mean over non-masked positions
118 mean = (emb * attention_mask).sum(dim=1) / attention_mask.sum(dim=1) # (b, d)
119 mean = mean.unsqueeze(1) # (b, 1, d)
120 # Compute squared differences from mean, only over non-masked positions
121 squared_diff = (emb - mean) ** 2 # (b, L, d)
122 # Sum squared differences over non-masked positions and divide by count
123 var = (squared_diff * attention_mask).sum(dim=1) / attention_mask.sum(dim=1) # (b, d)
124 return var
125
126 def cls_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: # (b, L, d) -> (b, d)
127 return emb[:, 0, :]
128
129 def __call__(
130 self,
131 emb: torch.Tensor,
132 attention_mask: Optional[torch.Tensor] = None,
133 attentions: Optional[torch.Tensor] = None
134 ) -> torch.Tensor: # [mean, max]
135 final_emb: List[torch.Tensor] = []
136 for pooling_type in self.pooling_types:
137 final_emb.append(self.pooling_options[pooling_type](emb=emb, attention_mask=attention_mask, attentions=attentions)) # (b, d)
138 return torch.cat(final_emb, dim=-1) # (b, n_pooling_types * d)
139
140
141class ProteinDataset(TorchDataset):
142 """Simple dataset for protein sequences."""
143 def __init__(self, sequences: List[str]) -> None:
144 self.sequences = sequences
145
146 def __len__(self) -> int:
147 return len(self.sequences)
148
149 def __getitem__(self, idx: int) -> str:
150 return self.sequences[idx]
151
152
153def build_collator(tokenizer: PreTrainedTokenizerBase) -> Callable[[List[str]], Dict[str, torch.Tensor]]:
154 def _collate_fn(sequences: List[str]) -> Dict[str, torch.Tensor]:
155 return tokenizer(sequences, return_tensors="pt", padding='longest')
156 return _collate_fn
157
158
159def parse_fasta(fasta_path: str) -> List[str]:
160 assert os.path.exists(fasta_path), f"FASTA file does not exist: {fasta_path}"
161 sequences = []
162 current_seq = []
163 with open(fasta_path, 'r') as f:
164 for line in f:
165 line = line.strip()
166 if not line:
167 continue
168 if line.startswith('>'):
169 if current_seq:
170 sequences.append(''.join(current_seq))
171 current_seq = []
172 else:
173 current_seq.append(line)
174 if current_seq:
175 sequences.append(''.join(current_seq))
176 return sequences
177
178
179class EmbeddingMixin:
180 def _embed(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
181 raise NotImplementedError
182
183 @property
184 def device(self) -> torch.device:
185 """Get the device of the model."""
186 return next(self.parameters()).device
187
188 def _read_sequences_from_db(self, db_path: str) -> Set[str]:
189 """Read sequences from SQLite database."""
190 sequences = []
191 with sqlite3.connect(db_path) as conn:
192 c = conn.cursor()
193 c.execute("SELECT sequence FROM embeddings")
194 while True:
195 row = c.fetchone()
196 if row is None:
197 break
198 sequences.append(row[0])
199 return set(sequences)
200
201 def _ensure_embeddings_table(self, conn: sqlite3.Connection) -> None:
202 cursor = conn.cursor()
203 cursor.execute(
204 "CREATE TABLE IF NOT EXISTS embeddings ("
205 "sequence TEXT PRIMARY KEY, "
206 "embedding BLOB NOT NULL, "
207 "shape TEXT, "
208 "dtype TEXT"
209 ")"
210 )
211 cursor.execute("PRAGMA table_info(embeddings)")
212 rows = cursor.fetchall()
213 column_names = [row[1] for row in rows]
214 if "shape" not in column_names:
215 cursor.execute("ALTER TABLE embeddings ADD COLUMN shape TEXT")
216 if "dtype" not in column_names:
217 cursor.execute("ALTER TABLE embeddings ADD COLUMN dtype TEXT")
218 conn.commit()
219
220 def load_embeddings_from_pth(self, save_path: str) -> Dict[str, torch.Tensor]:
221 assert os.path.exists(save_path), f"Embedding file does not exist: {save_path}"
222 payload = torch.load(save_path, map_location="cpu", weights_only=True)
223 assert isinstance(payload, dict), "Expected .pth embeddings file to contain a dictionary."
224 for sequence, tensor in payload.items():
225 assert isinstance(sequence, str), "Expected embedding dictionary keys to be sequences (str)."
226 assert isinstance(tensor, torch.Tensor), "Expected embedding dictionary values to be tensors."
227 return payload
228
229 def load_embeddings_from_db(self, db_path: str, sequences: Optional[List[str]] = None) -> Dict[str, torch.Tensor]:
230 assert os.path.exists(db_path), f"Embedding database does not exist: {db_path}"
231 loaded: Dict[str, torch.Tensor] = {}
232 with sqlite3.connect(db_path) as conn:
233 self._ensure_embeddings_table(conn)
234 cursor = conn.cursor()
235 if sequences is None:
236 cursor.execute("SELECT sequence, embedding, shape, dtype FROM embeddings")
237 else:
238 if len(sequences) == 0:
239 return loaded
240 placeholders = ",".join(["?"] * len(sequences))
241 cursor.execute(
242 f"SELECT sequence, embedding, shape, dtype FROM embeddings WHERE sequence IN ({placeholders})",
243 tuple(sequences),
244 )
245
246 rows = cursor.fetchall()
247 for row in rows:
248 sequence = row[0]
249 embedding_bytes = row[1]
250 shape_text = row[2]
251 dtype_text = row[3]
252 assert shape_text is not None, "Missing shape metadata in embeddings table."
253 assert dtype_text is not None, "Missing dtype metadata in embeddings table."
254 shape_values = [int(value) for value in shape_text.split(",") if len(value) > 0]
255 assert len(shape_values) > 0, f"Invalid shape metadata for sequence: {sequence}"
256 expected_size = int(np.prod(shape_values))
257 np_dtype = np.dtype(dtype_text)
258 array = np.frombuffer(embedding_bytes, dtype=np_dtype)
259 assert array.size == expected_size, f"Shape mismatch while reading sequence: {sequence}"
260 reshaped = array.copy().reshape(tuple(shape_values))
261 loaded[sequence] = torch.from_numpy(reshaped)
262 return loaded
263
264 def embed_dataset(
265 self,
266 sequences: Optional[List[str]] = None,
267 tokenizer: Optional[PreTrainedTokenizerBase] = None,
268 batch_size: int = 2,
269 max_len: int = 512,
270 truncate: bool = True,
271 full_embeddings: bool = False,
272 embed_dtype: torch.dtype = torch.float32,
273 pooling_types: List[str] = ['mean'],
274 num_workers: int = 0,
275 sql: bool = False,
276 save: bool = True,
277 sql_db_path: str = 'embeddings.db',
278 save_path: str = 'embeddings.pth',
279 fasta_path: Optional[str] = None,
280 **kwargs,
281 ) -> Optional[Dict[str, torch.Tensor]]:
282 """
283 Embed a dataset of protein sequences.
284
285 Supports two modes:
286 - Tokenizer mode (ESM2/ESM++): provide `tokenizer`, `_embed(input_ids, attention_mask)` is used.
287 - Sequence mode (E1): pass `tokenizer=None`, `_embed(sequences, return_attention_mask=True, **kwargs)` is used.
288
289 Sequences can be supplied as a list via `sequences`, parsed from a FASTA file via
290 `fasta_path`, or both (the two sources are combined). At least one must be provided.
291 """
292 if fasta_path is not None:
293 fasta_sequences = parse_fasta(fasta_path)
294 sequences = list(sequences or []) + fasta_sequences
295 assert sequences is not None and len(sequences) > 0, \
296 "Must provide at least one sequence via `sequences` or `fasta_path`."
297 sequences = list(set([seq[:max_len] if truncate else seq for seq in sequences]))
298 sequences = sorted(sequences, key=len, reverse=True)
299 hidden_size = self.config.hidden_size
300 pooler = Pooler(pooling_types) if not full_embeddings else None
301 tokenizer_mode = tokenizer is not None
302 if tokenizer_mode:
303 collate_fn = build_collator(tokenizer)
304 device = self.device
305 else:
306 collate_fn = None
307 device = None
308
309 def get_embeddings(residue_embeddings: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
310 assert isinstance(residue_embeddings, torch.Tensor)
311 if full_embeddings or residue_embeddings.ndim == 2:
312 return residue_embeddings
313 return pooler(residue_embeddings, attention_mask)
314
315 def iter_batches(to_embed: List[str]):
316 if tokenizer_mode:
317 assert collate_fn is not None
318 assert device is not None
319 dataset = ProteinDataset(to_embed)
320 dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, collate_fn=collate_fn, shuffle=False)
321 for i, batch in tqdm(enumerate(dataloader), total=len(dataloader), desc='Embedding batches'):
322 seqs = to_embed[i * batch_size:(i + 1) * batch_size]
323 input_ids = batch['input_ids'].to(device)
324 attention_mask = batch['attention_mask'].to(device)
325 residue_embeddings = self._embed(input_ids, attention_mask)
326 yield seqs, residue_embeddings, attention_mask
327 else:
328 for batch_start in tqdm(range(0, len(to_embed), batch_size), desc='Embedding batches'):
329 seqs = to_embed[batch_start:batch_start + batch_size]
330 batch_output = self._embed(seqs, return_attention_mask=True, **kwargs)
331 assert isinstance(batch_output, tuple), "Sequence mode _embed must return (last_hidden_state, attention_mask)."
332 assert len(batch_output) == 2, "Sequence mode _embed must return exactly two values."
333 residue_embeddings, attention_mask = batch_output
334 assert isinstance(attention_mask, torch.Tensor), "Sequence mode _embed must return attention_mask as a torch.Tensor."
335 yield seqs, residue_embeddings, attention_mask
336
337 if sql:
338 conn = sqlite3.connect(sql_db_path)
339 self._ensure_embeddings_table(conn)
340 c = conn.cursor()
341 already_embedded = self._read_sequences_from_db(sql_db_path)
342 to_embed = [seq for seq in sequences if seq not in already_embedded]
343 print(f"Found {len(already_embedded)} already embedded sequences in {sql_db_path}")
344 print(f"Embedding {len(to_embed)} new sequences")
345 if len(to_embed) > 0:
346 with torch.no_grad():
347 for i, (seqs, residue_embeddings, attention_mask) in enumerate(iter_batches(to_embed)):
348 embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype)
349 for seq, emb, mask in zip(seqs, embeddings, attention_mask):
350 if full_embeddings:
351 emb = emb[mask.bool()].reshape(-1, hidden_size)
352 emb_np = emb.cpu().numpy()
353 emb_shape = ",".join([str(dim) for dim in emb_np.shape])
354 emb_dtype = str(emb_np.dtype)
355 c.execute(
356 "INSERT OR REPLACE INTO embeddings (sequence, embedding, shape, dtype) VALUES (?, ?, ?, ?)",
357 (seq, emb_np.tobytes(), emb_shape, emb_dtype),
358 )
359 if tokenizer_mode and (i + 1) % 100 == 0:
360 conn.commit()
361 conn.commit()
362 conn.close()
363 return None
364
365 embeddings_dict = {}
366 if os.path.exists(save_path):
367 embeddings_dict = self.load_embeddings_from_pth(save_path)
368 to_embed = [seq for seq in sequences if seq not in embeddings_dict]
369 print(f"Found {len(embeddings_dict)} already embedded sequences in {save_path}")
370 print(f"Embedding {len(to_embed)} new sequences")
371 else:
372 to_embed = sequences
373 print(f"Embedding {len(to_embed)} new sequences")
374
375 if len(to_embed) > 0:
376 with torch.no_grad():
377 for seqs, residue_embeddings, attention_mask in iter_batches(to_embed):
378 embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype)
379 for seq, emb, mask in zip(seqs, embeddings, attention_mask):
380 if full_embeddings:
381 emb = emb[mask.bool()].reshape(-1, hidden_size)
382 embeddings_dict[seq] = emb.cpu()
383
384 if save:
385 torch.save(embeddings_dict, save_path)
386
387 return embeddings_dict
388
389
390if __name__ == "__main__":
391 # py -m pooler
392 pooler = Pooler(pooling_types=['max', 'parti'])
393 batch_size = 8
394 seq_len = 64
395 hidden_size = 128
396 num_layers = 12
397 emb = torch.randn(batch_size, seq_len, hidden_size)
398 attentions = torch.randn(batch_size, num_layers, seq_len, seq_len)
399 attention_mask = torch.ones(batch_size, seq_len)
400 y = pooler(emb=emb, attention_mask=attention_mask, attentions=attentions)
401 print(y.shape)
402