Remove ablang2 folder - repository now fully self-contained

ablang2/__init__.py

@@ -1 +0,0 @@
-from .pretrained import pretrained

ablang2/adapter.py

@@ -1,306 +0,0 @@
-from ablang2.pretrained_utils.restoration import AbRestore
-from ablang2.pretrained_utils.encodings import AbEncoding
-from ablang2.pretrained_utils.alignment import AbAlignment
-from ablang2.pretrained_utils.scores import AbScores
-import torch
-import numpy as np
-from ablang2.pretrained_utils.extra_utils import res_to_seq, res_to_list
-
-class HuggingFaceTokenizerAdapter:
-    def __init__(self, tokenizer, device):
-        self.tokenizer = tokenizer
-        self.device = device
-        self.pad_token_id = tokenizer.pad_token_id
-        self.mask_token_id = getattr(tokenizer, 'mask_token_id', None) or tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
-        self.vocab = tokenizer.get_vocab() if hasattr(tokenizer, 'get_vocab') else tokenizer.vocab
-        self.inv_vocab = {v: k for k, v in self.vocab.items()}
-        self.all_special_tokens = tokenizer.all_special_tokens
-
-    def __call__(self, seqs, pad=True, w_extra_tkns=False, device=None, mode=None):
-        tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
-        input_ids = tokens['input_ids'].to(self.device if device is None else device)
-        if mode == 'decode':
-            # seqs is a tensor of token ids
-            if isinstance(seqs, torch.Tensor):
-                seqs = seqs.cpu().numpy()
-            decoded = []
-            for i, seq in enumerate(seqs):
-                chars = [self.inv_vocab.get(int(t), '') for t in seq if self.inv_vocab.get(int(t), '') not in {'-', '*', '<', '>'} and self.inv_vocab.get(int(t), '') != '']
-                # Use res_to_seq for formatting, pass (sequence, length) tuple as in original code
-                # The length is not always available, so use len(chars) as fallback
-                formatted = res_to_seq([ ''.join(chars), len(chars) ], mode='restore')
-                decoded.append(formatted)
-            return decoded
-        return input_ids
-
-class HFAbRestore(AbRestore):
-    def __init__(self, hf_model, hf_tokenizer, spread=11, device='cpu', ncpu=1):
-        super().__init__(spread=spread, device=device, ncpu=ncpu)
-        self.used_device = device
-        self._hf_model = hf_model
-        self.tokenizer = HuggingFaceTokenizerAdapter(hf_tokenizer, device)
-
-    @property
-    def AbLang(self):
-        def model_call(x):
-            output = self._hf_model(x)
-            if hasattr(output, 'last_hidden_state'):
-                return output.last_hidden_state
-            return output
-        return model_call
-
-def add_angle_brackets(seq):
-    # Assumes input is 'VH|VL' or 'VH|' or '|VL'
-    if '|' in seq:
-        vh, vl = seq.split('|', 1)
-    else:
-        vh, vl = seq, ''
-    return f"<{vh}>|<{vl}>"
-
-class AbLang2PairedHuggingFaceAdapter(AbEncoding, AbRestore, AbAlignment, AbScores):
-    """
-    Adapter to use pretrained utilities with a HuggingFace-loaded ablang2_paired model and tokenizer.
-    Automatically uses CUDA if available, otherwise CPU.
-    """
-    def __init__(self, model, tokenizer, device=None, ncpu=1):
-        super().__init__()
-        if device is None:
-            self.used_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        else:
-            self.used_device = torch.device(device)
-        self.AbLang = model  # HuggingFace model instance
-        self.tokenizer = tokenizer
-        self.AbLang.to(self.used_device)
-        self.AbLang.eval()
-        # Always get AbRep from the underlying model
-        if hasattr(self.AbLang, 'model') and hasattr(self.AbLang.model, 'AbRep'):
-            self.AbRep = self.AbLang.model.AbRep
-        else:
-            raise AttributeError("Could not find AbRep in the HuggingFace model or its underlying model.")
-        self.ncpu = ncpu
-        self.spread = 11  # For compatibility with original utilities
-        # The following is no longer needed since all_special_tokens now returns IDs directly
-        # self.tokenizer.all_special_token_ids = [
-        #     self.tokenizer.convert_tokens_to_ids(tok) for tok in self.tokenizer.all_special_tokens
-        # ]
-        # self.tokenizer._all_special_tokens_str = self.tokenizer.all_special_tokens
-        # self.tokenizer.all_special_tokens = [
-        #     self.tokenizer.convert_tokens_to_ids(tok) for tok in self.tokenizer._all_special_tokens_str
-        # ]
-
-    def freeze(self):
-        self.AbLang.eval()
-
-    def unfreeze(self):
-        self.AbLang.train()
-
-    def _encode_sequences(self, seqs):
-        # Use HuggingFace-style padding and return PyTorch tensors
-        tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
-        tokens = extract_input_ids(tokens, self.used_device)
-        return self.AbRep(tokens).last_hidden_states.detach()
-
-    def _predict_logits(self, seqs):
-        tokens = self.tokenizer(seqs, padding=True, return_tensors='pt')
-        tokens = extract_input_ids(tokens, self.used_device)
-        output = self.AbLang(tokens)
-        if hasattr(output, 'last_hidden_state'):
-            return output.last_hidden_state.detach()
-        return output.detach()
-
-    def _preprocess_labels(self, labels):
-        labels = extract_input_ids(labels, self.used_device)
-        return labels
-
-    def __call__(self, seqs, mode='seqcoding', align=False, stepwise_masking=False, fragmented=False, batch_size=50):
-        """
-        Use different modes for different usecases, mimicking the original pretrained class.
-        """
-        from ablang2.pretrained import format_seq_input
-
-        valid_modes = [
-            'rescoding', 'seqcoding', 'restore', 'likelihood', 'probability',
-            'pseudo_log_likelihood', 'confidence'
-        ]
-        if mode not in valid_modes:
-            raise SyntaxError(f"Given mode doesn't exist. Please select one of the following: {valid_modes}.")
-
-        seqs, chain = format_seq_input(seqs, fragmented=fragmented)
-
-        if align:
-            numbered_seqs, seqs, number_alignment = self.number_sequences(
-                seqs, chain=chain, fragmented=fragmented
-            )
-        else:
-            numbered_seqs = None
-            number_alignment = None
-
-        subset_list = []
-        for subset in [seqs[x:x+batch_size] for x in range(0, len(seqs), batch_size)]:
-            subset_list.append(getattr(self, mode)(subset, align=align, stepwise_masking=stepwise_masking))
-
-        return self.reformat_subsets(
-            subset_list,
-            mode=mode,
-            align=align,
-            numbered_seqs=numbered_seqs,
-            seqs=seqs,
-            number_alignment=number_alignment,
-        )
-
-    def pseudo_log_likelihood(self, seqs, **kwargs):
-        """
-        Original (non-vectorized) pseudo log-likelihood computation matching notebook behavior.
-        """
-        # Format input: join VH and VL with '|'
-        formatted_seqs = []
-        for s in seqs:
-            if isinstance(s, (list, tuple)):
-                formatted_seqs.append('|'.join(s))
-            else:
-                formatted_seqs.append(s)
-
-        # Tokenize all sequences in batch
-        labels = self.tokenizer(
-            formatted_seqs, padding=True, return_tensors='pt'
-        )
-        labels = extract_input_ids(labels, self.used_device)
-
-        # Convert special tokens to IDs
-        if isinstance(self.tokenizer.all_special_tokens[0], int):
-            special_token_ids = set(self.tokenizer.all_special_tokens)
-        else:
-            special_token_ids = set(self.tokenizer.convert_tokens_to_ids(tok) for tok in self.tokenizer.all_special_tokens)
-        pad_token_id = self.tokenizer.pad_token_id
-
-        mask_token_id = getattr(self.tokenizer, 'mask_token_id', None)
-        if mask_token_id is None:
-            mask_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)
-
-        plls = []
-        with torch.no_grad():
-            for i, seq_label in enumerate(labels):
-                seq_pll = []
-                for j, token_id in enumerate(seq_label):
-                    if token_id.item() in special_token_ids or token_id.item() == pad_token_id:
-                        continue
-                    masked = seq_label.clone()
-                    masked[j] = mask_token_id
-                    logits = self.AbLang(masked.unsqueeze(0))
-                    if hasattr(logits, 'last_hidden_state'):
-                        logits = logits.last_hidden_state
-                    logits = logits[0, j]
-                    nll = torch.nn.functional.cross_entropy(
-                        logits.unsqueeze(0), token_id.unsqueeze(0), reduction="none"
-                    )
-                    seq_pll.append(-nll.item())
-                if seq_pll:
-                    plls.append(np.mean(seq_pll))
-                else:
-                    plls.append(float('nan'))
-        return np.array(plls)
-
-    def confidence(self, seqs, **kwargs):
-        """Confidence calculation - match original ablang2 implementation by excluding all special tokens from loss."""
-        # Format input: join VH and VL with '|'
-        formatted_seqs = []
-        for s in seqs:
-            if isinstance(s, (list, tuple)):
-                formatted_seqs.append('|'.join(s))
-            else:
-                formatted_seqs.append(s)
-        
-        plls = []
-        for seq in formatted_seqs:
-            tokens = self.tokenizer([seq], padding=True, return_tensors='pt')
-            input_ids = extract_input_ids(tokens, self.used_device)
-            
-            with torch.no_grad():
-                output = self.AbLang(input_ids)
-                if hasattr(output, 'last_hidden_state'):
-                    logits = output.last_hidden_state
-                else:
-                    logits = output
-                
-                # Get the sequence (remove batch dimension)
-                logits = logits[0]  # [seq_len, vocab_size]
-                input_ids = input_ids[0]  # [seq_len]
-                
-                # Exclude all special tokens (pad, mask, etc.)
-                if isinstance(self.tokenizer.all_special_tokens[0], int):
-                    special_token_ids = set(self.tokenizer.all_special_tokens)
-                else:
-                    special_token_ids = set(self.tokenizer.convert_tokens_to_ids(tok) for tok in self.tokenizer.all_special_tokens)
-                valid_mask = ~torch.isin(input_ids, torch.tensor(list(special_token_ids), device=input_ids.device))
-                
-                if valid_mask.sum() > 0:
-                    valid_logits = logits[valid_mask]
-                    valid_labels = input_ids[valid_mask]
-                    
-                    # Calculate cross-entropy loss
-                    nll = torch.nn.functional.cross_entropy(
-                        valid_logits,
-                        valid_labels,
-                        reduction="mean"
-                    )
-                    pll = -nll.item()
-                else:
-                    pll = 0.0
-                
-                plls.append(pll)
-        
-        return np.array(plls, dtype=np.float32)
-
-    def probability(self, seqs, align=False, stepwise_masking=False, **kwargs):
-        """
-        Probability of mutations - applies softmax to logits to get probabilities
-        """
-        # Format input: join VH and VL with '|'
-        formatted_seqs = []
-        for s in seqs:
-            if isinstance(s, (list, tuple)):
-                formatted_seqs.append('|'.join(s))
-            else:
-                formatted_seqs.append(s)
-
-        # Get logits
-        if stepwise_masking:
-            # For stepwise masking, we need to implement it similar to likelihood
-            # This is a simplified version - you might want to implement full stepwise masking
-            logits = self._predict_logits(formatted_seqs)
-        else:
-            logits = self._predict_logits(formatted_seqs)
-        
-        # Apply softmax to get probabilities
-        probs = logits.softmax(-1).cpu().numpy()
-        
-        if align:
-            return probs
-        else:
-            # Return residue-level probabilities (excluding special tokens)
-            return [res_to_list(state, seq) for state, seq in zip(probs, formatted_seqs)]
-
-    def restore(self, seqs, align=False, **kwargs):
-        hf_abrestore = HFAbRestore(self.AbLang, self.tokenizer, spread=self.spread, device=self.used_device, ncpu=self.ncpu)
-        restored = hf_abrestore.restore(seqs, align=align)
-        # Apply angle brackets formatting
-        if isinstance(restored, np.ndarray):
-            restored = np.array([add_angle_brackets(seq) for seq in restored])
-        else:
-            restored = [add_angle_brackets(seq) for seq in restored]
-        return restored
-
-def extract_input_ids(tokens, device):
-    if hasattr(tokens, 'input_ids'):
-        return tokens.input_ids.to(device)
-    elif isinstance(tokens, dict):
-        if 'input_ids' in tokens:
-            return tokens['input_ids'].to(device)
-        else:
-            for v in tokens.values():
-                if hasattr(v, 'ndim') or torch.is_tensor(v):
-                    return v.to(device)
-    elif torch.is_tensor(tokens):
-        return tokens.to(device)
-    else:
-        raise ValueError("Could not extract input_ids from tokenizer output")

ablang2/alignment.py

