File size: 3,205 Bytes
4c6c214 | 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 | 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):
# residue-only mean: drop CLS (position 0) and EOS (last non-pad
# position of each sequence) from the mask before averaging.
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]
|