Metadata-Version: 2.4
Name: msa-pairformer
Version: 1.0.2
Summary: MSA Pairformer model repository
Author-email: Yo Akiyama <yo_aki@mit.edu>
Maintainer-email: Yo Akiyama <yo_aki@mit.edu>
License: Modified MIT License (Coffee edition)
        
        Copyright (c) 2025 Yo Akiyama
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
        If we meet some day, and you think this work is valuable, you can buy me a 
        cup of coffee in return.
License-File: LICENSE
Requires-Python: >=3.10
Requires-Dist: biopython>=1.85
Requires-Dist: cuequivariance-ops-cu12>=0.7.0; sys_platform == 'linux'
Requires-Dist: cuequivariance-ops-torch-cu12>=0.7.0; sys_platform == 'linux'
Requires-Dist: cuequivariance-torch>=0.7.0; sys_platform == 'linux'
Requires-Dist: cuequivariance>=0.7.0; sys_platform == 'linux'
Requires-Dist: einops>=0.8.0
Requires-Dist: einx>=0.3.0
Requires-Dist: environs>=14.5.0
Requires-Dist: jaxtyping>=0.2.38
Requires-Dist: matplotlib>=3.5.1
Requires-Dist: numpy<2.0,>=1.26
Requires-Dist: scikit-learn>=0.23.2
Requires-Dist: scipy>=1.8.0
Requires-Dist: torch>=2.5.0
Requires-Dist: tqdm>=4.67.1
Description-Content-Type: text/markdown

# MSA Pairformer
<!-- ![Logo](msa_pairformer_logo.png) -->
<div align="left">
  <img src="msa_pairformer_logo.png" width="300" alt="Neural Network Logo">
</div>

