| import numpy as np |
| import torch |
|
|
| from extra_utils import res_to_seq, get_sequences_from_anarci |
|
|
|
|
| class AbRestore: |
| def __init__(self, spread = 11, device = 'cpu', ncpu = 1): |
| self.spread = spread |
| self.device = device |
| self.ncpu = ncpu |
| |
| def _initiate_abrestore(self, model, tokenizer): |
| self.AbLang = model |
| self.tokenizer = tokenizer |
|
|
| def restore(self, seqs, align = False, **kwargs): |
| """ |
| Restore sequences |
| """ |
| n_seqs = len(seqs) |
| |
| if align: |
| |
| seqs = self._sequence_aligning(seqs) |
| nr_seqs = len(seqs)//self.spread |
| |
| tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device) |
| predictions = self.AbLang(tokens)[:,:,1:21] |
|
|
| |
| tokens = tokens.reshape(nr_seqs, self.spread, -1) |
| predictions = predictions.reshape(nr_seqs, self.spread, -1, 20) |
| seqs = seqs.reshape(nr_seqs, -1) |
|
|
| |
| best_seq_idx = torch.argmax(torch.max(predictions, -1).values[:,:,1:2].mean(2), -1) |
|
|
| |
| tokens = tokens.gather(1, best_seq_idx.view(-1, 1).unsqueeze(1).repeat(1, 1, tokens.shape[-1])).squeeze(1) |
| predictions = predictions[range(predictions.shape[0]), best_seq_idx] |
| seqs = np.take_along_axis(seqs, best_seq_idx.view(-1, 1).cpu().numpy(), axis=1) |
|
|
| else: |
| tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device) |
| predictions = self.AbLang(tokens)[:,:,1:21] |
|
|
| predicted_tokens = torch.max(predictions, -1).indices + 1 |
| restored_tokens = torch.where(tokens==23, predicted_tokens, tokens) |
|
|
| restored_seqs = self.tokenizer(restored_tokens, mode="decode") |
|
|
| if n_seqs < len(restored_seqs): |
| restored_seqs = [f"{h}|{l}".replace('-','') for h,l in zip(restored_seqs[:n_seqs], restored_seqs[n_seqs:])] |
| seqs = [f"{h}|{l}" for h,l in zip(seqs[:n_seqs], seqs[n_seqs:])] |
| |
| return np.array([res_to_seq(seq, 'restore') for seq in np.c_[restored_seqs, np.vectorize(len)(seqs)]]) |
| |
| def _create_spread_of_sequences(self, seqs, chain = 'H'): |
| import pandas as pd |
| import anarci |
| |
| chain_idx = 0 if chain == 'H' else 1 |
| numbered_seqs = anarci.run_anarci( |
| pd.DataFrame([seq[chain_idx].replace('*', 'X') for seq in seqs]).reset_index().values.tolist(), |
| ncpu=self.ncpu, |
| scheme='imgt', |
| allowed_species=['human', 'mouse'], |
| ) |
| |
| anarci_data = pd.DataFrame( |
| [str(anarci[0][0]) if anarci else 'ANARCI_error' for anarci in numbered_seqs[1]], |
| columns=['anarci'] |
| ).astype('<U90') |
| |
| max_position = 128 if chain == 'H' else 127 |
| |
| seqs = anarci_data.apply( |
| lambda x: get_sequences_from_anarci( |
| x.anarci, |
| max_position, |
| self.spread |
| ), axis=1, result_type='expand' |
| ).to_numpy().reshape(-1) |
| |
| return seqs |
| |
| |
| def _sequence_aligning(self, seqs): |
|
|
| tmp_seqs = [pairs.replace(">", "").replace("<", "").split("|") for pairs in seqs] |
| |
| spread_heavy = [f"<{seq}>" for seq in self._create_spread_of_sequences(tmp_seqs, chain = 'H')] |
| spread_light = [f"<{seq}>" for seq in self._create_spread_of_sequences(tmp_seqs, chain = 'L')] |
| |
| return np.concatenate([np.array(spread_heavy),np.array(spread_light)]) |
