Added Tranception model

README.md

@@ -1,3 +1,133 @@
----
-license: mit
----
+# Tranception
+
+This is the official code repository for the paper "Tranception: protein fitness prediction with autoregressive transformers and inference-time retrieval". This project is a joint collaboration between the [Marks lab](https://www.deboramarkslab.com/) and the [OATML group](https://oatml.cs.ox.ac.uk/).
+
+## Abstract
+The ability to accurately model the fitness landscape of protein sequences is critical to a wide range of applications, from quantifying the effects of human variants on disease likelihood, to predicting immune-escape mutations in viruses and designing novel biotherapeutic proteins. Deep generative models of protein sequences trained on multiple sequence alignments have been the most successful approaches so far to address these tasks. The performance of these methods is however contingent on the availability of sufficiently deep and diverse alignments for reliable training. Their potential scope is thus limited by the fact many protein families are hard, if not impossible, to align. Large language models trained on massive quantities of non-aligned protein sequences from diverse families address these problems and show potential to eventually bridge the performance gap. We introduce Tranception, a novel transformer architecture leveraging autoregressive predictions and retrieval of homologous sequences at inference to achieve state-of-the-art fitness prediction performance. Given its markedly higher performance on multiple mutants, robustness to shallow alignments and ability to score indels, our approach offers significant gain of scope over existing approaches. To enable more rigorous model testing across a broader range of protein families, we develop ProteinGym -- an extensive set of multiplexed assays of variant effects, substantially increasing both the number and diversity of assays compared to existing benchmarks.
+
+## Setup
+You may download the Tranception repository and create a conda environment with the proper dependencies (as listed in `tranception_env.yml`) as follows:
+```
+git clone https://github.com/OATML-Markslab/Tranception.git
+conda env create -f tranception_env.yml
+```
+
+## Tranception
+Tranception is a novel autoregressive transformer architecture that was designed with two core principles in mind: 1) promoting specialization across attention heads 2) explicitly extracting patterns from contiguous subsequences.
+
+To download the *Tranception Large* model checkpoint (~3.1GB unzipped):
+```
+curl -o Tranception_Large.zip https://marks.hms.harvard.edu/ProteinGym/Tranception_Large.zip
+unzip Tranception_Large.zip
+rm Tranception_Large.zip
+```
+
+Tranception is also made available through the [Huggging Face hub](https://huggingface.co/OATML-Markslab/Tranception).
+
+When scoring with retrieval, we compute weighted pseudocounts at each position using sequence weights as per the procedure described in [Hopf et al.](https://www.nature.com/articles/nbt.3769).
+Weights for all proteins in the ProteinGym benchmarks may be downloaded as follows (~68M unzipped):
+```
+curl -o MSA_weights.zip https://marks.hms.harvard.edu/ProteinGym/MSA_weights.zip
+unzip MSA_weights.zip
+rm MSA_weights.zip
+```
+To compute sequence weights for new proteins, you may use the MSA_processing class under `tranception/utils/msa_utils.py`.
+
+The `examples` folder provides several bash scripts that may be used for scoring and evaluating Tranception on the ProteinGym benchmarks. We also provide a colab notebook illustrating how to load Tranception from the Hugging Face hub and then score the mutated sequences in the ProteinGym benchmarks with it.
+
+## ProteinGym
+ProteinGym is an extensive set of Deep Mutational Scanning (DMS) assays curated to enable thorough comparisons of various mutation effect predictors indifferent regimes. ProteinGym is comprised of two benchmarks: 1) a substitution benchmark which consists of the experimental characterisation of ∼1.5M missense variants across 87 DMS assays 2) an indel benchmark that includes ∼300k mutants across 7 DMS assays.
+
+Each processed file in each benchmark corresponds to a single DMS assay, and contains the following three variables:
+- mutant (str): 
+    - for the substitution benchmark, it describes the set of substitutions to apply on the reference sequence to obtain the mutated sequence (eg., A1P:D2N implies the amino acid 'A' at position 1 should be replaced by 'P', and 'D' at position 2 should be replaced by 'N')
+    - for the indel benchmark, it corresponds to the full mutated sequence
+- DMS_score (float): corresponds to the experimental measurement in the DMS assay. Across all assays, the higher the DMS_score value, the higher the fitness of the mutated protein
+- DMS_score_bin (int): indicates whether the DMS_score is above the fitness cutoff (1 is fit, 0 is not fit)
+
+Additionally, we provide reference files in the [ProteinGym folder](https://github.com/OATML-Markslab/Tranception/tree/main/ProteinGym) that give further details on each assay and contain in particular:
+- The UniProt_ID of the corresponding protein, along with taxon and MSA depth category
+- The target sequence (target_seq) used in the assay
+- Details on how the DMS_score was created from the raw files and how it was binarized 
+
+To download the substitution benchmark (~224M unzipped):
+```
+curl -o ProteinGym_substitutions.zip https://marks.hms.harvard.edu/ProteinGym/ProteinGym_substitutions.zip 
+unzip ProteinGym_substitutions.zip
+rm ProteinGym_substitutions.zip
+```
+
+Similarly, to download the indel benchmark (~86M unzipped):
+```
+curl -o ProteinGym_indels.zip https://marks.hms.harvard.edu/ProteinGym/ProteinGym_indels.zip
+unzip ProteinGym_indels.zip
+rm ProteinGym_indels.zip
+```
+
+## Fitness prediction performance
+
+The [proteingym folder](https://github.com/OATML-Markslab/Tranception/tree/main/ProteinGym) provides detailed performance files for Tranception and baselines on the two ProteinGym benchmarks. 
+
+We recommand to aggregate fitness pediction performance at the Uniprot ID level to avoid biasing results towards proteins for which several DMS assays are available in ProteinGym. The corresponding aggregated files are suffixed with "_Uniprot_level", while the non aggregated performance files are suffixed with "_DMS_level".
+Furthermore, to enable fair comparison with models trained multiple-sequence alignments (eg., EVE, DeepSequence, EVmutation), we only evaluate on the subset of mutations where position coverage is deemed high enough by these models to make a prediction. The corresponding files are preffixed with "All_models_". For comprehensiveness, we also provide performance files on all possible mutants available in ProteinGym, comparing only with the baselines that are able to score all mutants.
+Note that for the ProteinGym indel benchmark, baselines that are able to score indels do not have the aforementionned coverage constraints (ie., no distinction between "All_models_" and "All_mutants_") and there is at most one DMS per Uniprot_ID (ie., no difference between "_Uniprot_level" and "_DMS_level"). We thus only provide one set of performance metrics for that benchmark.
+
+### ProteinGym substitution benchmark - Leaderboard
+The table below provides the average Spearman's rank correlation between DMS experimental fitness measurements and fitness predictions from Tranception or other baselines on the ProteinGym substitution benchmark. Following the terminology introduced above, we report the performance at the "Uniprot" level for "All models".
+
+Rank | Model name | Spearman | Reference
+--- | --- | --- | --- |
+1 | Ensemble Tranception & EVE | 0.476 | [Notin et al.](https://arxiv.org/abs/2205.13760)
+2 | Tranception (w/ retrieval) | 0.451 | [Notin et al.](https://arxiv.org/abs/2205.13760)
+3 | EVE | 0.448 | [Frazer et al.](https://www.nature.com/articles/s41586-021-04043-8)
+4 | EVmutation | 0.427 | [Hopf et al.](https://www.nature.com/articles/nbt.3769)
+5 | MSA Transformer | 0.422 | [Rao et al.](https://proceedings.mlr.press/v139/rao21a.html)
+6 | DeepSequence | 0.415 | [Riesselman et al.](https://www.nature.com/articles/s41592-018-0138-4)
+7 | Tranception (no retrieval) | 0.406 | [Notin et al.](https://arxiv.org/abs/2205.13760)
+8 | Wavenet | 0.398 | [Shin et al.](https://www.nature.com/articles/s41467-021-22732-w)
+9 | Site Independent | 0.397 | [Hopf et al.](https://www.nature.com/articles/nbt.3769)
+10 | ESM-1v | 0.371 | [Meier et al.](https://proceedings.neurips.cc/paper/2021/hash/f51338d736f95dd42427296047067694-Abstract.html)
+
+### ProteinGym indel benchmark - Leaderboard
+The table below provides the average Spearman's rank correlation between DMS experimental fitness measurements and fitness predictions from Tranception or other baselines on the ProteinGym indel benchmark.
+
+Rank | Model name | Spearman | Reference
+--- | --- | --- | --- |
+[Notin et al.](https://arxiv.org/abs/2205.13760)
+[Notin et al.](https://arxiv.org/abs/2205.13760)
+[Shin et al.](https://www.nature.com/articles/s41467-021-22732-w)
+
+## Aggregated model scoring files
+The scores for all DMS assays in the ProteinGym substitution benchmark for Tranception and other baselines (eg., EVE, Wavenet, ESM-1v, MSA Transformer) may be downloaded as follows;
+```
+curl -o scores_all_models_proteingym_substitutions.zip https://marks.hms.harvard.edu/ProteinGym/scores_all_models_proteingym_substitutions.zip
+unzip scores_all_models_proteingym_substitutions.zip
+rm scores_all_models_proteingym_substitutions.zip
+```
+Similarly for the indel benchmark, all scoring files may be downloaded as follows:
+```
+curl -o scores_all_models_proteingym_indels.zip https://marks.hms.harvard.edu/ProteinGym/scores_all_models_proteingym_indels.zip
+unzip scores_all_models_proteingym_indels.zip
+rm scores_all_models_proteingym_indels.zip
+```
+
+## Multiple Sequence Alignments (MSAs)
+
+The MSAs used to train alignment-based methods or used at inference in Tranception with retrieval and MSA Transformer may be downloaded as follows (~2.2GB unzipped):
+```
+curl -o MSA_ProteinGym.zip https://marks.hms.harvard.edu/ProteinGym/MSA_ProteinGym.zip
+unzip MSA_ProteinGym.zip
+rm MSA_ProteinGym.zip
+```
+
+## License
+This project is available under the MIT license.
+
+## Reference
+If you use Tranception, ProteinGym or other files provided through this repository (eg., aggregated model scoring files) in your work, please cite the following paper:
+```
+Notin, P., Dias, M., Frazer, J., Marchena-Hurtado, J., Gomez, A., Marks, D.S., Gal, Y. (2022). Tranception: Protein Fitness Prediction with Autoregressive Transformers and Inference-time Retrieval. ICML.
+```
+
+## Links
+Pre-print: https://arxiv.org/abs/2205.13760

__init__.py

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

activations.py

@@ -0,0 +1,114 @@
+import math
+
+import torch
+from packaging import version
+from torch import nn
+
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def _gelu_python(x):
+    """
+    Original Implementation of the GELU activation function in Google BERT repo when initially created. For
+    information: OpenAI GPT's GELU is slightly different (and gives slightly different results): 0.5 * x * (1 +
+    torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) This is now written in C in nn.functional
+    Also see the Gaussian Error Linear Units paper: https://arxiv.org/abs/1606.08415
+    """
+    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+
+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))))
+
+
+if version.parse(torch.__version__) < version.parse("1.4"):
+    gelu = _gelu_python
+else:
+    gelu = nn.functional.gelu
+
+
+def gelu_fast(x):
+    return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x)))
+
+
+def quick_gelu(x):
+    return x * torch.sigmoid(1.702 * x)
+
+
+def _silu_python(x):
+    """
+    See Gaussian Error Linear Units (Hendrycks et al., https://arxiv.org/abs/1606.08415) where the SiLU (Sigmoid Linear
+    Unit) was originally introduced and coined, and see Sigmoid-Weighted Linear Units for Neural Network Function
+    Approximation in Reinforcement Learning (Elfwing et al., https://arxiv.org/abs/1702.03118) and Swish: a Self-Gated
+    Activation Function (Ramachandran et al., https://arxiv.org/abs/1710.05941v1) where the SiLU was experimented with
+    later.
+    """
+    return x * torch.sigmoid(x)
+
+
+if version.parse(torch.__version__) < version.parse("1.7"):
+    silu = _silu_python
+else:
+    silu = nn.functional.silu
+
+
+def _mish_python(x):
+    """
+    See Mish: A Self-Regularized Non-Monotonic Activation Function (Misra., https://arxiv.org/abs/1908.08681). Also
+    visit the official repository for the paper: https://github.com/digantamisra98/Mish
+    """
+    return x * torch.tanh(nn.functional.softplus(x))
+
+
+if version.parse(torch.__version__) < version.parse("1.9"):
+    mish = _mish_python
+else:
+    mish = nn.functional.mish
+
+
+def linear_act(x):
+    return x
+
+def squared_relu(x):
+    """
+    Squared ReLU variant that is fastest with Pytorch.
+    """
+    x = nn.functional.relu(x)
+    return x*x
+
+def squared_relu_xla(x):
+    """
+    Squared ReLU variant that is fastest with JAX.
+    """
+    x = nn.functional.relu(x)
+    return x**2
+
+tranception_ACT2FN = {
+    "relu": nn.functional.relu,
+    "silu": silu,
+    "swish": silu,
+    "gelu": gelu,
+    "tanh": torch.tanh,
+    "gelu_new": gelu_new,
+    "gelu_fast": gelu_fast,
+    "quick_gelu": quick_gelu,
+    "mish": mish,
+    "linear": linear_act,
+    "sigmoid": torch.sigmoid,
+    "squared_relu": squared_relu,
+    "squared_relu_xla": squared_relu_xla,
+}
+
+
+def get_activation(activation_string):
+    if activation_string in tranception_ACT2FN:
+        return tranception_ACT2FN[activation_string]
+    else:
+        raise KeyError(f"function {activation_string} not found in ACT2FN mapping {list(tranception_ACT2FN.keys())}")

config.json

@@ -0,0 +1,46 @@
+{
+  "MSA_end": null,
+  "MSA_filename": null,
+  "MSA_start": null,
+  "MSA_weight_file_name": null,
+  "_name_or_path": "Tranception_Large",
+  "activation_function": "squared_relu",
+  "architectures": [
+    "TranceptionLMHeadModel"
+  ],
+  "attention_mode": "tranception",
+  "attn_pdrop": 0.1,
+  "bos_token_id": 1,
+  "clustal_omega_location": null,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 2,
+  "full_protein_length": null,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "local_batch_size": 1,
+  "model_type": "tranception",
+  "n_ctx": 1024,
+  "n_embd": 1280,
+  "n_head": 20,
+  "n_inner": 5120,
+  "n_layer": 36,
+  "n_positions": 1024,
+  "position_embedding": "grouped_alibi",
+  "reorder_and_upcast_attn": false,
+  "resid_pdrop": 0.1,
+  "retrieval_aggregation_mode": null,
+  "retrieval_inference_weight": 0.6,
+  "scale_attn_by_inverse_layer_idx": false,
+  "scale_attn_weights": true,
+  "scoring_window": "optimal",
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "tokenizer": null,
+  "torch_dtype": "float32",
+  "transformers_version": "4.17.0",
+  "use_cache": true,
+  "vocab_size": 25
+}

config.py

@@ -0,0 +1,36 @@
+from transformers import GPT2Config
+
+class TranceptionConfig(GPT2Config):
+    """
+    Config subclass for Tranception model architecture.
+    """
+    def __init__(
+        self,
+        attention_mode="tranception",
+        position_embedding="grouped_alibi",
+        tokenizer=None,
+        retrieval_aggregation_mode=None,
+        retrieval_inference_weight=0.6,
+        MSA_filename=None,
+        MSA_weight_file_name=None,
+        MSA_start=None,
+        MSA_end=None,
+        full_protein_length=None,
+        clustal_omega_location=None,
+        scoring_window=None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.model_type="tranception"
+        self.attention_mode=attention_mode
+        self.position_embedding=position_embedding
+        self.tokenizer = tokenizer
+        self.retrieval_aggregation_mode = retrieval_aggregation_mode
+        self.retrieval_inference_weight = retrieval_inference_weight
+        self.MSA_filename = MSA_filename
+        self.MSA_weight_file_name = MSA_weight_file_name
+        self.MSA_start=MSA_start
+        self.MSA_end=MSA_end
+        self.full_protein_length = full_protein_length
+        self.clustal_omega_location = clustal_omega_location
+        self.scoring_window=scoring_window

model_pytorch.py

@@ -0,0 +1,917 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import math
+import os
+import pandas as pd
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss, NLLLoss
+import torch.nn.functional as F
+from transformers import GPT2PreTrainedModel
+
+from transformers.modeling_utils import (
+    Conv1D,
+    PreTrainedModel,
+    SequenceSummary,
+    find_pruneable_heads_and_indices,
+    prune_conv1d_layer,
+)
+from transformers.file_utils import (
+    ModelOutput,
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    replace_return_docstrings
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput
+)
+from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
+
+from tranception.activations import tranception_ACT2FN
+from tranception.config import TranceptionConfig
+from tranception.outputs import (
+    TranceptionCausalLMOutputWithCrossAttentions,
+)
+from tranception.utils import msa_utils
+from tranception.utils import scoring_utils
+
+def nanmean(v, *args, inplace=False, **kwargs):
+    if not inplace:
+        v = v.clone()
+    is_nan = torch.isnan(v)
+    v[is_nan] = 0
+    return v.sum(*args, **kwargs) / (~is_nan).float().sum(*args, **kwargs)
+
+def get_slopes(n, mode="standard_alibi", verbose=False):
+    """
+    Function to compute the m constant for each attention head. Code has been adapted from the official ALiBi codebase at:
+    https://github.com/ofirpress/attention_with_linear_biases/blob/master/fairseq/models/transformer.py
+    """
+    def get_slopes_power_of_2(n):
+        start = (2**(-2**-(math.log2(n)-3)))
+        ratio = start
+        return [start*ratio**i for i in range(n)]
+    if mode=="grouped_alibi":
+        n = n // 4
+    if math.log2(n).is_integer():
+        result = get_slopes_power_of_2(n)                   
+    else:
+        #Workaround when the number of heads is not a power of 2
+        closest_power_of_2 = 2**math.floor(math.log2(n))  
+        result = get_slopes_power_of_2(closest_power_of_2) + get_slopes(2*closest_power_of_2)[0::2][:n-closest_power_of_2]
+    if mode=="grouped_alibi":
+        result = result * 4
+        if verbose:
+            print("ALiBi slopes: {}".format(result))
+    return result
+
+class SpatialDepthWiseConvolution(nn.Module):
+    def __init__(self, head_dim: int, kernel_size: int = 3):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.conv = nn.Conv1d(in_channels=head_dim, out_channels=head_dim, kernel_size=(kernel_size,), padding=(kernel_size - 1,), groups=head_dim)
+    
+    def forward(self, x: torch.Tensor):
+        batch_size, heads, seq_len, head_dim = x.shape
+        x = x.permute(0, 1, 3, 2).contiguous()
+        x = x.view(batch_size * heads, head_dim, seq_len)
+        x = self.conv(x)
+        if self.kernel_size>1:
+            x = x[:, :, :-(self.kernel_size - 1)]
+        x = x.view(batch_size, heads, head_dim, seq_len)
+        x = x.permute(0, 1, 3, 2)
+        return x
+
+class TranceptionBlockAttention(nn.Module):
+    def __init__(self, config, is_cross_attention=False, SDWC_kernel_size=None):
+        super().__init__()
+
+        max_positions = config.max_position_embeddings
+        self.register_buffer(
+            "bias",
+            torch.tril(torch.ones((max_positions, max_positions), dtype=torch.uint8)).view(
+                1, 1, max_positions, max_positions
+            ),
+        )
+        self.register_buffer("masked_bias", torch.tensor(-1e4))
+
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        self.split_size = self.embed_dim
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
+            )
+
+        self.scale_attn_weights = config.scale_attn_weights
+        self.is_cross_attention = is_cross_attention
+
+        if self.is_cross_attention:
+            self.c_attn = Conv1D(2 * self.embed_dim, self.embed_dim)
+            self.q_attn = Conv1D(self.embed_dim, self.embed_dim)
+        else:
+            self.c_attn = Conv1D(3 * self.embed_dim, self.embed_dim)
+        self.c_proj = Conv1D(self.embed_dim, self.embed_dim)
+
+        self.attn_dropout = nn.Dropout(config.attn_pdrop)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+
+        self.pruned_heads = set()
+
+        self.attention_mode=config.attention_mode
+        
+        if self.attention_mode=="tranception":
+            assert self.num_heads%4==0, "Invalid number of heads. Tranception requires the number of heads to be a multiple of 4."
+            self.num_heads_per_kernel_size = self.num_heads // 4
+            self.query_depthwiseconv = nn.ModuleDict()
+            self.key_depthwiseconv = nn.ModuleDict()
+            self.value_depthwiseconv = nn.ModuleDict()
+            for kernel_idx, kernel in enumerate([3,5,7]):
+                self.query_depthwiseconv[str(kernel_idx)] = SpatialDepthWiseConvolution(self.head_dim,kernel)
+                self.key_depthwiseconv[str(kernel_idx)]   = SpatialDepthWiseConvolution(self.head_dim,kernel)
+                self.value_depthwiseconv[str(kernel_idx)] = SpatialDepthWiseConvolution(self.head_dim,kernel)
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(heads, self.num_heads, self.head_dim, self.pruned_heads)
+        index_attn = torch.cat([index, index + self.split_size, index + (2 * self.split_size)])
+
+        # Prune conv1d layers
+        self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1)
+        self.c_proj = prune_conv1d_layer(self.c_proj, index, dim=0)
+
+        # Update hyper params
+        self.split_size = (self.split_size // self.num_heads) * (self.num_heads - len(heads))
+        self.num_heads = self.num_heads - len(heads)
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def _attn(self, query, key, value, attention_mask=None, head_mask=None, alibi_bias=None):
+        attn_weights = torch.matmul(query, key.transpose(-1, -2))
+
+        if self.scale_attn_weights:
+            attn_weights = attn_weights / (float(value.size(-1)) ** 0.5)
+
+        if not self.is_cross_attention:
+            # if only "normal" attention layer implements causal mask
+            query_length, key_length = query.size(-2), key.size(-2)
+            causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
+            attn_weights = torch.where(causal_mask, attn_weights, self.masked_bias.to(attn_weights.dtype))
+
+        if alibi_bias is not None:
+            attn_weights = attn_weights + alibi_bias[:,:,:attn_weights.size(-1)]
+
+        if attention_mask is not None:
+            # Apply the attention mask
+            attn_weights = attn_weights + attention_mask
+
+        attn_weights = nn.Softmax(dim=-1)(attn_weights)
+        attn_weights = self.attn_dropout(attn_weights)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attn_weights = attn_weights * head_mask
+
+        attn_output = torch.matmul(attn_weights, value)
+
+        return attn_output, attn_weights
+
+    def _split_heads(self, tensor, num_heads, attn_head_size):
+        """
+        Splits hidden_size dim into attn_head_size and num_heads
+        """
+        new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
+        tensor = tensor.view(*new_shape)
+        return tensor.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)
+
+    def _merge_heads(self, tensor, num_heads, attn_head_size):
+        """
+        Merges attn_head_size dim and num_attn_heads dim into hidden_size
+        """
+        tensor = tensor.permute(0, 2, 1, 3).contiguous()
+        new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
+        return tensor.view(new_shape)
+
+    def forward(
+        self,
+        hidden_states,
+        layer_past=None,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        use_cache=False,
+        output_attentions=False,
+        alibi_bias=None,
+    ):
+        if encoder_hidden_states is not None:
+            if not hasattr(self, "q_attn"):
+                raise ValueError(
+                    "If class is used as cross attention, the weights `q_attn` have to be defined. "
+                    "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
+                )
+
+            query = self.q_attn(hidden_states)
+            key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
+            attention_mask = encoder_attention_mask
+        else:
+            query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
+
+        query = self._split_heads(query, self.num_heads, self.head_dim)
+        key = self._split_heads(key, self.num_heads, self.head_dim)
+        value = self._split_heads(value, self.num_heads, self.head_dim)
+
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            key = torch.cat((past_key, key), dim=-2)
+            value = torch.cat((past_value, value), dim=-2)
+
+        if use_cache is True:
+            present = (key, value)
+        else:
+            present = None
+        
+        if self.attention_mode=="tranception":
+            # We do not do anything on the first self.num_heads_per_kernel_size heads (kernel =1)
+            query_list=[query[:,:self.num_heads_per_kernel_size,:,:]]
+            key_list=[key[:,:self.num_heads_per_kernel_size,:,:]]
+            value_list=[value[:,:self.num_heads_per_kernel_size,:,:]]
+            for kernel_idx in range(3):
+                query_list.append(self.query_depthwiseconv[str(kernel_idx)](query[:,(kernel_idx+1)*self.num_heads_per_kernel_size:(kernel_idx+2)*self.num_heads_per_kernel_size,:,:]))
+                key_list.append(self.key_depthwiseconv[str(kernel_idx)](key[:,(kernel_idx+1)*self.num_heads_per_kernel_size:(kernel_idx+2)*self.num_heads_per_kernel_size,:,:]))
+                value_list.append(self.value_depthwiseconv[str(kernel_idx)](value[:,(kernel_idx+1)*self.num_heads_per_kernel_size:(kernel_idx+2)*self.num_heads_per_kernel_size,:,:]))
+            query=torch.cat(query_list, dim=1)
+            key=torch.cat(key_list, dim=1)
+            value=torch.cat(value_list, dim=1)
+        
+        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask, alibi_bias=alibi_bias)
+
+        attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
+        attn_output = self.c_proj(attn_output)
+        attn_output = self.resid_dropout(attn_output)
+
+        outputs = (attn_output, present)
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs  # a, present, (attentions)
+
+class TranceptionBlockMLP(nn.Module):
+    def __init__(self, intermediate_size, config):
+        super().__init__()
+        embed_dim = config.hidden_size
+        self.c_fc = Conv1D(intermediate_size, embed_dim)
+        self.c_proj = Conv1D(embed_dim, intermediate_size)
+        self.act = tranception_ACT2FN[config.activation_function]
+        self.dropout = nn.Dropout(config.resid_pdrop)
+    
+    def forward(self, hidden_states):
+        hidden_states = self.c_fc(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.c_proj(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        return hidden_states
+
+class TranceptionBlock(nn.Module):
+    def __init__(self, config, SDWC_kernel_size=None):
+        super().__init__()
+        hidden_size = config.hidden_size
+        inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
+
+        self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.attn = TranceptionBlockAttention(config, SDWC_kernel_size=SDWC_kernel_size)
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+
+        if config.add_cross_attention:
+            self.crossattention = TranceptionBlockAttention(config, is_cross_attention=True, SDWC_kernel_size=SDWC_kernel_size)
+            self.ln_cross_attn = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+
+        self.mlp = TranceptionBlockMLP(inner_dim, config)
+    
+    def forward(
+        self,
+        hidden_states,
+        layer_past=None,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        use_cache=False,
+        output_attentions=False,
+        alibi_bias=None,
+    ):
+        residual = hidden_states
+        hidden_states = self.ln_1(hidden_states)
+        attn_outputs = self.attn(
+            hidden_states,
+            layer_past=layer_past,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            alibi_bias=alibi_bias,
+        )
+        attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)
+        outputs = attn_outputs[1:]
+        # residual connection
+        hidden_states = attn_output + residual
+
+        if encoder_hidden_states is not None:
+            # add one self-attention block for cross-attention
+            if not hasattr(self, "crossattention"):
+                raise ValueError(
+                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
+                    "cross-attention layers by setting `config.add_cross_attention=True`"
+                )
+            residual = hidden_states
+            hidden_states = self.ln_cross_attn(hidden_states)
+            cross_attn_outputs = self.crossattention(
+                hidden_states,
+                attention_mask=attention_mask,
+                head_mask=head_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                output_attentions=output_attentions,
+            )
+            attn_output = cross_attn_outputs[0]
+            # residual connection
+            hidden_states = residual + attn_output
+            outputs = outputs + cross_attn_outputs[2:]  # add cross attentions if we output attention weights
+
+        residual = hidden_states
+        hidden_states = self.ln_2(hidden_states)
+
+        feed_forward_hidden_states = self.mlp(hidden_states)
+        
+        # residual connection
+        hidden_states = residual + feed_forward_hidden_states
+
+        if use_cache:
+            outputs = (hidden_states,) + outputs
+        else:
+            outputs = (hidden_states,) + outputs[1:]
+
+        return outputs  # hidden_states, present, (attentions, cross_attentions)
+
+class TranceptionModel(GPT2PreTrainedModel):
+    _keys_to_ignore_on_load_missing = ["attn.masked_bias"]
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embed_dim = config.hidden_size
+        self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
+        self.position_embedding = config.position_embedding if hasattr(config, "position_embedding") else "learned"
+        if self.position_embedding=="learned":
+            self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)
+            self.alibi = None
+        elif self.position_embedding=="grouped_alibi":
+            maxpos = config.n_positions
+            attn_heads = config.n_head
+            self.slopes = torch.Tensor(get_slopes(attn_heads, mode=self.position_embedding))
+            #The softmax operation is invariant to translation, and bias functions used are always linear. 
+            alibi = self.slopes.unsqueeze(1).unsqueeze(1) * torch.arange(maxpos).unsqueeze(0).unsqueeze(0).expand(attn_heads, -1, -1)
+            alibi = alibi.view(attn_heads, 1, maxpos)
+            self.register_buffer('alibi',alibi)
+
+        self.drop = nn.Dropout(config.embd_pdrop)
+        self.h = nn.ModuleList([TranceptionBlock(config) for _ in range(config.num_hidden_layers)])
+        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
+
+        self.init_weights()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+        self.gradient_checkpointing = False
+    
+    def parallelize(self, device_map=None, num_cores=None):
+        self.device_map = (
+                get_device_map(len(self.h), range(torch.cuda.device_count())) if device_map is None else device_map
+            )
+        device_prefix="cuda:"
+        assert_device_map(self.device_map, len(self.h))
+        self.model_parallel = True
+        self.first_device = "cpu" if "cpu" in self.device_map.keys() else device_prefix + str(min(self.device_map.keys()))
+        self.last_device = device_prefix + str(max(self.device_map.keys()))
+        self.wte = self.wte.to(self.first_device)
+        if self.position_embedding=="learned":
+            self.wpe = self.wpe.to(self.first_device)
+        for k, v in self.device_map.items():
+            print("k,v :"+str(k)+","+str(v))
+            for block in v:
+                cuda_device = device_prefix + str(k)
+                self.h[block] = self.h[block].to(cuda_device)
+        self.ln_f = self.ln_f.to(self.last_device)
+    
+    def deparallelize(self):
+        self.model_parallel = False
+        self.device_map = None
+        self.first_device = "cpu"
+        self.last_device = "cpu"
+        self.wte = self.wte.to("cpu")
+        if self.position_embedding=="learned":
+            self.wpe = self.wpe.to("cpu")
+        for index in range(len(self.h)):
+            self.h[index] = self.h[index].to("cpu")
+        self.ln_f = self.ln_f.to("cpu")
+        torch.cuda.empty_cache()
+
+    def get_input_embeddings(self):
+        return self.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.wte = new_embeddings
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
+        """
+        for layer, heads in heads_to_prune.items():
+            self.h[layer].attn.prune_heads(heads)
+
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            input_ids = input_ids.view(-1, input_shape[-1])
+            batch_size = input_ids.shape[0]
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size = inputs_embeds.shape[0]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids.view(-1, input_shape[-1])
+        if position_ids is not None:
+            position_ids = position_ids.view(-1, input_shape[-1])
+
+        if past_key_values is None:
+            past_length = 0
+            past_key_values = tuple([None] * len(self.h))
+        else:
+            past_length = past_key_values[0][0].size(-2)
+        if position_ids is None:
+            position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
+            position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])
+
+        # GPT2Attention mask.
+        if attention_mask is not None:
+            if batch_size <= 0:
+                raise ValueError("batch_size has to be defined and > 0")
+            attention_mask = attention_mask.view(batch_size, -1)
+            # We create a 3D attention mask from a 2D tensor mask.
+            # Sizes are [batch_size, 1, 1, to_seq_length]
+            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
+            # this attention mask is more simple than the triangular masking of causal attention
+            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
+            attention_mask = attention_mask[:, None, None, :]
+
+            # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+            # masked positions, this operation will create a tensor which is 0.0 for
+            # positions we want to attend and -10000.0 for masked positions.
+            # Since we are adding it to the raw scores before the softmax, this is
+            # effectively the same as removing these entirely.
+            attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+            attention_mask = (1.0 - attention_mask) * -10000.0
+
+        # If a 2D ou 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.config.add_cross_attention and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+            encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # head_mask has shape n_layer x batch x n_heads x N x N
+        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.wte(input_ids)
+        if self.position_embedding=="learned":
+            position_embeds = self.wpe(position_ids)
+            hidden_states = inputs_embeds + position_embeds
+        else:
+            hidden_states = inputs_embeds
+
+        if token_type_ids is not None:
+            token_type_embeds = self.wte(token_type_ids)
+            hidden_states = hidden_states + token_type_embeds
+
+        hidden_states = self.drop(hidden_states)
+
+        output_shape = input_shape + (hidden_states.size(-1),)
+
+        presents = () if use_cache else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
+        all_hidden_states = () if output_hidden_states else None
+        
+        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
+            # Model parallel
+            if self.model_parallel:
+                torch.cuda.set_device(hidden_states.device)
+                # Ensure layer_past is on same device as hidden_states (might not be correct)
+                if layer_past is not None:
+                    layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past)
+                # Ensure that attention_mask is always on the same device as hidden_states
+                if attention_mask is not None:
+                    attention_mask = attention_mask.to(hidden_states.device)
+                if isinstance(head_mask, torch.Tensor):
+                    head_mask = head_mask.to(hidden_states.device)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                if use_cache:
+                    logger.warning(
+                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                    )
+                    use_cache = False
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, use_cache, output_attentions)
+
+                    return custom_forward
+
+                outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    None,
+                    attention_mask,
+                    head_mask[i],
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                outputs = block(
+                    hidden_states,
+                    layer_past=layer_past,
+                    attention_mask=attention_mask,
+                    head_mask=head_mask[i],
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                    alibi_bias=self.alibi if hasattr(self, "alibi") else None
+                )
+
+            hidden_states = outputs[0]
+            
+            if use_cache is True:
+                presents = presents + (outputs[1],)
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
+                if self.config.add_cross_attention:
+                    all_cross_attentions = all_cross_attentions + (outputs[3 if use_cache else 2],)
+
+            if self.model_parallel:
+                device_prefix="cuda:"
+                for k, v in self.device_map.items():
+                    if i == v[-1] and device_prefix + str(k) != self.last_device:
+                        hidden_states = hidden_states.to(device_prefix + str(k + 1))
+
+        hidden_states = self.ln_f(hidden_states)
+
+        hidden_states = hidden_states.view(*output_shape)
+        # Add last hidden state
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, presents, all_hidden_states, all_self_attentions, all_cross_attentions, moe_loss]
+                if v is not None
+            )
+        
+        return BaseModelOutputWithPastAndCrossAttentions(
+                last_hidden_state=hidden_states,
+                past_key_values=presents,
+                hidden_states=all_hidden_states,
+                attentions=all_self_attentions,
+                cross_attentions=all_cross_attentions,
+            )
+
+class TranceptionLMHeadModel(GPT2PreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"attn.bias", r"lm_head.weight"]
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = TranceptionModel(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.config = config
+
+        self.init_weights()
+        
+        self.default_model_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+        
+        self.retrieval_aggregation_mode = config.retrieval_aggregation_mode if hasattr(config, "retrieval_aggregation_mode") else None
+        if self.retrieval_aggregation_mode is not None:
+            print("Model leverages both autoregressive and retrieval inference")
+            self.MSA_filename = config.MSA_filename if hasattr(config, "MSA_filename") else False
+            self.MSA_folder = '/'.join(self.MSA_filename.split(os.sep)[:-1])
+            self.MSA_name = self.MSA_filename.split(os.sep)[-1]
+            self.retrieval_inference_weight_LR = config.retrieval_inference_weight if hasattr(config, "retrieval_inference_weight") else 0.6
+            self.retrieval_inference_weight_RL = config.retrieval_inference_weight if hasattr(config, "retrieval_inference_weight") else 0.6
+            self.MSA_start=config.MSA_start
+            self.MSA_end=config.MSA_end
+            self.full_protein_length = config.full_protein_length if hasattr(config, "full_protein_length") else -1
+            
+            self.MSA_log_prior = torch.log(torch.tensor(
+                                                        msa_utils.get_msa_prior(
+                                                            MSA_data_file=self.MSA_filename, 
+                                                            MSA_weight_file_name=config.MSA_weight_file_name, 
+                                                            retrieval_aggregation_mode=self.retrieval_aggregation_mode,
+                                                            MSA_start=self.MSA_start,
+                                                            MSA_end=self.MSA_end,
+                                                            len_target_seq=self.full_protein_length, 
+                                                            vocab=config.tokenizer.get_vocab(), 
+                                                            verbose=False
+                                                        )
+                                            ).float().to(self.default_model_device))
+        else:
+            print("Model only uses autoregressive inference")
+
+    def parallelize(self, device_map=None, num_cores=None, num_pipelines=1):
+        self.num_pipelines=num_pipelines
+        self.device_map = (
+                get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
+                if device_map is None
+                else device_map
+            )
+        assert_device_map(self.device_map, len(self.transformer.h))
+        self.transformer.parallelize(self.device_map, num_cores=num_cores)
+        self.lm_head = self.lm_head.to(self.transformer.first_device)
+        self.model_parallel = True
+
+    def deparallelize(self):
+        self.transformer.deparallelize()
+        self.transformer = self.transformer.to("cpu")
+        self.lm_head = self.lm_head.to("cpu")
+        self.model_parallel = False
+        torch.cuda.empty_cache()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)
+        # only last token for inputs_ids if past is defined in kwargs
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        
+        return {
+                "input_ids": input_ids,
+                "past_key_values": past,
+                "use_cache": kwargs.get("use_cache"),
+                "position_ids": position_ids,
+                "attention_mask": attention_mask,
+                "token_type_ids": token_type_ids,
+                "flip": kwargs.get("flip", None),
+            }
+
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        flip=None,
+        start_slice=None,
+        end_slice=None,
+        full_raw_sequence=None,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+            ``labels = input_ids`` Indices are selected in ``[-100, 0, ..., config.vocab_size]`` All labels set to
+            ``-100`` are ignored (masked), the loss is only computed for labels in ``[0, ..., config.vocab_size]``
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+        hidden_states = transformer_outputs[0]
+            
+        # Set device for model parallelism
+        if self.model_parallel:
+            torch.cuda.set_device(self.transformer.first_device)
+            hidden_states = hidden_states.to(self.lm_head.weight.device)
+            self.MSA_log_prior = self.MSA_log_prior.to(self.lm_head.weight.device)
+
+        lm_logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            
+            if self.retrieval_aggregation_mode is not None:
+                batch_size = input_ids.size(0)
+                
+                if self.retrieval_aggregation_mode=="aggregate_indel":
+                    assert batch_size==1, "Aggregate indel is only supported for batch size of 1"
+                    truncated_sequence_text = full_raw_sequence[0][start_slice[0]:end_slice[0]]
+                    if len(truncated_sequence_text)!=shift_logits.shape[1]-1: # shift_logits only has one extra token compared to truncated_sequence_text (the BOS token)
+                        print("Tokenization error -- seq length: {} and shift_logits length - 1 : {}".format(len(full_raw_sequence),shift_logits.shape[1]-1))
+                    MSA_log_prior, MSA_start, MSA_end = msa_utils.update_retrieved_MSA_log_prior_indel(self, self.MSA_log_prior, self.MSA_start, self.MSA_end, full_raw_sequence[0])  
+                
+                elif self.retrieval_aggregation_mode=="aggregate_substitution":
+                    MSA_log_prior=self.MSA_log_prior
+                    MSA_start=self.MSA_start
+                    MSA_end=self.MSA_end
+                
+                shift_log_probas = torch.log_softmax(shift_logits,dim=-1)
+                fused_shift_log_probas = shift_log_probas.clone()
+                if flip is None:
+                    flip = torch.zeros(batch_size).to(fused_shift_log_probas.device)
+                flip = flip > 0
+                
+                for seq_index in range(batch_size):
+                    min_prior_slice = max(start_slice[seq_index], MSA_start) 
+                    max_prior_slice = min(end_slice[seq_index], MSA_end)
+                    
+                    if max_prior_slice <= min_prior_slice:
+                        print("Non overlapping region detected: min_prior_slice {} and max_prior_slice {}".format(min_prior_slice,max_prior_slice))
+                        continue
+                    
+                    slice_prior = MSA_log_prior[min_prior_slice:max_prior_slice,:].to(fused_shift_log_probas.device) 
+                    if flip[seq_index]:
+                        slice_prior = torch.flip(slice_prior,dims=(0,))
+                        min_logits_slice = max(0,end_slice[seq_index]-MSA_end) 
+                        max_logits_slice = min_logits_slice + (max_prior_slice-min_prior_slice)
+                        fused_shift_log_probas[seq_index,min_logits_slice:max_logits_slice,:] = (1-self.retrieval_inference_weight_RL)*shift_log_probas[seq_index,min_logits_slice:max_logits_slice,:] + self.retrieval_inference_weight_RL*slice_prior
+                    else:
+                        min_logits_slice = max(0, MSA_start-start_slice[seq_index]) 
+                        max_logits_slice = min_logits_slice + (max_prior_slice-min_prior_slice)
+                        fused_shift_log_probas[seq_index,min_logits_slice:max_logits_slice,:] = (1-self.retrieval_inference_weight_LR)*shift_log_probas[seq_index,min_logits_slice:max_logits_slice,:] + self.retrieval_inference_weight_LR*slice_prior
+                
+                if self.retrieval_aggregation_mode=="aggregate_indel":
+                    try:
+                        # If a given residue colume is an added zero-column, then we overwrite prior fusion and only predict based on the autoregressive transformer inference mode.
+                        inserted_retrieval_positions = [True if slice_prior[i].sum()==0 else False for i in range(len(slice_prior))]+[True] #Last True is for the end of sentence token
+                        fused_shift_log_probas[:,inserted_retrieval_positions,:]=shift_log_probas[:,inserted_retrieval_positions,:]
+                    except:
+                        print("Error when adding zero column(s) to account for insertion mutations.")
+                
+                loss_fct = NLLLoss(reduction='none')
+                loss = loss_fct(input=fused_shift_log_probas.view(-1, fused_shift_log_probas.size(-1)), target=shift_labels.view(-1)).view(fused_shift_log_probas.shape[0],fused_shift_log_probas.shape[1])
+                mask = attention_mask[..., 1:].float()
+                mask[mask==0]=float('nan')
+                loss *= mask
+                loss = nanmean(loss, dim=1).mean()
+            else:
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+                fused_shift_log_probas = None
+
+        if not return_dict:
+            output = (lm_logits,) + transformer_outputs[1:] 
+            return ((loss,) + output) if loss is not None else output
+        
+        return TranceptionCausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+            fused_shift_log_probas=fused_shift_log_probas
+        )
+
+
+    @staticmethod
+    def _reorder_cache(past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor) -> Tuple[Tuple[torch.Tensor]]:
+        """
+        This function is used to re-order the :obj:`past_key_values` cache if
+        :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+        called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
+            for layer_past in past
+        )
+    
+    def score_mutants(self, DMS_data, target_seq, scoring_mirror=True, batch_size_inference=10, num_workers=10, indel_mode=False):
+        """
+        Method to score mutants in an input DMS file.
+        DMS_data: (dataframe) Dataframe containing the list of mutant triplets (substitutions) or full mutated sequences (indels) for scoring.
+        target_seq: (string) Full reference sequence (wild type) that is mutated in the DMS assay.
+        scoring_mirror: (bool) Whether to score mutated sequences from both directions (Left->Right and Right->Left).
+        batch_size_inference: (int) Batch size for scoring.
+        num_workers: (int) Number of workers to be used in the data loader.
+        indel_mode: (bool) Flag to be used when scoring insertions and deletions. Otherwise assumes substitutions.
+        """
+        df = DMS_data.copy()
+        df['mutated_sequence'] = df['mutant'].apply(lambda x: scoring_utils.get_mutated_sequence(target_seq, x)) if not indel_mode else df['mutant']
+        if 'DMS_score' in df: del df['DMS_score'] 
+        if 'DMS_score_bin' in df: del df['DMS_score_bin'] 
+        df_left_to_right_slices = scoring_utils.get_sequence_slices(df, target_seq=target_seq, model_context_len = self.config.n_ctx - 2, indel_mode=indel_mode, scoring_window=self.config.scoring_window)
+        print("Scoring sequences from left to right")
+        scores_L_to_R = scoring_utils.get_tranception_scores_mutated_sequences(model=self, mutated_sequence_df=df_left_to_right_slices, batch_size_inference=batch_size_inference, score_var_name='avg_score_L_to_R', len_target_seq=len(target_seq), num_workers=num_workers, indel_mode=indel_mode)
+        if scoring_mirror: 
+            print("Scoring sequences from right to left")
+            df_right_to_left_slices = df_left_to_right_slices.copy()
+            df_right_to_left_slices['mutated_sequence'] = df_right_to_left_slices['mutated_sequence'].apply(lambda x: x[::-1])
+            scores_R_to_L = scoring_utils.get_tranception_scores_mutated_sequences(model=self, mutated_sequence_df=df_right_to_left_slices, batch_size_inference=batch_size_inference, score_var_name='avg_score_R_to_L', len_target_seq=len(target_seq), num_workers=num_workers, reverse=True, indel_mode=indel_mode)
+            all_scores = pd.merge(scores_L_to_R, scores_R_to_L, on='mutant', how='left',suffixes=('','_R_to_L'))
+            all_scores['avg_score'] = (all_scores['avg_score_L_to_R'] + all_scores['avg_score_R_to_L']) / 2.0
+        else:
+            all_scores = scores_L_to_R
+            all_scores['avg_score'] = all_scores['avg_score_L_to_R']
+        return all_scores
+
+    def encode_batch(self, protein_sequence, sequence_name="mutated_sequence"):
+        """
+        Method to process an input AA sequence batch (protein_sequence) and return a tokenized sequence (via the tokenizer associated to the model).
+        """
+        protein_sequence[sequence_name] = scoring_utils.sequence_replace(sequences=protein_sequence[sequence_name], char_to_replace='X', char_replacements='ACDEFGHIKLMNPQRSTVWY')
+        protein_sequence[sequence_name] = scoring_utils.sequence_replace(sequences=protein_sequence[sequence_name], char_to_replace='B', char_replacements='DN')
+        protein_sequence[sequence_name] = scoring_utils.sequence_replace(sequences=protein_sequence[sequence_name], char_to_replace='J', char_replacements='IL')
+        protein_sequence[sequence_name] = scoring_utils.sequence_replace(sequences=protein_sequence[sequence_name], char_to_replace='Z', char_replacements='EQ')
+        return self.config.tokenizer(list(protein_sequence[sequence_name]), add_special_tokens=True, truncation=True, padding=True, max_length=self.config.n_ctx)
+

outputs.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+
+from transformers.file_utils import ModelOutput
+
+@dataclass
+class TranceptionCausalLMOutputWithCrossAttentions(ModelOutput):
+    """
+    Class for Tranception causal language model (or autoregressive) outputs.
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (for next-token prediction).
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
+            shape `(batch_size, sequence_length, hidden_size)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+            Cross attentions weights after the attention softmax, used to compute the weighted average in the
+            cross-attention heads.
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `torch.FloatTensor` tuples of length `config.n_layers`, with each tuple containing the cached key,
+            value states of the self-attention and the cross-attention layers if model is used in encoder-decoder
+            setting. Only relevant if `config.is_decoder = True`.
+            Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
+            `past_key_values` input) to speed up sequential decoding.
+        fused_shift_log_probas (`torch.FloatTensor` of shape (batch_size, sequence_length, config.vocab_size), *optional*, returned when config.retrieval_aggregation_mode is not None.
+            log_probas for each residue position after aggregating autoregressive logits and retrieval logits.
+            
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
+    fused_shift_log_probas: Optional[torch.FloatTensor] = None

pytorch_model.bin

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

utils/__init__.py

@@ -0,0 +1 @@
+from . import scoring_utils, msa_utils

utils/dms_utils.py

@@ -0,0 +1,26 @@
+import pandas as pd
+import numpy as np
+
+def DMS_file_cleanup(DMS_filename, target_seq, start_idx=1, end_idx=None, DMS_mutant_column='mutant', DMS_phenotype_name='score', DMS_directionality=1, AA_vocab = "ACDEFGHIKLMNPQRSTVWY"):
+    """
+    Function to process the raw DMS assay data (eg., removing invalid mutants, aggregate silent mutations)
+    """
+    DMS_data = pd.read_csv(DMS_filename, low_memory=False)
+    end_idx = start_idx + len(target_seq) - 1 if end_idx is None else end_idx
+    DMS_data['mutant'] = DMS_data[DMS_mutant_column]
+    
+    DMS_data=DMS_data[DMS_data['mutant'].notnull()].copy()
+    DMS_data=DMS_data[DMS_data['mutant'].apply(lambda x: all([len(y)>=3 for y in x.split(":")]))].copy() #Mutant triplets should have at least 3 or more characters
+    DMS_data=DMS_data[DMS_data['mutant'].apply(lambda x: all([(y[0] in AA_vocab) and (y[1:-1].isnumeric()) and (y[-1] in AA_vocab) for y in x.split(":")]))].copy()
+    DMS_data=DMS_data[DMS_data['mutant'].apply(lambda x: all([int(y[1:-1])-start_idx >=0 and int(y[1:-1]) <= end_idx for y in x.split(":")]))].copy()
+    DMS_data=DMS_data[DMS_data['mutant'].apply(lambda x: all([y[0]==target_seq[int(y[1:-1])-start_idx] for y in x.split(":")]))].copy()
+    
+    DMS_data[DMS_phenotype_name]=pd.to_numeric(DMS_data[DMS_phenotype_name],errors='coerce')
+    DMS_data=DMS_data[np.isfinite(DMS_data[DMS_phenotype_name])]
+    DMS_data.dropna(subset = [DMS_phenotype_name], inplace=True)
+    DMS_data['DMS_score'] = DMS_data[DMS_phenotype_name] * DMS_directionality
+    DMS_data=DMS_data[['mutant','DMS_score']]
+    DMS_data = DMS_data.groupby('mutant').mean().reset_index()
+    
+    return DMS_data
+

utils/msa_utils.py

@@ -0,0 +1,361 @@
+import numpy as np
+import pandas as pd
+from collections import defaultdict
+import random
+import os
+import torch
+from Bio.Align.Applications import ClustalOmegaCommandline
+
+def filter_msa(msa_data, num_sequences_kept=3):
+    """
+    Helper function to filter an input MSA msa_data (obtained via process_msa_data) and keep only num_sequences_kept aligned sequences.
+    If the MSA already has fewer sequences than num_sequences_kept, we keep the MSA as is.
+    If filtering, we always keep the first sequence of the MSA (ie. the wild type) by default.
+    Sampling is done without replacement.
+    """
+    if len(list(msa_data.keys())) <= num_sequences_kept:
+        return  msa_data
+    filtered_msa = {}
+    wt_name = next(iter(msa_data)) 
+    filtered_msa[wt_name] = msa_data[wt_name]
+    del msa_data[wt_name]
+    sequence_names = list(msa_data.keys())
+    sequence_names_sampled  = random.sample(sequence_names,k=num_sequences_kept-1)
+    for seq in sequence_names_sampled:
+        filtered_msa[seq] = msa_data[seq]
+    return filtered_msa
+
+def process_msa_data(MSA_data_file):
+    """
+    Helper function that takes as input a path to a MSA file (expects a2m format) and returns a dict mapping sequence ID to the corresponding AA sequence.
+    """
+    msa_data = defaultdict(str)
+    sequence_name = ""
+    with open(MSA_data_file, "r") as msa_file:
+        for i, line in enumerate(msa_file):
+            line = line.rstrip()
+            if line.startswith(">"):
+                sequence_name = line
+            else:
+                msa_data[sequence_name] += line.upper()
+    return msa_data
+
+def get_one_hot_sequences_dict(msa_data,MSA_start,MSA_end,vocab):
+    vocab_size = len(vocab.keys())
+    num_sequences_msa = len(msa_data.keys())
+    one_hots = np.zeros((num_sequences_msa,MSA_end-MSA_start,vocab_size))
+    for i,seq_name in enumerate(msa_data.keys()):
+        sequence = msa_data[seq_name]
+        for j,letter in enumerate(sequence):
+            if letter in vocab: 
+                k = vocab[letter]
+                one_hots[i,j,k] = 1.0
+    return one_hots
+
+def one_hot(sequence_string,vocab):
+    one_hots = np.zeros((len(sequence_string),len(vocab.keys())))
+    for j,letter in enumerate(sequence_string):
+        if letter in vocab: 
+            k = vocab[letter]
+            one_hots[j,k] = 1.0
+    return one_hots.flatten()
+
+def get_msa_prior(MSA_data_file, MSA_weight_file_name, MSA_start, MSA_end, len_target_seq, vocab, retrieval_aggregation_mode="aggregate_substitution", filter_MSA=True, verbose=False):
+    """
+    Function to enable retrieval inference mode, via computation of (weighted) pseudocounts of AAs at each position of the retrieved MSA.
+    MSA_data_file: (string) path to MSA file (expects a2m format).
+    MSA_weight_file_name: (string) path to sequence weights in MSA.
+    MSA_start: (int) Sequence position that the MSA starts at (1-indexing).
+    MSA_end: (int) Sequence position that the MSA ends at (1-indexing).
+    len_target_seq: (int) Full length of sequence to be scored.
+    vocab: (dict) Vocabulary of the tokenizer.
+    retrieval_aggregation_mode: (string) Mode for retrieval inference (aggregate_substitution Vs aggregate_indel). If None, places a uniform prior over each token.
+    filter_MSA: (bool) Whether to filter out sequences with very low hamming similarity (< 0.2) to the reference sequence in the MSA (first sequence).
+    verbose: (bool) Whether to print to the console processing details along the way.
+    """
+    msa_data = process_msa_data(MSA_data_file)
+    vocab_size = len(vocab.keys())
+    if verbose: print("Target seq len is {}, MSA length is {}, start position is {}, end position is {} and vocab size is {}".format(len_target_seq,MSA_end-MSA_start,MSA_start,MSA_end,vocab_size))
+
+    if filter_MSA:
+        if verbose: print("Num sequences in MSA pre filtering: {}".format(len(msa_data.keys())))
+        list_sequence_names = list(msa_data.keys())
+        focus_sequence_name = list(msa_data.keys())[0]
+        ref_sequence_hot = one_hot(msa_data[focus_sequence_name],vocab)
+        for sequence_name in list_sequence_names:
+            seq_hot = one_hot(msa_data[sequence_name],vocab)
+            hamming_similarity_seq_ref = np.dot(ref_sequence_hot,seq_hot) / np.dot(ref_sequence_hot,ref_sequence_hot)
+            if hamming_similarity_seq_ref < 0.2:
+                del msa_data[sequence_name]
+        if verbose: print("Num sequences in MSA post filtering: {}".format(len(msa_data.keys())))
+
+    if MSA_weight_file_name is not None:
+        if verbose: print("Using weights in {} for sequences in MSA.".format(MSA_weight_file_name))
+        assert os.path.exists(MSA_weight_file_name), "Weights file not located on disk."
+        MSA_EVE = MSA_processing(
+                MSA_location=MSA_data_file,
+                use_weights=True,
+                weights_location=MSA_weight_file_name
+        )
+        #We scan through all sequences to see if we have a weight for them as per EVE pre-processing. We drop them otherwise.
+        dropped_sequences=0
+        list_sequence_names = list(msa_data.keys())
+        MSA_weight=[]
+        for sequence_name in list_sequence_names:
+            if sequence_name not in MSA_EVE.seq_name_to_sequence:
+                dropped_sequences +=1
+                del msa_data[sequence_name]
+            else:
+                MSA_weight.append(MSA_EVE.seq_name_to_weight[sequence_name])
+        if verbose: print("Dropped {} sequences from MSA due to absent sequence weights".format(dropped_sequences))
+    else:
+        MSA_weight = [1] *  len(list(msa_data.keys()))
+
+    if retrieval_aggregation_mode=="aggregate_substitution" or retrieval_aggregation_mode=="aggregate_indel":
+        one_hots = get_one_hot_sequences_dict(msa_data,MSA_start,MSA_end,vocab)
+        MSA_weight = np.expand_dims(np.array(MSA_weight),axis=(1,2))
+        base_rate = 1e-5
+        base_rates = np.ones_like(one_hots) * base_rate
+        weighted_one_hots = (one_hots + base_rates) * MSA_weight
+        MSA_weight_norm_counts = weighted_one_hots.sum(axis=-1).sum(axis=0)
+        MSA_weight_norm_counts = np.tile(MSA_weight_norm_counts.reshape(-1,1), (1,vocab_size))
+        one_hots_avg = weighted_one_hots.sum(axis=0) / MSA_weight_norm_counts
+        msa_prior = np.zeros((len_target_seq,vocab_size))
+        msa_prior[MSA_start:MSA_end,:]=one_hots_avg
+    else:
+        msa_prior = np.ones((len_target_seq,vocab_size)) / vocab_size
+    
+    if verbose:
+        for idx, position in enumerate(msa_prior):
+            if len(position)!=25:
+                print("Size error")
+            if not round(position.sum(),2)==1.0:
+                print("Position at index {} does not add up to 1: {}".format(idx, position.sum()))
+    
+    return msa_prior
+
+
+def update_retrieved_MSA_log_prior_indel(model, MSA_log_prior, MSA_start, MSA_end, full_raw_sequence):
+    """
+    Function to process MSA when scoring indels.
+    To identify positions to add / remove in the retrieved MSA, we append and align the sequence to be scored to the original MSA for that protein family with Clustal Omega.
+    If the original MSA is relatively deep (over 100k sequences), we sample (by default) 100k rows at random from that MSA to speed computations.
+    MSA sampling is performed only once (for the first sequence to be scored). Subsequent scoring use the same MSA sample.
+    """
+    if not os.path.isdir(model.MSA_folder + os.sep + "Sampled"):
+        os.mkdir(model.MSA_folder + os.sep + "Sampled")
+    sampled_MSA_location = model.MSA_folder + os.sep + "Sampled" + os.sep + "Sampled_" + model.MSA_filename.split(os.sep)[-1]
+    
+    if not os.path.exists(sampled_MSA_location):
+        msa_data = process_msa_data(model.MSA_filename)
+        msa_data_sampled = filter_msa(msa_data, num_sequences_kept=100000) #If MSA has less than 100k sequences, the sample is identical to original MSA
+        with open(sampled_MSA_location, 'w') as sampled_write_location:
+            for index, key in enumerate(msa_data_sampled):
+                key_name = ">REFERENCE_SEQUENCE" if index==0 else key
+                msa_data_sampled[key] = msa_data_sampled[key].upper()
+                msa_data_sampled[key] = msa_data_sampled[key].replace(".","-")
+                sampled_write_location.write(key_name+"\n"+"\n".join([msa_data_sampled[key][i:i+80] for i in range(0, len(msa_data_sampled[key]), 80)])+"\n")
+    
+    seq_to_align_location = model.MSA_folder + os.sep + "Sampled" + os.sep + "Seq_to_align_" + model.MSA_filename.split(os.sep)[-1]
+    sequence_text_split = [full_raw_sequence[i:i+80] for i in range(0, len(full_raw_sequence), 80)]
+    sequence_text_split_split_join = "\n".join([">SEQ_TO_SCORE"]+sequence_text_split)
+    os.system("echo '"+sequence_text_split_split_join+"' > "+seq_to_align_location)
+    
+    expanded_MSA_location = model.MSA_folder + os.sep + "Sampled" + os.sep + "Expanded_" + model.MSA_filename.split(os.sep)[-1]
+    clustalw_cline = ClustalOmegaCommandline(cmd=model.config.clustal_omega_location,
+                                            profile1=sampled_MSA_location,
+                                            profile2=seq_to_align_location,
+                                            outfile=expanded_MSA_location,
+                                            force=True)                                        
+    stdout, stderr = clustalw_cline()
+    msa_data = process_msa_data(expanded_MSA_location)
+    aligned_seqA, aligned_seqB = msa_data[">SEQ_TO_SCORE"], msa_data[">REFERENCE_SEQUENCE"]
+    try:
+        keep_column=[]
+        for column_index_pairwise_alignment in range(len(aligned_seqA)):
+            if aligned_seqA[column_index_pairwise_alignment]=="-" and aligned_seqB[column_index_pairwise_alignment]=="-":
+                continue
+            elif aligned_seqA[column_index_pairwise_alignment]=="-":
+                keep_column.append(False)
+            elif aligned_seqB[column_index_pairwise_alignment]=="-":
+                MSA_log_prior=torch.cat((MSA_log_prior[:column_index_pairwise_alignment], torch.zeros(MSA_log_prior.shape[1]).view(1,-1).cuda(), MSA_log_prior[column_index_pairwise_alignment:]),dim=0)
+                keep_column.append(True) #keep the zero column we just added
+            else:
+                keep_column.append(True)
+        MSA_log_prior = MSA_log_prior[keep_column]
+        MSA_end = MSA_start + len(MSA_log_prior)
+    except: 
+        print("Error when processing the following alignment: {}".format(expanded_MSA_location))
+    return MSA_log_prior, MSA_start, MSA_end
+
+class MSA_processing:
+    def __init__(self,
+        MSA_location="",
+        theta=0.2,
+        use_weights=True,
+        weights_location="./data/weights",
+        preprocess_MSA=True,
+        threshold_sequence_frac_gaps=0.5,
+        threshold_focus_cols_frac_gaps=0.3,
+        remove_sequences_with_indeterminate_AA_in_focus_cols=True
+        ):
+        
+        """
+        This MSA_processing class is directly borrowed from the EVE codebase: https://github.com/OATML-Markslab/EVE
+        
+        Parameters:
+        - msa_location: (path) Location of the MSA data. Constraints on input MSA format: 
+            - focus_sequence is the first one in the MSA data
+            - first line is structured as follows: ">focus_seq_name/start_pos-end_pos" (e.g., >SPIKE_SARS2/310-550)
+            - corespondding sequence data located on following line(s)
+            - then all other sequences follow with ">name" on first line, corresponding data on subsequent lines
+        - theta: (float) Sequence weighting hyperparameter. Generally: Prokaryotic and eukaryotic families =  0.2; Viruses = 0.01
+        - use_weights: (bool) If False, sets all sequence weights to 1. If True, checks weights_location -- if non empty uses that; 
+            otherwise compute weights from scratch and store them at weights_location
+        - weights_location: (path) Location to load from/save to the sequence weights
+        - preprocess_MSA: (bool) performs pre-processing of MSA to remove short fragments and positions that are not well covered.
+        - threshold_sequence_frac_gaps: (float, between 0 and 1) Threshold value to define fragments
+            - sequences with a fraction of gap characters above threshold_sequence_frac_gaps are removed
+            - default is set to 0.5 (i.e., fragments with 50% or more gaps are removed)
+        - threshold_focus_cols_frac_gaps: (float, between 0 and 1) Threshold value to define focus columns
+            - positions with a fraction of gap characters above threshold_focus_cols_pct_gaps will be set to lower case (and not included in the focus_cols)
+            - default is set to 0.3 (i.e., focus positions are the ones with 30% of gaps or less, i.e., 70% or more residue occupancy)
+        - remove_sequences_with_indeterminate_AA_in_focus_cols: (bool) Remove all sequences that have indeterminate AA (e.g., B, J, X, Z) at focus positions of the wild type
+        """
+        np.random.seed(2021)
+        self.MSA_location = MSA_location
+        self.weights_location = weights_location
+        self.theta = theta
+        self.alphabet = "ACDEFGHIKLMNPQRSTVWY"
+        self.use_weights = use_weights
+        self.preprocess_MSA = preprocess_MSA
+        self.threshold_sequence_frac_gaps = threshold_sequence_frac_gaps
+        self.threshold_focus_cols_frac_gaps = threshold_focus_cols_frac_gaps
+        self.remove_sequences_with_indeterminate_AA_in_focus_cols = remove_sequences_with_indeterminate_AA_in_focus_cols
+
+        self.gen_alignment()
+        
+    def gen_alignment(self, verbose=False):
+        """ Read training alignment and store basics in class instance """
+        self.aa_dict = {}
+        for i,aa in enumerate(self.alphabet):
+            self.aa_dict[aa] = i
+
+        self.seq_name_to_sequence = defaultdict(str)
+        name = ""
+        with open(self.MSA_location, "r") as msa_data:
+            for i, line in enumerate(msa_data):
+                line = line.rstrip()
+                if line.startswith(">"):
+                    name = line
+                    if i==0:
+                        self.focus_seq_name = name
+                else:
+                    self.seq_name_to_sequence[name] += line
+
+        
+        ## MSA pre-processing to remove inadequate columns and sequences
+        if self.preprocess_MSA:
+            msa_df = pd.DataFrame.from_dict(self.seq_name_to_sequence, orient='index', columns=['sequence'])
+            # Data clean up
+            msa_df.sequence = msa_df.sequence.apply(lambda x: x.replace(".","-")).apply(lambda x: ''.join([aa.upper() for aa in x]))
+            # Remove columns that would be gaps in the wild type
+            non_gap_wt_cols = [aa!='-' for aa in msa_df.sequence[self.focus_seq_name]]
+            msa_df['sequence'] = msa_df['sequence'].apply(lambda x: ''.join([aa for aa,non_gap_ind in zip(x, non_gap_wt_cols) if non_gap_ind]))
+            assert 0.0 <= self.threshold_sequence_frac_gaps <= 1.0,"Invalid fragment filtering parameter"
+            assert 0.0 <= self.threshold_focus_cols_frac_gaps <= 1.0,"Invalid focus position filtering parameter"
+            msa_array = np.array([list(seq) for seq in msa_df.sequence])
+            gaps_array = np.array(list(map(lambda seq: [aa=='-' for aa in seq], msa_array)))
+            # Identify fragments with too many gaps
+            seq_gaps_frac = gaps_array.mean(axis=1)
+            seq_below_threshold = seq_gaps_frac <= self.threshold_sequence_frac_gaps
+            if verbose: print("Proportion of sequences dropped due to fraction of gaps: "+str(round(float(1 - seq_below_threshold.sum()/seq_below_threshold.shape)*100,2))+"%")
+            # Identify focus columns
+            columns_gaps_frac = gaps_array[seq_below_threshold].mean(axis=0)
+            index_cols_below_threshold = columns_gaps_frac <= self.threshold_focus_cols_frac_gaps
+            if verbose: print("Proportion of non-focus columns removed: "+str(round(float(1 - index_cols_below_threshold.sum()/index_cols_below_threshold.shape)*100,2))+"%")
+            # Lower case non focus cols and filter fragment sequences
+            msa_df['sequence'] = msa_df['sequence'].apply(lambda x: ''.join([aa.upper() if upper_case_ind else aa.lower() for aa, upper_case_ind in zip(x, index_cols_below_threshold)]))
+            msa_df = msa_df[seq_below_threshold]
+            # Overwrite seq_name_to_sequence with clean version
+            self.seq_name_to_sequence = defaultdict(str)
+            for seq_idx in range(len(msa_df['sequence'])):
+                self.seq_name_to_sequence[msa_df.index[seq_idx]] = msa_df.sequence[seq_idx]
+
+        self.focus_seq = self.seq_name_to_sequence[self.focus_seq_name]
+        self.focus_cols = [ix for ix, s in enumerate(self.focus_seq) if s == s.upper() and s!='-'] 
+        self.focus_seq_trimmed = [self.focus_seq[ix] for ix in self.focus_cols]
+        self.seq_len = len(self.focus_cols)
+        self.alphabet_size = len(self.alphabet)
+
+        # Connect local sequence index with uniprot index (index shift inferred from 1st row of MSA)
+        focus_loc = self.focus_seq_name.split("/")[-1]
+        start,stop = focus_loc.split("-")
+        self.focus_start_loc = int(start)
+        self.focus_stop_loc = int(stop)
+        self.uniprot_focus_col_to_wt_aa_dict \
+            = {idx_col+int(start):self.focus_seq[idx_col] for idx_col in self.focus_cols} 
+        self.uniprot_focus_col_to_focus_idx \
+            = {idx_col+int(start):idx_col for idx_col in self.focus_cols} 
+
+        # Move all letters to CAPS; keeps focus columns only
+        self.raw_seq_name_to_sequence = self.seq_name_to_sequence.copy()
+        for seq_name,sequence in self.seq_name_to_sequence.items():
+            sequence = sequence.replace(".","-")
+            self.seq_name_to_sequence[seq_name] = [sequence[ix].upper() for ix in self.focus_cols]
+
+        # Remove sequences that have indeterminate AA (e.g., B, J, X, Z) in the focus columns
+        if self.remove_sequences_with_indeterminate_AA_in_focus_cols:
+            alphabet_set = set(list(self.alphabet))
+            seq_names_to_remove = []
+            for seq_name,sequence in self.seq_name_to_sequence.items():
+                for letter in sequence:
+                    if letter not in alphabet_set and letter != "-":
+                        seq_names_to_remove.append(seq_name)
+                        continue
+            seq_names_to_remove = list(set(seq_names_to_remove))
+            for seq_name in seq_names_to_remove:
+                del self.seq_name_to_sequence[seq_name]
+
+        # Encode the sequences
+        self.one_hot_encoding = np.zeros((len(self.seq_name_to_sequence.keys()),len(self.focus_cols),len(self.alphabet)))
+        if verbose: print("One-hot encoded sequences shape:" + str(self.one_hot_encoding.shape))
+        for i,seq_name in enumerate(self.seq_name_to_sequence.keys()):
+            sequence = self.seq_name_to_sequence[seq_name]
+            for j,letter in enumerate(sequence):
+                if letter in self.aa_dict: 
+                    k = self.aa_dict[letter]
+                    self.one_hot_encoding[i,j,k] = 1.0
+
+        if self.use_weights:
+            try:
+                self.weights = np.load(file=self.weights_location)
+                if verbose: print("Loaded sequence weights from disk")
+            except:
+                if verbose: print ("Computing sequence weights")
+                list_seq = self.one_hot_encoding
+                list_seq = list_seq.reshape((list_seq.shape[0], list_seq.shape[1] * list_seq.shape[2]))
+                def compute_weight(seq):
+                    number_non_empty_positions = np.dot(seq,seq)
+                    if number_non_empty_positions>0:
+                        denom = np.dot(list_seq,seq) / np.dot(seq,seq) 
+                        denom = np.sum(denom > 1 - self.theta) 
+                        return 1/denom
+                    else:
+                        return 0.0 #return 0 weight if sequence is fully empty
+                self.weights = np.array(list(map(compute_weight,list_seq)))
+                np.save(file=self.weights_location, arr=self.weights)
+        else:
+            # If not using weights, use an isotropic weight matrix
+            if verbose: print("Not weighting sequence data")
+            self.weights = np.ones(self.one_hot_encoding.shape[0])
+
+        self.Neff = np.sum(self.weights)
+        self.num_sequences = self.one_hot_encoding.shape[0]
+        self.seq_name_to_weight={}
+        for i,seq_name in enumerate(self.seq_name_to_sequence.keys()):
+            self.seq_name_to_weight[seq_name]=self.weights[i]
+
+        if verbose:
+            print ("Neff =",str(self.Neff))
+            print ("Data Shape =",self.one_hot_encoding.shape)

utils/scoring_utils.py

@@ -0,0 +1,192 @@
+import os
+import tqdm
+import re
+import numpy as np
+import pandas as pd
+
+import torch
+from torch.nn import CrossEntropyLoss, NLLLoss
+from torch.utils.data.sampler import Sampler, SequentialSampler
+
+from transformers import DataCollatorForLanguageModeling, PreTrainedTokenizerFast
+from datasets import Dataset
+
+AA_vocab = "ACDEFGHIKLMNPQRSTVWY"
+
+def get_mutated_sequence(focus_seq, mutant, start_idx=1, AA_vocab=AA_vocab):
+    """
+    Helper function that mutates an input sequence (focus_seq) via an input mutation triplet (substitutions only).
+    Mutation triplet are typically based on 1-indexing: start_idx is used for switching to 0-indexing.
+    """
+    mutated_seq = list(focus_seq)
+    for mutation in mutant.split(":"):
+        try:
+            from_AA, position, to_AA = mutation[0], int(mutation[1:-1]), mutation[-1]
+        except:
+            print("Issue with mutant: "+str(mutation))
+        relative_position = position - start_idx
+        assert (from_AA==focus_seq[relative_position]), "Invalid from_AA or mutant position: "+str(mutation)+" from_AA: "+str(from_AA) + " relative pos: "+str(relative_position) + " focus_seq: "+str(focus_seq)
+        assert (to_AA in AA_vocab) , "Mutant to_AA is invalid: "+str(mutation)
+        mutated_seq[relative_position] = to_AA
+    return "".join(mutated_seq)
+
+def nanmean(v, *args, inplace=False, **kwargs):
+    if not inplace:
+        v = v.clone()
+    is_nan = torch.isnan(v)
+    v[is_nan] = 0
+    return v.sum(*args, **kwargs) / (~is_nan).float().sum(*args, **kwargs)
+
+def nansum(v, *args, inplace=False, **kwargs):
+    if not inplace:
+        v = v.clone()
+    is_nan = torch.isnan(v)
+    v[is_nan] = 0
+    return v.sum(*args, **kwargs)
+
+def get_optimal_window(mutation_position_relative, seq_len_wo_special, model_window):
+    """
+    Helper function that selects an optimal sequence window that fits the maximum model context size.
+    If the sequence length is less than the maximum context size, the full sequence is returned.
+    """
+    half_model_window = model_window // 2
+    if seq_len_wo_special <= model_window:
+        return [0,seq_len_wo_special]
+    elif mutation_position_relative < half_model_window:
+        return [0,model_window]
+    elif mutation_position_relative >= seq_len_wo_special - half_model_window:
+        return [seq_len_wo_special - model_window, seq_len_wo_special]
+    else:
+        return [max(0,mutation_position_relative-half_model_window), min(seq_len_wo_special,mutation_position_relative+half_model_window)]
+
+def sequence_replace_single(sequence, char_to_replace, char_replacements):
+    char_replacements = list(char_replacements)
+    positions = [m.start() for m in re.finditer(char_to_replace, sequence)]
+    replacements = np.random.choice(a=char_replacements, size=len(positions), replace=True)
+    sequence=list(sequence)
+    for idx, position in enumerate(positions):
+        sequence[position]=replacements[idx]
+    return ''.join(sequence)
+
+def sequence_replace(sequences, char_to_replace, char_replacements):
+    """
+    Helper function that replaces all Amino Acids passsed in via char_to_replace (as a string of AAs) with Amino Acids sampled from char_replacements (also a string of eligible AAs).
+    """
+    return [sequence_replace_single(sequence, char_to_replace, char_replacements) for sequence in sequences]
+
+def get_tranception_scores_mutated_sequences(model, mutated_sequence_df, batch_size_inference, score_var_name, len_target_seq, num_workers=10, reverse=False, indel_mode=False):
+    """
+    Helper function that takes as input a set of mutated sequences (in a pandas dataframe) and returns scores for each mutation (delta log likelihood wrt wild type sequence).
+    """
+    scores = {} 
+    scores['mutant']=[]
+    scores['window_start']=[]
+    scores['window_end']=[]
+    scores['score']=[]
+    with torch.no_grad():
+        ds = Dataset.from_pandas(mutated_sequence_df)
+        ds.set_transform(model.encode_batch)
+        data_collator = DataCollatorForLanguageModeling(
+                        tokenizer=model.config.tokenizer,
+                        mlm=False)
+        sampler = SequentialSampler(ds)
+        ds_loader = torch.utils.data.DataLoader(ds, batch_size=batch_size_inference, sampler=sampler, collate_fn=data_collator, num_workers=num_workers, pin_memory=True, drop_last=False)
+        mutant_index=0
+        for encoded_batch in tqdm.tqdm(ds_loader):
+            full_batch_length = len(encoded_batch['input_ids'])
+            scores['mutant'] += list(mutated_sequence_df['mutant'][mutant_index:mutant_index+full_batch_length])
+            window_start = np.array(mutated_sequence_df['window_start'][mutant_index:mutant_index+full_batch_length])
+            scores['window_start'] += list(window_start)
+            window_end = np.array(mutated_sequence_df['window_end'][mutant_index:mutant_index+full_batch_length])
+            scores['window_end'] += list(window_end)
+            full_raw_sequence = np.array(mutated_sequence_df['full_raw_sequence'][mutant_index:mutant_index+full_batch_length])
+            for k, v in encoded_batch.items():
+                if isinstance(v, torch.Tensor):
+                    encoded_batch[k] = v.to(model.device)            
+            shift_labels = encoded_batch['labels'][..., 1:].contiguous()
+            if (hasattr(model.config,"retrieval_aggregation_mode")) and (model.config.retrieval_aggregation_mode is not None):
+                if reverse:
+                    encoded_batch['flip']=torch.tensor([1]*full_batch_length)
+                encoded_batch['start_slice']=window_start
+                encoded_batch['end_slice']=window_end
+                encoded_batch['full_raw_sequence'] = full_raw_sequence #only mutated_sequence is flipped if the scoring_mirror branch of score_mutants. No need to flip full_raw_sequence for MSA re-aligning
+                fused_shift_log_probas=model(**encoded_batch,return_dict=True).fused_shift_log_probas
+                loss_fct = NLLLoss(reduction='none')
+                loss = - loss_fct(input=fused_shift_log_probas.view(-1, fused_shift_log_probas.size(-1)), target=shift_labels.view(-1)).view(fused_shift_log_probas.shape[0],fused_shift_log_probas.shape[1])
+            else:
+                lm_logits=model(**encoded_batch,return_dict=True).logits
+                shift_logits = lm_logits[..., :-1, :].contiguous()
+                loss_fct = CrossEntropyLoss(reduction='none') 
+                loss = - loss_fct(input=shift_logits.view(-1, shift_logits.size(-1)), target=shift_labels.view(-1)).view(shift_logits.shape[0],shift_logits.shape[1])
+            mask = encoded_batch['attention_mask'][..., 1:].float()
+            mask[mask==0]=float('nan')
+            loss *= mask
+            loss =  nanmean(loss, dim=1)
+            scores_batch = list(loss.cpu().numpy())
+            full_batch_length = len(encoded_batch['input_ids'])
+            scores['score'] += scores_batch
+            mutant_index+=full_batch_length
+    scores = pd.DataFrame(scores)
+    scores_mutated_seq = scores[scores.mutant != 'wt']
+    scores_wt = scores[scores.mutant == 'wt']
+    delta_scores = pd.merge(scores_mutated_seq,scores_wt,how='left',on=['window_start'],suffixes=('','_wt'))
+    delta_scores[score_var_name]  = delta_scores['score'] - delta_scores['score_wt']
+    delta_scores=delta_scores[['mutant',score_var_name]].groupby('mutant').mean().reset_index()
+    return delta_scores
+
+def get_sequence_slices(df, target_seq, model_context_len, start_idx=1, scoring_window="optimal", indel_mode=False):
+    """
+    Helper function that takes as input a (pandas) dataframe df that contains a list of mutant triplets (substitutions) or full mutated sequences (indels) for scoring.
+    It returns a processed DMS in which sequences have been sliced to satisfy the maximum context window of the model.
+    df: (dataframe) Input dataframe to be processed
+    target_seq: (string) Full reference sequence (wild type) that is mutated in the DMS assay.
+    model_context_len: (int) Maximum context size for the model.
+    start_idx: (int) Integer to move to 0-indexing of positions (mutation triplet are typically based on 1-indexing).
+    scoring_window: (string) Method to slice sequences longer than maximum context size: 
+        - optimal selects a single window as large as possible via the get_optimal_window function (this is the default)
+        - sliding splits the full sequence in contiguous (non-overlapping) chunks that are of size equal to the max context (except the last chunk which may be shorter)
+    indel_mode: (bool) Flag to be used when scoring insertions and deletions. Otherwise assumes substitutions.
+    Note: when scoring indels for sequences that would be longer than the model max context length, it is preferable to use the "sliding" scoring_window. Use "optimal" otherwise.
+    """
+    len_target_seq = len(target_seq)
+    num_mutants = len(df['mutant'])
+    df=df.reset_index(drop=True)
+    if scoring_window=="optimal":
+        df['mutation_barycenter'] = df['mutant'].apply(lambda x: int(np.array([int(mutation[1:-1]) - start_idx for mutation in x.split(':')]).mean())) if not indel_mode else df['mutant'].apply(lambda x: len(x)//2)
+        df['scoring_optimal_window'] = df['mutation_barycenter'].apply(lambda x: get_optimal_window(x, len_target_seq, model_context_len)) if not indel_mode else df['mutant'].apply(lambda x: (0,len(x)))
+        df['full_raw_sequence'] = df['mutated_sequence']
+        df['mutated_sequence'] = [df['mutated_sequence'][index][df['scoring_optimal_window'][index][0]:df['scoring_optimal_window'][index][1]] for index in range(num_mutants)]
+        df['window_start'] = df['scoring_optimal_window'].map(lambda x: x[0]) 
+        df['window_end'] = df['scoring_optimal_window'].map(lambda x: x[1])
+        del df['scoring_optimal_window']
+        df_wt=df.copy()
+        df_wt['mutant'] = ['wt'] * num_mutants
+        df_wt['full_raw_sequence'] = [target_seq] * num_mutants
+        if indel_mode: # For indels, we set the wild type reference to be always the same (full length) sequence. We assume here that the length is lower than model context size (otherwise use "Sliding")
+            df_wt['mutation_barycenter'] = [len_target_seq // 2] * num_mutants
+            df_wt['window_end'] = df_wt['full_raw_sequence'].map(lambda x:len(x))
+        df_wt['mutated_sequence'] = [target_seq[df_wt['window_start'][index]:df_wt['window_end'][index]] for index in range(num_mutants)]
+        df = pd.concat([df,df_wt], axis=0)
+        df = df.drop_duplicates()
+    elif scoring_window=="sliding":
+        len_target_seq = len(target_seq)
+        num_windows = 1 + int( len_target_seq / model_context_len)
+        df_list=[]
+        start=0
+        for window_index in range(1, num_windows+1):
+            df_sliced = df.copy()
+            df_sliced['full_raw_sequence'] = df_sliced['mutated_sequence']
+            df_sliced['mutated_sequence'] = df_sliced['mutated_sequence'].map(lambda x: x[start:start+model_context_len]) 
+            df_sliced['window_start'] = [start] * num_mutants 
+            df_sliced['window_end']  =  df_sliced['full_raw_sequence'].map(lambda x: min(len(x), start+model_context_len)) 
+            df_sliced_wt = df_sliced.copy()
+            df_sliced_wt['mutant'] = ['wt'] * num_mutants
+            df_sliced_wt['full_raw_sequence'] = [target_seq] * num_mutants
+            df_sliced_wt['mutated_sequence'] = df_sliced_wt['full_raw_sequence'].map(lambda x: x[start:start+model_context_len])
+            df_sliced_wt['window_end'] = df_sliced_wt['full_raw_sequence'].map(lambda x: min(len(x), start+model_context_len)) #Need to adjust end index if WT and sequence are not same full length
+            df_list.append(df_sliced)
+            df_list.append(df_sliced_wt)
+            start += model_context_len
+        df_final = pd.concat(df_list,axis=0)
+        df = df_final.drop_duplicates()
+    return df.reset_index(drop=True)

utils/tokenizers/Basic_tokenizer

@@ -0,0 +1 @@
+{"version":"1.0","truncation":null,"padding":null,"added_tokens":[{"id":0,"special":true,"content":"[UNK]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false},{"id":1,"special":true,"content":"[CLS]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false},{"id":2,"special":true,"content":"[SEP]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false},{"id":3,"special":true,"content":"[PAD]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false},{"id":4,"special":true,"content":"[MASK]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false}],"normalizer":null,"pre_tokenizer":{"type":"Whitespace"},"post_processor":{"type":"TemplateProcessing","single":[{"SpecialToken":{"id":"[CLS]","type_id":0}},{"Sequence":{"id":"A","type_id":0}},{"SpecialToken":{"id":"[SEP]","type_id":0}}],"pair":[{"SpecialToken":{"id":"[CLS]","type_id":0}},{"Sequence":{"id":"A","type_id":0}},{"SpecialToken":{"id":"[SEP]","type_id":0}},{"Sequence":{"id":"B","type_id":1}},{"SpecialToken":{"id":"[SEP]","type_id":1}}],"special_tokens":{"[CLS]":{"id":"[CLS]","ids":[1],"tokens":["[CLS]"]},"[SEP]":{"id":"[SEP]","ids":[2],"tokens":["[SEP]"]}}},"decoder":null,"model":{"type":"BPE","dropout":null,"unk_token":"[UNK]","continuing_subword_prefix":null,"end_of_word_suffix":null,"fuse_unk":false,"vocab":{"[UNK]":0,"[CLS]":1,"[SEP]":2,"[PAD]":3,"[MASK]":4,"A":5,"C":6,"D":7,"E":8,"F":9,"G":10,"H":11,"I":12,"K":13,"L":14,"M":15,"N":16,"P":17,"Q":18,"R":19,"S":20,"T":21,"V":22,"W":23,"Y":24},"merges":[]}}