@@ -1,87 +0,0 @@
-from dataclasses import dataclass
-import numpy as np
-import torch
-
-from .extra_utils import paired_msa_numbering, unpaired_msa_numbering, create_alignment
-
-
-class AbAlignment:
-
-    def __init__(self, device = 'cpu', ncpu = 1):
-        
-        self.device = device
-        self.ncpu = ncpu
-        
-    def number_sequences(self, seqs, chain = 'H', fragmented = False):
-        if chain == 'HL':
-            numbered_seqs, seqs, number_alignment = paired_msa_numbering(seqs, fragmented = fragmented, n_jobs = self.ncpu)
-        else:
-            assert chain == 'HL', 'Currently "Align==True" only works for paired sequences. \nPlease use paired sequences or Align=False.'
-            numbered_seqs, seqs, number_alignment = unpaired_msa_numbering(
-                seqs, chain = chain, fragmented = fragmented, n_jobs = self.ncpu
-            )
-        
-        return numbered_seqs, seqs, number_alignment
-    
-    def align_encodings(self, encodings, numbered_seqs, seqs, number_alignment):
-        
-        aligned_encodings = np.concatenate(
-            [[
-                create_alignment(
-                    res_embed, numbered_seq, seq, number_alignment
-                ) for res_embed, numbered_seq, seq in zip(encodings, numbered_seqs, seqs)
-            ]], axis=0
-        )
-        return aligned_encodings
-        
-        
-    def reformat_subsets(
-        self, 
-        subset_list, 
-        mode = 'seqcoding', 
-        align = False,
-        numbered_seqs = None, 
-        seqs = None,
-        number_alignment = None,
-    ):
-        
-        if mode in [
-            'seqcoding', 
-            'restore',
-            'pseudo_log_likelihood',
-            'confidence'
-        ]:
-            return np.concatenate(subset_list)
-        elif align:
-            subset_list = [
-                self.align_encodings(
-                    subset, 
-                    numbered_seqs[num*len(subset):(num+1)*len(subset)],
-                    seqs[num*len(subset):(num+1)*len(subset)], 
-                    number_alignment
-                ) for num, subset in enumerate(subset_list)
-            ]       
-            
-            subset = np.concatenate(subset_list)
-            
-            return aligned_results(
-                aligned_seqs = [''.join(alist) for alist in subset[:,:,-1]],
-                aligned_embeds = subset[:,:,:-1].astype(float),
-                number_alignment=number_alignment.apply(lambda x: '{}{}'.format(*x[0]), axis=1).values
-            ) 
-    
-        elif not align:
-            return sum(subset_list, [])
-        else:
-            return np.concatenate(subset_list) # this needs to be changed
-        
-
-@dataclass
-class aligned_results():
-    """
-    Dataclass used to store output.
-    """
-    
-    aligned_seqs: None
-    aligned_embeds: None
-    number_alignment: None

ablang2/config.json

@@ -1,18 +0,0 @@
-{
-  "model_type": "ablang2-paired",
-  "vocab_size": 26,
-  "hidden_embed_size": 480,
-  "n_attn_heads": 20,
-  "n_encoder_blocks": 12,
-  "padding_tkn": 21,
-  "mask_tkn": 23,
-  "layer_norm_eps": 1e-12,
-  "a_fn": "swiglu",
-  "dropout": 0.0,
-  "tokenizer_class": "AbLang2PairedTokenizer",
-  "auto_map": {
-    "AutoConfig": "configuration_ablang2paired.AbLang2PairedConfig",
-    "AutoModel": "modeling_ablang2paired.AbLang2PairedHFModel",
-    "AutoTokenizer": ["tokenizer_ablang2paired.AbLang2PairedTokenizer", "tokenizer_ablang2paired.AbLang2PairedTokenizer"]
-  }
-}

ablang2/configuration_ablang2paired.py

@@ -1,31 +0,0 @@
-from transformers import PretrainedConfig
-
-class AbLang2PairedConfig(PretrainedConfig):
-    model_type = "ablang2-paired"
-
-    def __init__(
-        self,
-        vocab_size=26,
-        hidden_embed_size=480,
-        n_attn_heads=20,
-        n_encoder_blocks=12,
-        padding_tkn=21,
-        mask_tkn=23,
-        layer_norm_eps=1e-12,
-        a_fn="swiglu",
-        dropout=0.0,
-        **kwargs
-    ):
-        super().__init__(**kwargs)
-        self.vocab_size = vocab_size
-        self.hidden_embed_size = hidden_embed_size
-        self.hidden_size = hidden_embed_size  # Add this for Hugging Face compatibility
-        self.n_attn_heads = n_attn_heads
-        self.num_attention_heads = n_attn_heads  # Add this for Hugging Face compatibility
-        self.num_hidden_layers = n_encoder_blocks  # Add this for Hugging Face compatibility
-        self.n_encoder_blocks = n_encoder_blocks
-        self.padding_tkn = padding_tkn
-        self.mask_tkn = mask_tkn
-        self.layer_norm_eps = layer_norm_eps
-        self.a_fn = a_fn
-        self.dropout = dropout 

ablang2/encodings.py

@@ -1,97 +0,0 @@
-import numpy as np
-import torch
-
-from .extra_utils import res_to_list, res_to_seq
-
-
-class AbEncoding:
-
-    def __init__(self, device = 'cpu', ncpu = 1):
-        
-        self.device = device
-        self.ncpu = ncpu
-        
-    def _initiate_abencoding(self, model, tokenizer):
-        self.AbLang = model
-        self.tokenizer = tokenizer
-        
-    def _encode_sequences(self, seqs):
-        tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
-        with torch.no_grad():
-            return self.AbLang.AbRep(tokens).last_hidden_states
-        
-    def _predict_logits(self, seqs):
-        tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
-        with torch.no_grad():
-            return self.AbLang(tokens)
-        
-    def _predict_logits_with_step_masking(self, seqs):
-        
-        tokens = self.tokenizer(seqs, pad=True, w_extra_tkns=False, device=self.used_device)
-        
-        logits = []
-        for single_seq_tokens in tokens:
-            
-            tkn_len = len(single_seq_tokens)
-            masked_tokens = single_seq_tokens.repeat(tkn_len, 1)
-            for num in range(tkn_len):
-                masked_tokens[num, num] = self.tokenizer.mask_token
-            
-            with torch.no_grad():
-                logits_tmp = self.AbLang(masked_tokens)
-                       
-            logits_tmp = torch.stack([logits_tmp[num, num] for num in range(tkn_len)])
-
-            logits.append(logits_tmp)
-    
-        return torch.stack(logits, dim=0)        
-
-    def seqcoding(self, seqs, **kwargs):
-        """
-        Sequence specific representations
-        """
-        
-        encodings = self._encode_sequences(seqs).cpu().numpy()
-        
-        lens = np.vectorize(len)(seqs)
-        lens = np.tile(lens.reshape(-1,1,1), (encodings.shape[2], 1))
-
-        return np.apply_along_axis(res_to_seq, 2, np.c_[np.swapaxes(encodings,1,2), lens])
-        
-    def rescoding(self, seqs, align=False, **kwargs):
-        """
-        Residue specific representations.
-        """
-        encodings = self._encode_sequences(seqs).cpu().numpy()
-           
-        if align: return encodings
-            
-        else: return [res_to_list(state, seq) for state, seq in zip(encodings, seqs)]
-        
-    def likelihood(self, seqs, align=False, stepwise_masking=False, **kwargs):
-        """
-        Likelihood of mutations
-        """
-        if stepwise_masking:
-            logits = self._predict_logits_with_step_masking(seqs).cpu().numpy()
-        else:
-            logits = self._predict_logits(seqs).cpu().numpy()
-        
-        if align: return logits
-            
-        else: return [res_to_list(state, seq) for state, seq in zip(logits, seqs)]
-        
-    def probability(self, seqs, align=False, stepwise_masking=False, **kwargs):
-        """
-        Probability of mutations
-        """
-        if stepwise_masking:
-            logits = self._predict_logits_with_step_masking(seqs)
-        else:
-            logits = self._predict_logits(seqs)
-        probs = logits.softmax(-1).cpu().numpy()
-        
-        if align: return probs
-            
-        else: return [res_to_list(state, seq) for state, seq in zip(probs, seqs)]
-        

ablang2/environment.yaml

@@ -1,44 +0,0 @@
-name: AbLang
-channels:
-  - conda-forge
-  - pytorch
-  - bioconda
-  - defaults
-dependencies:
-  - python=3.10.18
-  - pip
-  - pytorch=2.5.1
-  - pytorch-cuda=12.4
-  - numpy=2.2.6
-  - pandas=2.3.1
-  - transformers=4.53.3
-  - anarci=2024.05.21
-  - jupyter=7.4.4
-  - notebook=7.4.4
-  - ipython=8.37.0
-  - ipykernel=6.29.5
-  - matplotlib-inline=0.1.7
-  - scikit-learn
-  - matplotlib
-  - seaborn
-  - biopython=1.85
-  - huggingface_hub=0.33.4
-  - tokenizers=0.21.3
-  - safetensors=0.5.3
-  - einops=0.8.1
-  - tqdm=4.67.1
-  - requests=2.32.4
-  - urllib3=2.5.0
-  - certifi=2025.7.14
-  - filelock=3.18.0
-  - fsspec=2025.3.0
-  - packaging=25.0
-  - regex=2024.11.6
-  - sympy=1.13.3
-  - networkx=3.4.2
-  - jinja2=3.1.6
-  - pyyaml=6.0.2
-  - typing_extensions=4.14.1
-  - pip:
-    - numba=0.61.2
-    - llvmlite=0.44.0 

ablang2/extra_utils.py

@@ -1,165 +0,0 @@
-import string, re
-import numpy as np
-
-
-def res_to_list(logits, seq):
-    return logits[:len(seq)]
-
-def res_to_seq(a, mode='mean'):
-    """
-    Function for how we go from n_values for each amino acid to n_values for each sequence.
-    
-    We leave out padding tokens.
-    """
-    
-    if mode=='sum':
-        return a[0:(int(a[-1]))].sum()
-    
-    elif mode=='mean':
-        return a[0:(int(a[-1]))].mean()
-    
-    elif mode=='restore':
-        return a[0][0:(int(a[-1]))]
-
-def get_number_alignment(numbered_seqs):
-    """
-    Creates a number alignment from the anarci results.
-    """
-    import pandas as pd
-
-    alist = [pd.DataFrame(aligned_seq, columns = [0,1,'resi']) for aligned_seq in numbered_seqs]
-    unsorted_alignment = pd.concat(alist).drop_duplicates(subset=0)
-    max_alignment = get_max_alignment()    
-    
-    return max_alignment.merge(unsorted_alignment.query("resi!='-'"), left_on=0, right_on=0)[[0,1]]
-
-def get_max_alignment():
-    """
-    Create maximum possible alignment for sorting
-    """
-    import pandas as pd
-
-    sortlist = [[("<", "")]]
-    for num in range(1, 128+1):
-        if num in [33,61,112]:
-            for char in string.ascii_uppercase[::-1]:
-                sortlist.append([(num, char)])
-
-            sortlist.append([(num,' ')])
-        else:
-            sortlist.append([(num,' ')])
-            for char in string.ascii_uppercase:
-                sortlist.append([(num, char)])
-                
-    return pd.DataFrame(sortlist + [[(">", "")]])
-
-
-def paired_msa_numbering(ab_seqs, fragmented = False, n_jobs = 10):
-    
-    import pandas as pd
-    
-    tmp_seqs = [pairs.replace(">", "").replace("<", "").split("|") for pairs in ab_seqs]
-
-    numbered_seqs_heavy, seqs_heavy, number_alignment_heavy = unpaired_msa_numbering(
-        [i[0] for i in tmp_seqs], 'H', fragmented = fragmented, n_jobs = n_jobs
-    )
-    numbered_seqs_light, seqs_light, number_alignment_light = unpaired_msa_numbering(
-        [i[1] for i in tmp_seqs], 'L', fragmented = fragmented, n_jobs = n_jobs
-    )
-    
-    number_alignment = pd.concat([
-        number_alignment_heavy, 
-        pd.DataFrame([[("|",""), "|"]]), 
-        number_alignment_light]
-    ).reset_index(drop=True)
-    
-    seqs = [f"{heavy}|{light}" for heavy, light in zip(seqs_heavy, seqs_light)]
-    numbered_seqs = [
-        heavy + [(("|",""), "|", "|")] + light for heavy, light in zip(numbered_seqs_heavy, numbered_seqs_light)
-    ]
-    
-    return numbered_seqs, seqs, number_alignment
-
-
-def unpaired_msa_numbering(seqs, chain = 'H', fragmented = False, n_jobs = 10):
-    
-    numbered_seqs = number_with_anarci(seqs, chain = chain, fragmented = fragmented, n_jobs = n_jobs)
-    number_alignment = get_number_alignment(numbered_seqs)
-    number_alignment[1] = chain   
-    
-    seqs = [''.join([i[2] for i in numbered_seq]).replace('-','') for numbered_seq in numbered_seqs]    
-    return numbered_seqs, seqs, number_alignment
-
-
-def number_with_anarci(seqs, chain = 'H', fragmented = False, n_jobs = 1):
-    
-    import anarci
-    import pandas as pd
-
-    anarci_out = anarci.run_anarci(
-        pd.DataFrame(seqs).reset_index().values.tolist(), 
-        ncpu=n_jobs, 
-        scheme='imgt',
-        allowed_species=['human', 'mouse'],
-    )
-    
-    numbered_seqs = []
-    for onarci in anarci_out[1]:            
-        numbered_seq = []
-        for i in onarci[0][0]:
-            if i[1] != '-':
-                numbered_seq.append((i[0], chain, i[1]))
-            
-        if fragmented:
-            numbered_seqs.append(numbered_seq) 
-        else:
-            numbered_seqs.append([(("<",""), chain, "<")] + numbered_seq + [((">",""), chain, ">")])
-    
-    return numbered_seqs
-
-
-def create_alignment(res_embeds, numbered_seqs, seq, number_alignment):
-    
-    import pandas as pd
-    
-    datadf = pd.DataFrame(numbered_seqs)
-    sequence_alignment = number_alignment.merge(datadf, how='left', on=[0, 1]).fillna('-')[2]
-
-    idxs = np.where(sequence_alignment.values == '-')[0]
-    idxs = [idx-num for num, idx in enumerate(idxs)]
-    
-    aligned_embeds = pd.DataFrame(np.insert(res_embeds[:len(seq)], idxs , 0, axis=0))
-    
-    return pd.concat([aligned_embeds, sequence_alignment], axis=1).values
-
-
-def get_spread_sequences(seq, spread, start_position):
-    """
-    Test sequences which are 8 positions shorter (position 10 + max CDR1 gap of 7) up to 2 positions longer (possible insertions).
-    """
-    spread_sequences = []
-
-    for diff in range(start_position-8, start_position+2+1):
-        spread_sequences.append('*'*diff+seq)
-    
-    return np.array(spread_sequences)
-
-def get_sequences_from_anarci(out_anarci, max_position, spread):
-    """
-    Ensures correct masking on each side of sequence
-    """
-    
-    if out_anarci == 'ANARCI_error':
-        return np.array(['ANARCI-ERR']*spread)
-    
-    end_position = int(re.search(r'\d+', out_anarci[::-1]).group()[::-1])
-    # Fixes ANARCI error of poor numbering of the CDR1 region
-    start_position = int(re.search(r'\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+',
-                                   out_anarci).group().split(',')[0]) - 1
-    
-    sequence = "".join(re.findall(r"(?i)[A-Z*]", "".join(re.findall(r'\),\s\'[A-Z*]', out_anarci))))
-
-    sequence_j = ''.join(sequence).replace('-','').replace('X','*') + '*'*(max_position-int(end_position))
-
-    return get_spread_sequences(sequence_j, spread, start_position)
-

