| import torch |
| from transformers import AutoTokenizer, AutoModel |
|
|
|
|
| def batch_sequences_by_length(sequences, max_tokens=8000): |
| """Groups sequences into batches whose padded token budget (n_seqs x |
| max_len_in_batch) stays under max_tokens. Sequences are sorted |
| longest-first so the first entry in each batch sets the padding length.""" |
| lengths = torch.tensor([len(s) for s in sequences]) |
| sort_idx = torch.argsort(lengths, descending=True) |
|
|
| sorted_seqs = [sequences[i] for i in sort_idx] |
| sorted_lengths = lengths[sort_idx] |
|
|
| batches, reorder_map = [], [] |
| current_batch, current_indices, batch_max_len = [], [], 0 |
|
|
| for seq, length, orig_idx in zip(sorted_seqs, sorted_lengths, sort_idx): |
| l_val = length.item() |
| new_max = max(batch_max_len, l_val) |
| if (len(current_batch) + 1) * new_max <= max_tokens: |
| current_batch.append(seq) |
| current_indices.append(orig_idx.item()) |
| batch_max_len = new_max |
| else: |
| if current_batch: |
| batches.append(current_batch) |
| reorder_map.extend(current_indices) |
| current_batch = [seq] |
| current_indices = [orig_idx.item()] |
| batch_max_len = l_val |
|
|
| if current_batch: |
| batches.append(current_batch) |
| reorder_map.extend(current_indices) |
|
|
| restore_indices = torch.argsort(torch.tensor(reorder_map)) |
| return batches, restore_indices |
|
|
|
|
| class ESM2Embedder: |
| def __init__(self, model_name="facebook/esm2_t12_35M_UR50D", device="cpu", max_aa=8000): |
| self.model_name = model_name |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
| self.model = AutoModel.from_pretrained(model_name).to(device) |
| self.model.eval() |
| self.device = device |
| self.max_aa = max_aa |
|
|
| def _pool(self, last_hidden_state, attention_mask): |
| |
| |
| mask = attention_mask.clone() |
| mask[:, 0] = 0 |
| eos_pos = attention_mask.sum(dim=1) - 1 |
| mask[torch.arange(mask.size(0)), eos_pos] = 0 |
| mask = mask.unsqueeze(-1).float() |
| return (last_hidden_state * mask).sum(1) / mask.sum(1) |
|
|
| def _embed_batch(self, batch): |
| if len(batch) == 0: |
| return [] |
| inputs = self.tokenizer( |
| batch, return_tensors="pt", padding=True, truncation=False, |
| ).to(self.device) |
| with torch.no_grad(): |
| outputs = self.model(**inputs) |
| emb = self._pool(outputs.last_hidden_state, inputs["attention_mask"]).float().cpu() |
| return [emb] |
|
|
| def get_embeddings(self, sequences, max_aa=None): |
| limit = max_aa if max_aa is not None else self.max_aa |
| if len(sequences) == 1: |
| batches = [sequences] |
| restore_idx = torch.arange(1) |
| else: |
| batches, restore_idx = batch_sequences_by_length(sequences, limit) |
|
|
| all_results = [] |
| for batch in batches: |
| all_results.extend(self._embed_batch(batch)) |
|
|
| combined = torch.cat(all_results, dim=0) |
| return combined[restore_idx] |
|
|