[![Contact prediction](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/yoakiyama/MSA_Pairformer/blob/main/MSA_Pairformer_with_MMseqs2.ipynb/)

- [MSA Pairformer](#MSA-Pairformer)
  - [Installation ](#installation-)
  - [MSA Pairformer](#MSA-Pairformer--)
    - [Quickstart ](#quickstart--)
    - [Example Usage](#example-usage--)
  - [Licenses  ](#licenses--)
  - [Citations  ](#citations--)

This repository contains the latest release of MSA Pairformer and Google Colab notebooks for relevant analyses. Here, you will find how to use MSA Pairformer to embed protein sequences, predict residue-residue interactions in monomers and at the interface of protein-protein interactions, and perform zero-shot variant effect prediction.

## Installation <a name="installation"></a>

To get started with MSA Pairformer, install the python library using pip:

```bash
pip install msa-pairformer
```

or download this Github repository and install manually
```bash
git clone git@github.com:yoakiyama/MSA_Pairformer.git
pip install -e .
```

### Installing hhsuite (for filtering MSAs)

We use hhfilter (part of hhsuite) as the default option for subsampling MSAs. To install hhsuite, please follow these directions:

```
curl -fsSL https://github.com/soedinglab/hh-suite/releases/download/v3.3.0/hhsuite-3.3.0-SSE2-Linux.tar.gz
tar xz -C hhsuite
export PATH="$PATH:/path/to/hhsuite/bin:/path/to/hhsuite/scripts"
```
Alternatively, in python:
```
def _setup_tools():
  """Download and compile C++ tools."""

  # Install HHsuite
  hhsuite_path = "hhsuite"
  if not os.path.isdir(hhsuite_path):
      print("Installing HHsuite...")
      os.makedirs(hhsuite_path, exist_ok=True)
      url = "https://github.com/soedinglab/hh-suite/releases/download/v3.3.0/hhsuite-3.3.0-SSE2-Linux.tar.gz"
      os.system(f"curl -fsSL {url} | tar xz -C {hhsuite_path}/")
  os.environ['PATH'] += f":{hhsuite_path}/bin:{hhsuite_path}/scripts"
```

## MSA Pairformer <a name="MSA-Pairformer"></a>

[MSA Pairformer](https://www.cell.com/cell/fulltext/S0092-8674(26)00749-X) is an MSA-based protein language model that can model the coevolution of interacting proteins. In this repository, we provide the model source code, a Google Colab notebook for quick experimentation, and notebooks and scripts to reproduce the results of our manuscript. We are excited to deliver this tool to the community and to see all of its applications as a tool to study and engineer biology.

### Getting started with MSA Pairformer <a name="getting-started"></a>
The model's weights can be downloaded from Huggingface under [HuggingFace/yakiyama/MSA-Pairformer](https://huggingface.co/yakiyama/MSA-Pairformer/).
```py
import torch
import numpy as np
from huggingface_hub import login
from MSA_Pairformer.model import MSAPairformer
from MSA_Pairformer.dataset import MSA, prepare_msa_masks, aa2tok_d

# Use the GPU if available
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {torch.cuda.get_device_name(device)}")

# This function will allow you to login to huggingface via an API key
login()

# Download model weights and load model
# As long as the cache doesn't get cleared, you won't need to re-download the weights whenever you re-run this
model = MSAPairformer.from_pretrained(device=device)

# You can also save the downloaded weights to a specified directory in your filesystem.
# Saving the model weights like so will allow you to load the model without re-downloading if your cache gets cleared.
# Once you run this code once, you can re-run and it will automatically load the weights
save_model_dir = "model_weights"
model = MSAPairformer.from_pretrained(weights_dir=save_model_dir, device=device)
```

Here, we walk through how to run MSA Pairformer on the phenylalanyl tRNA synthetase pheS and pheT dimer (PDB ID: 1B70).
```py
# Subsample MSA using hhfilter and greedy diversification
msa_file = "data/1B70_A_1B70_B.fas"
max_msa_depth = 512
max_length = 10240
chain_break_idx = 265
np.random.seed(42)
msa_obj = MSA(
    msa_file_path=msa_file,
    max_seqs=max_msa_depth,
    max_length=max_length,
    max_tokens=np.inf,
    diverse_select_method="hhfilter",
    hhfilter_kwargs={"binary": "hhfilter"}
)
# Prepare MSA and mask tensors
msa_tokenized_t = msa_obj.diverse_tokenized_msa
msa_onehot_t = torch.nn.functional.one_hot(msa_tokenized_t, num_classes=len(aa2tok_d)).unsqueeze(0).float().to(device)
mask, msa_mask, full_mask, pairwise_mask = prepare_msa_masks(msa_obj.diverse_tokenized_msa.unsqueeze(0))
mask, msa_mask, full_mask, pairwise_mask = mask.to(device), msa_mask.to(device), full_mask.to(device), pairwise_mask.to(device)

# Run MSA Pairformer to generate embeddings and predict contacts
with torch.no_grad():
    with torch.amp.autocast(dtype=torch.bfloat16, device_type="cuda"):
        res = model(
            msa=msa_onehot_t.to(torch.bfloat16),
            mask=mask,
            msa_mask=msa_mask,
            full_mask=full_mask,
            pairwise_mask=pairwise_mask,
            complex_chain_break_indices=[[chain_break_idx]],
            return_seq_weights=True,
            return_pairwise_repr_layer_idx=None,
            return_msa_repr_layer_idx=None,
            return_cb_contacts=True,
            return_confind_contacts=True
        )
# res is a dictionary with the following keys: final_msa_repr, final_pairwise_repr, msa_repr_d, pairwise_repr_d, seq_weights_list_d, predicted_cb_contacts, predicted_confind_contacts

# Just predict Cb-Cb
with torch.no_grad():
    with torch.amp.autocast(dtype=torch.bfloat16, device_type="cuda"):
        res = model.predict_cb_contacts(
            msa=msa_onehot_t.to(torch.bfloat16),
            mask=mask,
            msa_mask=msa_mask,
            full_mask=full_mask,
            pairwise_mask=pairwise_mask,
            complex_chain_break_indices=[[chain_break_idx]],
            return_seq_weights=True,
        )
# Just predict ConFind contacts
with torch.no_grad():
    with torch.amp.autocast(dtype=torch.bfloat16, device_type="cuda"):
        res = model.predict_confind_contacts(
            msa=msa_onehot_t.to(torch.bfloat16),
            mask=mask,
            msa_mask=msa_mask,
            full_mask=full_mask,
            pairwise_mask=pairwise_mask,
            complex_chain_break_indices=[[chain_break_idx]],
            return_seq_weights=True,
        )
```

That's it -- you've generated embeddings and predicted contacts using MSA Pairformer!

## Licenses <a name="licenses"></a>

MSA Pairformer code and model weights are released under a permissive, slightly modified ☕️ MIT license. It can be freely used for both academic and commercial purposes.

## Citation <a name="citation"></a>
If you use MSA Pairformer in your work, please use the following citation
```
@article {Akiyama2026,
	author = {Akiyama, Yo and Zhang, Zhidian and Tang, Olivia and Kim, Rachel Seongeun and Mirdita, Milot and Steinegger, Martin and Ovchinnikov, Sergey},
	title = {Expanding the scope of protein language modeling to protein-protein interactions with MSA Pairformer},
	year = {2026},
	doi = {10.1016/j.cell.2026.06.029},
	publisher = {Cell Press},
	URL = {https://doi.org/10.1016/j.cell.2026.06.029},
	journal = {Cell}
}

```