ablang2/hparams.json

@@ -1 +0,0 @@
-{"name": "AbLang-2", "n_encoder_blocks": 12, "hidden_embed_size": 480, "n_attn_heads": 20, "a_fn": "swiglu", "layer_norm_eps": 1e-12, "pad_tkn": 21, "start_tkn": 0, "end_tkn": 22, "sep_tkn": 25, "mask_tkn": 23, "vocab_size": 26}

ablang2/load_model.py

@@ -1,119 +0,0 @@
-import os, subprocess, json, argparse,requests
-import torch
-
-list_of_models = {
-    "ablang1-heavy":["https://opig.stats.ox.ac.uk/data/downloads/ablang-heavy.tar.gz", "amodel.pt"], 
-    "ablang1-light":["https://opig.stats.ox.ac.uk/data/downloads/ablang-light.tar.gz", "amodel.pt"],
-    "ablang2-paired":["https://zenodo.org/records/10185169/files/ablang2-weights.tar.gz", "model.pt"],
-    "tcrlang-paired":["https://zenodo.org/records/11208211/files/tcrlang-weights.tar.gz", "model.pt"],
-}
-ablang1_models = ["ablang1-heavy", "ablang1-light"]
-ablang2_models = ["ablang2-paired", "tcrlang-paired"]
-
-
-def load_model(model_to_use = "ablang2-paired", random_init = False, device = 'cpu'):
-
-    if model_to_use in ablang1_models:
-        AbLang, tokenizer, hparams = fetch_ablang1(
-            model_to_use, 
-            random_init=random_init, 
-            device=device
-        )
-    elif model_to_use in ablang2_models:
-        AbLang, tokenizer, hparams = fetch_ablang2(
-            model_to_use, 
-            random_init=random_init, 
-            device=device
-        )
-    elif "ABLANG-" in model_to_use:
-        AbLang, tokenizer, hparams = fetch_ablang2(
-            model_to_use, 
-            random_init=random_init, 
-            device=device
-        )
-    else: 
-        assert False, f"The selected model to use ({model_to_use}) does not exist.\
-        Please select a valid model."   
-
-    return AbLang, tokenizer, hparams
-    
-    
-def download_model(model_to_use = "ablang2-paired"):
-    """
-    If not already downloaded, download model inside environment.
-    """
-
-    local_model_folder = os.path.join(os.path.dirname(__file__), "model-weights-{}".format(model_to_use))
-    os.makedirs(local_model_folder, exist_ok = True)
-
-    file_w_weights, file_model = list_of_models[model_to_use] # modify list of models
-
-    if not os.path.isfile(os.path.join(local_model_folder, file_model)):
-        print("Downloading model ...")
-        tmp_file = os.path.join(local_model_folder, "tmp.tar.gz")
-
-        with open(tmp_file,'wb') as f: f.write(requests.get(file_w_weights).content)
-
-        subprocess.run(["tar", "-zxvf", tmp_file, "-C", local_model_folder], check = True) 
-        os.remove(tmp_file)
-    
-    return local_model_folder
-    
-        
-def fetch_ablang1(model_to_use, random_init=False, device='cpu'):
-    
-    from .models.ablang1 import model as ablang_1_model
-    from .models.ablang1 import tokenizers as ablang_1_tokenizer
-    
-    local_model_folder = download_model(model_to_use)
-        
-    with open(os.path.join(local_model_folder, 'hparams.json'), 'r', encoding='utf-8') as f:
-        hparams = argparse.Namespace(**json.load(f))    
-
-    AbLang = ablang_1_model.AbLang(hparams)
-    if not random_init:
-        AbLang.load_state_dict(
-            torch.load(
-                os.path.join(local_model_folder, 'amodel.pt'),
-                map_location=torch.device(device)
-            )
-        )
-    tokenizer = ablang_1_tokenizer.ABtokenizer(os.path.join(local_model_folder, 'vocab.json'))
-        
-    return AbLang, tokenizer, hparams
-
-
-def fetch_ablang2(model_to_use, random_init=False, device='cpu'):
-    
-    from .models.ablang2 import ablang
-    from .models.ablang2 import tokenizers
-    
-    if model_to_use in ablang2_models:
-        local_model_folder = download_model(model_to_use)
-    else:
-        local_model_folder = model_to_use
-    
-    with open(os.path.join(local_model_folder, 'hparams.json'), 'r', encoding='utf-8') as f:
-        hparams = argparse.Namespace(**json.load(f))    
-        
-    AbLang = ablang.AbLang(
-        vocab_size = hparams.vocab_size,
-        hidden_embed_size = hparams.hidden_embed_size,
-        n_attn_heads = hparams.n_attn_heads,
-        n_encoder_blocks = hparams.n_encoder_blocks,
-        padding_tkn = hparams.pad_tkn,
-        mask_tkn = hparams.mask_tkn,
-        layer_norm_eps = hparams.layer_norm_eps,
-        a_fn = hparams.a_fn,
-    )
-
-    if not random_init:
-        AbLang.load_state_dict(
-            torch.load(
-                os.path.join(local_model_folder, 'model.pt'), 
-                map_location=torch.device(device)
-            )
-        )
-    tokenizer = tokenizers.ABtokenizer()
-    
-    return AbLang, tokenizer, hparams

ablang2/model.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:56d6f07862a6f824f88c8707bbc03e4026c9db762be2d3041e9767e2e6f86386
-size 179314477

ablang2/modeling_ablang2paired.py

@@ -1,81 +0,0 @@
-import torch
-import os
-from torch import nn
-from transformers import PreTrainedModel
-from ablang2.models.ablang2.ablang import AbLang as AbLang2
-from ablang2_paired.configuration_ablang2paired import AbLang2PairedConfig
-
-class AbLang2PairedHFModel(PreTrainedModel):
-    config_class = AbLang2PairedConfig
-    model_type = "ablang2-paired"
-
-    def __init__(self, config: AbLang2PairedConfig):
-        super().__init__(config)
-        self.model = AbLang2(
-            vocab_size=config.vocab_size,
-            hidden_embed_size=config.hidden_embed_size,
-            n_attn_heads=config.n_attn_heads,
-            n_encoder_blocks=config.n_encoder_blocks,
-            padding_tkn=config.padding_tkn,
-            mask_tkn=config.mask_tkn,
-            layer_norm_eps=config.layer_norm_eps,
-            a_fn=config.a_fn,
-            dropout=config.dropout,
-        )
-
-    def forward(self, input_ids=None, x=None, attention_mask=None, **kwargs):
-        # Handle both Hugging Face format (input_ids) and original format (x)
-        if input_ids is not None:
-            x = input_ids
-        elif x is None:
-            raise ValueError("Either input_ids or x must be provided")
-        
-        # Get the output from the underlying model
-        output = self.model(x, attention_mask)
-        
-        # Return as a simple object with last_hidden_state attribute
-        class ModelOutput:
-            def __init__(self, last_hidden_state):
-                self.last_hidden_state = last_hidden_state
-        
-        return ModelOutput(output)
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-        # Check if we have custom weights
-        model_path = pretrained_model_name_or_path
-        custom_weights_path = os.path.join(model_path, "model.pt")
-        
-        if os.path.exists(custom_weights_path):
-            # Load config
-            config = kwargs.get("config")
-            if config is None:
-                from transformers import AutoConfig
-                config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
-            
-            # Create model with only the config argument
-            model = cls(config)
-            
-            # Load custom weights
-            state_dict = torch.load(custom_weights_path, map_location="cpu", weights_only=True)
-            model.model.load_state_dict(state_dict)
-            
-            # Move model to appropriate device (GPU if available, otherwise CPU)
-            device = kwargs.get("device", None)
-            if device is None:
-                device = "cuda" if torch.cuda.is_available() else "cpu"
-            model = model.to(device)
-            
-            return model
-        else:
-            # Fall back to standard Hugging Face loading
-            return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
-
-    def save_pretrained(self, save_directory, **kwargs):
-        os.makedirs(save_directory, exist_ok=True)
-        # Save custom weights
-        torch.save(self.model.state_dict(), f"{save_directory}/model.pt")
-        # Save config
-        self.config.save_pretrained(save_directory)
-        # Call parent method for any additional saving
-        super().save_pretrained(save_directory, **kwargs) 

ablang2/models/ablang1/__init__.py

@@ -1,3 +0,0 @@
-from .tokenizers import ABtokenizer
-from .model import AbLang, AbRep, AbHead
-from .pretrained import pretrained

ablang2/models/ablang1/embedding.py

@@ -1,36 +0,0 @@
-import torch
-
-
-class AbEmbeddings(torch.nn.Module):
-    """
-    Residue embedding and Positional embedding
-    """
-    
-    def __init__(self, hparams):
-        super().__init__()
-        self.pad_token_id = hparams.pad_token_id
-        
-        self.AAEmbeddings = torch.nn.Embedding(hparams.vocab_size, hparams.hidden_size, padding_idx=self.pad_token_id)
-        self.PositionEmbeddings = torch.nn.Embedding(hparams.max_position_embeddings, hparams.hidden_size, padding_idx=0) # here padding_idx is always 0
-        
-        self.LayerNorm = torch.nn.LayerNorm(hparams.hidden_size, eps=hparams.layer_norm_eps)
-        self.Dropout = torch.nn.Dropout(hparams.hidden_dropout_prob)
-
-    def forward(self, src):
-        
-        inputs_embeds = self.AAEmbeddings(src)
-        
-        position_ids = self.create_position_ids_from_input_ids(src, self.pad_token_id)   
-        position_embeddings = self.PositionEmbeddings(position_ids)
-
-        embeddings = inputs_embeds + position_embeddings
-
-        return self.Dropout(self.LayerNorm(embeddings))
-        
-    def create_position_ids_from_input_ids(self, input_ids, padding_idx):
-        """
-        Replace non-padding symbols with their position numbers. Padding idx will get position 0, which will be ignored later on.
-        """
-        mask = input_ids.ne(padding_idx).int()
-        
-        return torch.cumsum(mask, dim=1).long() * mask

ablang2/models/ablang1/encoderblocks.py

@@ -1,141 +0,0 @@
-import math
-from typing import List, Optional, Tuple
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-#from fairseq.modules.multihead_attention import MultiheadAttention
-from .fairseq_mha import MultiheadAttention
-
-from .extra_fns import ACT2FN
-
-
-@dataclass
-class AbRepOutput():
-    """
-    Dataclass used to store AbRep output.
-    """
-
-    last_hidden_states: torch.FloatTensor
-    all_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-    attentions: Optional[Tuple[torch.FloatTensor]] = None
-
-        
-class EncoderBlocks(torch.nn.Module):
-    """
-    Wrapper for multiple EncoderBlocks (or a single).
-    """
-    def __init__(self, hparams):
-        super().__init__()
-        self.hparams = hparams
-        self.Layers = nn.ModuleList([EncoderBlock(hparams) for _ in range(hparams.num_hidden_layers)])
-
-    def forward(self, hidden_states, attention_mask=None, output_attentions=False, output_hidden_states=False):
-        
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attentions = () if output_attentions else None
-        
-        for num_block, a_EncoderBlock in enumerate(self.Layers):
-            
-            hidden_states, attentions = a_EncoderBlock(hidden_states, attention_mask, output_attentions)
-            #print(attentions)
-            
-            if output_hidden_states: 
-                all_hidden_states = all_hidden_states + (hidden_states,) # Takes out each hidden states after each EncoderBlock
-            
-            if output_attentions: 
-                all_self_attentions = all_self_attentions + (attentions,) # Takes out attention layers for analysis
-                  
-        return AbRepOutput(last_hidden_states=hidden_states, all_hidden_states=all_hidden_states, attentions=all_self_attentions)
-    
-    
-class EncoderBlock(torch.nn.Module):
-    """
-    Single EncoderBlock.
-    
-    An EncoderBlock consists of a MultiHeadAttention and a IntermediateLayer.
-    """
-    def __init__(self, hparams):
-        super().__init__()
-        
-        self.MultiHeadAttention = ThirdMultiHeadAttention(hparams)
-        self.MHADropout = nn.Dropout(hparams.hidden_dropout_prob)
-        self.MHALayerNorm = nn.LayerNorm(hparams.hidden_size, eps=hparams.layer_norm_eps)
-        
-        self.IntermediateLayer = IntermediateLayer(hparams)
-        
-    def forward(self, hidden_states, attention_mask=None, output_attentions=False):
-
-        MHAoutput, attentions = self.MultiHeadAttention(hidden_states, attention_mask, output_attentions=output_attentions)
-         
-        output = self.MHADropout(MHAoutput)
-        output = self.MHALayerNorm(output + hidden_states) # HIDDEN_STATES ARE ADDED FOR RESIDUAL BLOCK EFFECT
-        
-        output = self.IntermediateLayer(output) # INTERMEDIATELAYER HAS RESIDUAL BLOCK EFFECT INTERNALLY
-
-        #outputs = (layer_output,) + self_attention_outputs[1:]  # if output_attentions=False then 1: is empty
-        
-        return output, attentions
-    
-    
-class ThirdMultiHeadAttention(torch.nn.Module):
-    """
-    New MultiHeadAttention which can return the weights of the individual heads.
-    """
-    
-    def __init__(self, hparams):
-        super().__init__()
-                
-        self.Attention = MultiheadAttention(hparams.hidden_size, hparams.num_attention_heads, dropout=hparams.attention_probs_dropout_prob, self_attention=True)
-        
-    def forward(self, hidden_states, attention_mask=None, output_attentions=False):
-                
-        hidden_states = torch.transpose(hidden_states, 0, 1)
-        
-        # static_kv is only True because there is currently a bug which doesn't return the head weights unaveraged unless its true
-        attn_output, attn_weights = self.Attention(hidden_states, hidden_states, hidden_states, key_padding_mask=attention_mask, static_kv=True, 
-                                                   need_weights=output_attentions, need_head_weights=output_attentions)
-
-        return torch.transpose(attn_output, 0, 1), attn_weights
-
-
-class OldMultiHeadAttention(torch.nn.Module):
-    """
-    MultiHeadAttention contains a Scaled Dot Product Attention and a Linear Layer.    
-    """
-    def __init__(self, config):
-        super().__init__()
-        self.Attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, config.attention_probs_dropout_prob)
-
-    def forward(self, hidden_states, attention_mask=None, output_attentions=False):  
-        
-        hidden_states = torch.transpose(hidden_states, 0, 1)
-        output, attentions = self.Attention(hidden_states, hidden_states, hidden_states, key_padding_mask=attention_mask, need_weights=output_attentions) 
-        
-        attention_output = torch.transpose(output, 0, 1)
-
-        return attention_output, attentions
-    
-
-class IntermediateLayer(nn.Module):
-    """
-    Contains an expanding layer, while also functioning as a residual block ending with a drop-norm layer
-    """
-    def __init__(self, config):
-        super().__init__()
-        self.expand_dense = nn.Linear(config.hidden_size, config.intermediate_size)
-        self.intermediate_act_fn = ACT2FN[config.hidden_act]
-        
-        self.dense_dense = nn.Linear(config.intermediate_size, config.hidden_size)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
-    def forward(self, hidden_states):
-        output = self.expand_dense(hidden_states)
-        output = self.intermediate_act_fn(output)
-               
-        output = self.dense_dense(output)
-        output = self.dropout(output)
-        output = self.LayerNorm(output + hidden_states)
-                
-        return output 

ablang2/models/ablang1/extra_fns.py

@@ -1,26 +0,0 @@
-import torch
-import math
-
-
-def gelu_new(x):
-    """
-    Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT). Also see
-    the Gaussian Error Linear Units paper: https://arxiv.org/abs/1606.08415
-    """
-    return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
-
-def gelu_fast(x):
-    return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x)))
-
-def mish(x):
-    return x * torch.tanh(torch.nn.functional.softplus(x))
-
-ACT2FN = {
-    "relu": torch.nn.functional.relu,
-    "gelu": torch.nn.functional.gelu,
-    "tanh": torch.tanh,
-    "gelu_new": gelu_new,
-    "gelu_fast": gelu_fast,
-    "mish": mish,
-    "sigmoid": torch.sigmoid,
-}

ablang2/models/ablang1/fairseq_mha.py

@@ -1,1306 +0,0 @@
-import math
-from typing import Dict, List, Optional, Tuple
-import uuid
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor, nn
-from torch.nn import Parameter
-
-_xformers_available = False
-
-# TODO: move this into xformers?
-# TODO: uint8 input type should just output a bool
-def _mask_for_xformers(mask: Tensor, to_dtype: Optional[torch.dtype] = None):
-    """
-    call to pytorch multihead accepts three mask types:
-        - ByteTensor where non-zero means to mask
-        - FloatTensor which is an additive mask
-        - BoolTensor where True means to mask
-    xFormers currently accepts boolean and additive maks. For boolean masks
-    the values have opposite meaning. For a BoolTensor True mean to keep the value.
-    """
-    float_types = [torch.float, torch.float16]
-    # If an input mask is a float it is an additive mask. Otherwise it is either uint8 or bool.
-    additive = mask.dtype in float_types
-    # If to_dype is not specified, keep same dtype as mask.
-    to_dtype = mask.dtype if to_dtype is None else to_dtype
-    to_additive = to_dtype in float_types
-
-    if additive:
-        if to_additive:
-            return mask.to(to_dtype)
-        mask = mask < 0
-
-    if to_additive:
-        # return additive mask
-        new_mask = torch.zeros_like(mask, dtype=to_dtype)
-        new_mask = new_mask.masked_fill_(mask, -float("inf"))
-        return new_mask
-
-    # In xFormers True is value to keep rather than value to mask
-    mask = ~mask.to(torch.bool)
-    mask = mask.to(to_dtype)
-    return mask
-
-class FairseqDecoder(nn.Module):
-    """Base class for decoders."""
-
-    def __init__(self, dictionary):
-        super().__init__()
-        self.dictionary = dictionary
-        self.onnx_trace = False
-        self.adaptive_softmax = None
-
-    def forward(self, prev_output_tokens, encoder_out=None, **kwargs):
-        """
-        Args:
-            prev_output_tokens (LongTensor): shifted output tokens of shape
-                `(batch, tgt_len)`, for teacher forcing
-            encoder_out (dict, optional): output from the encoder, used for
-                encoder-side attention
-
-        Returns:
-            tuple:
-                - the decoder's output of shape `(batch, tgt_len, vocab)`
-                - a dictionary with any model-specific outputs
-        """
-        x, extra = self.extract_features(
-            prev_output_tokens, encoder_out=encoder_out, **kwargs
-        )
-        x = self.output_layer(x)
-        return x, extra
-
-    def extract_features(self, prev_output_tokens, encoder_out=None, **kwargs):
-        """
-        Returns:
-            tuple:
-                - the decoder's features of shape `(batch, tgt_len, embed_dim)`
-                - a dictionary with any model-specific outputs
-        """
-        raise NotImplementedError
-
-    def output_layer(self, features, **kwargs):
-        """
-        Project features to the default output size, e.g., vocabulary size.
-
-        Args:
-            features (Tensor): features returned by *extract_features*.
-        """
-        raise NotImplementedError
-
-    def get_normalized_probs(
-        self,
-        net_output: Tuple[Tensor, Optional[Dict[str, List[Optional[Tensor]]]]],
-        log_probs: bool,
-        sample: Optional[Dict[str, Tensor]] = None,
-    ):
-        """Get normalized probabilities (or log probs) from a net's output."""
-        return self.get_normalized_probs_scriptable(net_output, log_probs, sample)
-
-    # TorchScript doesn't support super() method so that the scriptable Subclass
-    # can't access the base class model in Torchscript.
-    # Current workaround is to add a helper function with different name and
-    # call the helper function from scriptable Subclass.
-    def get_normalized_probs_scriptable(
-        self,
-        net_output: Tuple[Tensor, Optional[Dict[str, List[Optional[Tensor]]]]],
-        log_probs: bool,
-        sample: Optional[Dict[str, Tensor]] = None,
-    ):
-        """Get normalized probabilities (or log probs) from a net's output."""
-
-        if hasattr(self, "adaptive_softmax") and self.adaptive_softmax is not None:
-            if sample is not None:
-                assert "target" in sample
-                target = sample["target"]
-            else:
-                target = None
-            out = self.adaptive_softmax.get_log_prob(net_output[0], target=target)
-            return out.exp_() if not log_probs else out
-
-        logits = net_output[0]
-        if log_probs:
-            return log_softmax(logits, dim=-1, onnx_trace=self.onnx_trace)
-        else:
-            return softmax(logits, dim=-1, onnx_trace=self.onnx_trace)
-
-    def max_positions(self):
-        """Maximum input length supported by the decoder."""
-        return 1e6  # an arbitrary large number
-
-    def upgrade_state_dict_named(self, state_dict, name):
-        """Upgrade old state dicts to work with newer code."""
-        return state_dict
-
-    def prepare_for_onnx_export_(self):
-        self.onnx_trace = True
-
-
-class FairseqIncrementalState(object):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.init_incremental_state()
-
-    def init_incremental_state(self):
-        self._incremental_state_id = str(uuid.uuid4())
-
-    def _get_full_incremental_state_key(self, key: str) -> str:
-        return "{}.{}".format(self._incremental_state_id, key)
-
-    def get_incremental_state(
-        self,
-        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
-        key: str,
-    ) -> Optional[Dict[str, Optional[Tensor]]]:
-        """Helper for getting incremental state for an nn.Module."""
-        full_key = self._get_full_incremental_state_key(key)
-        if incremental_state is None or full_key not in incremental_state:
-            return None
-        return incremental_state[full_key]
-
-    def set_incremental_state(
-        self,
-        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
-        key: str,
-        value: Dict[str, Optional[Tensor]],
-    ) -> Optional[Dict[str, Dict[str, Optional[Tensor]]]]:
-        """Helper for setting incremental state for an nn.Module."""
-        if incremental_state is not None:
-            full_key = self._get_full_incremental_state_key(key)
-            incremental_state[full_key] = value
-        return incremental_state
-
-
-def with_incremental_state(cls):
-    cls.__bases__ = (FairseqIncrementalState,) + tuple(
-        b for b in cls.__bases__ if b != FairseqIncrementalState
-    )
-    return cls
-        
-        
-@with_incremental_state
-class FairseqIncrementalDecoder(FairseqDecoder):
-    """Base class for incremental decoders.
-
-    Incremental decoding is a special mode at inference time where the Model
-    only receives a single timestep of input corresponding to the previous
-    output token (for teacher forcing) and must produce the next output
-    *incrementally*. Thus the model must cache any long-term state that is
-    needed about the sequence, e.g., hidden states, convolutional states, etc.
-
-    Compared to the standard :class:`FairseqDecoder` interface, the incremental
-    decoder interface allows :func:`forward` functions to take an extra keyword
-    argument (*incremental_state*) that can be used to cache state across
-    time-steps.
-
-    The :class:`FairseqIncrementalDecoder` interface also defines the
-    :func:`reorder_incremental_state` method, which is used during beam search
-    to select and reorder the incremental state based on the selection of beams.
-
-    To learn more about how incremental decoding works, refer to `this blog
-    <http://www.telesens.co/2019/04/21/understanding-incremental-decoding-in-fairseq/>`_.
-    """
-
-    def __init__(self, dictionary):
-        super().__init__(dictionary)
-
-    def forward(
-        self, prev_output_tokens, encoder_out=None, incremental_state=None, **kwargs
-    ):
-        """
-        Args:
-            prev_output_tokens (LongTensor): shifted output tokens of shape
-                `(batch, tgt_len)`, for teacher forcing
-            encoder_out (dict, optional): output from the encoder, used for
-                encoder-side attention
-            incremental_state (dict, optional): dictionary used for storing
-                state during :ref:`Incremental decoding`
-
-        Returns:
-            tuple:
-                - the decoder's output of shape `(batch, tgt_len, vocab)`
-                - a dictionary with any model-specific outputs
-        """
-        raise NotImplementedError
-
-    def extract_features(
-        self, prev_output_tokens, encoder_out=None, incremental_state=None, **kwargs
-    ):
-        """
-        Returns:
-            tuple:
-                - the decoder's features of shape `(batch, tgt_len, embed_dim)`
-                - a dictionary with any model-specific outputs
-        """
-        raise NotImplementedError
-
-    def reorder_incremental_state(
-        self,
-        incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
-        new_order: Tensor,
-    ):
-        """Reorder incremental state.
-
-        This will be called when the order of the input has changed from the
-        previous time step. A typical use case is beam search, where the input
-        order changes between time steps based on the selection of beams.
-        """
-        pass
-
-    def reorder_incremental_state_scripting(
-        self,
-        incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
-        new_order: Tensor,
-    ):
-        """Main entry point for reordering the incremental state.
-
-        Due to limitations in TorchScript, we call this function in
-        :class:`fairseq.sequence_generator.SequenceGenerator` instead of
-        calling :func:`reorder_incremental_state` directly.
-        """
-        for module in self.modules():
-            if hasattr(module, "reorder_incremental_state"):
-                result = module.reorder_incremental_state(incremental_state, new_order)
-                if result is not None:
-                    incremental_state = result
-
-    def set_beam_size(self, beam_size):
-        """Sets the beam size in the decoder and all children."""
-        if getattr(self, "_beam_size", -1) != beam_size:
-            seen = set()
-
-            def apply_set_beam_size(module):
-                if (
-                    module != self
-                    and hasattr(module, "set_beam_size")
-                    and module not in seen
-                ):
-                    seen.add(module)
-                    module.set_beam_size(beam_size)
-
-            self.apply(apply_set_beam_size)
-            self._beam_size = beam_size
-
-
-            
-            
-            
-            
-class MultiheadAttention(FairseqIncrementalDecoder):
-    """Multi-headed attention.
-
-    See "Attention Is All You Need" for more details.
-    """
-
-    def __init__(
-        self,
-        embed_dim,
-        num_heads,
-        kdim=None,
-        vdim=None,
-        dropout=0.0,
-        bias=True,
-        add_bias_kv=False,
-        add_zero_attn=False,
-        self_attention=False,
-        encoder_decoder_attention=False,
-        dictionary=None,
-        q_noise=0.0,
-        qn_block_size=8,
-        # TODO: pass in config rather than string.
-        # config defined in xformers.components.attention.AttentionConfig
-        xformers_att_config: Optional[str] = None,
-        xformers_blocksparse_layout: Optional[
-            torch.Tensor
-        ] = None,  # This should be part of the config
-        xformers_blocksparse_blocksize: Optional[
-            int
-        ] = 16,  # This should be part of the config
-    ):
-        super().__init__(dictionary)
-
-        #xformers_att_config = utils.eval_str_dict(xformers_att_config)
-        self.use_xformers = False #xformers_att_config is not None
-        if self.use_xformers and not _xformers_available:
-            raise ImportError("\n\n  Please install xFormers.")
-        self.embed_dim = embed_dim
-        self.kdim = kdim if kdim is not None else embed_dim
-        self.vdim = vdim if vdim is not None else embed_dim
-        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
-
-        self.num_heads = num_heads
-        self.dropout_module = FairseqDropout(
-            dropout, module_name=self.__class__.__name__
-        )
-
-        self.head_dim = embed_dim // num_heads
-        assert (
-            self.head_dim * num_heads == self.embed_dim
-        ), "embed_dim must be divisible by num_heads"
-        self.scaling = self.head_dim**-0.5
-
-        self.self_attention = self_attention
-        self.encoder_decoder_attention = encoder_decoder_attention
-
-        assert not self.self_attention or self.qkv_same_dim, (
-            "Self-attention requires query, key and " "value to be of the same size"
-        )
-
-        self.k_proj = quant_noise(
-            nn.Linear(self.kdim, embed_dim, bias=bias), q_noise, qn_block_size
-        )
-        self.v_proj = quant_noise(
-            nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size
-        )
-        self.q_proj = quant_noise(
-            nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size
-        )
-
-        self.out_proj = quant_noise(
-            nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size
-        )
-
-        if add_bias_kv:
-            self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))
-            self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))
-        else:
-            self.bias_k = self.bias_v = None
-
-        self.add_zero_attn = add_zero_attn
-        self.beam_size = 1
-        self.reset_parameters()
-
-        if self.use_xformers:
-            xformers_att_config["dropout"] = xformers_att_config.get("dropout", dropout)
-            xformers_att_config["num_heads"] = xformers_att_config.get(
-                "num_heads", num_heads
-            )
-
-            if xformers_blocksparse_layout is not None:
-                # Could be part of a single config passed only once
-                xformers_att_config["block_size"] = xformers_blocksparse_blocksize
-                xformers_att_config["layout"] = xformers_blocksparse_layout
-                xformers_att_config["name"] = "blocksparse"
-
-            self.attention = build_attention(xformers_att_config)
-
-        self.onnx_trace = False
-        self.skip_embed_dim_check = False
-        self.init_incremental_state()
-
-    def prepare_for_onnx_export_(self):
-        self.onnx_trace = True
-
-    def reset_parameters(self):
-        if self.qkv_same_dim:
-            # Empirically observed the convergence to be much better with
-            # the scaled initialization
-            nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
-            nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
-            nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
-        else:
-            nn.init.xavier_uniform_(self.k_proj.weight)
-            nn.init.xavier_uniform_(self.v_proj.weight)
-            nn.init.xavier_uniform_(self.q_proj.weight)
-
-        nn.init.xavier_uniform_(self.out_proj.weight)
-        if self.out_proj.bias is not None:
-            nn.init.constant_(self.out_proj.bias, 0.0)
-        if self.bias_k is not None:
-            nn.init.xavier_normal_(self.bias_k)
-        if self.bias_v is not None:
-            nn.init.xavier_normal_(self.bias_v)
-
-    def _get_reserve_head_index(self, num_heads_to_keep: int):
-        k_proj_heads_norm = []
-        q_proj_heads_norm = []
-        v_proj_heads_norm = []
-
-        for i in range(self.num_heads):
-            start_idx = i * self.head_dim
-            end_idx = (i + 1) * self.head_dim
-            k_proj_heads_norm.append(
-                torch.sum(
-                    torch.abs(
-                        self.k_proj.weight[
-                            start_idx:end_idx,
-                        ]
-                    )
-                ).tolist()
-                + torch.sum(torch.abs(self.k_proj.bias[start_idx:end_idx])).tolist()
-            )
-            q_proj_heads_norm.append(
-                torch.sum(
-                    torch.abs(
-                        self.q_proj.weight[
-                            start_idx:end_idx,
-                        ]
-                    )
-                ).tolist()
-                + torch.sum(torch.abs(self.q_proj.bias[start_idx:end_idx])).tolist()
-            )
-            v_proj_heads_norm.append(
-                torch.sum(
-                    torch.abs(
-                        self.v_proj.weight[
-                            start_idx:end_idx,
-                        ]
-                    )
-                ).tolist()
-                + torch.sum(torch.abs(self.v_proj.bias[start_idx:end_idx])).tolist()
-            )
-
-        heads_norm = []
-        for i in range(self.num_heads):
-            heads_norm.append(
-                k_proj_heads_norm[i] + q_proj_heads_norm[i] + v_proj_heads_norm[i]
-            )
-
-        sorted_head_index = sorted(
-            range(self.num_heads), key=lambda k: heads_norm[k], reverse=True
-        )
-        reserve_head_index = []
-        for i in range(num_heads_to_keep):
-            start = sorted_head_index[i] * self.head_dim
-            end = (sorted_head_index[i] + 1) * self.head_dim
-            reserve_head_index.append((start, end))
-        return reserve_head_index
-
-    def _adaptive_prune_heads(self, reserve_head_index: List[Tuple[int, int]]):
-        new_q_weight = []
-        new_q_bias = []
-        new_k_weight = []
-        new_k_bias = []
-        new_v_weight = []
-        new_v_bias = []
-        new_out_proj_weight = []
-
-        for ele in reserve_head_index:
-            start_idx, end_idx = ele
-            new_q_weight.append(
-                self.q_proj.weight[
-                    start_idx:end_idx,
-                ]
-            )
-            new_q_bias.append(self.q_proj.bias[start_idx:end_idx])
-
-            new_k_weight.append(
-                self.k_proj.weight[
-                    start_idx:end_idx,
-                ]
-            )
-
-            new_k_bias.append(self.k_proj.bias[start_idx:end_idx])
-
-            new_v_weight.append(
-                self.v_proj.weight[
-                    start_idx:end_idx,
-                ]
-            )
-            new_v_bias.append(self.v_proj.bias[start_idx:end_idx])
-
-            new_out_proj_weight.append(self.out_proj.weight[:, start_idx:end_idx])
-
-        new_q_weight = torch.cat(new_q_weight).detach()
-        new_k_weight = torch.cat(new_k_weight).detach()
-        new_v_weight = torch.cat(new_v_weight).detach()
-        new_out_proj_weight = torch.cat(new_out_proj_weight, dim=-1).detach()
-        new_q_weight.requires_grad = True
-        new_k_weight.requires_grad = True
-        new_v_weight.requires_grad = True
-        new_out_proj_weight.requires_grad = True
-
-        new_q_bias = torch.cat(new_q_bias).detach()
-        new_q_bias.requires_grad = True
-
-        new_k_bias = torch.cat(new_k_bias).detach()
-        new_k_bias.requires_grad = True
-
-        new_v_bias = torch.cat(new_v_bias).detach()
-        new_v_bias.requires_grad = True
-
-        self.q_proj.weight = torch.nn.Parameter(new_q_weight)
-        self.q_proj.bias = torch.nn.Parameter(new_q_bias)
-
-        self.k_proj.weight = torch.nn.Parameter(new_k_weight)
-        self.k_proj.bias = torch.nn.Parameter(new_k_bias)
-
-        self.v_proj.weight = torch.nn.Parameter(new_v_weight)
-        self.v_proj.bias = torch.nn.Parameter(new_v_bias)
-
-        self.out_proj.weight = torch.nn.Parameter(new_out_proj_weight)
-
-        self.num_heads = len(reserve_head_index)
-        self.embed_dim = self.head_dim * self.num_heads
-        self.q_proj.out_features = self.embed_dim
-        self.k_proj.out_features = self.embed_dim
-        self.v_proj.out_features = self.embed_dim
-
-    def _set_skip_embed_dim_check(self):
-        self.skip_embed_dim_check = True
-
-    def _pad_masks(
-        self,
-        key_padding_mask: Optional[Tensor],
-        attn_mask: Optional[Tensor],
-    ) -> Tuple[Optional[Tensor], Optional[Tensor]]:
-        if attn_mask is not None:
-            shape = attn_mask.size()[:-1] + torch.Size([1])
-            attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(shape)], dim=-1)
-        if key_padding_mask is not None:
-            shape = key_padding_mask.size()[:-1] + torch.Size([1])
-            key_padding_mask = torch.cat(
-                [
-                    key_padding_mask,
-                    key_padding_mask.new_zeros(shape),
-                ],
-                dim=-1,
-            )
-        return key_padding_mask, attn_mask
-
-    def _add_bias(
-        self,
-        k: Tensor,
-        v: Tensor,
-        key_padding_mask: Optional[Tensor],
-        attn_mask: Optional[Tensor],
-        bsz: int,
-    ) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:
-        assert self.bias_k is not None
-        assert self.bias_v is not None
-        k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
-        v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
-        key_padding_mask, attn_mask = self._pad_masks(
-            key_padding_mask=key_padding_mask, attn_mask=attn_mask
-        )
-        return k, v, key_padding_mask, attn_mask
-
-    def _append_zero_attn(
-        self,
-        k: Tensor,
-        v: Tensor,
-        key_padding_mask: Optional[Tensor],
-        attn_mask: Optional[Tensor],
-    ) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:
-        zero_attn_shape = k.size()[:-2] + torch.Size([1]) + k.size()[-1:]
-        k = torch.cat(
-            [k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=-2
-        )
-        v = torch.cat(
-            [v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=-2
-        )
-        key_padding_mask, attn_mask = self._pad_masks(
-            key_padding_mask=key_padding_mask, attn_mask=attn_mask
-        )
-        return k, v, key_padding_mask, attn_mask
-
-    def _xformers_attn_forward(
-        self,
-        query,
-        key: Optional[Tensor],
-        value: Optional[Tensor],
-        key_padding_mask: Optional[Tensor] = None,
-        need_weights: bool = True,
-        attn_mask: Optional[Tensor] = None,
-    ) -> Tuple[Tensor, Optional[Tensor]]:
-
-        tgt_len, bsz, embed_dim = query.size()
-
-        if key_padding_mask is not None:
-            assert key_padding_mask.size(0) == bsz
-            assert key_padding_mask.size(1) == tgt_len
-
-        if self.self_attention:
-            key = query
-            value = query
-        elif self.encoder_decoder_attention:
-            value = key
-
-        q = self.q_proj(query)
-        k = self.k_proj(key)
-        v = self.v_proj(value)
-
-        if self.bias_k is not None:
-            assert self.bias_v is not None
-            k, v, attn_mask, key_padding_mask = self._add_bias(
-                k, v, attn_mask, key_padding_mask, bsz
-            )
-
-        def fold_heads(x):
-            return (
-                x.contiguous()
-                .view(-1, bsz * self.num_heads, self.head_dim)
-                .transpose(0, 1)
-            )
-
-        def split_heads(x):
-            return (
-                x.contiguous()
-                .view(-1, bsz, self.num_heads, self.head_dim)
-                .transpose(0, 1)
-                .transpose(1, 2)
-            )
-
-        massage = split_heads if self.attention.requires_head_dimension else fold_heads
-        q = massage(q)
-        if k is not None:
-            k = massage(k)
-        if v is not None:
-            v = massage(v)
-
-        if self.add_zero_attn:
-            k, v, key_padding_mask, attn_mask = self._append_zero_attn(
-                k=k, v=v, key_padding_mask=key_padding_mask, attn_mask=attn_mask
-            )
-
-        kwargs = {}
-
-        if attn_mask is not None and self.attention.supports_attention_mask:
-            attn_mask = _mask_for_xformers(attn_mask, to_dtype=q.dtype)
-            kwargs["att_mask"] = attn_mask
-
-        if key_padding_mask is not None:
-            key_padding_mask = _mask_for_xformers(key_padding_mask, to_dtype=torch.bool)
-            if not self.attention.requires_separate_masks:
-                attn_mask = maybe_merge_masks(
-                    attn_mask,
-                    key_padding_mask,
-                    batch_size=bsz,
-                    src_len=k.size(-2),
-                    tgt_len=q.size(-2),
-                    num_heads=self.num_heads,
-                )
-                key_padding_mask = None
-                kwargs["att_mask"] = attn_mask
-            if self.attention.supports_key_padding_mask:
-                kwargs["key_padding_mask"] = key_padding_mask
-
-        y = self.attention(q, k, v, **kwargs)
-
-        y = (
-            y.view(bsz, self.num_heads, tgt_len, self.head_dim)
-            .transpose(1, 2)
-            .flatten(start_dim=2, end_dim=3)
-            .transpose(0, 1)
-        )
-        assert list(y.size()) == [tgt_len, bsz, embed_dim]
-
-        # Dropout not needed because already applied in attention.
-        # It is applied to the attention weights before matmul with v.
-        y = self.out_proj(y)
-
-        # TODO: support returning attention weights if needed.
-        return y, None
-
-    def forward(
-        self,
-        query: Tensor,
-        key: Optional[Tensor],
-        value: Optional[Tensor],
-        key_padding_mask: Optional[Tensor] = None,
-        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
-        need_weights: bool = True,
-        static_kv: bool = False,
-        attn_mask: Optional[Tensor] = None,
-        before_softmax: bool = False,
-        need_head_weights: bool = False,
-    ) -> Tuple[Tensor, Optional[Tensor]]:
-        """Input shape: Time x Batch x Channel
-
-        Args:
-            key_padding_mask (ByteTensor, optional): mask to exclude
-                keys that are pads, of shape `(batch, src_len)`, where
-                padding elements are indicated by 1s.
-            need_weights (bool, optional): return the attention weights,
-                averaged over heads (default: False).
-            attn_mask (ByteTensor, optional): typically used to
-                implement causal attention, where the mask prevents the
-                attention from looking forward in time (default: None).
-            before_softmax (bool, optional): return the raw attention
-                weights and values before the attention softmax.
-            need_head_weights (bool, optional): return the attention
-                weights for each head. Implies *need_weights*. Default:
-                return the average attention weights over all heads.
-        """
-        if need_head_weights:
-            need_weights = True
-
-        is_tpu = query.device.type == "xla"
-
-        tgt_len, bsz, embed_dim = query.size()
-        src_len = tgt_len
-        if not self.skip_embed_dim_check:
-            assert (
-                embed_dim == self.embed_dim
-            ), f"query dim {embed_dim} != {self.embed_dim}"
-        assert list(query.size()) == [tgt_len, bsz, embed_dim]
-        if key is not None:
-            src_len, key_bsz, _ = key.size()
-            if not torch.jit.is_scripting():
-                assert value is not None
-                assert src_len, key_bsz == value.shape[:2]
-
-        if (
-            not self.onnx_trace
-            and not is_tpu  # don't use PyTorch version on TPUs
-            and incremental_state is None
-            and not static_kv
-            # A workaround for quantization to work. Otherwise JIT compilation
-            # treats bias in linear module as method.
-            and not torch.jit.is_scripting()
-            # The Multihead attention implemented in pytorch forces strong dimension check
-            # for input embedding dimention and K,Q,V projection dimension.
-            # Since pruning will break the dimension check and it is not easy to modify the pytorch API,
-            # it is preferred to bypass the pytorch MHA when we need to skip embed_dim_check
-            and not self.skip_embed_dim_check
-        ):
-            assert key is not None and value is not None
-
-            if self.use_xformers:
-                return self._xformers_attn_forward(
-                    query, key, value, key_padding_mask, need_weights, attn_mask
-                )
-
-            else:
-                return F.multi_head_attention_forward(
-                    query,
-                    key,
-                    value,
-                    self.embed_dim,
-                    self.num_heads,
-                    torch.empty([0]),
-                    torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)),
-                    self.bias_k,
-                    self.bias_v,
-                    self.add_zero_attn,
-                    self.dropout_module.p,
-                    self.out_proj.weight,
-                    self.out_proj.bias,
-                    self.training or self.dropout_module.apply_during_inference,
-                    key_padding_mask.bool() if key_padding_mask is not None else None,
-                    need_weights,
-                    attn_mask,
-                    use_separate_proj_weight=True,
-                    q_proj_weight=self.q_proj.weight,
-                    k_proj_weight=self.k_proj.weight,
-                    v_proj_weight=self.v_proj.weight,
-                )
-
-        if incremental_state is not None:
-            saved_state = self._get_input_buffer(incremental_state)
-            if saved_state is not None and "prev_key" in saved_state:
-                # previous time steps are cached - no need to recompute
-                # key and value if they are static
-                if static_kv:
-                    assert self.encoder_decoder_attention and not self.self_attention
-                    key = value = None
-        else:
-            saved_state = None
-
-        if self.self_attention:
-            q = self.q_proj(query)
-            k = self.k_proj(query)
-            v = self.v_proj(query)
-        elif self.encoder_decoder_attention:
-            # encoder-decoder attention
-            q = self.q_proj(query)
-            if key is None:
-                assert value is None
-                k = v = None
-            else:
-                if self.beam_size > 1 and bsz == key.size(1):
-                    # key is [T, bsz*beam_size, C], reduce to [T, bsz, C]
-                    key = key.view(key.size(0), -1, self.beam_size, key.size(2))[
-                        :, :, 0, :
-                    ]
-                    if key_padding_mask is not None:
-                        key_padding_mask = key_padding_mask.view(
-                            -1, self.beam_size, key_padding_mask.size(1)
-                        )[:, 0, :]
-                k = self.k_proj(key)
-                v = self.v_proj(key)
-
-        else:
-            assert key is not None and value is not None
-            q = self.q_proj(query)
-            k = self.k_proj(key)
-            v = self.v_proj(value)
-        q *= self.scaling
-
-        if self.bias_k is not None:
-            assert self.bias_v is not None
-            k, v, attn_mask, key_padding_mask = self._add_bias(
-                k, v, attn_mask, key_padding_mask, bsz
-            )
-
-        q = (
-            q.contiguous()
-            .view(tgt_len, bsz * self.num_heads, self.head_dim)
-            .transpose(0, 1)
-        )
-        kv_bsz = bsz  # need default value for scripting
-        if k is not None:
-            kv_bsz = k.size(1)
-            k = (
-                k.contiguous()
-                .view(-1, kv_bsz * self.num_heads, self.head_dim)
-                .transpose(0, 1)
-            )
-        if v is not None:
-            v = (
-                v.contiguous()
-                .view(-1, kv_bsz * self.num_heads, self.head_dim)
-                .transpose(0, 1)
-            )
-
-        if saved_state is not None:
-            # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
-            if "prev_key" in saved_state:
-                _prev_key = saved_state["prev_key"]
-                assert _prev_key is not None
-                kv_bsz = _prev_key.size(0)
-                prev_key = _prev_key.view(kv_bsz * self.num_heads, -1, self.head_dim)
-                if static_kv:
-                    k = prev_key
-                else:
-                    assert k is not None
-                    k = torch.cat([prev_key, k], dim=1)
-                src_len = k.size(1)
-            if "prev_value" in saved_state:
-                _prev_value = saved_state["prev_value"]
-                assert _prev_value is not None
-                assert kv_bsz == _prev_value.size(0)
-                prev_value = _prev_value.view(
-                    kv_bsz * self.num_heads, -1, self.head_dim
-                )
-                if static_kv:
-                    v = prev_value
-                else:
-                    assert v is not None
-                    v = torch.cat([prev_value, v], dim=1)
-            prev_key_padding_mask: Optional[Tensor] = None
-            if "prev_key_padding_mask" in saved_state:
-                prev_key_padding_mask = saved_state["prev_key_padding_mask"]
-            assert k is not None and v is not None
-            key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(
-                key_padding_mask=key_padding_mask,
-                prev_key_padding_mask=prev_key_padding_mask,
-                batch_size=kv_bsz,
-                src_len=k.size(1),
-                static_kv=static_kv,
-            )
-
-            saved_state["prev_key"] = k.view(kv_bsz, self.num_heads, -1, self.head_dim)
-            saved_state["prev_value"] = v.view(
-                kv_bsz, self.num_heads, -1, self.head_dim
-            )
-            saved_state["prev_key_padding_mask"] = key_padding_mask
-            # In this branch incremental_state is never None
-            assert incremental_state is not None
-            incremental_state = self._set_input_buffer(incremental_state, saved_state)
-        assert k is not None
-        assert k.size(1) == src_len
-
-        # This is part of a workaround to get around fork/join parallelism
-        # not supporting Optional types.
-        if key_padding_mask is not None and key_padding_mask.dim() == 0:
-            key_padding_mask = None
-
-        if key_padding_mask is not None:
-            assert key_padding_mask.size(0) == kv_bsz
-            assert key_padding_mask.size(1) == src_len
-
-        if self.add_zero_attn:
-            assert v is not None
-            src_len += 1
-            k, v, key_padding_mask, attn_mask = self._append_zero_attn(
-                k=k, v=v, key_padding_mask=key_padding_mask, attn_mask=attn_mask
-            )
-
-        if self.encoder_decoder_attention and bsz != kv_bsz:
-            attn_weights = torch.einsum(
-                "bxhtd,bhsd->bxhts",
-                q.view((kv_bsz, -1, self.num_heads) + q.size()[1:]),
-                k.view((kv_bsz, self.num_heads) + k.size()[1:]),
-            )
-            attn_weights = attn_weights.reshape((-1,) + attn_weights.size()[-2:])
-        else:
-            attn_weights = torch.bmm(q, k.transpose(1, 2))
-        attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
-
-        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
-
-        if attn_mask is not None:
-            attn_mask = attn_mask.unsqueeze(0)
-            if self.onnx_trace:
-                attn_mask = attn_mask.repeat(attn_weights.size(0), 1, 1)
-            attn_weights += attn_mask
-
-        if key_padding_mask is not None:
-            # don't attend to padding symbols
-            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
-            if not is_tpu:
-                attn_weights = attn_weights.view(
-                    kv_bsz, -1, self.num_heads, tgt_len, src_len
-                )
-                attn_weights = attn_weights.masked_fill(
-                    key_padding_mask.unsqueeze(1)
-                    .unsqueeze(2)
-                    .unsqueeze(3)
-                    .to(torch.bool),
-                    float("-inf"),
-                )
-            else:
-                attn_weights = attn_weights.transpose(0, 2)
-                attn_weights = attn_weights.masked_fill(key_padding_mask, float("-inf"))
-                attn_weights = attn_weights.transpose(0, 2)
-            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
-
-        if before_softmax:
-            return attn_weights, v
-
-        attn_weights_float = softmax(
-            attn_weights, dim=-1, onnx_trace=self.onnx_trace
-        )
-        attn_weights = attn_weights_float.type_as(attn_weights)
-        attn_probs = self.dropout_module(attn_weights)
-
-        assert v is not None
-        attn: Optional[Tensor] = None
-        if self.encoder_decoder_attention and bsz != kv_bsz:
-            attn = torch.einsum(
-                "bxhts,bhsd->bxhtd",
-                attn_probs.view(
-                    (
-                        kv_bsz,
-                        -1,
-                        self.num_heads,
-                    )
-                    + attn_probs.size()[1:]
-                ),
-                v.view(
-                    (
-                        kv_bsz,
-                        self.num_heads,
-                    )
-                    + v.size()[1:]
-                ),
-            )
-            attn = attn.reshape((-1,) + attn.size()[-2:])
-        else:
-            attn = torch.bmm(attn_probs, v)
-        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
-        if self.onnx_trace and attn.size(1) == 1:
-            # when ONNX tracing a single decoder step (sequence length == 1)
-            # the transpose is a no-op copy before view, thus unnecessary
-            attn = attn.contiguous().view(tgt_len, bsz, self.embed_dim)
-        else:
-            attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, self.embed_dim)
-        attn = self.out_proj(attn)
-        attn_weights: Optional[Tensor] = None
-        if need_weights:
-            attn_weights = attn_weights_float.view(
-                bsz, self.num_heads, tgt_len, src_len
-            ).transpose(1, 0)
-            if not need_head_weights:
-                # average attention weights over heads
-                attn_weights = attn_weights.mean(dim=0)
-
-        return attn, attn_weights
-
-    @staticmethod
-    def _append_prev_key_padding_mask(
-        key_padding_mask: Optional[Tensor],
-        prev_key_padding_mask: Optional[Tensor],
-        batch_size: int,
-        src_len: int,
-        static_kv: bool,
-    ) -> Optional[Tensor]:
-        # saved key padding masks have shape (bsz, seq_len)
-        if prev_key_padding_mask is not None and static_kv:
-            new_key_padding_mask = prev_key_padding_mask
-        elif prev_key_padding_mask is not None and key_padding_mask is not None:
-            new_key_padding_mask = torch.cat(
-                [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1
-            )
-        # During incremental decoding, as the padding token enters and
-        # leaves the frame, there will be a time when prev or current
-        # is None
-        elif prev_key_padding_mask is not None:
-            if src_len > prev_key_padding_mask.size(1):
-                filler = torch.zeros(
-                    (batch_size, src_len - prev_key_padding_mask.size(1)),
-                    device=prev_key_padding_mask.device,
-                )
-                new_key_padding_mask = torch.cat(
-                    [prev_key_padding_mask.float(), filler.float()], dim=1
-                )
-            else:
-                new_key_padding_mask = prev_key_padding_mask.float()
-        elif key_padding_mask is not None:
-            if src_len > key_padding_mask.size(1):
-                filler = torch.zeros(
-                    (batch_size, src_len - key_padding_mask.size(1)),
-                    device=key_padding_mask.device,
-                )
-                new_key_padding_mask = torch.cat(
-                    [filler.float(), key_padding_mask.float()], dim=1
-                )
-            else:
-                new_key_padding_mask = key_padding_mask.float()
-        else:
-            new_key_padding_mask = prev_key_padding_mask
-        return new_key_padding_mask
-
-    @torch.jit.export
-    def reorder_incremental_state(
-        self,
-        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
-        new_order: Tensor,
-    ):
-        """Reorder buffered internal state (for incremental generation)."""
-        input_buffer = self._get_input_buffer(incremental_state)
-        if input_buffer is not None:
-            for k in input_buffer.keys():
-                input_buffer_k = input_buffer[k]
-                if input_buffer_k is not None:
-                    if self.encoder_decoder_attention:
-                        if input_buffer_k.size(0) * self.beam_size == new_order.size(0):
-                            return incremental_state
-                        elif self.beam_size > 1:
-                            input_buffer[k] = input_buffer_k.index_select(
-                                0,
-                                new_order.reshape(-1, self.beam_size)[:, 0]
-                                // self.beam_size,
-                            )
-                        else:
-                            input_buffer[k] = input_buffer_k.index_select(0, new_order)
-                    else:
-                        input_buffer[k] = input_buffer_k.index_select(0, new_order)
-            incremental_state = self._set_input_buffer(incremental_state, input_buffer)
-        return incremental_state
-
-    def set_beam_size(self, beam_size):
-        """Used for effiecient beamable enc-dec attention"""
-        self.beam_size = beam_size
-
-    def _get_input_buffer(
-        self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]
-    ) -> Dict[str, Optional[Tensor]]:
-        result = self.get_incremental_state(incremental_state, "attn_state")
-        if result is not None:
-            return result
-        else:
-            empty_result: Dict[str, Optional[Tensor]] = {}
-            return empty_result
-
-    def _set_input_buffer(
-        self,
-        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
-        buffer: Dict[str, Optional[Tensor]],
-    ):
-        return self.set_incremental_state(incremental_state, "attn_state", buffer)
-
-    def apply_sparse_mask(self, attn_weights, tgt_len: int, src_len: int, bsz: int):
-        return attn_weights
-
-    def upgrade_state_dict_named(self, state_dict, name):
-        prefix = name + "." if name != "" else ""
-        items_to_add = {}
-        keys_to_remove = []
-        for k in state_dict.keys():
-            if k.endswith(prefix + "in_proj_weight"):
-                # in_proj_weight used to be q + k + v with same dimensions
-                dim = int(state_dict[k].shape[0] / 3)
-                items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim]
-                items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim : 2 * dim]
-                items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim :]
-
-                keys_to_remove.append(k)
-
-                k_bias = prefix + "in_proj_bias"
-                if k_bias in state_dict.keys():
-                    dim = int(state_dict[k].shape[0] / 3)
-                    items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim]
-                    items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][
-                        dim : 2 * dim
-                    ]
-                    items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :]
-
-                    keys_to_remove.append(prefix + "in_proj_bias")
-
-        for k in keys_to_remove:
-            del state_dict[k]
-
-        for key, value in items_to_add.items():
-            state_dict[key] = value
-            
-            
-            
-            
-            
-            
-            
-            
-            
-            
-            
-class FairseqDropout(nn.Module):
-    def __init__(self, p, module_name=None):
-        super().__init__()
-        self.p = p
-        self.module_name = module_name
-        self.apply_during_inference = False
-
-    def forward(self, x, inplace: bool = False):
-        if self.p > 0 and (self.training or self.apply_during_inference):
-            return F.dropout(x, p=self.p, training=True, inplace=inplace)
-        else:
-            return x
-
-    def make_generation_fast_(
-        self,
-        name: str,
-        retain_dropout: bool = False,
-        retain_dropout_modules: Optional[List[str]] = None,
-        **kwargs
-    ):
-        if retain_dropout:
-            if retain_dropout_modules is not None and self.module_name is None:
-                logger.warning(
-                    "Cannot enable dropout during inference for module {} "
-                    "because module_name was not set".format(name)
-                )
-            elif (
-                retain_dropout_modules is None  # if None, apply to all modules
-                or self.module_name in retain_dropout_modules
-            ):
-                logger.info(
-                    "Enabling dropout during inference for module: {}".format(name)
-                )
-                self.apply_during_inference = True
-            else:
-                logger.info("Disabling dropout for module: {}".format(name))
-                
-                
-def quant_noise(module, p, block_size):
-    """
-    Wraps modules and applies quantization noise to the weights for
-    subsequent quantization with Iterative Product Quantization as
-    described in "Training with Quantization Noise for Extreme Model Compression"
-
-    Args:
-        - module: nn.Module
-        - p: amount of Quantization Noise
-        - block_size: size of the blocks for subsequent quantization with iPQ
-
-    Remarks:
-        - Module weights must have the right sizes wrt the block size
-        - Only Linear, Embedding and Conv2d modules are supported for the moment
-        - For more detail on how to quantize by blocks with convolutional weights,
-          see "And the Bit Goes Down: Revisiting the Quantization of Neural Networks"
-        - We implement the simplest form of noise here as stated in the paper
-          which consists in randomly dropping blocks
-    """
-
-    # if no quantization noise, don't register hook
-    if p <= 0:
-        return module
-
-    # supported modules
-    assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d))
-
-    # test whether module.weight has the right sizes wrt block_size
-    is_conv = module.weight.ndim == 4
-
-    # 2D matrix
-    if not is_conv:
-        assert (
-            module.weight.size(1) % block_size == 0
-        ), "Input features must be a multiple of block sizes"
-
-    # 4D matrix
-    else:
-        # 1x1 convolutions
-        if module.kernel_size == (1, 1):
-            assert (
-                module.in_channels % block_size == 0
-            ), "Input channels must be a multiple of block sizes"
-        # regular convolutions
-        else:
-            k = module.kernel_size[0] * module.kernel_size[1]
-            assert k % block_size == 0, "Kernel size must be a multiple of block size"
-
-    def _forward_pre_hook(mod, input):
-        # no noise for evaluation
-        if mod.training:
-            if not is_conv:
-                # gather weight and sizes
-                weight = mod.weight
-                in_features = weight.size(1)
-                out_features = weight.size(0)
-
-                # split weight matrix into blocks and randomly drop selected blocks
-                mask = torch.zeros(
-                    in_features // block_size * out_features, device=weight.device
-                )
-                mask.bernoulli_(p)
-                mask = mask.repeat_interleave(block_size, -1).view(-1, in_features)
-
-            else:
-                # gather weight and sizes
-                weight = mod.weight
-                in_channels = mod.in_channels
-                out_channels = mod.out_channels
-
-                # split weight matrix into blocks and randomly drop selected blocks
-                if mod.kernel_size == (1, 1):
-                    mask = torch.zeros(
-                        int(in_channels // block_size * out_channels),
-                        device=weight.device,
-                    )
-                    mask.bernoulli_(p)
-                    mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels)
-                else:
-                    mask = torch.zeros(
-                        weight.size(0), weight.size(1), device=weight.device
-                    )
-                    mask.bernoulli_(p)
-                    mask = (
-                        mask.unsqueeze(2)
-                        .unsqueeze(3)
-                        .repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1])
-                    )
-
-            # scale weights and apply mask
-            mask = mask.to(
-                torch.bool
-            )  # x.bool() is not currently supported in TorchScript
-            s = 1 / (1 - p)
-            mod.weight.data = s * weight.masked_fill(mask, 0)
-
-    module.register_forward_pre_hook(_forward_pre_hook)
-    return module
-
-
-
-        
-        
-        
-
-
-
-
-
-def softmax(x, dim: int, onnx_trace: bool = False):
-    if onnx_trace:
-        return F.softmax(x.float(), dim=dim)
-    else:
-        return F.softmax(x, dim=dim, dtype=torch.float32)
-    
-def log_softmax(x, dim: int, onnx_trace: bool = False):
-    if onnx_trace:
-        return F.log_softmax(x.float(), dim=dim)
-    else:
-        return F.log_softmax(x, dim=dim, dtype=torch.float32)

ablang2/models/ablang1/model.py

@@ -1,102 +0,0 @@
-import torch
-
-from .extra_fns import ACT2FN
-from .encoderblocks import EncoderBlocks
-from .embedding import AbEmbeddings
-
-
-class AbLang(torch.nn.Module):
-    """
-    Pretraining model includes Abrep and the head model used for training.
-    """
-    def __init__(self, hparams):
-        super().__init__()
-        self.hparams = hparams
-        
-        self.AbRep = AbRep(self.hparams)       
-        self.AbHead = AbHead(self.hparams)
-        
-    def forward(self, x, attention_mask=None):
-        
-        representations = self.AbRep(x, attention_mask)
-        
-        output = self.AbHead(representations.last_hidden_states)
-        
-        return output
-    
-    def get_aa_embeddings(self):
-        "This function is used to extract the trained aa_embeddings."
-        return self.AbRep.AbEmbeddings.aa_embeddings#().weight.detach()
-
-    
-class AbRep(torch.nn.Module):
-    """
-    This is the AbRep model.
-    """
-    def __init__(self, hparams):
-        super().__init__()
-        self.hparams = hparams
-        
-        self.AbEmbeddings = AbEmbeddings(self.hparams)    
-        self.EncoderBlocks = EncoderBlocks(self.hparams)
-        
-        self.init_weights()
-        
-    def forward(self, src, attention_mask=None, output_attentions=False):
-        
-        attention_mask = torch.zeros(*src.shape, device=src.device).masked_fill(src == self.hparams.pad_token_id, 1)
-
-        src = self.AbEmbeddings(src)
-        
-        output = self.EncoderBlocks(src, attention_mask=attention_mask, output_attentions=output_attentions)
-        
-        return output
-    
-    def _init_weights(self, module):
-        """ Initialize the weights """
-        if isinstance(module, (torch.nn.Linear, torch.nn.Embedding)):
-            module.weight.data.normal_(mean=0.0, std=self.hparams.initializer_range)
-        elif isinstance(module, torch.nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-        if isinstance(module, torch.nn.Linear) and module.bias is not None:
-            module.bias.data.zero_()
-            
-    def init_weights(self):
-        """
-        Initializes and prunes weights if needed.
-        """
-        # Initialize weights
-        self.apply(self._init_weights)
-    
-
-class AbHead(torch.nn.Module):
-    """
-    Head for masked sequence prediction.
-    """
-
-    def __init__(self, hparams):
-        super().__init__()
-        self.hparams = hparams
-        self.dense = torch.nn.Linear(self.hparams.hidden_size, self.hparams.hidden_size)
-        self.layer_norm = torch.nn.LayerNorm(self.hparams.hidden_size, eps=self.hparams.layer_norm_eps)
-
-        self.decoder = torch.nn.Linear(self.hparams.hidden_size, self.hparams.vocab_size, bias=False)
-        self.bias = torch.nn.Parameter(torch.zeros(self.hparams.vocab_size))
-        
-        self.activation = ACT2FN[self.hparams.hidden_act]
-        
-        ## self.init_weights() - need to have a function doing this
-
-        self.decoder.bias = self.bias # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
-
-    def forward(self, features, **kwargs):
-        x = self.dense(features)
-
-        x = self.activation(x)
-        x = self.layer_norm(x)
-
-        # project back to size of vocabulary with bias
-        x = self.decoder(x)
-
-        return x

ablang2/models/ablang1/pretrained.py

@@ -1,358 +0,0 @@
-import os, json, argparse, string, subprocess, re
-from dataclasses import dataclass
-
-from numba import jit
-from numba.typed import Dict, List
-from numba.types import unicode_type, DictType
-
-import numpy as np
-import torch
-import requests
-
-from . import tokenizers, model
-
-
-class pretrained:
-    """
-    Initializes AbLang for heavy or light chains.    
-    """
-    
-    def __init__(self, chain="heavy", model_folder="download", random_init=False, ncpu=7, device='cpu'):
-        super().__init__()
-        
-        self.used_device = torch.device(device)
-        
-        if model_folder == "download":
-            # Download model and save to specific place - if already downloaded do not download again
-            model_folder = os.path.join(os.path.dirname(__file__), "model-weights-{}".format(chain))
-            os.makedirs(model_folder, exist_ok = True)
-            
-            if not os.path.isfile(os.path.join(model_folder, "amodel.pt")):
-                print("Downloading model ...")
-                
-                url = "https://opig.stats.ox.ac.uk/data/downloads/ablang-{}.tar.gz".format(chain)
-                tmp_file = os.path.join(model_folder, "tmp.tar.gz")
-
-                with open(tmp_file,'wb') as f: f.write(requests.get(url).content)
-                
-                subprocess.run(["tar", "-zxvf", tmp_file, "-C", model_folder], check = True) 
-                
-                os.remove(tmp_file)
-
-        self.hparams_file = os.path.join(model_folder, 'hparams.json')
-        self.model_file = os.path.join(model_folder, 'amodel.pt')
-        
-        with open(self.hparams_file, 'r', encoding='utf-8') as f:
-            self.hparams = argparse.Namespace(**json.load(f))    
-
-        self.AbLang = model.AbLang(self.hparams)
-        self.AbLang.to(self.used_device)
-        
-        if not random_init:
-            self.AbLang.load_state_dict(torch.load(self.model_file, map_location=self.used_device))
-        
-        self.tokenizer = tokenizers.ABtokenizer(os.path.join(model_folder, 'vocab.json'))
-        self.AbRep = self.AbLang.AbRep
-        
-        self.ncpu = ncpu
-        self.spread = 11 # Based on get_spread_sequences function
-        if chain == 'heavy':
-            self.max_position = 128
-        else:
-            self.max_position = 127
-        
-        
-    def freeze(self):
-        self.AbLang.eval()
-        
-    def unfreeze(self):
-        self.AbLang.train()
-        
-    def __call__(self, sequence, mode='seqcoding', align=False, splitSize=50):
-        """
-        Mode: sequence, residue, restore or likelihood.
-        """
-        if not mode in ['rescoding', 'seqcoding', 'restore', 'likelihood']:
-            raise SyntaxError("Given mode doesn't exist.")
-        
-        if isinstance(sequence, str): sequence = [sequence]
-        
-        
-        if align and mode=='restore':
-            sequence = self.sequence_aligning(sequence)
-            splitSize = ((splitSize//self.spread)+1)*self.spread
-        
-        aList = []
-        for sequence_part in [sequence[x:x+splitSize] for x in range(0, len(sequence), splitSize)]:
-            aList.append(getattr(self, mode)(sequence_part, align))
-        
-        if mode == 'rescoding':
-            if align==True:
-                return aList
-            
-            return sum(aList, [])
-        
-        return np.concatenate(aList)
-    
-    def seqcoding(self, seqs, align=False):
-        """
-        Sequence specific representations
-        """
-        
-        tokens = self.tokenizer(seqs, pad=True, device=self.used_device)
-
-        residue_states = self.AbRep(tokens).last_hidden_states
-        
-        if torch.is_tensor(residue_states): residue_states = residue_states.cpu().detach().numpy()
-        
-        lens = np.vectorize(len)(seqs)
-        
-        lens = np.tile(lens.reshape(-1,1,1), (residue_states.shape[2], 1))
-        
-        seq_codings = np.apply_along_axis(res_to_seq, 2, np.c_[np.swapaxes(residue_states,1,2), lens])
-        
-        del lens
-        del residue_states
-        
-        return seq_codings
-    
-    def restore(self, seqs, align=False):
-        """
-        Restore sequences
-        """
-
-        if align:
-            nr_seqs = len(seqs)//self.spread
-
-            tokens = self.tokenizer(seqs, pad=True, device=self.used_device)
-            predictions = self.AbLang(tokens)[:,:,1:21]
-            
-            # Reshape
-            tokens = tokens.reshape(nr_seqs, self.spread, -1)
-            predictions = predictions.reshape(nr_seqs, self.spread, -1, 20)
-            seqs = seqs.reshape(nr_seqs, -1)
-            
-            # Find index of best predictions
-            best_seq_idx = torch.argmax(torch.max(predictions, -1).values[:,:,1:2].mean(2), -1)           
-            
-            # Select best predictions           
-            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, 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, encode=False)
-
-        return np.array([res_to_seq(seq, 'reconstruct') for seq in np.c_[restored_seqs, np.vectorize(len)(seqs)]])
-    
-    def likelihood(self, seqs, align=False):
-        """
-        Possible Mutations
-        """
-        
-        tokens = self.tokenizer(seqs, pad=True, device=self.used_device)
-        
-        predictions = self.AbLang(tokens)[:,:,1:21]
-        
-        if torch.is_tensor(predictions): predictions = predictions.cpu().detach().numpy()
-
-        return predictions
-    
-    def rescoding(self, seqs, align=False):
-        """
-        Residue specific representations.
-        """
-           
-        if align:
-            
-            import pandas as pd
-            import anarci
-            
-            anarci_out = anarci.run_anarci(pd.DataFrame(seqs).reset_index().values.tolist(), ncpu=7, scheme='imgt')
-            number_alignment = get_number_alignment(anarci_out)
-            
-            seqs = np.array([''.join([i[1] for i in onarci[0][0]]).replace('-','') for onarci in anarci_out[1]])
-            
-            tokens = self.tokenizer(seqs, pad=True, device=self.used_device)
-            residue_states = self.AbRep(tokens).last_hidden_states
-            
-            if torch.is_tensor(residue_states): residue_states = residue_states.cpu().detach().numpy()
-            
-            residue_output = np.array([create_alignment(res_embed, oanarci, seq, number_alignment) for res_embed, oanarci, seq in zip(residue_states, anarci_out[1], seqs)])
-            del residue_states
-            del tokens
-            
-            return output(aligned_embeds=residue_output, number_alignment=number_alignment.apply(lambda x: '{}{}'.format(*x[0]), axis=1).values)
-            
-        else:
-            
-            tokens = self.tokenizer(seqs, pad=True, device=self.used_device)
-            residue_states = self.AbRep(tokens).last_hidden_states
-        
-            if torch.is_tensor(residue_states): residue_states = residue_states.cpu().detach().numpy()
-            
-            residue_output = [res_to_list(state, seq) for state, seq in zip(residue_states, seqs)]
-
-            return residue_output
-        
-    def sequence_aligning(self, seqs):
-        
-        import pandas as pd
-        import anarci
-
-        anarci_out = anarci.run_anarci(
-            pd.DataFrame([seq.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 anarci_out[1]], columns=['anarci']).astype('<U90')
-
-        seqs = anarci_data.apply(lambda x: get_sequences_from_anarci(x.anarci, 
-                                                                     self.max_position, 
-                                                                     self.spread), axis=1, result_type='expand').to_numpy().reshape(-1)
-        
-        return seqs
-        
-        
-            
-            
-
-@dataclass
-class output():
-    """
-    Dataclass used to store output.
-    """
-
-    aligned_embeds: None
-    number_alignment: None
-    
-    
-def res_to_list(state, seq):
-    return state[1:1+len(seq)]
-
-def res_to_seq(a, mode='mean'):
-    """
-    Function for how we go from n_values for each amino acid to n_values for each sequence.
-    
-    We leave out the start, end and padding tokens.
-    """
-    if mode=='sum':
-        return a[1:(1+int(a[-1]))].sum()
-    
-    elif mode=='mean':
-        return a[1:(1+int(a[-1]))].mean()
-    
-    elif mode=='reconstruct':
-        
-        return a[0][1:(1+int(a[-1]))]
-
-def get_number_alignment(oanarci):
-    """
-    Creates a number alignment from the anarci results.
-    """
-    
-    import pandas as pd
-    
-    alist = []
-    
-    for aligned_seq in oanarci[1]:
-        alist.append(pd.DataFrame(aligned_seq[0][0])[0])
-
-    unsorted_alignment = pd.concat(alist).drop_duplicates()
-    max_alignment = get_max_alignment()
-    
-    return max_alignment.merge(unsorted_alignment.to_frame(), left_on=0, right_on=0)
-
-def get_max_alignment():
-    """
-    Create maximum possible alignment for sorting
-    """
-    
-    import pandas as pd
-
-    sortlist = []
-
-    for num in range(1, 128+1):
-
-        if num==112:
-            for char in string.ascii_uppercase[::-1]:
-                sortlist.append([(num, char)])
-
-            sortlist.append([(num,' ')])
-
-        else:
-            sortlist.append([(num,' ')])
-            for char in string.ascii_uppercase:
-                sortlist.append([(num, char)])
-                
-    return pd.DataFrame(sortlist)
-
-
-def create_alignment(res_embeds, oanarci, seq, number_alignment):
-    
-    import pandas as pd
-
-    datadf = pd.DataFrame(oanarci[0][0])
-
-    sequence_alignment = number_alignment.merge(datadf, how='left', on=0).fillna('-')[1]
-        
-    idxs = np.where(sequence_alignment.values == '-')[0]
-
-    idxs = [idx-num for num, idx in enumerate(idxs)]
-
-    aligned_embeds = pd.DataFrame(np.insert(res_embeds[1:1+len(seq)], idxs , 0, axis=0))
-
-    return pd.concat([aligned_embeds, sequence_alignment], axis=1).values
-
-def turn_into_numba(anarcis):
-    """
-    Turns the nested anarci dictionary into a numba item, allowing us to use numba on it.
-    """
-    
-    anarci_list = List.empty_list(unicode_type)
-    [anarci_list.append(str(anarci)) for anarci in anarcis]
-
-    return anarci_list
-
-@jit(nopython=True)
-def get_spread_sequences(seq, spread, start_position, numbaList):
-    """
-    Test sequences which are 8 positions shorter (position 10 + max CDR1 gap of 7) up to 2 positions longer (possible insertions).
-    """
-
-    for diff in range(start_position-8, start_position+2+1):
-        numbaList.append('*'*diff+seq)
-    
-    return numbaList
-
-def get_sequences_from_anarci(out_anarci, max_position, spread):
-    """
-    Ensures correct masking on each side of sequence
-    """
-    
-    if out_anarci == 'ANARCI_error':
-        return np.array(['ANARCI-ERR']*spread)
-    
-    end_position = int(re.search(r'\d+', out_anarci[::-1]).group()[::-1])
-    # Fixes ANARCI error of poor numbering of the CDR1 region
-    start_position = int(re.search(r'\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+\'\),\s\(\(\d+,\s\'.\'\),\s\'[^-]+',
-                                   out_anarci).group().split(',')[0]) - 1
-    
-    sequence = "".join(re.findall(r"(?i)[A-Z*]", "".join(re.findall(r'\),\s\'[A-Z*]', out_anarci))))
-
-    sequence_j = ''.join(sequence).replace('-','').replace('X','*') + '*'*(max_position-int(end_position))
-    
-    numba_list = List.empty_list(unicode_type)
-
-    spread_seqs = np.array(get_spread_sequences(sequence_j, spread, start_position, numba_list))
-
-    return spread_seqs
-
-

ablang2/models/ablang1/tokenizers.py

@@ -1,50 +0,0 @@
-import json
-import torch
-
-class ABtokenizer():
-    """
-    Tokenizer for proteins. Both aa to token and token to aa.
-    """
-    
-    def __init__(self, vocab_dir):
-        self.set_vocabs(vocab_dir)
-        self.pad_token = self.vocab_to_token['-']
-        
-    def __call__(self, sequenceList, encode=True, pad=False, device='cpu'):
-        #assert isinstance(sequenceList, list)
-        
-        if encode: 
-            data = [self.encode(seq, device=device) for seq in sequenceList]
-            if pad: return torch.nn.utils.rnn.pad_sequence(data, batch_first=True, padding_value=self.pad_token)
-            else: return data
-        
-        else: return [self.decode(token) for token in sequenceList]
-    
-    def set_vocabs(self, vocab_dir):
-        with open(vocab_dir, encoding="utf-8") as vocab_handle:
-            self.vocab_to_token=json.load(vocab_handle)
-            
-        self.vocab_to_aa = {v: k for k, v in self.vocab_to_token.items()}
-     
-    def encode(self, sequence, device='cpu'):
-        try:
-            encoded = [self.vocab_to_token["<"]]+[self.vocab_to_token[resn] for resn in sequence]+[self.vocab_to_token[">"]]
-        except KeyError as e:
-            
-            wrong_aa = e.args
-            
-            e.args = (f"Following character(s) not accepted in sequences: {wrong_aa}. \
-Please only use amino acids (MRHKDESTNQCGPAVIFYWL) or the mask token (*).",)
-            raise 
-        
-        return torch.tensor(encoded, dtype=torch.long, device=device)
-        # Start and Stop token should probably not be added here, but instead earlier
-    
-    def decode(self, seqtokens):
-        
-        if torch.is_tensor(seqtokens): seqtokens = seqtokens.cpu().numpy()
-
-        return ''.join([self.vocab_to_aa[token] for token in seqtokens])
-    
-
-