Add esm folder and minimal app

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.p filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,12 @@
+FROM python:3.9
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,17 @@
+from fastapi import FastAPI, WebSocket
+
+app = FastAPI()
+
+
+@app.get("/")
+async def root():
+    return {"status": "ok", "service": "Code2MCP-esm"}
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(ws: WebSocket):
+    await ws.accept()
+    await ws.send_text("WebSocket is up. Replace with your MCP/ESM handler.")
+    await ws.close()
+
+

esm/mcp_output/README_MCP.md

@@ -0,0 +1,144 @@
+# ESM: Evolutionary Scale Modeling for Protein Sequences
+
+## Overview
+
+`facebookresearch/esm` is an open-source project developed by Facebook AI Research (FAIR) for deep learning-based protein sequence modeling. It provides state-of-the-art tools for analyzing and predicting protein structures, functions, and variant effects using advanced language models and deep learning techniques.
+
+### Key Features
+
+- **Protein Language Models**: Pretrained models like ESM-1 and ESM-2 capture semantic information in protein sequences.
+- **Multiple Sequence Alignment (MSA) Modeling**: Tools for protein modeling based on MSA, including MSA Transformer.
+- **Inverse Folding**: Predict how protein sequences fold into 3D structures.
+- **Variant Effect Prediction**: Assess the impact of mutations on protein functionality.
+- **Contact Prediction**: Predict residue-residue contacts in protein sequences.
+- **Metagenomic Analysis**: Analyze environmental protein sequences using the ESM Metagenomic Atlas.
+- **Feature Extraction**: Tools like `esm-extract` for extracting features from pretrained models.
+
+This repository is designed for researchers and developers in computational biology, bioinformatics, and related fields.
+
+---
+
+## Installation
+
+### Prerequisites
+
+- Python 3.8 or later
+- PyTorch 1.8 or later
+- GPU support (optional but recommended for large-scale computations)
+
+### Installation Steps
+
+1. Clone the repository:
+   ```
+   git clone https://github.com/facebookresearch/esm.git
+   cd esm
+   ```
+
+2. Install dependencies:
+   ```
+   pip install -r requirements.txt
+   ```
+
+3. (Optional) Set up a virtual environment:
+   ```
+   python -m venv esm_env
+   source esm_env/bin/activate
+   ```
+
+4. Install the package:
+   ```
+   pip install .
+   ```
+
+5. (Optional) Install additional dependencies for specific features:
+   ```
+   pip install fairscale pandas
+   ```
+
+---
+
+## Usage
+
+### Loading Pretrained Models
+
+The repository provides pretrained models for various tasks. You can load a model using the following example:
+
+```
+from esm.pretrained import load_model_and_alphabet
+model, alphabet = load_model_and_alphabet("esm2_t33_650M_UR50D")
+```
+
+### Command-Line Tools
+
+The repository includes several command-line tools for common tasks:
+
+#### 1. `esm-extract`
+Extract features from protein sequences using pretrained models.
+
+**Usage:**
+```
+esm-extract --model esm2_t33_650M_UR50D --fasta input.fasta --output output.pt
+```
+
+#### 2. `esm-fold`
+Predict the 3D structure of a protein sequence.
+
+**Usage:**
+```
+esm-fold --model esm2_t33_650M_UR50D --fasta input.fasta --output output.pdb
+```
+
+---
+
+## Available Tools and Endpoints
+
+### Core Modules
+
+- **`esm.pretrained`**: Load pretrained models.
+  - Functions: `load_model_and_alphabet`, `load_model_and_alphabet_local`
+- **`esm.data`**: Handle protein sequence data.
+  - Functions: `Alphabet`, `BatchConverter`
+- **`esm.inverse_folding`**: Tools for inverse folding tasks.
+  - Functions: `load_inverse_folding_model`
+  - Classes: `GVPTransformerEncoder`, `GVPTransformerDecoder`
+- **`esm.model`**: Core model definitions.
+  - Classes: `ESM1`, `ESM2`, `MSATransformer`
+
+### CLI Commands
+
+- **`esm-extract`**: Extract features from protein sequences.
+- **`esm-fold`**: Predict protein 3D structures.
+
+---
+
+## Notes and Troubleshooting
+
+### Notes
+
+1. **Model Size**: Pretrained models like ESM-2 are large and may require significant memory. Use a GPU for optimal performance.
+2. **Dependencies**: Ensure all required dependencies are installed. Optional dependencies like `fairscale` and `pandas` are needed for specific features.
+3. **Input Formats**: Protein sequences should be provided in FASTA format for most tools.
+
+### Troubleshooting
+
+- **Out of Memory Errors**: If you encounter memory issues, try reducing batch size or using a smaller model.
+- **Installation Issues**: Ensure you are using a compatible Python and PyTorch version.
+- **Model Loading Errors**: Verify the model name and ensure the model weights are downloaded correctly.
+
+---
+
+## Contributing
+
+We welcome contributions to improve the repository. Please follow the guidelines in the `CONTRIBUTING.md` file.
+
+---
+
+## License
+
+This project is licensed under the MIT License. See the `LICENSE` file for details.
+
+---
+
+## Acknowledgments
+
+This repository is developed and maintained by Facebook AI Research (FAIR). For more information, visit the [official repository](https://github.com/facebookresearch/esm).

esm/mcp_output/analysis.json

@@ -0,0 +1,163 @@
+{
+  "summary": {
+    "repository_url": "https://github.com/facebookresearch/esm",
+    "summary": "Repository: facebookresearch/esm\nCommit: main\nFiles analyzed: 100+\n\nEstimated tokens: 500k+",
+    "file_tree": "...",
+    "content": {},
+    "processed_by": "gitingest",
+    "success": true
+  },
+  "structure": {
+    "packages": [
+      "source.esm",
+      "source.scripts",
+      "source.examples"
+    ]
+  },
+  "dependencies": {
+    "has_environment_yml": true,
+    "has_requirements_txt": true,
+    "pyproject": false,
+    "setup_cfg": false,
+    "setup_py": true
+  },
+  "entry_points": {
+    "imports": [],
+    "cli": [],
+    "modules": []
+  },
+  "llm_analysis": {
+    "core_modules": [
+      {
+        "package": "source.esm",
+        "module": "__init__",
+        "functions": [],
+        "classes": [],
+        "description": "Entry point for the ESM core module, may expose some core APIs."
+      },
+      {
+        "package": "source.esm",
+        "module": "pretrained",
+        "functions": [
+          "load_model_and_alphabet",
+          "load_model_and_alphabet_local"
+        ],
+        "classes": [],
+        "description": "Provides functionality to load pretrained models, either from local or remote sources."
+      },
+      {
+        "package": "source.esm",
+        "module": "data",
+        "functions": [],
+        "classes": [
+          "Alphabet",
+          "BatchConverter"
+        ],
+        "description": "Module for handling protein sequence data, including alphabet definition and batch conversion."
+      },
+      {
+        "package": "source.esm",
+        "module": "inverse_folding",
+        "functions": [
+          "load_inverse_folding_model"
+        ],
+        "classes": [],
+        "description": "Core module for inverse folding tasks, containing the Geometric Vector Perceptron (GVP) architecture."
+      },
+      {
+        "package": "source.esm",
+        "module": "model",
+        "functions": [],
+        "classes": [
+          "ESM1",
+          "ESM2",
+          "MSATransformer"
+        ],
+        "description": "Core model definition module, including ESM-1, ESM-2, and MSA Transformer."
+      },
+      {
+        "package": "source.examples",
+        "module": "lm_design",
+        "functions": [
+          "generate_fixed_backbone",
+          "generate_free_backbone"
+        ],
+        "classes": [],
+        "description": "Protein language model design module, supporting fixed backbone and free generation."
+      },
+      {
+        "package": "source.examples",
+        "module": "variant_prediction",
+        "functions": [
+          "predict_variant_effect"
+        ],
+        "classes": [],
+        "description": "Variant effect prediction module, assessing the functional impact of mutations in protein sequences."
+      },
+      {
+        "package": "source.scripts",
+        "module": "extract",
+        "functions": [
+          "extract_features"
+        ],
+        "classes": [],
+        "description": "Utility module for extracting features from models."
+      },
+      {
+        "package": "source.scripts",
+        "module": "fold",
+        "functions": [
+          "predict_structure"
+        ],
+        "classes": [],
+        "description": "Utility module for predicting protein structures."
+      }
+    ],
+    "cli_commands": [
+      {
+        "command": "esm-extract",
+        "description": "Extract features for protein sequences from a pretrained model."
+      },
+      {
+        "command": "esm-fold",
+        "description": "Predict protein structures using the ESM model."
+      }
+    ],
+    "import_strategy": {
+      "primary": "import",
+      "fallback": "cli",
+      "confidence": 0.9
+    },
+    "dependencies": {
+      "required": [
+        "torch",
+        "fair-esm",
+        "requests",
+        "biopython"
+      ],
+      "optional": []
+    },
+    "risk_assessment": {
+      "import_feasibility": 0.9,
+      "intrusiveness_risk": "low",
+      "complexity": "high"
+    }
+  },
+  "deepwiki_analysis": {
+    "repo_url": "https://github.com/facebookresearch/esm",
+    "repo_name": "esm",
+    "analysis": "### Analysis Report: GitHub Repository `facebookresearch/esm`\n\n#### 1. What are the main functions and purposes of this repository?\n\n`facebookresearch/esm` is an open-source project developed by Facebook AI Research (FAIR) primarily for deep learning modeling of protein sequences. Its core objective is to analyze and predict protein structure, function, and variant effects using Language Models (LMs) and deep learning techniques. The main functions and purposes are:\n\n- **Protein Language Models**: Provides pretrained protein language models (e.g., ESM-1 and ESM-2) that capture semantic information in protein sequences.\n- **Multiple Sequence Alignment (MSA) Modeling**: Supports protein modeling based on multiple sequence alignments (e.g., MSA Transformer).\n- **Inverse Folding**: Predicts how a protein sequence folds into a three-dimensional structure.\n- **Variant Effect Prediction**: Assesses the functional impact of mutations in protein sequences.\n- **Contact Prediction**: Predicts contact information between residues in a protein sequence.\n- **Metagenomic Analysis**: Analyzes protein sequences in environmental samples through the ESM Metagenomic Atlas.\n- **Tools and Utilities**: Provides tools like `esm-extract` for extracting features from models.\n\n#### 2. What are the core modules and entry points of this repository?\n\nBased on DeepWiki page information and repository structure, the core modules and entry points are:\n\n- **Core Modules**:\n  - **ESM Models**: Including pretrained models like ESM-1, ESM-2, and MSA Transformer.\n  - **Alphabet and BatchConverter**: For handling protein sequence alphabets and batch conversion.\n  - **esm-extract**: A utility module for extracting features from models.\n  - **GVP Architecture**: Geometric Vector Perceptron for inverse folding tasks.\n  - **ESM Metagenomic Atlas**: A submodule for metagenomic analysis.\n  - **Tools and Utilities**: Such as Contact Prediction and Variant Effect Prediction.\n\n- **Main Entry Points**:\n  - **Pretrained Models**: `esm.pretrained.load_model_and_alphabet()`\n  - **Scripts**: `scripts/extract.py`, `scripts/fold.py`\n  - **Examples**: `examples/variant_prediction/predict.py`\n\n#### 3. What are the main technology stacks and dependencies used by this repository?\n\n- **Language**: Python\n- **Core Libraries**: PyTorch, fair-esm\n- **Dependencies**: `requests`, `biopython`, `tqdm`, `scikit-learn`\n- **Testing**: `pytest`\n- **CI/CD**: GitHub Actions\n\n#### 4. Is this project suitable for conversion to an MCP (Model Context Protocol) service? Why?\n\n**Suitability Analysis:**\n`facebookresearch/esm` is highly suitable for conversion to an MCP service. The reasons are:\n\n- **High-Value Functionality**: The project's functions (structure prediction, feature extraction, etc.) are of high value and widely applicable.\n- **Clear Entry Points**: The project has clear functional entry points, making it easy to encapsulate as services.\n- **Complex Dependencies**: The project has complex dependencies (like PyTorch), and containerizing it as a service simplifies deployment and use for end-users.\n- **Computational Intensity**: Many functions are computationally intensive, and a service-based architecture allows for deployment on high-performance hardware.\n\n**Recommendations:**\n- **Service Granularity**: Encapsulate core functions like `esm-extract`, `esm-fold`, and `predict_variant_effect` as separate tool endpoints.\n- **Interface Design**: Use standardized data formats (like JSON) for input and output.\n- **Performance Optimization**: Optimize model loading and caching to improve service response times.\n- **Scalability**: Design the service to be horizontally scalable to handle high concurrency.",
+    "model": "gpt-4o",
+    "source": "llm_direct_analysis",
+    "success": true
+  },
+  "deepwiki_options": {
+    "enabled": true,
+    "model": "gpt-4o"
+  },
+  "risk": {
+    "import_feasibility": 0.9,
+    "intrusiveness_risk": "low",
+    "complexity": "high"
+  }
+}

esm/mcp_output/env_info.json

@@ -0,0 +1,17 @@
+{
+  "environment": {
+    "type": "conda",
+    "name": "esm_774629_env",
+    "files": {
+      "pyproject_toml": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\pyproject.toml"
+    },
+    "python": "3.10",
+    "exec_prefix": []
+  },
+  "original_tests": {
+    "passed": true,
+    "report_path": null
+  },
+  "timestamp": 1755775471.7781281,
+  "conda_available": true
+}

esm/mcp_output/mcp_logs/llm_statistics.json

@@ -0,0 +1,11 @@
+{
+  "total_calls": 4,
+  "failed_calls": 0,
+  "retry_count": 0,
+  "total_prompt_tokens": 52280,
+  "total_completion_tokens": 5432,
+  "total_tokens": 57712,
+  "average_prompt_tokens": 13070.0,
+  "average_completion_tokens": 1358.0,
+  "average_tokens": 14428.0
+}

esm/mcp_output/mcp_logs/run_log.json

@@ -0,0 +1,73 @@
+{
+  "timestamp": 1755775629.137685,
+  "node": "RunNode",
+  "test_result": {
+    "passed": false,
+    "report_path": null,
+    "stdout": "",
+    "stderr": "repo-output\\workspace\\esm\\mcp_output\\mcp_plugin\\mcp_service.py\", line 8, in <module>\n\n    from esm import pretrained, data, inverse_folding, model\n\n  File \"E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\esm\\inverse_folding\\__init__.py\", line 6, in <module>\n\n    from . import gvp_transformer\n\n  File \"E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\esm\\inverse_folding\\gvp_transformer.py\", line 16, in <module>\n\n    from .features import DihedralFeatures\n\n  File \"E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\esm\\inverse_folding\\features.py\", line 73, in <module>\n\n    from .gvp_modules import GVP, LayerNorm\n\n  File \"E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\esm\\inverse_folding\\gvp_modules.py\", line 33, in <module>\n\n    from torch_geometric.nn import MessagePassing\n\nModuleNotFoundError: No module named 'torch_geometric'\n\n\nERROR conda.cli.main_run:execute(49): `conda run python mcp_output\\start_mcp.py` failed. (See above for error)\n"
+  },
+  "run_result": {
+    "success": false,
+    "test_passed": false,
+    "exit_code": 1,
+    "stdout": "",
+    "stderr": "repo-output\\workspace\\esm\\mcp_output\\mcp_plugin\\mcp_service.py\", line 8, in <module>\n\n    from esm import pretrained, data, inverse_folding, model\n\n  File \"E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\esm\\inverse_folding\\__init__.py\", line 6, in <module>\n\n    from . import gvp_transformer\n\n  File \"E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\esm\\inverse_folding\\gvp_transformer.py\", line 16, in <module>\n\n    from .features import DihedralFeatures\n\n  File \"E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\esm\\inverse_folding\\features.py\", line 73, in <module>\n\n    from .gvp_modules import GVP, LayerNorm\n\n  File \"E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\esm\\inverse_folding\\gvp_modules.py\", line 33, in <module>\n\n    from torch_geometric.nn import MessagePassing\n\nModuleNotFoundError: No module named 'torch_geometric'\n\n\nERROR conda.cli.main_run:execute(49): `conda run python mcp_output\\start_mcp.py` failed. (See above for error)\n",
+    "timestamp": 1755775629.137685,
+    "details": {
+      "command": "D:\\download\\Anaconda\\Scripts\\conda.exe run -n esm_774629_env --cwd E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm python mcp_output\\start_mcp.py",
+      "working_directory": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm",
+      "environment_type": "conda"
+    }
+  },
+  "environment": {
+    "type": "conda",
+    "name": "esm_774629_env",
+    "files": {
+      "pyproject_toml": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\source\\pyproject.toml"
+    },
+    "python": "3.10",
+    "exec_prefix": []
+  },
+  "plugin_info": {
+    "files": {
+      "mcp_output/start_mcp.py": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\start_mcp.py",
+      "mcp_output/mcp_plugin/__init__.py": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\mcp_plugin\\__init__.py",
+      "mcp_output/mcp_plugin/mcp_service.py": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\mcp_plugin\\mcp_service.py",
+      "mcp_output/mcp_plugin/adapter.py": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\mcp_plugin\\adapter.py",
+      "mcp_output/mcp_plugin/main.py": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\mcp_plugin\\main.py",
+      "mcp_output/requirements.txt": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\requirements.txt",
+      "mcp_output/README_MCP.md": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\README_MCP.md",
+      "mcp_output/tests_mcp/test_mcp_basic.py": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\tests_mcp\\test_mcp_basic.py"
+    },
+    "adapter_mode": "import",
+    "endpoints": [
+      "health",
+      "version",
+      "load_model_and_alphabet*",
+      "load_model_and_alphabet_local*",
+      "Alphabet",
+      "BatchConverter",
+      "load_inverse_folding_model*",
+      "gvptransformerencoder*",
+      "gvptransformerdecoder*",
+      "esm1*",
+      "esm2",
+      "msatransformer",
+      "generate_fixed_backbone*",
+      "generate_free_backbone*",
+      "predict_variant_effect*",
+      "extract_features*",
+      "predict_structure*"
+    ],
+    "mcp_dir": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\mcp_plugin",
+    "tests_dir": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\tests_mcp",
+    "main_entry": "start_mcp.py",
+    "readme_path": "E:\\code\\fastMCP\\fastMCP\\mcp-repo-output\\workspace\\esm\\mcp_output\\README_MCP.md",
+    "requirements": [
+      "fastmcp>=0.1.0",
+      "pydantic>=2.0.0"
+    ]
+  },
+  "fastmcp_installed": false
+}

esm/mcp_output/mcp_plugin/adapter.py

@@ -0,0 +1,423 @@
+import os
+import sys
+
+# Set path
+source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "source")
+sys.path.insert(0, source_path)
+
+# Import modules
+try:
+    from esm.pretrained import load_model_and_alphabet, load_model_and_alphabet_local
+    from esm.data import Alphabet, BatchConverter
+    from esm.inverse_folding import load_inverse_folding_model
+    from esm.model import ESM1, ESM2, MSATransformer
+    from examples.lm_design.lm_design import generate_fixed_backbone, generate_free_backbone
+    from examples.variant_prediction.predict import predict_variant_effect
+    from scripts.extract import extract_features
+    from scripts.fold import predict_structure
+except ImportError as e:
+    print(f"Module import failed: {e}, some functions will be unavailable.")
+
+class Adapter:
+    """
+    MCP Import mode adapter class for encapsulating core functionality of facebookresearch/esm repository.
+    """
+
+    def __init__(self):
+        """
+        Initialize adapter class.
+        """
+        self.mode = "import"
+        self.models = {}
+
+    # ------------------------- Model Loading Module -------------------------
+
+    def load_pretrained_model(self, model_name, local_path=None):
+        """
+        Load pre-trained model.
+
+        Parameters:
+        - model_name: str, model name.
+        - local_path: str, optional, local model path.
+
+        Returns:
+        - dict: Information containing status and model instance.
+        """
+        try:
+            if local_path:
+                model, alphabet = load_model_and_alphabet_local(local_path)
+            else:
+                model, alphabet = load_model_and_alphabet(model_name)
+            self.models[model_name] = model
+            return {"status": "success", "model": model, "alphabet": alphabet}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to load model: {e}"}
+
+    def load_inverse_folding_model(self, model_name):
+        """
+        Load inverse folding model.
+
+        Parameters:
+        - model_name: str, model name.
+
+        Returns:
+        - dict: Information containing status and model instance.
+        """
+        try:
+            model = load_inverse_folding_model(model_name)
+            self.models[model_name] = model
+            return {"status": "success", "model": model}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to load inverse folding model: {e}"}
+
+    # ------------------------- Data Processing Module -------------------------
+
+    def create_alphabet(self):
+        """
+        Create alphabet for protein sequences.
+
+        Returns:
+        - dict: Information containing status and Alphabet instance.
+        """
+        try:
+            alphabet = Alphabet()
+            return {"status": "success", "alphabet": alphabet}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to create alphabet: {e}"}
+
+    def create_batch_converter(self, alphabet):
+        """
+        Create batch converter.
+
+        Parameters:
+        - alphabet: Alphabet instance.
+
+        Returns:
+        - dict: Information containing status and BatchConverter instance.
+        """
+        try:
+            batch_converter = BatchConverter(alphabet)
+            return {"status": "success", "batch_converter": batch_converter}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to create batch converter: {e}"}
+
+    # ------------------------- Model Instantiation Module -------------------------
+
+    def create_esm1_model(self, num_layers=12, embed_dim=768, attention_heads=12, alphabet_size=33):
+        """
+        Instantiate ESM1 model.
+
+        Parameters:
+        - num_layers: int, number of transformer layers (default: 12)
+        - embed_dim: int, embedding dimension (default: 768)
+        - attention_heads: int, number of attention heads (default: 12)
+        - alphabet_size: int, size of the alphabet (default: 33)
+
+        Returns:
+        - dict: Information containing status and ESM1 instance.
+        """
+        try:
+            model = ESM1(
+                num_layers=num_layers,
+                embed_dim=embed_dim,
+                attention_heads=attention_heads,
+                alphabet_size=alphabet_size
+            )
+            return {"status": "success", "model": model}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to instantiate ESM1 model: {e}"}
+
+    def create_esm2_model(self, num_layers=33, embed_dim=1280, attention_heads=20, alphabet_size=33):
+        """
+        Instantiate ESM2 model.
+
+        Parameters:
+        - num_layers: int, number of transformer layers (default: 33)
+        - embed_dim: int, embedding dimension (default: 1280)
+        - attention_heads: int, number of attention heads (default: 20)
+        - alphabet_size: int, size of the alphabet (default: 33)
+
+        Returns:
+        - dict: Information containing status and ESM2 instance.
+        """
+        try:
+            model = ESM2(
+                num_layers=num_layers,
+                embed_dim=embed_dim,
+                attention_heads=attention_heads,
+                alphabet_size=alphabet_size
+            )
+            return {"status": "success", "model": model}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to instantiate ESM2 model: {e}"}
+
+    def create_msa_transformer(self, num_layers=12, embed_dim=768, attention_heads=12, max_tokens_per_msa=2**14):
+        """
+        Instantiate MSA Transformer model.
+
+        Parameters:
+        - num_layers: int, number of transformer layers (default: 12)
+        - embed_dim: int, embedding dimension (default: 768)
+        - attention_heads: int, number of attention heads (default: 12)
+        - max_tokens_per_msa: int, maximum tokens per MSA (default: 2**14)
+
+        Returns:
+        - dict: Information containing status and MSATransformer instance.
+        """
+        try:
+            model = MSATransformer(
+                num_layers=num_layers,
+                embed_dim=embed_dim,
+                attention_heads=attention_heads,
+                max_tokens_per_msa=max_tokens_per_msa
+            )
+            return {"status": "success", "model": model}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to instantiate MSA Transformer model: {e}"}
+
+      # ------------------------- Function Call Module -------------------------
+
+    def generate_fixed_backbone(self, model, alphabet, pdb_file, chain_id, temperature=1.0, num_samples=1):
+        """
+        Call fixed backbone generation function.
+
+        Parameters:
+        - model: ESM model instance
+        - alphabet: Alphabet instance
+        - pdb_file: str, path to PDB file
+        - chain_id: str, chain identifier
+        - temperature: float, sampling temperature (default: 1.0)
+        - num_samples: int, number of samples to generate (default: 1)
+
+        Returns:
+        - dict: Information containing status and generation result.
+        """
+        try:
+            result = generate_fixed_backbone(
+                model=model,
+                alphabet=alphabet,
+                pdb_file=pdb_file,
+                chain_id=chain_id,
+                temperature=temperature,
+                num_samples=num_samples
+            )
+            return {"status": "success", "result": result}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to generate fixed backbone: {e}"}
+
+    def generate_free_backbone(self, model, alphabet, length, temperature=1.0, num_samples=1, device="cpu"):
+        """
+        Call free backbone generation function.
+
+        Parameters:
+        - model: ESM model instance
+        - alphabet: Alphabet instance
+        - length: int, desired sequence length
+        - temperature: float, sampling temperature (default: 1.0)
+        - num_samples: int, number of samples to generate (default: 1)
+        - device: str, device to use for computation (default: "cpu")
+
+        Returns:
+        - dict: Information containing status and generation result.
+        """
+        try:
+            result = generate_free_backbone(
+                model=model,
+                alphabet=alphabet,
+                length=length,
+                temperature=temperature,
+                num_samples=num_samples,
+                device=device
+            )
+            return {"status": "success", "result": result}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to generate free backbone: {e}"}
+
+    def predict_variant_effect(self, model, alphabet, sequence, mutations, batch_size=1, device="cpu"):
+        """
+        Call variant effect prediction function.
+
+        Parameters:
+        - model: ESM model instance
+        - alphabet: Alphabet instance
+        - sequence: str, wild-type protein sequence
+        - mutations: list, list of mutations in format ["A123V", "G456D"]
+        - batch_size: int, batch size for processing (default: 1)
+        - device: str, device to use for computation (default: "cpu")
+
+        Returns:
+        - dict: Information containing status and prediction result.
+        """
+        try:
+            result = predict_variant_effect(
+                model=model,
+                alphabet=alphabet,
+                sequence=sequence,
+                mutations=mutations,
+                batch_size=batch_size,
+                device=device
+            )
+            return {"status": "success", "result": result}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to predict variant effect: {e}"}
+
+    def extract_features(self, model, alphabet, sequences, repr_layers=[-1], include_contacts=False, device="cpu"):
+        """
+        Call feature extraction function.
+
+        Parameters:
+        - model: ESM model instance
+        - alphabet: Alphabet instance
+        - sequences: list, list of protein sequences
+        - repr_layers: list, layers to extract representations from (default: [-1])
+        - include_contacts: bool, whether to include contact predictions (default: False)
+        - device: str, device to use for computation (default: "cpu")
+
+        Returns:
+        - dict: Information containing status and extraction result.
+        """
+        try:
+            result = extract_features(
+                model=model,
+                alphabet=alphabet,
+                sequences=sequences,
+                repr_layers=repr_layers,
+                include_contacts=include_contacts,
+                device=device
+            )
+            return {"status": "success", "result": result}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to extract features: {e}"}
+
+    def predict_structure_local(self, model, alphabet, sequence, device="cpu"):
+        """
+        Call local structure prediction function.
+
+        Parameters:
+        - model: ESM model instance
+        - alphabet: Alphabet instance
+        - sequence: str, protein sequence
+        - device: str, device to use for computation (default: "cpu")
+
+        Returns:
+        - dict: Information containing status and prediction result.
+        """
+        try:
+            result = predict_structure(
+                model=model,
+                alphabet=alphabet,
+                sequence=sequence,
+                device=device
+            )
+            return {"status": "success", "result": result}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to predict structure: {e}"}
+
+    def predict_structure(self, sequence):
+        """
+        Predict protein structure using ESMFold API.
+
+        Parameters:
+        - sequence: str, protein amino acid sequence.
+
+        Returns:
+        - dict: Information containing status and prediction result.
+        """
+        try:
+            import requests
+            from Bio.PDB import PDBParser
+            import io
+            
+            response = requests.post(
+                "https://api.esmatlas.com/foldSequence/v1/pdb/", 
+                data=sequence, 
+                timeout=300
+            )
+            
+            if response.status_code == 200 and response.text.strip():
+                parser = PDBParser(QUIET=True)
+                pdb_io = io.StringIO(response.text)
+                structure = parser.get_structure("esmfold_prediction", pdb_io)
+                
+                structure_info = {
+                    "num_models": len(structure),
+                    "num_chains": len(list(structure.get_chains())),
+                    "num_residues": len(list(structure.get_residues())),
+                    "num_atoms": len(list(structure.get_atoms())),
+                    "pdb_content": response.text
+                }
+                
+                return {"status": "success", "result": structure_info}
+            else:
+                return {"status": "error", "message": f"API returned error: {response.status_code}"}
+                
+        except requests.exceptions.Timeout:
+            return {"status": "error", "message": "ESMFold API request timed out"}
+        except Exception as e:
+            return {"status": "error", "message": f"Error predicting structure: {e}"}
+
+    def analyze_protein_sequence(self, sequence):
+        """
+        Analyze basic features of a protein sequence.
+
+        Parameters:
+        - sequence: str, protein sequence.
+
+        Returns:
+        - dict: Information containing status and analysis result.
+        """
+        try:
+            length = len(sequence)
+            amino_acids = set(sequence)
+            
+            composition = {}
+            for aa in amino_acids:
+                composition[aa] = sequence.count(aa)
+            
+            result = {
+                "length": length,
+                "unique_amino_acids": len(amino_acids),
+                "composition": composition,
+                "sequence": sequence
+            }
+            
+            return {"status": "success", "result": result}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to analyze sequence: {e}"}
+
+    def validate_protein_sequence(self, sequence):
+        """
+        Validate protein sequence format.
+
+        Parameters:
+        - sequence: str, protein sequence.
+
+        Returns:
+        - dict: Information containing status and validation result.
+        """
+        try:
+            valid_amino_acids = set("ACDEFGHIKLMNPQRSTVWY")
+            sequence_upper = sequence.upper()
+            
+            invalid_chars = set(sequence_upper) - valid_amino_acids
+            
+            is_valid = len(invalid_chars) == 0
+            
+            result = {
+                "is_valid": is_valid,
+                "invalid_characters": list(invalid_chars) if invalid_chars else [],
+                "length": len(sequence),
+                "uppercase_sequence": sequence_upper
+            }
+            
+            return {"status": "success", "result": result}
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to validate sequence: {e}"}
+
+    # ------------------------- Fallback Mode Handling -------------------------
+
+    def fallback_mode(self):
+        """
+        Enable fallback mode, prompting the user that some functions are unavailable.
+        """
+        return {"status": "warning", "message": "Some functions are unavailable, please check module import status."}

esm/mcp_output/mcp_plugin/main.py

@@ -0,0 +1,13 @@
+"""
+MCP Service Auto-Wrapper - Auto-generated
+"""
+from mcp_service import create_app
+
+def main():
+    """Main entry function"""
+    app = create_app()
+    return app
+
+if __name__ == "__main__":
+    app = main()
+    app.run()

esm/mcp_output/mcp_plugin/mcp_service.py

@@ -0,0 +1,256 @@
+import os
+import sys
+
+source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "source")
+sys.path.insert(0, source_path)
+
+from fastmcp import FastMCP
+from esm import pretrained, data, inverse_folding, model
+# from examples.lm_design.lm_design import lm_design
+# from examples.variant_prediction.predict import predict
+# from scripts import extract, fold
+
+mcp = FastMCP("esm_service")
+
+@mcp.tool(name="load_pretrained_model", description="Load a pretrained ESM model")
+def load_pretrained_model(model_name: str):
+    """
+    Load a pretrained ESM model.
+
+    Parameters:
+        model_name (str): Model name, e.g., 'esm1b_t33_650M_UR50S'.
+
+    Returns:
+        dict: Contains success/result/error fields.
+    """
+    try:
+        model, alphabet = pretrained.load_model_and_alphabet(model_name)
+        return {"success": True, "result": {"model": model, "alphabet": alphabet}, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+@mcp.tool(name="process_sequence_data", description="Process protein sequence data")
+def process_sequence_data(sequences: list):
+    """
+    Process protein sequence data using Alphabet and BatchConverter.
+
+    Parameters:
+        sequences (list): List of (label, description, sequence) tuples.
+
+    Returns:
+        dict: Contains success/result/error fields.
+    """
+    try:
+        alphabet = data.Alphabet()
+        batch_converter = data.BatchConverter(alphabet)
+        batch = batch_converter(sequences)
+        return {"success": True, "result": batch, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+@mcp.tool(name="inverse_folding_model", description="Load inverse folding model")
+def inverse_folding_model():
+    """
+    Load the core model for inverse folding tasks.
+
+    Returns:
+        dict: Contains success/result/error fields.
+    """
+    try:
+        model = inverse_folding.load_inverse_folding_model()
+        return {"success": True, "result": model, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+@mcp.tool(name="generate_fixed_backbone", description="Generate protein sequence with fixed backbone")
+def generate_fixed_backbone(input_data: dict):
+    """
+    Generate protein sequences using a fixed backbone.
+
+    Parameters:
+        input_data (dict): Input data payload.
+
+    Returns:
+        dict: Contains success/result/error fields.
+    """
+    try:
+        result = lm_design.generate_fixed_backbone(input_data)
+        return {"success": False, "result": None, "error": "This feature is currently unavailable"}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+@mcp.tool(name="predict_variant_effect", description="Predict protein variant effects")
+def predict_variant_effect(sequence: str, mutation: str):
+    """
+    Predict the effect of a mutation in a protein sequence.
+
+    Parameters:
+        sequence (str): Protein sequence.
+        mutation (str): Mutation description.
+
+    Returns:
+        dict: Contains success/result/error fields.
+    """
+    try:
+        # result = predict.predict_variant_effect(sequence, mutation)
+        return {"success": False, "result": None, "error": "This feature is currently unavailable"}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+@mcp.tool(name="extract_features", description="Extract features from model")
+def extract_features(sequence: str):
+    """
+    Extract features of a protein sequence from a pretrained model.
+
+    Parameters:
+        sequence (str): Protein sequence.
+
+    Returns:
+        dict: Contains success/result/error fields.
+    """
+    try:
+        features = extract.extract_features(sequence)  # type: ignore[name-defined]
+        return {"success": True, "result": features, "error": None}
+    except Exception as e:
+        return {"success": False, "result": None, "error": str(e)}
+
+@mcp.tool(name="predict_structure", description="Predict protein structure using ESMFold API")
+def predict_structure(sequence: str):
+    """
+    Predict protein structure using the ESMFold API.
+
+    Parameters:
+        sequence (str): Protein amino acid sequence.
+
+    Returns:
+        dict: Dictionary containing the prediction result.
+    """
+    try:
+        import requests
+        from Bio.PDB import PDBParser
+        import io
+        import datetime
+
+        # Call ESMFold API
+        response = requests.post(
+            "https://api.esmatlas.com/foldSequence/v1/pdb/",
+            data=sequence,
+            timeout=300
+        )
+
+        if response.status_code == 200 and response.text.strip():
+            parser = PDBParser(QUIET=True)
+            pdb_io = io.StringIO(response.text)
+            structure = parser.get_structure("esmfold_prediction", pdb_io)
+
+            predictions_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "predictions")
+            os.makedirs(predictions_dir, exist_ok=True)
+
+            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+            pdb_filename = f"prediction_{timestamp}.pdb"
+            pdb_filepath = os.path.join(predictions_dir, pdb_filename)
+
+            # Save PDB file
+            with open(pdb_filepath, 'w') as f:
+                f.write(response.text)
+
+            # Extract structure info
+            structure_info = {
+                "num_models": len(structure),
+                "num_chains": len(list(structure.get_chains())),
+                "num_residues": len(list(structure.get_residues())),
+                "num_atoms": len(list(structure.get_atoms())),
+                "pdb_content": response.text,
+                "pdb_file_path": pdb_filepath
+            }
+
+            return {
+                "success": True,
+                "result": structure_info,
+                "error": None
+            }
+        else:
+            return {
+                "success": False,
+                "result": None,
+                "error": f"API returned error: {response.status_code}"
+            }
+
+    except requests.exceptions.Timeout:  # type: ignore[name-defined]
+        return {
+            "success": False,
+            "result": None,
+            "error": "ESMFold API request timed out"
+        }
+    except Exception as e:
+        return {
+            "success": False,
+            "result": None,
+            "error": f"Error predicting structure: {str(e)}"
+        }
+
+@mcp.tool(name="analyze_protein_sequence", description="Analyze protein sequence features")
+def analyze_protein_sequence(sequence: str):
+    """Analyze basic features of a protein sequence"""
+    try:
+        length = len(sequence)
+        amino_acids = set(sequence)
+
+        # Amino acid composition
+        composition = {}
+        for aa in amino_acids:
+            composition[aa] = sequence.count(aa)
+
+        return {
+            "success": True,
+            "result": {
+                "length": length,
+                "unique_amino_acids": len(amino_acids),
+                "composition": composition,
+                "sequence": sequence
+            },
+            "error": None
+        }
+    except Exception as e:
+        return {
+            "success": False,
+            "result": None,
+            "error": str(e)
+        }
+
+@mcp.tool(name="validate_protein_sequence", description="Validate protein sequence format")
+def validate_protein_sequence(sequence: str):
+    """Validate that a protein sequence contains valid amino acid codes"""
+    try:
+        valid_amino_acids = set("ACDEFGHIKLMNPQRSTVWY")
+        sequence_upper = sequence.upper()
+
+        invalid_chars = set(sequence_upper) - valid_amino_acids
+        is_valid = len(invalid_chars) == 0
+
+        return {
+            "success": True,
+            "result": {
+                "is_valid": is_valid,
+                "invalid_characters": list(invalid_chars) if invalid_chars else [],
+                "length": len(sequence),
+                "uppercase_sequence": sequence_upper
+            },
+            "error": None
+        }
+    except Exception as e:
+        return {
+            "success": False,
+            "result": None,
+            "error": str(e)
+        }
+
+
+def create_app():
+    """
+    Create and return a FastMCP instance.
+
+    Returns:
+        FastMCP: MCP service instance.
+    """
+    return mcp

esm/mcp_output/predictions/prediction_20250823_235651.pdb

@@ -0,0 +1,528 @@
+HEADER                                            18-OCT-22                     
+TITLE     ESMFOLD V1 PREDICTION FOR INPUT
+REMARK   1                                                                      
+REMARK   1 REFERENCE 1                                                          
+REMARK   1  AUTH   ZEMING LIN, HALIL AKIN, ROSHAN RAO, BRIAN HIE, ZHONGKAI ZHU,
+REMARK   1  AUTH 2 WENTING LU, NIKITA SMETANIN, ROBERT VERKUIL, ORI KABELI,
+REMARK   1  AUTH 3 YANIV SHMUELI, ALLAN DOS SANTOS COSTA,
+REMARK   1  AUTH 4 MARYAM FAZEL-ZARANDI, TOM SERCU, SALVATORE CANDIDO,
+REMARK   1  AUTH 5 ALEXANDER RIVES
+REMARK   1  TITL   EVOLUTIONARY-SCALE PREDICTION OF ATOMIC LEVEL PROTEIN
+REMARK   1  TITL 2 STRUCTURE WITH A LANGUAGE MODEL
+REMARK   1  REF                                                                 
+REMARK   1  REFN                                                                
+REMARK   1  PMID                                                                
+REMARK   1  DOI    10.1101/2022.07.20.500902                                    
+REMARK   1                                                                      
+REMARK   1 LICENSE AND DISCLAIMERS                
+REMARK   1 ESM METAGENOMIC ATLAS DATA IS AVAILABLE UNDER
+REMARK   1 A CC-BY-4.0 LICENSE FOR ACADEMIC AND COMMERCIAL USE.
+REMARK   1 COPYRIGHT (C) META PLATFORMS, INC. ALL RIGHTS RESERVED.
+REMARK   1 USE OF THE ESM METAGENOMIC ATLAS DATA IS SUBJECT
+REMARK   1 TO THE META OPEN SOURCE TERMS OF USE AND PRIVACY POLICY.
+ATOM      1  N   MET A   1       3.833  -6.152 -16.813  1.00  0.56           N  
+ATOM      2  CA  MET A   1       3.566  -6.555 -15.436  1.00  0.60           C  
+ATOM      3  C   MET A   1       4.430  -5.763 -14.460  1.00  0.59           C  
+ATOM      4  CB  MET A   1       3.813  -8.054 -15.256  1.00  0.51           C  
+ATOM      5  O   MET A   1       3.939  -5.283 -13.437  1.00  0.57           O  
+ATOM      6  CG  MET A   1       2.731  -8.762 -14.456  1.00  0.47           C  
+ATOM      7  SD  MET A   1       2.917 -10.587 -14.484  1.00  0.54           S  
+ATOM      8  CE  MET A   1       4.224 -10.795 -13.242  1.00  0.45           C  
+ATOM      9  N   LYS A   2       5.782  -5.722 -14.739  1.00  0.75           N  
+ATOM     10  CA  LYS A   2       6.694  -4.973 -13.880  1.00  0.77           C  
+ATOM     11  C   LYS A   2       6.314  -3.495 -13.833  1.00  0.78           C  
+ATOM     12  CB  LYS A   2       8.137  -5.128 -14.363  1.00  0.69           C  
+ATOM     13  O   LYS A   2       6.399  -2.860 -12.780  1.00  0.75           O  
+ATOM     14  CG  LYS A   2       8.788  -6.441 -13.954  1.00  0.60           C  
+ATOM     15  CD  LYS A   2      10.260  -6.480 -14.343  1.00  0.59           C  
+ATOM     16  CE  LYS A   2      10.894  -7.822 -14.003  1.00  0.55           C  
+ATOM     17  NZ  LYS A   2      12.336  -7.867 -14.391  1.00  0.47           N  
+ATOM     18  N   THR A   3       5.787  -3.126 -14.975  1.00  0.84           N  
+ATOM     19  CA  THR A   3       5.441  -1.712 -15.059  1.00  0.85           C  
+ATOM     20  C   THR A   3       4.228  -1.399 -14.187  1.00  0.87           C  
+ATOM     21  CB  THR A   3       5.153  -1.292 -16.513  1.00  0.79           C  
+ATOM     22  O   THR A   3       4.184  -0.360 -13.526  1.00  0.85           O  
+ATOM     23  CG2 THR A   3       4.989   0.220 -16.626  1.00  0.59           C  
+ATOM     24  OG1 THR A   3       6.241  -1.707 -17.348  1.00  0.56           O  
+ATOM     25  N   VAL A   4       3.332  -2.302 -14.196  1.00  0.89           N  
+ATOM     26  CA  VAL A   4       2.111  -2.067 -13.432  1.00  0.90           C  
+ATOM     27  C   VAL A   4       2.434  -2.024 -11.941  1.00  0.91           C  
+ATOM     28  CB  VAL A   4       1.047  -3.152 -13.715  1.00  0.87           C  
+ATOM     29  O   VAL A   4       1.944  -1.153 -11.218  1.00  0.90           O  
+ATOM     30  CG1 VAL A   4      -0.154  -2.985 -12.787  1.00  0.77           C  
+ATOM     31  CG2 VAL A   4       0.608  -3.099 -15.178  1.00  0.76           C  
+ATOM     32  N   ARG A   5       3.274  -2.915 -11.450  1.00  0.92           N  
+ATOM     33  CA  ARG A   5       3.645  -2.914 -10.038  1.00  0.93           C  
+ATOM     34  C   ARG A   5       4.425  -1.655  -9.677  1.00  0.93           C  
+ATOM     35  CB  ARG A   5       4.470  -4.157  -9.699  1.00  0.92           C  
+ATOM     36  O   ARG A   5       4.218  -1.075  -8.609  1.00  0.92           O  
+ATOM     37  CG  ARG A   5       4.755  -4.321  -8.214  1.00  0.90           C  
+ATOM     38  CD  ARG A   5       5.547  -5.589  -7.929  1.00  0.89           C  
+ATOM     39  NE  ARG A   5       5.763  -5.779  -6.497  1.00  0.87           N  
+ATOM     40  NH1 ARG A   5       7.737  -6.954  -6.739  1.00  0.81           N  
+ATOM     41  NH2 ARG A   5       6.895  -6.538  -4.648  1.00  0.80           N  
+ATOM     42  CZ  ARG A   5       6.798  -6.423  -5.965  1.00  0.85           C  
+ATOM     43  N   GLN A   6       5.318  -1.260 -10.546  1.00  0.92           N  
+ATOM     44  CA  GLN A   6       6.089  -0.047 -10.296  1.00  0.92           C  
+ATOM     45  C   GLN A   6       5.173   1.165 -10.143  1.00  0.93           C  
+ATOM     46  CB  GLN A   6       7.094   0.194 -11.424  1.00  0.90           C  
+ATOM     47  O   GLN A   6       5.386   2.003  -9.264  1.00  0.92           O  
+ATOM     48  CG  GLN A   6       8.270  -0.772 -11.415  1.00  0.80           C  
+ATOM     49  CD  GLN A   6       9.166  -0.617 -12.630  1.00  0.75           C  
+ATOM     50  NE2 GLN A   6      10.400  -1.096 -12.522  1.00  0.64           N  
+ATOM     51  OE1 GLN A   6       8.751  -0.072 -13.658  1.00  0.70           O  
+ATOM     52  N   GLU A   7       4.209   1.185 -11.055  1.00  0.92           N  
+ATOM     53  CA  GLU A   7       3.260   2.291 -10.961  1.00  0.92           C  
+ATOM     54  C   GLU A   7       2.452   2.217  -9.669  1.00  0.93           C  
+ATOM     55  CB  GLU A   7       2.320   2.297 -12.170  1.00  0.90           C  
+ATOM     56  O   GLU A   7       2.168   3.244  -9.049  1.00  0.92           O  
+ATOM     57  CG  GLU A   7       2.993   2.712 -13.470  1.00  0.81           C  
+ATOM     58  CD  GLU A   7       3.663   4.074 -13.390  1.00  0.76           C  
+ATOM     59  OE1 GLU A   7       3.045   5.025 -12.860  1.00  0.71           O  
+ATOM     60  OE2 GLU A   7       4.816   4.192 -13.863  1.00  0.68           O  
+ATOM     61  N   ARG A   8       2.161   1.024  -9.290  1.00  0.94           N  
+ATOM     62  CA  ARG A   8       1.415   0.847  -8.049  1.00  0.94           C  
+ATOM     63  C   ARG A   8       2.247   1.276  -6.844  1.00  0.95           C  
+ATOM     64  CB  ARG A   8       0.974  -0.609  -7.889  1.00  0.94           C  
+ATOM     65  O   ARG A   8       1.748   1.966  -5.953  1.00  0.94           O  
+ATOM     66  CG  ARG A   8       0.090  -0.856  -6.676  1.00  0.93           C  
+ATOM     67  CD  ARG A   8      -0.399  -2.296  -6.618  1.00  0.91           C  
+ATOM     68  NE  ARG A   8       0.707  -3.234  -6.450  1.00  0.90           N  
+ATOM     69  NH1 ARG A   8       0.132  -4.557  -8.255  1.00  0.83           N  
+ATOM     70  NH2 ARG A   8       1.970  -5.075  -6.987  1.00  0.82           N  
+ATOM     71  CZ  ARG A   8       0.934  -4.287  -7.231  1.00  0.88           C  
+ATOM     72  N   LEU A   9       3.502   0.910  -6.829  1.00  0.94           N  
+ATOM     73  CA  LEU A   9       4.402   1.277  -5.741  1.00  0.95           C  
+ATOM     74  C   LEU A   9       4.528   2.793  -5.628  1.00  0.94           C  
+ATOM     75  CB  LEU A   9       5.784   0.654  -5.956  1.00  0.94           C  
+ATOM     76  O   LEU A   9       4.464   3.345  -4.528  1.00  0.94           O  
+ATOM     77  CG  LEU A   9       5.867  -0.869  -5.836  1.00  0.93           C  
+ATOM     78  CD1 LEU A   9       7.227  -1.365  -6.316  1.00  0.90           C  
+ATOM     79  CD2 LEU A   9       5.608  -1.308  -4.399  1.00  0.90           C  
+ATOM     80  N   LYS A  10       4.686   3.490  -6.747  1.00  0.94           N  
+ATOM     81  CA  LYS A  10       4.773   4.947  -6.761  1.00  0.94           C  
+ATOM     82  C   LYS A  10       3.489   5.580  -6.231  1.00  0.94           C  
+ATOM     83  CB  LYS A  10       5.061   5.454  -8.175  1.00  0.93           C  
+ATOM     84  O   LYS A  10       3.534   6.594  -5.531  1.00  0.94           O  
+ATOM     85  CG  LYS A  10       6.475   5.166  -8.660  1.00  0.86           C  
+ATOM     86  CD  LYS A  10       6.688   5.659 -10.085  1.00  0.81           C  
+ATOM     87  CE  LYS A  10       8.032   5.206 -10.639  1.00  0.73           C  
+ATOM     88  NZ  LYS A  10       8.191   5.574 -12.077  1.00  0.64           N  
+ATOM     89  N   SER A  11       2.412   4.973  -6.576  1.00  0.95           N  
+ATOM     90  CA  SER A  11       1.124   5.485  -6.118  1.00  0.95           C  
+ATOM     91  C   SER A  11       0.985   5.357  -4.605  1.00  0.95           C  
+ATOM     92  CB  SER A  11      -0.022   4.745  -6.808  1.00  0.94           C  
+ATOM     93  O   SER A  11       0.476   6.266  -3.945  1.00  0.95           O  
+ATOM     94  OG  SER A  11      -0.073   5.069  -8.187  1.00  0.85           O  
+ATOM     95  N   ILE A  12       1.404   4.270  -4.118  1.00  0.95           N  
+ATOM     96  CA  ILE A  12       1.342   4.069  -2.674  1.00  0.96           C  
+ATOM     97  C   ILE A  12       2.158   5.149  -1.968  1.00  0.95           C  
+ATOM     98  CB  ILE A  12       1.851   2.666  -2.276  1.00  0.95           C  
+ATOM     99  O   ILE A  12       1.680   5.776  -1.019  1.00  0.95           O  
+ATOM    100  CG1 ILE A  12       0.873   1.587  -2.754  1.00  0.94           C  
+ATOM    101  CG2 ILE A  12       2.067   2.579  -0.762  1.00  0.94           C  
+ATOM    102  CD1 ILE A  12       1.390   0.165  -2.589  1.00  0.93           C  
+ATOM    103  N   VAL A  13       3.373   5.365  -2.383  1.00  0.95           N  
+ATOM    104  CA  VAL A  13       4.255   6.350  -1.765  1.00  0.95           C  
+ATOM    105  C   VAL A  13       3.625   7.738  -1.859  1.00  0.95           C  
+ATOM    106  CB  VAL A  13       5.653   6.353  -2.424  1.00  0.94           C  
+ATOM    107  O   VAL A  13       3.621   8.492  -0.883  1.00  0.94           O  
+ATOM    108  CG1 VAL A  13       6.485   7.530  -1.919  1.00  0.92           C  
+ATOM    109  CG2 VAL A  13       6.371   5.032  -2.155  1.00  0.92           C  
+ATOM    110  N   ARG A  14       3.008   8.094  -3.002  1.00  0.94           N  
+ATOM    111  CA  ARG A  14       2.369   9.390  -3.204  1.00  0.94           C  
+ATOM    112  C   ARG A  14       1.185   9.570  -2.261  1.00  0.94           C  
+ATOM    113  CB  ARG A  14       1.911   9.543  -4.656  1.00  0.93           C  
+ATOM    114  O   ARG A  14       1.009  10.640  -1.674  1.00  0.94           O  
+ATOM    115  CG  ARG A  14       3.035   9.869  -5.626  1.00  0.83           C  
+ATOM    116  CD  ARG A  14       2.503  10.221  -7.009  1.00  0.77           C  
+ATOM    117  NE  ARG A  14       2.183   9.025  -7.783  1.00  0.72           N  
+ATOM    118  NH1 ARG A  14       1.327  10.176  -9.594  1.00  0.52           N  
+ATOM    119  NH2 ARG A  14       1.383   7.884  -9.609  1.00  0.47           N  
+ATOM    120  CZ  ARG A  14       1.632   9.031  -8.994  1.00  0.68           C  
+ATOM    121  N   ILE A  15       0.432   8.498  -2.193  1.00  0.95           N  
+ATOM    122  CA  ILE A  15      -0.742   8.563  -1.329  1.00  0.95           C  
+ATOM    123  C   ILE A  15      -0.307   8.788   0.117  1.00  0.95           C  
+ATOM    124  CB  ILE A  15      -1.596   7.280  -1.438  1.00  0.95           C  
+ATOM    125  O   ILE A  15      -0.849   9.654   0.807  1.00  0.94           O  
+ATOM    126  CG1 ILE A  15      -2.264   7.198  -2.816  1.00  0.93           C  
+ATOM    127  CG2 ILE A  15      -2.640   7.229  -0.319  1.00  0.93           C  
+ATOM    128  CD1 ILE A  15      -2.880   5.841  -3.126  1.00  0.92           C  
+ATOM    129  N   LEU A  16       0.689   8.051   0.584  1.00  0.95           N  
+ATOM    130  CA  LEU A  16       1.143   8.158   1.966  1.00  0.95           C  
+ATOM    131  C   LEU A  16       1.813   9.505   2.215  1.00  0.95           C  
+ATOM    132  CB  LEU A  16       2.113   7.022   2.301  1.00  0.95           C  
+ATOM    133  O   LEU A  16       1.694  10.071   3.304  1.00  0.94           O  
+ATOM    134  CG  LEU A  16       1.521   5.612   2.331  1.00  0.95           C  
+ATOM    135  CD1 LEU A  16       2.609   4.587   2.634  1.00  0.94           C  
+ATOM    136  CD2 LEU A  16       0.397   5.525   3.357  1.00  0.93           C  
+ATOM    137  N   GLU A  17       2.509  10.095   1.227  1.00  0.94           N  
+ATOM    138  CA  GLU A  17       3.177  11.387   1.356  1.00  0.94           C  
+ATOM    139  C   GLU A  17       2.165  12.522   1.484  1.00  0.93           C  
+ATOM    140  CB  GLU A  17       4.100  11.637   0.161  1.00  0.92           C  
+ATOM    141  O   GLU A  17       2.413  13.505   2.184  1.00  0.92           O  
+ATOM    142  CG  GLU A  17       5.412  10.868   0.225  1.00  0.86           C  
+ATOM    143  CD  GLU A  17       6.272  11.044  -1.016  1.00  0.81           C  
+ATOM    144  OE1 GLU A  17       5.739  11.459  -2.070  1.00  0.78           O  
+ATOM    145  OE2 GLU A  17       7.489  10.764  -0.934  1.00  0.76           O  
+ATOM    146  N   ARG A  18       1.030  12.289   0.890  1.00  0.93           N  
+ATOM    147  CA  ARG A  18       0.060  13.378   0.835  1.00  0.93           C  
+ATOM    148  C   ARG A  18      -0.916  13.301   2.003  1.00  0.92           C  
+ATOM    149  CB  ARG A  18      -0.705  13.351  -0.490  1.00  0.91           C  
+ATOM    150  O   ARG A  18      -1.588  14.284   2.323  1.00  0.91           O  
+ATOM    151  CG  ARG A  18       0.139  13.731  -1.697  1.00  0.83           C  
+ATOM    152  CD  ARG A  18      -0.656  13.640  -2.992  1.00  0.79           C  
+ATOM    153  NE  ARG A  18       0.139  14.059  -4.143  1.00  0.73           N  
+ATOM    154  NH1 ARG A  18      -1.546  13.730  -5.688  1.00  0.56           N  
+ATOM    155  NH2 ARG A  18       0.507  14.490  -6.368  1.00  0.51           N  
+ATOM    156  CZ  ARG A  18      -0.302  14.092  -5.397  1.00  0.71           C  
+ATOM    157  N   SER A  19      -0.853  12.101   2.564  1.00  0.92           N  
+ATOM    158  CA  SER A  19      -1.857  11.925   3.607  1.00  0.92           C  
+ATOM    159  C   SER A  19      -1.372  12.481   4.942  1.00  0.91           C  
+ATOM    160  CB  SER A  19      -2.215  10.446   3.762  1.00  0.91           C  
+ATOM    161  O   SER A  19      -0.232  12.239   5.344  1.00  0.90           O  
+ATOM    162  OG  SER A  19      -3.219  10.272   4.748  1.00  0.85           O  
+ATOM    163  N   LYS A  20      -2.211  13.277   5.567  1.00  0.89           N  
+ATOM    164  CA  LYS A  20      -1.915  13.785   6.903  1.00  0.89           C  
+ATOM    165  C   LYS A  20      -2.272  12.759   7.974  1.00  0.89           C  
+ATOM    166  CB  LYS A  20      -2.667  15.092   7.158  1.00  0.87           C  
+ATOM    167  O   LYS A  20      -1.768  12.824   9.097  1.00  0.87           O  
+ATOM    168  CG  LYS A  20      -2.206  16.252   6.287  1.00  0.79           C  
+ATOM    169  CD  LYS A  20      -2.955  17.535   6.623  1.00  0.76           C  
+ATOM    170  CE  LYS A  20      -2.516  18.689   5.732  1.00  0.68           C  
+ATOM    171  NZ  LYS A  20      -3.289  19.935   6.018  1.00  0.61           N  
+ATOM    172  N   GLU A  21      -3.098  11.887   7.683  1.00  0.92           N  
+ATOM    173  CA  GLU A  21      -3.549  10.825   8.577  1.00  0.92           C  
+ATOM    174  C   GLU A  21      -3.148   9.451   8.047  1.00  0.92           C  
+ATOM    175  CB  GLU A  21      -5.066  10.892   8.770  1.00  0.90           C  
+ATOM    176  O   GLU A  21      -2.881   9.294   6.854  1.00  0.92           O  
+ATOM    177  CG  GLU A  21      -5.548  12.188   9.406  1.00  0.84           C  
+ATOM    178  CD  GLU A  21      -7.025  12.165   9.769  1.00  0.80           C  
+ATOM    179  OE1 GLU A  21      -7.741  11.233   9.339  1.00  0.78           O  
+ATOM    180  OE2 GLU A  21      -7.468  13.088  10.489  1.00  0.74           O  
+ATOM    181  N   PRO A  22      -3.069   8.392   8.977  1.00  0.94           N  
+ATOM    182  CA  PRO A  22      -2.802   7.027   8.516  1.00  0.95           C  
+ATOM    183  C   PRO A  22      -3.802   6.555   7.462  1.00  0.95           C  
+ATOM    184  CB  PRO A  22      -2.921   6.193   9.794  1.00  0.94           C  
+ATOM    185  O   PRO A  22      -4.990   6.878   7.544  1.00  0.94           O  
+ATOM    186  CG  PRO A  22      -2.711   7.168  10.907  1.00  0.93           C  
+ATOM    187  CD  PRO A  22      -3.262   8.500  10.485  1.00  0.91           C  
+ATOM    188  N   VAL A  23      -3.318   5.865   6.476  1.00  0.95           N  
+ATOM    189  CA  VAL A  23      -4.137   5.301   5.408  1.00  0.95           C  
+ATOM    190  C   VAL A  23      -4.274   3.793   5.603  1.00  0.95           C  
+ATOM    191  CB  VAL A  23      -3.543   5.607   4.015  1.00  0.95           C  
+ATOM    192  O   VAL A  23      -3.273   3.074   5.652  1.00  0.95           O  
+ATOM    193  CG1 VAL A  23      -4.477   5.115   2.910  1.00  0.93           C  
+ATOM    194  CG2 VAL A  23      -3.276   7.103   3.866  1.00  0.93           C  
+ATOM    195  N   SER A  24      -5.474   3.345   5.655  1.00  0.96           N  
+ATOM    196  CA  SER A  24      -5.672   1.925   5.923  1.00  0.96           C  
+ATOM    197  C   SER A  24      -5.316   1.077   4.706  1.00  0.96           C  
+ATOM    198  CB  SER A  24      -7.119   1.653   6.338  1.00  0.95           C  
+ATOM    199  O   SER A  24      -5.386   1.552   3.571  1.00  0.95           O  
+ATOM    200  OG  SER A  24      -7.981   1.698   5.214  1.00  0.91           O  
+ATOM    201  N   GLY A  25      -4.862  -0.173   5.000  1.00  0.95           N  
+ATOM    202  CA  GLY A  25      -4.628  -1.110   3.914  1.00  0.95           C  
+ATOM    203  C   GLY A  25      -5.840  -1.304   3.021  1.00  0.95           C  
+ATOM    204  O   GLY A  25      -5.706  -1.426   1.802  1.00  0.95           O  
+ATOM    205  N   ALA A  26      -7.029  -1.338   3.592  1.00  0.95           N  
+ATOM    206  CA  ALA A  26      -8.269  -1.507   2.839  1.00  0.95           C  
+ATOM    207  C   ALA A  26      -8.487  -0.346   1.873  1.00  0.95           C  
+ATOM    208  CB  ALA A  26      -9.456  -1.633   3.790  1.00  0.94           C  
+ATOM    209  O   ALA A  26      -8.901  -0.551   0.729  1.00  0.95           O  
+ATOM    210  N   GLN A  27      -8.247   0.886   2.350  1.00  0.95           N  
+ATOM    211  CA  GLN A  27      -8.373   2.066   1.501  1.00  0.95           C  
+ATOM    212  C   GLN A  27      -7.389   2.013   0.335  1.00  0.95           C  
+ATOM    213  CB  GLN A  27      -8.151   3.341   2.317  1.00  0.94           C  
+ATOM    214  O   GLN A  27      -7.757   2.295  -0.807  1.00  0.95           O  
+ATOM    215  CG  GLN A  27      -8.308   4.623   1.511  1.00  0.86           C  
+ATOM    216  CD  GLN A  27      -8.034   5.870   2.330  1.00  0.81           C  
+ATOM    217  NE2 GLN A  27      -7.923   7.010   1.657  1.00  0.73           N  
+ATOM    218  OE1 GLN A  27      -7.923   5.809   3.559  1.00  0.78           O  
+ATOM    219  N   LEU A  28      -6.152   1.593   0.619  1.00  0.96           N  
+ATOM    220  CA  LEU A  28      -5.160   1.467  -0.444  1.00  0.96           C  
+ATOM    221  C   LEU A  28      -5.574   0.395  -1.448  1.00  0.96           C  
+ATOM    222  CB  LEU A  28      -3.787   1.131   0.142  1.00  0.95           C  
+ATOM    223  O   LEU A  28      -5.474   0.602  -2.659  1.00  0.95           O  
+ATOM    224  CG  LEU A  28      -3.083   2.253   0.906  1.00  0.95           C  
+ATOM    225  CD1 LEU A  28      -1.839   1.718   1.608  1.00  0.93           C  
+ATOM    226  CD2 LEU A  28      -2.721   3.397  -0.035  1.00  0.92           C  
+ATOM    227  N   ALA A  29      -6.050  -0.750  -0.922  1.00  0.96           N  
+ATOM    228  CA  ALA A  29      -6.463  -1.859  -1.778  1.00  0.96           C  
+ATOM    229  C   ALA A  29      -7.599  -1.441  -2.708  1.00  0.96           C  
+ATOM    230  CB  ALA A  29      -6.886  -3.056  -0.931  1.00  0.95           C  
+ATOM    231  O   ALA A  29      -7.574  -1.741  -3.903  1.00  0.95           O  
+ATOM    232  N   GLU A  30      -8.555  -0.688  -2.175  1.00  0.96           N  
+ATOM    233  CA  GLU A  30      -9.687  -0.188  -2.949  1.00  0.95           C  
+ATOM    234  C   GLU A  30      -9.235   0.828  -3.995  1.00  0.95           C  
+ATOM    235  CB  GLU A  30     -10.735   0.440  -2.026  1.00  0.94           C  
+ATOM    236  O   GLU A  30      -9.618   0.736  -5.163  1.00  0.94           O  
+ATOM    237  CG  GLU A  30     -12.004   0.876  -2.744  1.00  0.85           C  
+ATOM    238  CD  GLU A  30     -13.067   1.424  -1.804  1.00  0.79           C  
+ATOM    239  OE1 GLU A  30     -12.798   1.542  -0.587  1.00  0.77           O  
+ATOM    240  OE2 GLU A  30     -14.177   1.736  -2.288  1.00  0.74           O  
+ATOM    241  N   GLU A  31      -8.407   1.749  -3.572  1.00  0.95           N  
+ATOM    242  CA  GLU A  31      -7.963   2.822  -4.456  1.00  0.94           C  
+ATOM    243  C   GLU A  31      -7.129   2.277  -5.612  1.00  0.94           C  
+ATOM    244  CB  GLU A  31      -7.159   3.864  -3.674  1.00  0.92           C  
+ATOM    245  O   GLU A  31      -7.197   2.791  -6.730  1.00  0.93           O  
+ATOM    246  CG  GLU A  31      -6.777   5.089  -4.493  1.00  0.78           C  
+ATOM    247  CD  GLU A  31      -6.323   6.263  -3.642  1.00  0.72           C  
+ATOM    248  OE1 GLU A  31      -6.684   6.319  -2.444  1.00  0.68           O  
+ATOM    249  OE2 GLU A  31      -5.603   7.135  -4.176  1.00  0.66           O  
+ATOM    250  N   LEU A  32      -6.421   1.249  -5.368  1.00  0.94           N  
+ATOM    251  CA  LEU A  32      -5.480   0.768  -6.373  1.00  0.94           C  
+ATOM    252  C   LEU A  32      -6.005  -0.494  -7.050  1.00  0.94           C  
+ATOM    253  CB  LEU A  32      -4.114   0.489  -5.738  1.00  0.94           C  
+ATOM    254  O   LEU A  32      -5.323  -1.081  -7.893  1.00  0.92           O  
+ATOM    255  CG  LEU A  32      -3.386   1.694  -5.140  1.00  0.92           C  
+ATOM    256  CD1 LEU A  32      -2.156   1.238  -4.363  1.00  0.89           C  
+ATOM    257  CD2 LEU A  32      -2.998   2.682  -6.234  1.00  0.89           C  
+ATOM    258  N   SER A  33      -7.187  -0.923  -6.692  1.00  0.95           N  
+ATOM    259  CA  SER A  33      -7.876  -2.070  -7.275  1.00  0.94           C  
+ATOM    260  C   SER A  33      -7.028  -3.333  -7.174  1.00  0.94           C  
+ATOM    261  CB  SER A  33      -8.228  -1.798  -8.738  1.00  0.93           C  
+ATOM    262  O   SER A  33      -6.863  -4.057  -8.159  1.00  0.93           O  
+ATOM    263  OG  SER A  33      -9.149  -0.726  -8.842  1.00  0.85           O  
+ATOM    264  N   VAL A  34      -6.509  -3.606  -5.951  1.00  0.94           N  
+ATOM    265  CA  VAL A  34      -5.794  -4.840  -5.646  1.00  0.94           C  
+ATOM    266  C   VAL A  34      -6.250  -5.380  -4.292  1.00  0.94           C  
+ATOM    267  CB  VAL A  34      -4.264  -4.625  -5.645  1.00  0.94           C  
+ATOM    268  O   VAL A  34      -7.013  -4.724  -3.580  1.00  0.94           O  
+ATOM    269  CG1 VAL A  34      -3.771  -4.252  -7.042  1.00  0.91           C  
+ATOM    270  CG2 VAL A  34      -3.878  -3.547  -4.634  1.00  0.91           C  
+ATOM    271  N   SER A  35      -5.841  -6.605  -4.013  1.00  0.95           N  
+ATOM    272  CA  SER A  35      -6.185  -7.200  -2.725  1.00  0.95           C  
+ATOM    273  C   SER A  35      -5.372  -6.579  -1.594  1.00  0.95           C  
+ATOM    274  CB  SER A  35      -5.960  -8.712  -2.756  1.00  0.94           C  
+ATOM    275  O   SER A  35      -4.314  -5.993  -1.833  1.00  0.95           O  
+ATOM    276  OG  SER A  35      -4.575  -9.014  -2.769  1.00  0.87           O  
+ATOM    277  N   ARG A  36      -5.839  -6.717  -0.444  1.00  0.96           N  
+ATOM    278  CA  ARG A  36      -5.113  -6.257   0.736  1.00  0.96           C  
+ATOM    279  C   ARG A  36      -3.765  -6.959   0.857  1.00  0.96           C  
+ATOM    280  CB  ARG A  36      -5.941  -6.489   2.001  1.00  0.93           C  
+ATOM    281  O   ARG A  36      -2.785  -6.360   1.306  1.00  0.95           O  
+ATOM    282  CG  ARG A  36      -7.147  -5.571   2.126  1.00  0.76           C  
+ATOM    283  CD  ARG A  36      -7.872  -5.769   3.449  1.00  0.72           C  
+ATOM    284  NE  ARG A  36      -7.044  -5.369   4.583  1.00  0.69           N  
+ATOM    285  NH1 ARG A  36      -8.302  -6.492   6.163  1.00  0.58           N  
+ATOM    286  NH2 ARG A  36      -6.450  -5.302   6.800  1.00  0.52           N  
+ATOM    287  CZ  ARG A  36      -7.267  -5.722   5.846  1.00  0.64           C  
+ATOM    288  N   GLN A  37      -3.757  -8.247   0.479  1.00  0.96           N  
+ATOM    289  CA  GLN A  37      -2.512  -9.005   0.531  1.00  0.96           C  
+ATOM    290  C   GLN A  37      -1.446  -8.374  -0.361  1.00  0.96           C  
+ATOM    291  CB  GLN A  37      -2.749 -10.458   0.117  1.00  0.94           C  
+ATOM    292  O   GLN A  37      -0.273  -8.312   0.013  1.00  0.96           O  
+ATOM    293  CG  GLN A  37      -1.503 -11.329   0.189  1.00  0.79           C  
+ATOM    294  CD  GLN A  37      -1.003 -11.522   1.609  1.00  0.72           C  
+ATOM    295  NE2 GLN A  37       0.254 -11.931   1.744  1.00  0.59           N  
+ATOM    296  OE1 GLN A  37      -1.740 -11.305   2.576  1.00  0.68           O  
+ATOM    297  N   VAL A  38      -1.860  -7.916  -1.492  1.00  0.96           N  
+ATOM    298  CA  VAL A  38      -0.934  -7.261  -2.410  1.00  0.96           C  
+ATOM    299  C   VAL A  38      -0.384  -5.989  -1.769  1.00  0.96           C  
+ATOM    300  CB  VAL A  38      -1.612  -6.928  -3.758  1.00  0.95           C  
+ATOM    301  O   VAL A  38       0.810  -5.699  -1.875  1.00  0.96           O  
+ATOM    302  CG1 VAL A  38      -0.732  -5.996  -4.589  1.00  0.93           C  
+ATOM    303  CG2 VAL A  38      -1.919  -8.209  -4.531  1.00  0.92           C  
+ATOM    304  N   ILE A  39      -1.219  -5.306  -1.044  1.00  0.96           N  
+ATOM    305  CA  ILE A  39      -0.801  -4.074  -0.385  1.00  0.96           C  
+ATOM    306  C   ILE A  39       0.250  -4.387   0.677  1.00  0.96           C  
+ATOM    307  CB  ILE A  39      -2.002  -3.337   0.250  1.00  0.96           C  
+ATOM    308  O   ILE A  39       1.263  -3.692   0.781  1.00  0.96           O  
+ATOM    309  CG1 ILE A  39      -2.970  -2.857  -0.838  1.00  0.94           C  
+ATOM    310  CG2 ILE A  39      -1.522  -2.166   1.112  1.00  0.94           C  
+ATOM    311  CD1 ILE A  39      -2.336  -1.930  -1.865  1.00  0.91           C  
+ATOM    312  N   VAL A  40       0.033  -5.399   1.497  1.00  0.96           N  
+ATOM    313  CA  VAL A  40       0.963  -5.793   2.550  1.00  0.96           C  
+ATOM    314  C   VAL A  40       2.317  -6.146   1.940  1.00  0.97           C  
+ATOM    315  CB  VAL A  40       0.420  -6.987   3.368  1.00  0.96           C  
+ATOM    316  O   VAL A  40       3.362  -5.740   2.454  1.00  0.96           O  
+ATOM    317  CG1 VAL A  40       1.503  -7.549   4.287  1.00  0.89           C  
+ATOM    318  CG2 VAL A  40      -0.806  -6.565   4.177  1.00  0.89           C  
+ATOM    319  N   GLN A  41       2.280  -6.826   0.823  1.00  0.96           N  
+ATOM    320  CA  GLN A  41       3.510  -7.216   0.142  1.00  0.96           C  
+ATOM    321  C   GLN A  41       4.236  -5.998  -0.424  1.00  0.96           C  
+ATOM    322  CB  GLN A  41       3.211  -8.215  -0.977  1.00  0.95           C  
+ATOM    323  O   GLN A  41       5.463  -5.909  -0.348  1.00  0.95           O  
+ATOM    324  CG  GLN A  41       2.779  -9.586  -0.477  1.00  0.86           C  
+ATOM    325  CD  GLN A  41       2.333 -10.506  -1.598  1.00  0.79           C  
+ATOM    326  NE2 GLN A  41       2.103 -11.773  -1.270  1.00  0.69           N  
+ATOM    327  OE1 GLN A  41       2.196 -10.082  -2.750  1.00  0.77           O  
+ATOM    328  N   ASP A  42       3.536  -5.124  -0.987  1.00  0.96           N  
+ATOM    329  CA  ASP A  42       4.119  -3.918  -1.568  1.00  0.96           C  
+ATOM    330  C   ASP A  42       4.728  -3.027  -0.487  1.00  0.96           C  
+ATOM    331  CB  ASP A  42       3.066  -3.140  -2.359  1.00  0.95           C  
+ATOM    332  O   ASP A  42       5.806  -2.461  -0.677  1.00  0.95           O  
+ATOM    333  CG  ASP A  42       2.757  -3.764  -3.709  1.00  0.94           C  
+ATOM    334  OD1 ASP A  42       3.553  -4.599  -4.191  1.00  0.91           O  
+ATOM    335  OD2 ASP A  42       1.711  -3.416  -4.297  1.00  0.92           O  
+ATOM    336  N   ILE A  43       4.006  -2.921   0.639  1.00  0.96           N  
+ATOM    337  CA  ILE A  43       4.512  -2.107   1.738  1.00  0.96           C  
+ATOM    338  C   ILE A  43       5.806  -2.713   2.276  1.00  0.96           C  
+ATOM    339  CB  ILE A  43       3.470  -1.976   2.871  1.00  0.96           C  
+ATOM    340  O   ILE A  43       6.770  -1.994   2.548  1.00  0.96           O  
+ATOM    341  CG1 ILE A  43       2.293  -1.106   2.416  1.00  0.94           C  
+ATOM    342  CG2 ILE A  43       4.117  -1.404   4.136  1.00  0.93           C  
+ATOM    343  CD1 ILE A  43       2.657   0.352   2.174  1.00  0.91           C  
+ATOM    344  N   ALA A  44       5.851  -4.052   2.411  1.00  0.96           N  
+ATOM    345  CA  ALA A  44       7.075  -4.731   2.830  1.00  0.96           C  
+ATOM    346  C   ALA A  44       8.218  -4.451   1.859  1.00  0.96           C  
+ATOM    347  CB  ALA A  44       6.837  -6.234   2.947  1.00  0.96           C  
+ATOM    348  O   ALA A  44       9.353  -4.213   2.278  1.00  0.96           O  
+ATOM    349  N   TYR A  45       7.872  -4.425   0.639  1.00  0.95           N  
+ATOM    350  CA  TYR A  45       8.877  -4.165  -0.387  1.00  0.95           C  
+ATOM    351  C   TYR A  45       9.363  -2.723  -0.320  1.00  0.95           C  
+ATOM    352  CB  TYR A  45       8.312  -4.461  -1.780  1.00  0.95           C  
+ATOM    353  O   TYR A  45      10.568  -2.464  -0.380  1.00  0.95           O  
+ATOM    354  CG  TYR A  45       9.316  -4.279  -2.893  1.00  0.93           C  
+ATOM    355  CD1 TYR A  45      10.439  -5.098  -2.985  1.00  0.89           C  
+ATOM    356  CD2 TYR A  45       9.143  -3.288  -3.853  1.00  0.88           C  
+ATOM    357  CE1 TYR A  45      11.367  -4.933  -4.008  1.00  0.88           C  
+ATOM    358  CE2 TYR A  45      10.065  -3.113  -4.880  1.00  0.88           C  
+ATOM    359  OH  TYR A  45      12.088  -3.771  -5.963  1.00  0.79           O  
+ATOM    360  CZ  TYR A  45      11.172  -3.940  -4.949  1.00  0.86           C  
+ATOM    361  N   LEU A  46       8.521  -1.815  -0.237  1.00  0.95           N  
+ATOM    362  CA  LEU A  46       8.894  -0.409  -0.123  1.00  0.95           C  
+ATOM    363  C   LEU A  46       9.792  -0.181   1.088  1.00  0.95           C  
+ATOM    364  CB  LEU A  46       7.645   0.470  -0.019  1.00  0.95           C  
+ATOM    365  O   LEU A  46      10.747   0.595   1.021  1.00  0.95           O  
+ATOM    366  CG  LEU A  46       6.847   0.669  -1.309  1.00  0.95           C  
+ATOM    367  CD1 LEU A  46       5.518   1.356  -1.010  1.00  0.93           C  
+ATOM    368  CD2 LEU A  46       7.656   1.476  -2.319  1.00  0.93           C  
+ATOM    369  N   ARG A  47       9.531  -0.912   2.192  1.00  0.96           N  
+ATOM    370  CA  ARG A  47      10.387  -0.815   3.370  1.00  0.96           C  
+ATOM    371  C   ARG A  47      11.791  -1.332   3.072  1.00  0.96           C  
+ATOM    372  CB  ARG A  47       9.781  -1.591   4.541  1.00  0.95           C  
+ATOM    373  O   ARG A  47      12.781  -0.740   3.506  1.00  0.95           O  
+ATOM    374  CG  ARG A  47       8.562  -0.923   5.158  1.00  0.93           C  
+ATOM    375  CD  ARG A  47       8.006  -1.730   6.323  1.00  0.91           C  
+ATOM    376  NE  ARG A  47       6.966  -0.996   7.039  1.00  0.90           N  
+ATOM    377  NH1 ARG A  47       5.819  -2.880   7.724  1.00  0.83           N  
+ATOM    378  NH2 ARG A  47       5.061  -0.797   8.306  1.00  0.81           N  
+ATOM    379  CZ  ARG A  47       5.951  -1.559   7.688  1.00  0.87           C  
+ATOM    380  N   SER A  48      11.822  -2.364   2.346  1.00  0.96           N  
+ATOM    381  CA  SER A  48      13.124  -2.927   2.003  1.00  0.96           C  
+ATOM    382  C   SER A  48      13.929  -1.969   1.131  1.00  0.95           C  
+ATOM    383  CB  SER A  48      12.957  -4.266   1.283  1.00  0.95           C  
+ATOM    384  O   SER A  48      15.159  -2.041   1.091  1.00  0.94           O  
+ATOM    385  OG  SER A  48      12.598  -4.066  -0.073  1.00  0.90           O  
+ATOM    386  N   LEU A  49      13.228  -1.054   0.452  1.00  0.95           N  
+ATOM    387  CA  LEU A  49      13.904  -0.075  -0.393  1.00  0.94           C  
+ATOM    388  C   LEU A  49      14.342   1.137   0.423  1.00  0.94           C  
+ATOM    389  CB  LEU A  49      12.989   0.368  -1.537  1.00  0.94           C  
+ATOM    390  O   LEU A  49      14.979   2.049  -0.107  1.00  0.92           O  
+ATOM    391  CG  LEU A  49      12.691  -0.678  -2.612  1.00  0.89           C  
+ATOM    392  CD1 LEU A  49      11.737  -0.109  -3.656  1.00  0.83           C  
+ATOM    393  CD2 LEU A  49      13.984  -1.156  -3.265  1.00  0.83           C  
+ATOM    394  N   GLY A  50      13.868   1.196   1.683  1.00  0.94           N  
+ATOM    395  CA  GLY A  50      14.344   2.263   2.549  1.00  0.94           C  
+ATOM    396  C   GLY A  50      13.258   3.253   2.927  1.00  0.94           C  
+ATOM    397  O   GLY A  50      13.514   4.216   3.652  1.00  0.93           O  
+ATOM    398  N   TYR A  51      12.070   3.128   2.433  1.00  0.94           N  
+ATOM    399  CA  TYR A  51      10.982   4.007   2.846  1.00  0.94           C  
+ATOM    400  C   TYR A  51      10.598   3.753   4.299  1.00  0.94           C  
+ATOM    401  CB  TYR A  51       9.760   3.812   1.942  1.00  0.94           C  
+ATOM    402  O   TYR A  51      10.466   2.601   4.721  1.00  0.93           O  
+ATOM    403  CG  TYR A  51       9.971   4.290   0.526  1.00  0.93           C  
+ATOM    404  CD1 TYR A  51       9.767   5.625   0.183  1.00  0.92           C  
+ATOM    405  CD2 TYR A  51      10.373   3.408  -0.472  1.00  0.91           C  
+ATOM    406  CE1 TYR A  51       9.958   6.069  -1.121  1.00  0.91           C  
+ATOM    407  CE2 TYR A  51      10.567   3.842  -1.779  1.00  0.91           C  
+ATOM    408  OH  TYR A  51      10.548   5.607  -3.386  1.00  0.87           O  
+ATOM    409  CZ  TYR A  51      10.358   5.172  -2.093  1.00  0.91           C  
+ATOM    410  N   ASN A  52      10.429   4.810   5.078  1.00  0.95           N  
+ATOM    411  CA  ASN A  52      10.061   4.716   6.487  1.00  0.95           C  
+ATOM    412  C   ASN A  52       8.547   4.669   6.670  1.00  0.95           C  
+ATOM    413  CB  ASN A  52      10.656   5.885   7.276  1.00  0.94           C  
+ATOM    414  O   ASN A  52       7.960   5.581   7.255  1.00  0.94           O  
+ATOM    415  CG  ASN A  52      10.661   5.638   8.772  1.00  0.86           C  
+ATOM    416  ND2 ASN A  52      10.958   6.677   9.543  1.00  0.77           N  
+ATOM    417  OD1 ASN A  52      10.401   4.522   9.230  1.00  0.75           O  
+ATOM    418  N   ILE A  53       7.934   3.537   6.187  1.00  0.95           N  
+ATOM    419  CA  ILE A  53       6.498   3.323   6.325  1.00  0.96           C  
+ATOM    420  C   ILE A  53       6.209   2.590   7.633  1.00  0.95           C  
+ATOM    421  CB  ILE A  53       5.927   2.530   5.128  1.00  0.95           C  
+ATOM    422  O   ILE A  53       6.695   1.477   7.849  1.00  0.94           O  
+ATOM    423  CG1 ILE A  53       6.259   3.239   3.810  1.00  0.94           C  
+ATOM    424  CG2 ILE A  53       4.416   2.335   5.281  1.00  0.94           C  
+ATOM    425  CD1 ILE A  53       5.911   2.431   2.568  1.00  0.93           C  
+ATOM    426  N   VAL A  54       5.441   3.118   8.475  1.00  0.95           N  
+ATOM    427  CA  VAL A  54       5.115   2.530   9.770  1.00  0.95           C  
+ATOM    428  C   VAL A  54       3.650   2.103   9.791  1.00  0.95           C  
+ATOM    429  CB  VAL A  54       5.400   3.514  10.927  1.00  0.94           C  
+ATOM    430  O   VAL A  54       2.779   2.826   9.299  1.00  0.95           O  
+ATOM    431  CG1 VAL A  54       4.993   2.903  12.267  1.00  0.78           C  
+ATOM    432  CG2 VAL A  54       6.876   3.906  10.942  1.00  0.78           C  
+ATOM    433  N   ALA A  55       3.436   0.986  10.239  1.00  0.95           N  
+ATOM    434  CA  ALA A  55       2.086   0.462  10.428  1.00  0.95           C  
+ATOM    435  C   ALA A  55       1.527   0.864  11.790  1.00  0.94           C  
+ATOM    436  CB  ALA A  55       2.080  -1.057  10.281  1.00  0.94           C  
+ATOM    437  O   ALA A  55       2.203   0.723  12.811  1.00  0.93           O  
+ATOM    438  N   THR A  56       0.347   1.432  11.827  1.00  0.93           N  
+ATOM    439  CA  THR A  56      -0.403   1.797  13.024  1.00  0.93           C  
+ATOM    440  C   THR A  56      -1.730   1.046  13.080  1.00  0.93           C  
+ATOM    441  CB  THR A  56      -0.665   3.314  13.077  1.00  0.92           C  
+ATOM    442  O   THR A  56      -2.129   0.405  12.106  1.00  0.92           O  
+ATOM    443  CG2 THR A  56       0.585   4.102  12.702  1.00  0.87           C  
+ATOM    444  OG1 THR A  56      -1.717   3.643  12.161  1.00  0.88           O  
+ATOM    445  N   PRO A  57      -2.436   0.978  14.221  1.00  0.92           N  
+ATOM    446  CA  PRO A  57      -3.759   0.351  14.257  1.00  0.92           C  
+ATOM    447  C   PRO A  57      -4.713   0.927  13.214  1.00  0.91           C  
+ATOM    448  CB  PRO A  57      -4.253   0.651  15.675  1.00  0.91           C  
+ATOM    449  O   PRO A  57      -5.663   0.255  12.803  1.00  0.89           O  
+ATOM    450  CG  PRO A  57      -3.010   0.760  16.497  1.00  0.89           C  
+ATOM    451  CD  PRO A  57      -1.939   1.398  15.660  1.00  0.88           C  
+ATOM    452  N   ARG A  58      -4.402   2.198  12.767  1.00  0.94           N  
+ATOM    453  CA  ARG A  58      -5.325   2.830  11.831  1.00  0.94           C  
+ATOM    454  C   ARG A  58      -4.823   2.702  10.396  1.00  0.93           C  
+ATOM    455  CB  ARG A  58      -5.525   4.305  12.186  1.00  0.92           C  
+ATOM    456  O   ARG A  58      -5.517   3.089   9.454  1.00  0.92           O  
+ATOM    457  CG  ARG A  58      -6.256   4.529  13.500  1.00  0.86           C  
+ATOM    458  CD  ARG A  58      -6.551   6.003  13.738  1.00  0.82           C  
+ATOM    459  NE  ARG A  58      -5.333   6.759  14.014  1.00  0.77           N  
+ATOM    460  NH1 ARG A  58      -6.405   8.786  14.289  1.00  0.64           N  
+ATOM    461  NH2 ARG A  58      -4.128   8.650  14.509  1.00  0.59           N  
+ATOM    462  CZ  ARG A  58      -5.291   8.063  14.270  1.00  0.75           C  
+ATOM    463  N   GLY A  59      -3.702   2.248  10.206  1.00  0.95           N  
+ATOM    464  CA  GLY A  59      -3.182   2.097   8.857  1.00  0.95           C  
+ATOM    465  C   GLY A  59      -1.700   2.406   8.749  1.00  0.95           C  
+ATOM    466  O   GLY A  59      -0.945   2.190   9.699  1.00  0.95           O  
+ATOM    467  N   TYR A  60      -1.248   2.930   7.637  1.00  0.96           N  
+ATOM    468  CA  TYR A  60       0.157   3.167   7.325  1.00  0.96           C  
+ATOM    469  C   TYR A  60       0.457   4.660   7.265  1.00  0.96           C  
+ATOM    470  CB  TYR A  60       0.531   2.505   5.996  1.00  0.96           C  
+ATOM    471  O   TYR A  60      -0.380   5.451   6.824  1.00  0.95           O  
+ATOM    472  CG  TYR A  60       0.342   1.008   5.988  1.00  0.95           C  
+ATOM    473  CD1 TYR A  60       1.278   0.165   6.583  1.00  0.94           C  
+ATOM    474  CD2 TYR A  60      -0.772   0.433   5.386  1.00  0.94           C  
+ATOM    475  CE1 TYR A  60       1.109  -1.216   6.577  1.00  0.94           C  
+ATOM    476  CE2 TYR A  60      -0.952  -0.946   5.374  1.00  0.94           C  
+ATOM    477  OH  TYR A  60      -0.180  -3.127   5.962  1.00  0.90           O  
+ATOM    478  CZ  TYR A  60      -0.008  -1.761   5.971  1.00  0.93           C  
+ATOM    479  N   VAL A  61       1.559   5.020   7.726  1.00  0.95           N  
+ATOM    480  CA  VAL A  61       2.043   6.393   7.628  1.00  0.95           C  
+ATOM    481  C   VAL A  61       3.484   6.399   7.124  1.00  0.95           C  
+ATOM    482  CB  VAL A  61       1.950   7.126   8.985  1.00  0.93           C  
+ATOM    483  O   VAL A  61       4.263   5.496   7.439  1.00  0.94           O  
+ATOM    484  CG1 VAL A  61       0.504   7.172   9.476  1.00  0.75           C  
+ATOM    485  CG2 VAL A  61       2.847   6.448  10.019  1.00  0.76           C  
+ATOM    486  N   LEU A  62       3.751   7.326   6.213  1.00  0.95           N  
+ATOM    487  CA  LEU A  62       5.136   7.577   5.829  1.00  0.95           C  
+ATOM    488  C   LEU A  62       5.806   8.537   6.807  1.00  0.94           C  
+ATOM    489  CB  LEU A  62       5.205   8.146   4.410  1.00  0.94           C  
+ATOM    490  O   LEU A  62       5.536   9.740   6.785  1.00  0.92           O  
+ATOM    491  CG  LEU A  62       6.604   8.327   3.819  1.00  0.89           C  
+ATOM    492  CD1 LEU A  62       7.315   6.982   3.717  1.00  0.83           C  
+ATOM    493  CD2 LEU A  62       6.525   9.001   2.453  1.00  0.82           C  
+ATOM    494  N   ALA A  63       6.595   7.972   7.715  1.00  0.91           N  
+ATOM    495  CA  ALA A  63       7.245   8.760   8.759  1.00  0.91           C  
+ATOM    496  C   ALA A  63       8.462   9.500   8.211  1.00  0.88           C  
+ATOM    497  CB  ALA A  63       7.653   7.864   9.926  1.00  0.87           C  
+ATOM    498  O   ALA A  63       9.150   9.000   7.317  1.00  0.84           O  
+ATOM    499  N   GLY A  64       8.844  10.801   8.768  1.00  0.79           N  
+ATOM    500  CA  GLY A  64      10.029  11.551   8.385  1.00  0.77           C  
+ATOM    501  C   GLY A  64       9.839  12.359   7.115  1.00  0.76           C  
+ATOM    502  O   GLY A  64      10.811  12.831   6.522  1.00  0.70           O  
+ATOM    503  N   GLY A  65       8.383  12.399   6.587  1.00  0.54           N  
+ATOM    504  CA  GLY A  65       8.133  13.333   5.501  1.00  0.54           C  
+ATOM    505  C   GLY A  65       7.712  14.709   5.983  1.00  0.53           C  
+ATOM    506  O   GLY A  65       7.142  14.845   7.067  1.00  0.51           O  

esm/mcp_output/predictions/prediction_20250830_220641.pdb

@@ -0,0 +1,489 @@
+HEADER                                            18-OCT-22                     
+TITLE     ESMFOLD V1 PREDICTION FOR INPUT
+REMARK   1                                                                      
+REMARK   1 REFERENCE 1                                                          
+REMARK   1  AUTH   ZEMING LIN, HALIL AKIN, ROSHAN RAO, BRIAN HIE, ZHONGKAI ZHU,
+REMARK   1  AUTH 2 WENTING LU, NIKITA SMETANIN, ROBERT VERKUIL, ORI KABELI,
+REMARK   1  AUTH 3 YANIV SHMUELI, ALLAN DOS SANTOS COSTA,
+REMARK   1  AUTH 4 MARYAM FAZEL-ZARANDI, TOM SERCU, SALVATORE CANDIDO,
+REMARK   1  AUTH 5 ALEXANDER RIVES
+REMARK   1  TITL   EVOLUTIONARY-SCALE PREDICTION OF ATOMIC LEVEL PROTEIN
+REMARK   1  TITL 2 STRUCTURE WITH A LANGUAGE MODEL
+REMARK   1  REF                                                                 
+REMARK   1  REFN                                                                
+REMARK   1  PMID                                                                
+REMARK   1  DOI    10.1101/2022.07.20.500902                                    
+REMARK   1                                                                      
+REMARK   1 LICENSE AND DISCLAIMERS                
+REMARK   1 ESM METAGENOMIC ATLAS DATA IS AVAILABLE UNDER
+REMARK   1 A CC-BY-4.0 LICENSE FOR ACADEMIC AND COMMERCIAL USE.
+REMARK   1 COPYRIGHT (C) META PLATFORMS, INC. ALL RIGHTS RESERVED.
+REMARK   1 USE OF THE ESM METAGENOMIC ATLAS DATA IS SUBJECT
+REMARK   1 TO THE META OPEN SOURCE TERMS OF USE AND PRIVACY POLICY.
+ATOM      1  N   MET A   1      12.955  22.762   2.808  1.00  0.40           N  
+ATOM      2  CA  MET A   1      13.442  21.402   3.023  1.00  0.43           C  
+ATOM      3  C   MET A   1      12.281  20.416   3.108  1.00  0.41           C  
+ATOM      4  CB  MET A   1      14.285  21.328   4.297  1.00  0.37           C  
+ATOM      5  O   MET A   1      11.322  20.643   3.847  1.00  0.40           O  
+ATOM      6  CG  MET A   1      15.524  20.457   4.162  1.00  0.36           C  
+ATOM      7  SD  MET A   1      16.674  20.646   5.579  1.00  0.46           S  
+ATOM      8  CE  MET A   1      16.455  19.037   6.387  1.00  0.35           C  
+ATOM      9  N   LYS A   2      11.743  19.862   2.050  1.00  0.44           N  
+ATOM     10  CA  LYS A   2      10.680  18.862   2.091  1.00  0.48           C  
+ATOM     11  C   LYS A   2      10.854  17.924   3.282  1.00  0.45           C  
+ATOM     12  CB  LYS A   2      10.648  18.058   0.791  1.00  0.41           C  
+ATOM     13  O   LYS A   2      11.967  17.481   3.573  1.00  0.44           O  
+ATOM     14  CG  LYS A   2       9.807  18.691  -0.308  1.00  0.40           C  
+ATOM     15  CD  LYS A   2       9.743  17.804  -1.544  1.00  0.44           C  
+ATOM     16  CE  LYS A   2       8.931  18.452  -2.657  1.00  0.42           C  
+ATOM     17  NZ  LYS A   2       8.883  17.595  -3.880  1.00  0.37           N  
+ATOM     18  N   THR A   3      10.260  18.153   4.498  1.00  0.59           N  
+ATOM     19  CA  THR A   3      10.394  17.359   5.714  1.00  0.62           C  
+ATOM     20  C   THR A   3      10.444  15.870   5.386  1.00  0.57           C  
+ATOM     21  CB  THR A   3       9.235  17.633   6.691  1.00  0.52           C  
+ATOM     22  O   THR A   3      10.034  15.454   4.300  1.00  0.52           O  
+ATOM     23  CG2 THR A   3       9.446  18.945   7.440  1.00  0.42           C  
+ATOM     24  OG1 THR A   3       8.007  17.707   5.957  1.00  0.44           O  
+ATOM     25  N   VAL A   4      11.392  14.978   5.807  1.00  0.55           N  
+ATOM     26  CA  VAL A   4      11.478  13.527   5.688  1.00  0.55           C  
+ATOM     27  C   VAL A   4      10.083  12.942   5.476  1.00  0.55           C  
+ATOM     28  CB  VAL A   4      12.139  12.893   6.932  1.00  0.50           C  
+ATOM     29  O   VAL A   4       9.905  12.019   4.678  1.00  0.54           O  
+ATOM     30  CG1 VAL A   4      12.097  11.368   6.849  1.00  0.40           C  
+ATOM     31  CG2 VAL A   4      13.578  13.384   7.080  1.00  0.41           C  
+ATOM     32  N   ARG A   5       9.045  13.367   6.280  1.00  0.59           N  
+ATOM     33  CA  ARG A   5       7.679  12.859   6.208  1.00  0.59           C  
+ATOM     34  C   ARG A   5       7.124  12.979   4.793  1.00  0.58           C  
+ATOM     35  CB  ARG A   5       6.775  13.607   7.191  1.00  0.55           C  
+ATOM     36  O   ARG A   5       6.507  12.043   4.280  1.00  0.57           O  
+ATOM     37  CG  ARG A   5       5.520  12.841   7.578  1.00  0.51           C  
+ATOM     38  CD  ARG A   5       4.719  13.573   8.645  1.00  0.53           C  
+ATOM     39  NE  ARG A   5       3.461  12.892   8.938  1.00  0.46           N  
+ATOM     40  NH1 ARG A   5       2.718  14.472  10.450  1.00  0.36           N  
+ATOM     41  NH2 ARG A   5       1.430  12.635   9.979  1.00  0.32           N  
+ATOM     42  CZ  ARG A   5       2.539  13.335   9.788  1.00  0.47           C  
+ATOM     43  N   GLN A   6       7.225  14.159   4.296  1.00  0.54           N  
+ATOM     44  CA  GLN A   6       6.739  14.443   2.950  1.00  0.53           C  
+ATOM     45  C   GLN A   6       7.437  13.564   1.916  1.00  0.54           C  
+ATOM     46  CB  GLN A   6       6.942  15.920   2.607  1.00  0.49           C  
+ATOM     47  O   GLN A   6       6.804  13.084   0.974  1.00  0.53           O  
+ATOM     48  CG  GLN A   6       5.903  16.844   3.227  1.00  0.46           C  
+ATOM     49  CD  GLN A   6       6.171  18.309   2.935  1.00  0.48           C  
+ATOM     50  NE2 GLN A   6       5.132  19.132   3.031  1.00  0.42           N  
+ATOM     51  OE1 GLN A   6       7.302  18.697   2.626  1.00  0.54           O  
+ATOM     52  N   GLU A   7       8.751  13.356   2.208  1.00  0.53           N  
+ATOM     53  CA  GLU A   7       9.522  12.566   1.253  1.00  0.51           C  
+ATOM     54  C   GLU A   7       9.048  11.115   1.229  1.00  0.52           C  
+ATOM     55  CB  GLU A   7      11.015  12.627   1.585  1.00  0.49           C  
+ATOM     56  O   GLU A   7       8.959  10.504   0.162  1.00  0.52           O  
+ATOM     57  CG  GLU A   7      11.700  13.898   1.106  1.00  0.45           C  
+ATOM     58  CD  GLU A   7      13.193  13.920   1.394  1.00  0.47           C  
+ATOM     59  OE1 GLU A   7      13.697  12.984   2.054  1.00  0.53           O  
+ATOM     60  OE2 GLU A   7      13.863  14.882   0.956  1.00  0.48           O  
+ATOM     61  N   ARG A   8       8.906  10.459   2.342  1.00  0.54           N  
+ATOM     62  CA  ARG A   8       8.501   9.057   2.381  1.00  0.53           C  
+ATOM     63  C   ARG A   8       7.127   8.866   1.748  1.00  0.54           C  
+ATOM     64  CB  ARG A   8       8.491   8.540   3.821  1.00  0.52           C  
+ATOM     65  O   ARG A   8       6.894   7.883   1.042  1.00  0.53           O  
+ATOM     66  CG  ARG A   8       9.859   8.124   4.336  1.00  0.50           C  
+ATOM     67  CD  ARG A   8       9.777   7.520   5.731  1.00  0.51           C  
+ATOM     68  NE  ARG A   8      11.087   7.083   6.207  1.00  0.45           N  
+ATOM     69  NH1 ARG A   8      10.307   6.243   8.213  1.00  0.37           N  
+ATOM     70  NH2 ARG A   8      12.540   6.130   7.708  1.00  0.33           N  
+ATOM     71  CZ  ARG A   8      11.308   6.486   7.375  1.00  0.48           C  
+ATOM     72  N   LEU A   9       6.231   9.830   2.211  1.00  0.55           N  
+ATOM     73  CA  LEU A   9       4.899   9.731   1.624  1.00  0.54           C  
+ATOM     74  C   LEU A   9       4.969   9.810   0.103  1.00  0.55           C  
+ATOM     75  CB  LEU A   9       3.992  10.840   2.163  1.00  0.52           C  
+ATOM     76  O   LEU A   9       4.234   9.107  -0.594  1.00  0.55           O  
+ATOM     77  CG  LEU A   9       3.348  10.585   3.527  1.00  0.50           C  
+ATOM     78  CD1 LEU A   9       2.845  11.893   4.128  1.00  0.45           C  
+ATOM     79  CD2 LEU A   9       2.211   9.577   3.402  1.00  0.47           C  
+ATOM     80  N   LEU A  10       5.770  10.725  -0.331  1.00  0.49           N  
+ATOM     81  CA  LEU A  10       5.935  10.895  -1.770  1.00  0.48           C  
+ATOM     82  C   LEU A  10       6.431   9.605  -2.417  1.00  0.49           C  
+ATOM     83  CB  LEU A  10       6.911  12.036  -2.067  1.00  0.46           C  
+ATOM     84  O   LEU A  10       5.980   9.237  -3.503  1.00  0.49           O  
+ATOM     85  CG  LEU A  10       6.296  13.336  -2.587  1.00  0.43           C  
+ATOM     86  CD1 LEU A  10       6.864  14.532  -1.829  1.00  0.38           C  
+ATOM     87  CD2 LEU A  10       6.540  13.481  -4.085  1.00  0.41           C  
+ATOM     88  N   LYS A  11       7.442   8.965  -1.722  1.00  0.51           N  
+ATOM     89  CA  LYS A  11       8.010   7.762  -2.324  1.00  0.50           C  
+ATOM     90  C   LYS A  11       6.960   6.663  -2.454  1.00  0.51           C  
+ATOM     91  CB  LYS A  11       9.197   7.258  -1.501  1.00  0.48           C  
+ATOM     92  O   LYS A  11       6.940   5.930  -3.445  1.00  0.51           O  
+ATOM     93  CG  LYS A  11      10.514   7.947  -1.828  1.00  0.46           C  
+ATOM     94  CD  LYS A  11      11.665   7.365  -1.017  1.00  0.49           C  
+ATOM     95  CE  LYS A  11      12.973   8.092  -1.300  1.00  0.43           C  
+ATOM     96  NZ  LYS A  11      14.104   7.528  -0.504  1.00  0.40           N  
+ATOM     97  N   ILE A  12       6.178   6.565  -1.403  1.00  0.53           N  
+ATOM     98  CA  ILE A  12       5.147   5.534  -1.450  1.00  0.52           C  
+ATOM     99  C   ILE A  12       4.183   5.820  -2.599  1.00  0.53           C  
+ATOM    100  CB  ILE A  12       4.377   5.443  -0.114  1.00  0.51           C  
+ATOM    101  O   ILE A  12       3.739   4.899  -3.289  1.00  0.54           O  
+ATOM    102  CG1 ILE A  12       5.283   4.883   0.988  1.00  0.45           C  
+ATOM    103  CG2 ILE A  12       3.117   4.587  -0.274  1.00  0.46           C  
+ATOM    104  CD1 ILE A  12       4.733   5.067   2.396  1.00  0.45           C  
+ATOM    105  N   SER A  13       3.703   7.056  -2.575  1.00  0.53           N  
+ATOM    106  CA  SER A  13       2.823   7.481  -3.659  1.00  0.52           C  
+ATOM    107  C   SER A  13       3.421   7.144  -5.020  1.00  0.53           C  
+ATOM    108  CB  SER A  13       2.549   8.982  -3.570  1.00  0.50           C  
+ATOM    109  O   SER A  13       2.699   6.764  -5.945  1.00  0.53           O  
+ATOM    110  OG  SER A  13       1.700   9.275  -2.474  1.00  0.49           O  
+ATOM    111  N   LEU A  14       4.831   7.388  -5.129  1.00  0.48           N  
+ATOM    112  CA  LEU A  14       5.477   7.149  -6.415  1.00  0.47           C  
+ATOM    113  C   LEU A  14       5.426   5.669  -6.781  1.00  0.48           C  
+ATOM    114  CB  LEU A  14       6.930   7.629  -6.382  1.00  0.44           C  
+ATOM    115  O   LEU A  14       5.233   5.321  -7.948  1.00  0.47           O  
+ATOM    116  CG  LEU A  14       7.205   9.002  -6.998  1.00  0.42           C  
+ATOM    117  CD1 LEU A  14       8.185   9.785  -6.130  1.00  0.36           C  
+ATOM    118  CD2 LEU A  14       7.740   8.854  -8.418  1.00  0.40           C  
+ATOM    119  N   VAL A  15       5.762   4.833  -5.746  1.00  0.52           N  
+ATOM    120  CA  VAL A  15       5.724   3.409  -6.064  1.00  0.51           C  
+ATOM    121  C   VAL A  15       4.335   3.031  -6.573  1.00  0.52           C  
+ATOM    122  CB  VAL A  15       6.097   2.543  -4.840  1.00  0.49           C  
+ATOM    123  O   VAL A  15       4.204   2.238  -7.509  1.00  0.51           O  
+ATOM    124  CG1 VAL A  15       5.908   1.060  -5.150  1.00  0.44           C  
+ATOM    125  CG2 VAL A  15       7.536   2.823  -4.408  1.00  0.46           C  
+ATOM    126  N   LEU A  16       3.392   3.725  -5.916  1.00  0.55           N  
+ATOM    127  CA  LEU A  16       2.029   3.405  -6.323  1.00  0.54           C  
+ATOM    128  C   LEU A  16       1.741   3.935  -7.724  1.00  0.55           C  
+ATOM    129  CB  LEU A  16       1.021   3.986  -5.328  1.00  0.52           C  
+ATOM    130  O   LEU A  16       0.957   3.342  -8.468  1.00  0.55           O  
+ATOM    131  CG  LEU A  16       0.971   3.325  -3.950  1.00  0.50           C  
+ATOM    132  CD1 LEU A  16       0.096   4.140  -3.003  1.00  0.46           C  
+ATOM    133  CD2 LEU A  16       0.458   1.893  -4.061  1.00  0.47           C  
+ATOM    134  N   SER A  17       2.343   5.204  -7.894  1.00  0.49           N  
+ATOM    135  CA  SER A  17       2.073   5.808  -9.195  1.00  0.48           C  
+ATOM    136  C   SER A  17       2.809   5.072 -10.309  1.00  0.49           C  
+ATOM    137  CB  SER A  17       2.475   7.284  -9.195  1.00  0.45           C  
+ATOM    138  O   SER A  17       2.409   5.139 -11.474  1.00  0.48           O  
+ATOM    139  OG  SER A  17       3.880   7.422  -9.077  1.00  0.43           O  
+ATOM    140  N   GLU A  18       4.061   4.645  -9.888  1.00  0.51           N  
+ATOM    141  CA  GLU A  18       4.832   3.986 -10.938  1.00  0.50           C  
+ATOM    142  C   GLU A  18       4.292   2.588 -11.224  1.00  0.51           C  
+ATOM    143  CB  GLU A  18       6.311   3.910 -10.552  1.00  0.47           C  
+ATOM    144  O   GLU A  18       4.769   1.907 -12.134  1.00  0.50           O  
+ATOM    145  CG  GLU A  18       6.999   5.266 -10.482  1.00  0.45           C  
+ATOM    146  CD  GLU A  18       8.505   5.185 -10.674  1.00  0.47           C  
+ATOM    147  OE1 GLU A  18       9.058   4.062 -10.671  1.00  0.48           O  
+ATOM    148  OE2 GLU A  18       9.138   6.254 -10.827  1.00  0.42           O  
+ATOM    149  N   LEU A  19       3.449   2.137 -10.297  1.00  0.52           N  
+ATOM    150  CA  LEU A  19       2.839   0.883 -10.722  1.00  0.52           C  
+ATOM    151  C   LEU A  19       2.051   1.073 -12.014  1.00  0.52           C  
+ATOM    152  CB  LEU A  19       1.921   0.334  -9.627  1.00  0.50           C  
+ATOM    153  O   LEU A  19       1.400   2.103 -12.204  1.00  0.52           O  
+ATOM    154  CG  LEU A  19       2.610  -0.333  -8.436  1.00  0.49           C  
+ATOM    155  CD1 LEU A  19       1.682  -0.348  -7.226  1.00  0.47           C  
+ATOM    156  CD2 LEU A  19       3.050  -1.748  -8.796  1.00  0.48           C  
+ATOM    157  N   PRO A  20       2.509   0.425 -12.947  1.00  0.52           N  
+ATOM    158  CA  PRO A  20       1.748   0.593 -14.187  1.00  0.50           C  
+ATOM    159  C   PRO A  20       0.239   0.494 -13.972  1.00  0.52           C  
+ATOM    160  CB  PRO A  20       2.250  -0.555 -15.067  1.00  0.49           C  
+ATOM    161  O   PRO A  20      -0.239  -0.467 -13.363  1.00  0.51           O  
+ATOM    162  CG  PRO A  20       3.024  -1.436 -14.141  1.00  0.46           C  
+ATOM    163  CD  PRO A  20       3.173  -0.730 -12.824  1.00  0.46           C  
+ATOM    164  N   LEU A  21      -0.387   1.673 -13.561  1.00  0.49           N  
+ATOM    165  CA  LEU A  21      -1.840   1.543 -13.546  1.00  0.49           C  
+ATOM    166  C   LEU A  21      -2.334   0.805 -14.786  1.00  0.49           C  
+ATOM    167  CB  LEU A  21      -2.501   2.921 -13.461  1.00  0.47           C  
+ATOM    168  O   LEU A  21      -3.422   0.225 -14.776  1.00  0.48           O  
+ATOM    169  CG  LEU A  21      -2.547   3.567 -12.075  1.00  0.45           C  
+ATOM    170  CD1 LEU A  21      -2.511   5.087 -12.197  1.00  0.41           C  
+ATOM    171  CD2 LEU A  21      -3.791   3.116 -11.316  1.00  0.42           C  
+ATOM    172  N   GLU A  22      -1.284   0.823 -15.746  1.00  0.49           N  
+ATOM    173  CA  GLU A  22      -1.690   0.325 -17.056  1.00  0.49           C  
+ATOM    174  C   GLU A  22      -1.183  -1.095 -17.289  1.00  0.49           C  
+ATOM    175  CB  GLU A  22      -1.184   1.252 -18.164  1.00  0.45           C  
+ATOM    176  O   GLU A  22      -1.105  -1.553 -18.431  1.00  0.47           O  
+ATOM    177  CG  GLU A  22      -1.870   2.611 -18.191  1.00  0.43           C  
+ATOM    178  CD  GLU A  22      -1.638   3.375 -19.484  1.00  0.46           C  
+ATOM    179  OE1 GLU A  22      -0.921   2.862 -20.372  1.00  0.44           O  
+ATOM    180  OE2 GLU A  22      -2.176   4.498 -19.610  1.00  0.40           O  
+ATOM    181  N   SER A  23      -0.507  -1.685 -16.243  1.00  0.48           N  
+ATOM    182  CA  SER A  23      -0.409  -3.087 -16.636  1.00  0.47           C  
+ATOM    183  C   SER A  23      -1.772  -3.771 -16.589  1.00  0.48           C  
+ATOM    184  CB  SER A  23       0.575  -3.830 -15.732  1.00  0.45           C  
+ATOM    185  O   SER A  23      -2.460  -3.724 -15.568  1.00  0.46           O  
+ATOM    186  OG  SER A  23       0.373  -3.480 -14.374  1.00  0.42           O  
+ATOM    187  N   LYS A  24      -2.783  -3.245 -17.333  1.00  0.48           N  
+ATOM    188  CA  LYS A  24      -3.816  -4.243 -17.597  1.00  0.47           C  
+ATOM    189  C   LYS A  24      -3.258  -5.658 -17.476  1.00  0.48           C  
+ATOM    190  CB  LYS A  24      -4.420  -4.035 -18.987  1.00  0.46           C  
+ATOM    191  O   LYS A  24      -2.311  -6.020 -18.177  1.00  0.46           O  
+ATOM    192  CG  LYS A  24      -5.249  -2.766 -19.119  1.00  0.44           C  
+ATOM    193  CD  LYS A  24      -6.028  -2.741 -20.427  1.00  0.46           C  
+ATOM    194  CE  LYS A  24      -6.820  -1.449 -20.582  1.00  0.40           C  
+ATOM    195  NZ  LYS A  24      -7.573  -1.413 -21.871  1.00  0.37           N  
+ATOM    196  N   PRO A  25      -3.101  -6.132 -16.224  1.00  0.48           N  
+ATOM    197  CA  PRO A  25      -2.777  -7.561 -16.224  1.00  0.47           C  
+ATOM    198  C   PRO A  25      -3.419  -8.311 -17.389  1.00  0.49           C  
+ATOM    199  CB  PRO A  25      -3.336  -8.049 -14.886  1.00  0.46           C  
+ATOM    200  O   PRO A  25      -4.457  -7.887 -17.905  1.00  0.48           O  
+ATOM    201  CG  PRO A  25      -4.170  -6.915 -14.383  1.00  0.44           C  
+ATOM    202  CD  PRO A  25      -3.916  -5.715 -15.250  1.00  0.45           C  
+ATOM    203  N   GLU A  26      -2.595  -8.749 -18.537  1.00  0.52           N  
+ATOM    204  CA  GLU A  26      -3.262  -9.728 -19.391  1.00  0.52           C  
+ATOM    205  C   GLU A  26      -4.559 -10.221 -18.756  1.00  0.52           C  
+ATOM    206  CB  GLU A  26      -2.334 -10.911 -19.678  1.00  0.49           C  
+ATOM    207  O   GLU A  26      -4.681 -10.258 -17.530  1.00  0.50           O  
+ATOM    208  CG  GLU A  26      -1.229 -10.598 -20.678  1.00  0.48           C  
+ATOM    209  CD  GLU A  26      -0.782 -11.811 -21.477  1.00  0.50           C  
+ATOM    210  OE1 GLU A  26      -1.195 -12.945 -21.142  1.00  0.51           O  
+ATOM    211  OE2 GLU A  26      -0.012 -11.627 -22.446  1.00  0.47           O  
+ATOM    212  N   PRO A  27      -5.776  -9.853 -19.347  1.00  0.50           N  
+ATOM    213  CA  PRO A  27      -6.879 -10.502 -18.634  1.00  0.49           C  
+ATOM    214  C   PRO A  27      -6.418 -11.681 -17.781  1.00  0.50           C  
+ATOM    215  CB  PRO A  27      -7.800 -10.971 -19.763  1.00  0.47           C  
+ATOM    216  O   PRO A  27      -5.769 -12.599 -18.290  1.00  0.49           O  
+ATOM    217  CG  PRO A  27      -7.065 -10.637 -21.020  1.00  0.45           C  
+ATOM    218  CD  PRO A  27      -5.782  -9.947 -20.654  1.00  0.45           C  
+ATOM    219  N   VAL A  28      -5.550 -11.387 -16.664  1.00  0.49           N  
+ATOM    220  CA  VAL A  28      -5.461 -12.521 -15.750  1.00  0.48           C  
+ATOM    221  C   VAL A  28      -6.814 -13.223 -15.665  1.00  0.49           C  
+ATOM    222  CB  VAL A  28      -4.999 -12.081 -14.343  1.00  0.47           C  
+ATOM    223  O   VAL A  28      -7.852 -12.571 -15.529  1.00  0.48           O  
+ATOM    224  CG1 VAL A  28      -4.676 -13.297 -13.476  1.00  0.45           C  
+ATOM    225  CG2 VAL A  28      -3.787 -11.157 -14.445  1.00  0.46           C  
+ATOM    226  N   GLN A  29      -7.108 -14.082 -16.749  1.00  0.53           N  
+ATOM    227  CA  GLN A  29      -8.211 -15.035 -16.676  1.00  0.53           C  
+ATOM    228  C   GLN A  29      -8.987 -14.878 -15.372  1.00  0.53           C  
+ATOM    229  CB  GLN A  29      -7.692 -16.467 -16.810  1.00  0.48           C  
+ATOM    230  O   GLN A  29      -8.402 -14.913 -14.287  1.00  0.51           O  
+ATOM    231  CG  GLN A  29      -7.428 -16.894 -18.247  1.00  0.46           C  
+ATOM    232  CD  GLN A  29      -7.062 -18.362 -18.366  1.00  0.48           C  
+ATOM    233  NE2 GLN A  29      -6.750 -18.799 -19.581  1.00  0.41           N  
+ATOM    234  OE1 GLN A  29      -7.058 -19.097 -17.373  1.00  0.50           O  
+ATOM    235  N   GLY A  30      -9.845 -13.707 -15.208  1.00  0.58           N  
+ATOM    236  CA  GLY A  30     -11.064 -13.567 -14.428  1.00  0.57           C  
+ATOM    237  C   GLY A  30     -10.991 -12.454 -13.400  1.00  0.58           C  
+ATOM    238  O   GLY A  30      -9.936 -12.218 -12.807  1.00  0.55           O  
+ATOM    239  N   ALA A  31     -11.469 -11.174 -13.727  1.00  0.65           N  
+ATOM    240  CA  ALA A  31     -11.865 -10.061 -12.868  1.00  0.64           C  
+ATOM    241  C   ALA A  31     -11.715 -10.426 -11.394  1.00  0.65           C  
+ATOM    242  CB  ALA A  31     -13.303  -9.645 -13.166  1.00  0.61           C  
+ATOM    243  O   ALA A  31     -11.286  -9.600 -10.584  1.00  0.65           O  
+ATOM    244  N   ALA A  32     -12.034 -11.668 -11.139  1.00  0.68           N  
+ATOM    245  CA  ALA A  32     -11.937 -12.088  -9.744  1.00  0.67           C  
+ATOM    246  C   ALA A  32     -10.490 -12.054  -9.260  1.00  0.68           C  
+ATOM    247  CB  ALA A  32     -12.521 -13.487  -9.567  1.00  0.64           C  
+ATOM    248  O   ALA A  32     -10.211 -11.605  -8.146  1.00  0.67           O  
+ATOM    249  N   LEU A  33      -9.591 -12.548 -10.047  1.00  0.67           N  
+ATOM    250  CA  LEU A  33      -8.192 -12.557  -9.631  1.00  0.66           C  
+ATOM    251  C   LEU A  33      -7.660 -11.135  -9.484  1.00  0.67           C  
+ATOM    252  CB  LEU A  33      -7.339 -13.332 -10.638  1.00  0.63           C  
+ATOM    253  O   LEU A  33      -6.890 -10.849  -8.564  1.00  0.65           O  
+ATOM    254  CG  LEU A  33      -5.950 -13.761 -10.161  1.00  0.59           C  
+ATOM    255  CD1 LEU A  33      -6.019 -15.128  -9.488  1.00  0.54           C  
+ATOM    256  CD2 LEU A  33      -4.967 -13.783 -11.326  1.00  0.55           C  
+ATOM    257  N   GLN A  34      -8.043 -10.287 -10.357  1.00  0.68           N  
+ATOM    258  CA  GLN A  34      -7.625  -8.892 -10.272  1.00  0.68           C  
+ATOM    259  C   GLN A  34      -8.075  -8.262  -8.957  1.00  0.69           C  
+ATOM    260  CB  GLN A  34      -8.176  -8.093 -11.454  1.00  0.65           C  
+ATOM    261  O   GLN A  34      -7.304  -7.552  -8.307  1.00  0.68           O  
+ATOM    262  CG  GLN A  34      -7.476  -6.759 -11.674  1.00  0.61           C  
+ATOM    263  CD  GLN A  34      -7.961  -6.040 -12.920  1.00  0.59           C  
+ATOM    264  NE2 GLN A  34      -7.359  -4.893 -13.214  1.00  0.48           N  
+ATOM    265  OE1 GLN A  34      -8.868  -6.513 -13.612  1.00  0.55           O  
+ATOM    266  N   ALA A  35      -9.311  -8.425  -8.724  1.00  0.72           N  
+ATOM    267  CA  ALA A  35      -9.858  -7.851  -7.498  1.00  0.71           C  
+ATOM    268  C   ALA A  35      -9.115  -8.371  -6.270  1.00  0.73           C  
+ATOM    269  CB  ALA A  35     -11.348  -8.161  -7.384  1.00  0.70           C  
+ATOM    270  O   ALA A  35      -8.849  -7.616  -5.332  1.00  0.73           O  
+ATOM    271  N   GLU A  36      -8.800  -9.602  -6.322  1.00  0.75           N  
+ATOM    272  CA  GLU A  36      -8.080 -10.212  -5.208  1.00  0.75           C  
+ATOM    273  C   GLU A  36      -6.689  -9.603  -5.050  1.00  0.76           C  
+ATOM    274  CB  GLU A  36      -7.971 -11.726  -5.403  1.00  0.73           C  
+ATOM    275  O   GLU A  36      -6.256  -9.313  -3.933  1.00  0.75           O  
+ATOM    276  CG  GLU A  36      -7.373 -12.458  -4.210  1.00  0.68           C  
+ATOM    277  CD  GLU A  36      -7.447 -13.972  -4.336  1.00  0.65           C  
+ATOM    278  OE1 GLU A  36      -8.001 -14.471  -5.342  1.00  0.64           O  
+ATOM    279  OE2 GLU A  36      -6.948 -14.664  -3.421  1.00  0.59           O  
+ATOM    280  N   LEU A  37      -5.986  -9.435  -6.129  1.00  0.74           N  
+ATOM    281  CA  LEU A  37      -4.628  -8.906  -6.060  1.00  0.73           C  
+ATOM    282  C   LEU A  37      -4.630  -7.469  -5.549  1.00  0.74           C  
+ATOM    283  CB  LEU A  37      -3.957  -8.970  -7.435  1.00  0.71           C  
+ATOM    284  O   LEU A  37      -3.797  -7.100  -4.718  1.00  0.73           O  
+ATOM    285  CG  LEU A  37      -3.472 -10.348  -7.888  1.00  0.66           C  
+ATOM    286  CD1 LEU A  37      -3.167 -10.337  -9.382  1.00  0.61           C  
+ATOM    287  CD2 LEU A  37      -2.244 -10.772  -7.090  1.00  0.61           C  
+ATOM    288  N   LEU A  38      -5.626  -6.678  -6.066  1.00  0.73           N  
+ATOM    289  CA  LEU A  38      -5.729  -5.296  -5.610  1.00  0.73           C  
+ATOM    290  C   LEU A  38      -6.038  -5.238  -4.118  1.00  0.75           C  
+ATOM    291  CB  LEU A  38      -6.811  -4.551  -6.397  1.00  0.71           C  
+ATOM    292  O   LEU A  38      -5.501  -4.391  -3.401  1.00  0.75           O  
+ATOM    293  CG  LEU A  38      -6.433  -4.102  -7.809  1.00  0.66           C  
+ATOM    294  CD1 LEU A  38      -7.679  -3.690  -8.586  1.00  0.60           C  
+ATOM    295  CD2 LEU A  38      -5.428  -2.956  -7.755  1.00  0.60           C  
+ATOM    296  N   SER A  39      -6.898  -6.106  -3.777  1.00  0.77           N  
+ATOM    297  CA  SER A  39      -7.237  -6.146  -2.358  1.00  0.77           C  
+ATOM    298  C   SER A  39      -6.013  -6.466  -1.507  1.00  0.78           C  
+ATOM    299  CB  SER A  39      -8.333  -7.179  -2.098  1.00  0.75           C  
+ATOM    300  O   SER A  39      -5.816  -5.871  -0.445  1.00  0.78           O  
+ATOM    301  OG  SER A  39      -8.631  -7.259  -0.715  1.00  0.68           O  
+ATOM    302  N   GLN A  40      -5.194  -7.402  -1.959  1.00  0.76           N  
+ATOM    303  CA  GLN A  40      -4.005  -7.786  -1.205  1.00  0.76           C  
+ATOM    304  C   GLN A  40      -3.014  -6.629  -1.113  1.00  0.76           C  
+ATOM    305  CB  GLN A  40      -3.334  -9.003  -1.843  1.00  0.73           C  
+ATOM    306  O   GLN A  40      -2.433  -6.384  -0.054  1.00  0.76           O  
+ATOM    307  CG  GLN A  40      -4.050 -10.317  -1.563  1.00  0.66           C  
+ATOM    308  CD  GLN A  40      -3.515 -11.466  -2.398  1.00  0.63           C  
+ATOM    309  NE2 GLN A  40      -4.359 -12.462  -2.645  1.00  0.54           N  
+ATOM    310  OE1 GLN A  40      -2.354 -11.457  -2.820  1.00  0.60           O  
+ATOM    311  N   VAL A  41      -2.792  -5.944  -2.237  1.00  0.75           N  
+ATOM    312  CA  VAL A  41      -1.862  -4.820  -2.237  1.00  0.74           C  
+ATOM    313  C   VAL A  41      -2.341  -3.752  -1.256  1.00  0.75           C  
+ATOM    314  CB  VAL A  41      -1.704  -4.213  -3.649  1.00  0.73           C  
+ATOM    315  O   VAL A  41      -1.550  -3.215  -0.478  1.00  0.74           O  
+ATOM    316  CG1 VAL A  41      -0.947  -2.888  -3.586  1.00  0.65           C  
+ATOM    317  CG2 VAL A  41      -0.989  -5.196  -4.575  1.00  0.66           C  
+ATOM    318  N   ARG A  42      -3.621  -3.448  -1.361  1.00  0.76           N  
+ATOM    319  CA  ARG A  42      -4.187  -2.474  -0.434  1.00  0.76           C  
+ATOM    320  C   ARG A  42      -3.930  -2.881   1.013  1.00  0.77           C  
+ATOM    321  CB  ARG A  42      -5.690  -2.314  -0.673  1.00  0.74           C  
+ATOM    322  O   ARG A  42      -3.588  -2.042   1.849  1.00  0.77           O  
+ATOM    323  CG  ARG A  42      -6.034  -1.420  -1.855  1.00  0.69           C  
+ATOM    324  CD  ARG A  42      -7.539  -1.280  -2.036  1.00  0.65           C  
+ATOM    325  NE  ARG A  42      -7.867  -0.358  -3.120  1.00  0.59           N  
+ATOM    326  NH1 ARG A  42     -10.131  -0.814  -3.100  1.00  0.49           N  
+ATOM    327  NH2 ARG A  42      -9.276   0.705  -4.589  1.00  0.45           N  
+ATOM    328  CZ  ARG A  42      -9.091  -0.158  -3.601  1.00  0.59           C  
+ATOM    329  N   GLN A  43      -4.215  -4.159   1.316  1.00  0.78           N  
+ATOM    330  CA  GLN A  43      -4.020  -4.652   2.675  1.00  0.78           C  
+ATOM    331  C   GLN A  43      -2.557  -4.540   3.096  1.00  0.78           C  
+ATOM    332  CB  GLN A  43      -4.490  -6.103   2.793  1.00  0.76           C  
+ATOM    333  O   GLN A  43      -2.260  -4.175   4.235  1.00  0.77           O  
+ATOM    334  CG  GLN A  43      -4.504  -6.630   4.222  1.00  0.70           C  
+ATOM    335  CD  GLN A  43      -5.452  -5.861   5.123  1.00  0.66           C  
+ATOM    336  NE2 GLN A  43      -4.992  -5.534   6.326  1.00  0.60           N  
+ATOM    337  OE1 GLN A  43      -6.588  -5.562   4.740  1.00  0.66           O  
+ATOM    338  N   ASP A  44      -1.615  -4.885   2.249  1.00  0.76           N  
+ATOM    339  CA  ASP A  44      -0.192  -4.815   2.566  1.00  0.75           C  
+ATOM    340  C   ASP A  44       0.231  -3.383   2.883  1.00  0.75           C  
+ATOM    341  CB  ASP A  44       0.643  -5.366   1.409  1.00  0.73           C  
+ATOM    342  O   ASP A  44       1.022  -3.152   3.800  1.00  0.74           O  
+ATOM    343  CG  ASP A  44       0.532  -6.874   1.262  1.00  0.68           C  
+ATOM    344  OD1 ASP A  44       0.048  -7.544   2.199  1.00  0.65           O  
+ATOM    345  OD2 ASP A  44       0.935  -7.396   0.200  1.00  0.66           O  
+ATOM    346  N   ILE A  45      -0.307  -2.445   2.112  1.00  0.76           N  
+ATOM    347  CA  ILE A  45      -0.006  -1.042   2.371  1.00  0.75           C  
+ATOM    348  C   ILE A  45      -0.514  -0.652   3.758  1.00  0.76           C  
+ATOM    349  CB  ILE A  45      -0.625  -0.122   1.295  1.00  0.73           C  
+ATOM    350  O   ILE A  45       0.201  -0.011   4.531  1.00  0.75           O  
+ATOM    351  CG1 ILE A  45       0.082  -0.322  -0.050  1.00  0.67           C  
+ATOM    352  CG2 ILE A  45      -0.561   1.343   1.735  1.00  0.67           C  
+ATOM    353  CD1 ILE A  45      -0.609   0.363  -1.221  1.00  0.64           C  
+ATOM    354  N   ALA A  46      -1.744  -0.973   3.967  1.00  0.76           N  
+ATOM    355  CA  ALA A  46      -2.318  -0.646   5.269  1.00  0.75           C  
+ATOM    356  C   ALA A  46      -1.481  -1.235   6.401  1.00  0.76           C  
+ATOM    357  CB  ALA A  46      -3.757  -1.149   5.358  1.00  0.74           C  
+ATOM    358  O   ALA A  46      -1.239  -0.572   7.413  1.00  0.75           O  
+ATOM    359  N   ASN A  47      -1.100  -2.513   6.275  1.00  0.77           N  
+ATOM    360  CA  ASN A  47      -0.287  -3.170   7.293  1.00  0.76           C  
+ATOM    361  C   ASN A  47       1.046  -2.456   7.493  1.00  0.75           C  
+ATOM    362  CB  ASN A  47      -0.053  -4.638   6.927  1.00  0.74           C  
+ATOM    363  O   ASN A  47       1.503  -2.291   8.626  1.00  0.74           O  
+ATOM    364  CG  ASN A  47      -1.313  -5.474   7.032  1.00  0.70           C  
+ATOM    365  ND2 ASN A  47      -1.301  -6.646   6.410  1.00  0.68           N  
+ATOM    366  OD1 ASN A  47      -2.291  -5.068   7.666  1.00  0.70           O  
+ATOM    367  N   SER A  48       1.642  -2.096   6.364  1.00  0.75           N  
+ATOM    368  CA  SER A  48       2.925  -1.406   6.457  1.00  0.73           C  
+ATOM    369  C   SER A  48       2.787  -0.078   7.194  1.00  0.73           C  
+ATOM    370  CB  SER A  48       3.508  -1.167   5.064  1.00  0.71           C  
+ATOM    371  O   SER A  48       3.641   0.278   8.009  1.00  0.70           O  
+ATOM    372  OG  SER A  48       3.811  -2.397   4.427  1.00  0.65           O  
+ATOM    373  N   LEU A  49       1.734   0.640   6.845  1.00  0.72           N  
+ATOM    374  CA  LEU A  49       1.508   1.919   7.510  1.00  0.72           C  
+ATOM    375  C   LEU A  49       1.271   1.721   9.003  1.00  0.72           C  
+ATOM    376  CB  LEU A  49       0.315   2.645   6.884  1.00  0.70           C  
+ATOM    377  O   LEU A  49       1.763   2.499   9.823  1.00  0.71           O  
+ATOM    378  CG  LEU A  49       0.534   3.228   5.487  1.00  0.66           C  
+ATOM    379  CD1 LEU A  49      -0.788   3.717   4.904  1.00  0.61           C  
+ATOM    380  CD2 LEU A  49       1.554   4.360   5.533  1.00  0.62           C  
+ATOM    381  N   ASN A  50       0.475   0.748   9.327  1.00  0.74           N  
+ATOM    382  CA  ASN A  50       0.188   0.475  10.732  1.00  0.73           C  
+ATOM    383  C   ASN A  50       1.452   0.099  11.499  1.00  0.73           C  
+ATOM    384  CB  ASN A  50      -0.860  -0.633  10.860  1.00  0.70           C  
+ATOM    385  O   ASN A  50       1.622   0.493  12.654  1.00  0.71           O  
+ATOM    386  CG  ASN A  50      -2.278  -0.113  10.732  1.00  0.65           C  
+ATOM    387  ND2 ASN A  50      -3.214  -1.008  10.440  1.00  0.63           N  
+ATOM    388  OD1 ASN A  50      -2.530   1.084  10.892  1.00  0.63           O  
+ATOM    389  N   ALA A  51       2.296  -0.751  10.857  1.00  0.70           N  
+ATOM    390  CA  ALA A  51       3.518  -1.202  11.518  1.00  0.69           C  
+ATOM    391  C   ALA A  51       4.408  -0.020  11.892  1.00  0.68           C  
+ATOM    392  CB  ALA A  51       4.279  -2.175  10.622  1.00  0.67           C  
+ATOM    393  O   ALA A  51       5.038  -0.020  12.952  1.00  0.67           O  
+ATOM    394  N   VAL A  52       4.525   0.951  11.010  1.00  0.67           N  
+ATOM    395  CA  VAL A  52       5.352   2.120  11.288  1.00  0.66           C  
+ATOM    396  C   VAL A  52       4.723   2.942  12.411  1.00  0.65           C  
+ATOM    397  CB  VAL A  52       5.540   2.995  10.029  1.00  0.63           C  
+ATOM    398  O   VAL A  52       5.431   3.495  13.256  1.00  0.64           O  
+ATOM    399  CG1 VAL A  52       6.261   4.296  10.378  1.00  0.57           C  
+ATOM    400  CG2 VAL A  52       6.309   2.226   8.956  1.00  0.58           C  
+ATOM    401  N   ALA A  53       3.415   3.071  12.338  1.00  0.62           N  
+ATOM    402  CA  ALA A  53       2.746   3.869  13.362  1.00  0.60           C  
+ATOM    403  C   ALA A  53       2.961   3.272  14.750  1.00  0.61           C  
+ATOM    404  CB  ALA A  53       1.253   3.980  13.059  1.00  0.58           C  
+ATOM    405  O   ALA A  53       2.919   3.988  15.753  1.00  0.61           O  
+ATOM    406  N   THR A  54       3.105   1.936  14.815  1.00  0.61           N  
+ATOM    407  CA  THR A  54       3.158   1.297  16.126  1.00  0.61           C  
+ATOM    408  C   THR A  54       4.591   1.253  16.648  1.00  0.61           C  
+ATOM    409  CB  THR A  54       2.583  -0.131  16.075  1.00  0.57           C  
+ATOM    410  O   THR A  54       4.835   0.805  17.770  1.00  0.59           O  
+ATOM    411  CG2 THR A  54       1.101  -0.114  15.712  1.00  0.51           C  
+ATOM    412  OG1 THR A  54       3.295  -0.892  15.092  1.00  0.54           O  
+ATOM    413  N   ARG A  55       5.532   1.631  15.809  1.00  0.65           N  
+ATOM    414  CA  ARG A  55       6.903   1.567  16.303  1.00  0.65           C  
+ATOM    415  C   ARG A  55       7.118   2.544  17.453  1.00  0.65           C  
+ATOM    416  CB  ARG A  55       7.895   1.859  15.175  1.00  0.61           C  
+ATOM    417  O   ARG A  55       6.621   3.672  17.418  1.00  0.63           O  
+ATOM    418  CG  ARG A  55       8.075   0.707  14.199  1.00  0.59           C  
+ATOM    419  CD  ARG A  55       9.132   1.017  13.148  1.00  0.59           C  
+ATOM    420  NE  ARG A  55       9.264  -0.067  12.180  1.00  0.52           N  
+ATOM    421  NH1 ARG A  55      10.992   0.916  11.004  1.00  0.41           N  
+ATOM    422  NH2 ARG A  55      10.179  -1.128  10.359  1.00  0.36           N  
+ATOM    423  CZ  ARG A  55      10.145  -0.090  11.183  1.00  0.55           C  
+ATOM    424  N   PRO A  56       7.605   2.129  18.545  1.00  0.62           N  
+ATOM    425  CA  PRO A  56       7.968   3.039  19.634  1.00  0.61           C  
+ATOM    426  C   PRO A  56       8.765   4.249  19.151  1.00  0.62           C  
+ATOM    427  CB  PRO A  56       8.815   2.161  20.559  1.00  0.56           C  
+ATOM    428  O   PRO A  56       9.656   4.109  18.310  1.00  0.59           O  
+ATOM    429  CG  PRO A  56       8.612   0.769  20.056  1.00  0.53           C  
+ATOM    430  CD  PRO A  56       8.067   0.841  18.659  1.00  0.55           C  
+ATOM    431  N   GLY A  57       8.356   5.535  19.374  1.00  0.58           N  
+ATOM    432  CA  GLY A  57       8.945   6.814  19.010  1.00  0.58           C  
+ATOM    433  C   GLY A  57       8.368   7.395  17.733  1.00  0.58           C  
+ATOM    434  O   GLY A  57       8.738   8.497  17.323  1.00  0.57           O  
+ATOM    435  N   TYR A  58       7.679   6.470  16.960  1.00  0.54           N  
+ATOM    436  CA  TYR A  58       7.100   7.135  15.799  1.00  0.54           C  
+ATOM    437  C   TYR A  58       6.095   8.199  16.224  1.00  0.53           C  
+ATOM    438  CB  TYR A  58       6.422   6.116  14.878  1.00  0.50           C  
+ATOM    439  O   TYR A  58       6.065   9.295  15.659  1.00  0.52           O  
+ATOM    440  CG  TYR A  58       6.176   6.630  13.480  1.00  0.49           C  
+ATOM    441  CD1 TYR A  58       4.946   7.177  13.124  1.00  0.47           C  
+ATOM    442  CD2 TYR A  58       7.173   6.568  12.512  1.00  0.48           C  
+ATOM    443  CE1 TYR A  58       4.714   7.650  11.837  1.00  0.47           C  
+ATOM    444  CE2 TYR A  58       6.953   7.038  11.222  1.00  0.48           C  
+ATOM    445  OH  TYR A  58       5.498   8.043   9.618  1.00  0.44           O  
+ATOM    446  CZ  TYR A  58       5.722   7.577  10.894  1.00  0.46           C  
+ATOM    447  N   LEU A  59       5.154   7.928  17.166  1.00  0.50           N  
+ATOM    448  CA  LEU A  59       4.166   8.868  17.683  1.00  0.50           C  
+ATOM    449  C   LEU A  59       4.547   9.344  19.081  1.00  0.50           C  
+ATOM    450  CB  LEU A  59       2.777   8.224  17.711  1.00  0.47           C  
+ATOM    451  O   LEU A  59       3.832  10.147  19.684  1.00  0.49           O  
+ATOM    452  CG  LEU A  59       2.130   7.949  16.353  1.00  0.45           C  
+ATOM    453  CD1 LEU A  59       0.903   7.059  16.521  1.00  0.42           C  
+ATOM    454  CD2 LEU A  59       1.756   9.257  15.663  1.00  0.44           C  
+ATOM    455  N   ALA A  60       5.718   9.319  19.497  1.00  0.46           N  
+ATOM    456  CA  ALA A  60       6.022   9.769  20.853  1.00  0.46           C  
+ATOM    457  C   ALA A  60       7.279  10.633  20.876  1.00  0.46           C  
+ATOM    458  CB  ALA A  60       6.186   8.572  21.787  1.00  0.42           C  
+ATOM    459  O   ALA A  60       8.383  10.140  20.637  1.00  0.44           O  
+ATOM    460  N   GLY A  61       7.357  11.746  20.071  1.00  0.45           N  
+ATOM    461  CA  GLY A  61       8.144  12.780  20.724  1.00  0.45           C  
+ATOM    462  C   GLY A  61       7.607  14.179  20.487  1.00  0.46           C  
+ATOM    463  O   GLY A  61       7.660  14.688  19.366  1.00  0.44           O  
+ATOM    464  N   GLY A  62       6.464  14.557  21.219  1.00  0.33           N  
+ATOM    465  CA  GLY A  62       6.288  15.850  21.860  1.00  0.36           C  
+ATOM    466  C   GLY A  62       7.363  16.852  21.487  1.00  0.33           C  
+ATOM    467  O   GLY A  62       8.470  16.469  21.102  1.00  0.32           O  

esm/mcp_output/requirements.txt

@@ -0,0 +1,4 @@
+fastmcp>=0.1.0
+pydantic>=2.0.0
+requests
+biopython

esm/mcp_output/start_mcp.py

@@ -0,0 +1,34 @@
+
+"""
+MCP Service Startup Entry Point
+"""
+import sys
+import os
+
+project_root = os.path.dirname(os.path.abspath(__file__))
+mcp_plugin_dir = os.path.join(project_root, "mcp_plugin")
+if mcp_plugin_dir not in sys.path:
+    sys.path.insert(0, mcp_plugin_dir)
+
+# Set path to point to source directory
+source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "source")
+sys.path.insert(0, source_path)
+
+from mcp_service import create_app
+
+def main():
+    """Start FastMCP Service"""
+    app = create_app()
+    # Use environment variable to configure port, default 8000
+    port = int(os.environ.get("MCP_PORT", "8000"))
+    
+    # Select transport mode based on environment variable
+    transport = os.environ.get("MCP_TRANSPORT", "stdio")
+    if transport == "http":
+        app.run(transport="http", host="0.0.0.0", port=port)
+    else:
+        # Default to STDIO mode
+        app.run()
+
+if __name__ == "__main__":
+    main()

esm/mcp_output/tests_mcp/test_mcp_basic.py

@@ -0,0 +1,49 @@
+"""
+MCP Service Basic Tests
+"""
+import sys
+import os
+
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+mcp_plugin_dir = os.path.join(project_root, "mcp_plugin")
+if mcp_plugin_dir not in sys.path:
+    sys.path.insert(0, mcp_plugin_dir)
+
+source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "source")
+sys.path.insert(0, source_path)
+
+def test_import_mcp_service():
+    """Test that the MCP service can be imported correctly"""
+    try:
+        from mcp_service import create_app
+        app = create_app()
+        assert app is not None
+        print("MCP service imported successfully")
+        return True
+    except Exception as e:
+        print(f"Failed to import MCP service: {e}")
+        return False
+
+def test_adapter_init():
+    """Test that the adapter can be initialized correctly"""
+    try:
+        from adapter import Adapter
+        adapter = Adapter()
+        assert adapter is not None
+        print("Adapter initialized successfully")
+        return True
+    except Exception as e:
+        print(f"Failed to initialize adapter: {e}")
+        return False
+
+if __name__ == "__main__":
+    print("Running MCP service basic tests...")
+    test1 = test_import_mcp_service()
+    test2 = test_adapter_init()
+    
+    if test1 and test2:
+        print("All basic tests passed")
+        sys.exit(0)
+    else:
+        print("Some tests failed")
+        sys.exit(1)

esm/mcp_output/tests_smoke/test_smoke.py

@@ -0,0 +1,29 @@
+import importlib, sys
+import os
+
+# Add current directory to Python path
+sys.path.insert(0, os.getcwd())
+
+source_dir = os.path.join(os.getcwd(), "source")
+if os.path.exists(source_dir):
+    sys.path.insert(0, source_dir)
+
+
+try:
+    importlib.import_module("esm")
+    print("OK - Successfully imported esm")
+except ImportError as e:
+    print(f"Failed to import esm: {e}")
+    fallback_packages = []
+
+    fallback_packages = ['esm']
+
+    for pkg in fallback_packages:
+        try:
+            importlib.import_module(pkg)
+            print(f"OK - Successfully imported {pkg}")
+            break
+        except ImportError:
+            continue
+    else:
+        print("All import attempts failed")

esm/source/.flake8

@@ -0,0 +1,10 @@
+[flake8]
+max-line-length = 99
+ignore = E203,W503
+exclude =
+    .git,
+    __pycache__,
+    build,
+    dist,
+    experimental
+    third_party

esm/source/.git-blame-ignore-revs

@@ -0,0 +1,2 @@
+# Migrate code style to Black
+8bc7e948cd9bf0b6d1f2113e221ef548ef663377

esm/source/.github/ISSUE_TEMPLATE/bug.md

@@ -0,0 +1,27 @@
+---
+name: "[Bug Report]"
+about: "Create a bug report. For other questions: see Discussions tab."
+
+---
+
+NOTE: if this is not a bug report, please use the [GitHub Discussions](https://github.com/facebookresearch/esm/discussions) for support questions (How do I do X?), feature requests, ideas, showcasing new applications, etc.
+
+
+**Bug description**
+Please enter a clear and concise description of what the bug is.
+
+**Reproduction steps**
+Enter steps to reproduce the behavior.
+
+**Expected behavior**
+Give a clear and concise description of what you expected to happen.
+
+**Logs**
+Please paste the command line output:
+
+```
+Output goes here
+```
+
+**Additional context**
+Add any other context about the problem here. (like proxy settings, network setup, overall goals, etc.)

esm/source/.gitignore

@@ -0,0 +1,31 @@
+# tensor dumps
+*.pt
+# Compiler Output #
+###################
+*.py[cod]
+*.so
+*.o
+*.exe
+*.class
+
+# Folders #
+###########
+bin/
+build/
+dist/
+local/
+tmp/
+__pycache__/
+*.egg-info/
+.idea/
+.ipynb_checkpoints/
+.vscode/
+esm/dev
+
+# Junk #
+########
+.DS_Store*
+.*.swp
+*.swp
+*.log
+*~

esm/source/CODE_OF_CONDUCT.rst

@@ -0,0 +1,6 @@
+Code of Conduct
+===============
+
+Facebook has adopted a Code of Conduct that we expect project participants to adhere to. Please `read the full text`__ so that you can understand what actions will and will not be tolerated.
+
+__ https://code.facebook.com/codeofconduct

esm/source/CONTRIBUTING.md

@@ -0,0 +1,31 @@
+# Contributing to esm
+We want to make contributing to this project as easy and transparent as
+possible.
+
+## Pull Requests
+We actively welcome your pull requests.
+
+1. Fork the repo and create your branch from `master`.
+2. If you've added code that should be tested, add tests.
+3. If you've changed APIs, update the documentation.
+4. Ensure the test suite passes.
+5. Make sure your code lints.
+6. If you haven't already, complete the Contributor License Agreement ("CLA").
+
+## Contributor License Agreement ("CLA")
+In order to accept your pull request, we need you to submit a CLA. You only need
+to do this once to work on any of Facebook's open source projects.
+
+Complete your CLA here: <https://code.facebook.com/cla>
+
+## Issues
+We use GitHub issues to track public bugs. Please ensure your description is
+clear and has sufficient instructions to be able to reproduce the issue.
+
+Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
+disclosure of security bugs. In those cases, please go through the process
+outlined on that page and do not file a public issue.
+
+## License
+By contributing to icp-block-mdp, you agree that your contributions will be licensed
+under the LICENSE file in the root directory of this source tree.

esm/source/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+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.

esm/source/README.md

@@ -0,0 +1,795 @@
+# Evolutionary Scale Modeling
+
+[![atlas](https://user-images.githubusercontent.com/3605224/199301187-a9e38b3f-71a7-44be-94f4-db0d66143c53.png)](https://esmatlas.com)
+
+***Update April 2023:*** Code for the two simultaneous preprints on protein design is now released! Code for "Language models generalize beyond natural proteins" is under [examples/lm-design/](examples/lm-design/). Code for "A high-level programming language for generative protein design" is under [examples/protein-programming-language/](examples/protein-programming-language/).
+
+This repository contains code and pre-trained weights for **Transformer protein language models** from the Meta Fundamental AI Research Protein Team (FAIR), including our state-of-the-art [**ESM-2** and **ESMFold**](#esmfold), as well as [**MSA Transformer**](https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1), [**ESM-1v**](#zs_variant) for predicting variant effects and [**ESM-IF1**](#invf) for inverse folding.
+Transformer protein language models were introduced in the [2019 preprint](https://doi.org/10.1101/622803) of the paper ["Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences"](https://doi.org/10.1073/pnas.2016239118).
+ESM-2 outperforms all tested single-sequence protein language models across a range of structure prediction tasks.
+ESMFold harnesses the ESM-2 language model to generate accurate structure predictions end to end directly from the sequence of a protein.
+
+In November 2022, we released `v0` of the [ESM Metagenomic Atlas](https://esmatlas.com), an open atlas of 617 million predicted metagenomic protein structures.
+The Atlas was updated in March 2023 in collaboration with EBI. The new `v2023_02` adds another 150 million predicted structures to the Atlas, as well as pre-computed ESM2 embeddings.
+Bulk download, blog post and the resources provided on the Atlas website are documented [on this README](#atlas).
+
+In December 2022, we released two simultaneous preprints on protein design.
+* "Language models generalize beyond natural proteins" ([PAPER](https://doi.org/10.1101/2022.12.21.521521), [CODE](examples/lm-design/)) uses ESM2 to design de novo proteins. The code and data associated with the preprint can be found [here](examples/lm-design/).
+* "A high-level programming language for generative protein design" ([PAPER](https://doi.org/10.1101/2022.12.21.521526), [CODE](examples/protein-programming-language/)) uses ESMFold to design proteins according to a high-level programming language.
+
+
+
+<details><summary><b>Citation</b></summary>
+For ESM2, ESMFold and ESM Atlas:
+```bibtex
+@article{lin2023evolutionary,
+title = {Evolutionary-scale prediction of atomic-level protein structure with a language model},
+author = {Zeming Lin  and Halil Akin  and Roshan Rao  and Brian Hie  and Zhongkai Zhu  and Wenting Lu  and Nikita Smetanin  and Robert Verkuil  and Ori Kabeli  and Yaniv Shmueli  and Allan dos Santos Costa  and Maryam Fazel-Zarandi  and Tom Sercu  and Salvatore Candido  and Alexander Rives },
+journal = {Science},
+volume = {379},
+number = {6637},
+pages = {1123-1130},
+year = {2023},
+doi = {10.1126/science.ade2574},
+URL = {https://www.science.org/doi/abs/10.1126/science.ade2574},
+note={Earlier versions as preprint: bioRxiv 2022.07.20.500902},
+}
+```
+
+For transformer protein language models:
+```bibtex
+@article{rives2021biological,
+  title={Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences},
+  author={Rives, Alexander and Meier, Joshua and Sercu, Tom and Goyal, Siddharth and Lin, Zeming and Liu, Jason and Guo, Demi and Ott, Myle and Zitnick, C Lawrence and Ma, Jerry and others},
+  journal={Proceedings of the National Academy of Sciences},
+  volume={118},
+  number={15},
+  pages={e2016239118},
+  year={2021},
+  publisher={National Acad Sciences},
+  note={bioRxiv 10.1101/622803},
+  doi={10.1073/pnas.2016239118},
+  url={https://www.pnas.org/doi/full/10.1073/pnas.2016239118},
+}
+```
+</details>
+
+<details open><summary><b>Table of contents</b></summary>
+
+- [Main models you should use](#main-models)
+- [Usage](#usage)
+  - [Quick Start](#quickstart)
+  - [Getting Started with this repository](#repostart)
+  - [ESMFold Structure Prediction](#esmfold)
+  - [Compute embeddings in bulk from FASTA](#bulk_fasta)
+  - [CPU offloading for inference with large models](#fsdp)
+  - [Zero-shot variant prediction](#zs_variant)
+  - [Inverse folding](#invf)
+- [ESM Metagenomic Atlas](#atlas)
+- [Notebooks](#notebooks)
+- [Available Models and Datasets](#available)
+  - [Pre-trained Models](#available-models)
+  - [ESM Structural Split Dataset](#available-esmssd)
+  - [Pre-training Dataset Split](#available-pretraining-split)
+  - [Comparison to related works](#perf_related)
+- [Citations](#citations)
+- [License](#license)
+</details>
+
+<details><summary><b>What's New</b></summary>
+
+- April 2023: Code for the protein design preprints released under [examples/lm-design/](examples/lm-design/).
+- March 2023: We release an update to the ESM Metagenomic Atlas, `v2023_02`. See [website](https://esmatlas.com/) and [bulk download details](#atlas).
+- December 2022: The Meta Fundamental AI Research Protein Team (FAIR) released two simultaneous preprints on protein design:
+["Language models generalize beyond natural proteins" (Verkuil, Kabeli, et al., 2022)](https://doi.org/10.1101/2022.12.21.521521), and ["A high-level programming language for generative protein design" (Hie, Candido, et al., 2022)](https://doi.org/10.1101/2022.12.21.521521).
+- November 2022: ESM Metagenomic Atlas, a repository of 600M+ metagenomics structures released, see [website](https://esmatlas.com/) and [bulk download details](#atlas)
+- November 2022: ESMFold - new end-to-end structure prediction model released (see [Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574))
+- August 2022: ESM-2 - new SOTA Language Models released (see [Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574))
+- April 2022: New inverse folding model ESM-IF1 released, trained on CATH and UniRef50 predicted structures.
+- August 2021: Added flexibility to tokenizer to allow for spaces and special tokens (like `<mask>`) in sequence.
+- July 2021: New pre-trained model ESM-1v released, trained on UniRef90 (see [Meier et al. 2021](https://doi.org/10.1101/2021.07.09.450648)).
+- July 2021: New MSA Transformer released, with a minor fix in the row positional embeddings (`ESM-MSA-1b`).
+- Feb 2021: MSA Transformer added (see [Rao et al. 2021](https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1)). Example usage in [notebook](#notebooks).
+- Dec 2020: [Self-Attention Contacts](#notebooks) for all pre-trained models (see [Rao et al. 2020](https://doi.org/10.1101/2020.12.15.422761))
+- Dec 2020: Added new pre-trained model [ESM-1b](#perf_related) (see [Rives et al. 2019](https://doi.org/10.1101/622803) Appendix B)
+- Dec 2020: [ESM Structural Split Dataset](#available-esmssd) (see [Rives et al. 2019](https://doi.org/10.1101/622803) Appendix A.10)
+
+</details>
+
+## Main models you should use <a name="main-models"></a>
+
+| Shorthand | `esm.pretrained.`           | Dataset | Description  |
+|-----------|-----------------------------|---------|--------------|
+| ESM-2    | `esm2_t36_3B_UR50D()` `esm2_t48_15B_UR50D()`       | UR50 (sample UR90)  | SOTA general-purpose protein language model. Can be used to predict structure, function and other protein properties directly from individual sequences. Released with [Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574) (Aug 2022 update). |
+| ESMFold   | `esmfold_v1()`         | PDB + UR50 | End-to-end single sequence 3D structure predictor (Nov 2022 update). |
+| ESM-MSA-1b| `esm_msa1b_t12_100M_UR50S()` |  UR50 + MSA  | MSA Transformer language model. Can be used to extract embeddings from an MSA. Enables SOTA inference of structure. Released with [Rao et al. 2021](https://www.biorxiv.org/content/10.1101/2021.02.12.430858v2) (ICML'21 version, June 2021).  |
+| ESM-1v    | `esm1v_t33_650M_UR90S_1()` ... `esm1v_t33_650M_UR90S_5()`| UR90  | Language model specialized for prediction of variant effects. Enables SOTA zero-shot prediction of the functional effects of sequence variations. Same architecture as ESM-1b, but trained on UniRef90. Released with [Meier et al. 2021](https://doi.org/10.1101/2021.07.09.450648). |
+| ESM-IF1  | `esm_if1_gvp4_t16_142M_UR50()` | CATH + UR50 | Inverse folding model. Can be used to design sequences for given structures, or to predict functional effects of sequence variation for given structures. Enables SOTA fixed backbone sequence design. Released with [Hsu et al. 2022](https://doi.org/10.1101/2022.04.10.487779). |
+
+For a complete list of available models, with details and release notes, see [Pre-trained Models](#available-models).
+
+
+## Usage <a name="usage"></a>
+
+### Quick start <a name="quickstart"></a>
+
+An easy way to get started is to load ESM or ESMFold through the [HuggingFace transformers library](https://huggingface.co/docs/transformers/model_doc/esm),
+which has simplified the ESMFold dependencies and provides a standardized API and tools to work with state-of-the-art pretrained models.
+
+Alternatively, [ColabFold](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/ESMFold.ipynb) has integrated ESMFold so that you can 
+easily run it directly in the browser on a Google Colab instance.
+
+We also provide an API which you can access through curl or on [the ESM Metagenomic Atlas web page](https://esmatlas.com/resources?action=fold).
+```
+curl -X POST --data "KVFGRCELAAAMKRHGLDNYRGYSLGNWVCAAKFESNFNTQATNRNTDGSTDYGILQINSRWWCNDGRTPGSRNLCNIPCSALLSSDITASVNCAKKIVSDGNGMNAWVAWRNRCKGTDVQAWIRGCRL" https://api.esmatlas.com/foldSequence/v1/pdb/
+```
+
+For ESM-MSA-1b, ESM-IF1, or any of the other models you can use the original implementation from our repo directly via the instructions below.
+
+### Getting started with this repo <a name="repostart"></a>
+
+As a prerequisite, you must have PyTorch installed to use this repository.
+
+You can use this one-liner for installation, using the latest release of esm:
+
+```bash
+pip install fair-esm  # latest release, OR:
+pip install git+https://github.com/facebookresearch/esm.git  # bleeding edge, current repo main branch
+```
+
+To use the ESMFold model, make sure you start from an environment with python <= 3.9 and pytorch installed.
+Then add the `[esmfold]` option to your pip install, which will install the dependencies for OpenFold
+automatically. Openfold installation requires `nvcc`.
+
+```bash
+pip install "fair-esm[esmfold]"
+# OpenFold and its remaining dependency
+pip install 'dllogger @ git+https://github.com/NVIDIA/dllogger.git'
+pip install 'openfold @ git+https://github.com/aqlaboratory/openfold.git@4b41059694619831a7db195b7e0988fc4ff3a307'
+```
+
+**NOTE**: If openfold installation fails, please double check that `nvcc` is available and that a cuda-compatable version of PyTorch has been installed.
+
+Alternatively, we provide the `esmfold` conda environment, which can be built via `conda env create -f environment.yml`.
+
+We also support PyTorch Hub, which removes the need to clone and/or install this repository yourself:
+
+```python
+import torch
+model, alphabet = torch.hub.load("facebookresearch/esm:main", "esm2_t33_650M_UR50D")
+```
+
+After pip install, you can load and use a pretrained model as follows:
+
+```python
+import torch
+import esm
+
+# Load ESM-2 model
+model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
+batch_converter = alphabet.get_batch_converter()
+model.eval()  # disables dropout for deterministic results
+
+# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
+data = [
+    ("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),
+    ("protein2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
+    ("protein2 with mask","KALTARQQEVFDLIRD<mask>ISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
+    ("protein3",  "K A <mask> I S Q"),
+]
+batch_labels, batch_strs, batch_tokens = batch_converter(data)
+batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)
+
+# Extract per-residue representations (on CPU)
+with torch.no_grad():
+    results = model(batch_tokens, repr_layers=[33], return_contacts=True)
+token_representations = results["representations"][33]
+
+# Generate per-sequence representations via averaging
+# NOTE: token 0 is always a beginning-of-sequence token, so the first residue is token 1.
+sequence_representations = []
+for i, tokens_len in enumerate(batch_lens):
+    sequence_representations.append(token_representations[i, 1 : tokens_len - 1].mean(0))
+
+# Look at the unsupervised self-attention map contact predictions
+import matplotlib.pyplot as plt
+for (_, seq), tokens_len, attention_contacts in zip(data, batch_lens, results["contacts"]):
+    plt.matshow(attention_contacts[: tokens_len, : tokens_len])
+    plt.title(seq)
+    plt.show()
+```
+
+
+### ESMFold Structure Prediction <a name="esmfold"></a>
+
+After installing with the `[esmfold]` option, you can use the ESMFold structure prediction model as follows:
+
+```python
+import torch
+import esm
+
+model = esm.pretrained.esmfold_v1()
+model = model.eval().cuda()
+
+# Optionally, uncomment to set a chunk size for axial attention. This can help reduce memory.
+# Lower sizes will have lower memory requirements at the cost of increased speed.
+# model.set_chunk_size(128)
+
+sequence = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
+# Multimer prediction can be done with chains separated by ':'
+
+with torch.no_grad():
+    output = model.infer_pdb(sequence)
+
+with open("result.pdb", "w") as f:
+    f.write(output)
+
+import biotite.structure.io as bsio
+struct = bsio.load_structure("result.pdb", extra_fields=["b_factor"])
+print(struct.b_factor.mean())  # this will be the pLDDT
+# 88.3
+```
+
+
+Besides `esm.pretrained.esmfold_v1()` which is the best performing model we recommend using, we
+also provide `esm.pretrained.esmfold_v0()` which was used for the experiments in
+[Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574).
+
+We also provide a command line interface (`esm-fold`) that efficiently predicts structures in bulk from a FASTA file using ESMFold:
+```
+usage: esm-fold [-h] -i FASTA -o PDB [--num-recycles NUM_RECYCLES]
+                [--max-tokens-per-batch MAX_TOKENS_PER_BATCH]
+                [--chunk-size CHUNK_SIZE] [--cpu-only] [--cpu-offload]
+
+optional arguments:
+  -h, --help            show this help message and exit
+  -i FASTA, --fasta FASTA
+                        Path to input FASTA file
+  -o PDB, --pdb PDB     Path to output PDB directory
+  --num-recycles NUM_RECYCLES
+                        Number of recycles to run. Defaults to number used in
+                        training (4).
+  --max-tokens-per-batch MAX_TOKENS_PER_BATCH
+                        Maximum number of tokens per gpu forward-pass. This
+                        will group shorter sequences together for batched
+                        prediction. Lowering this can help with out of memory
+                        issues, if these occur on short sequences.
+  --chunk-size CHUNK_SIZE
+                        Chunks axial attention computation to reduce memory
+                        usage from O(L^2) to O(L). Equivalent to running a for
+                        loop over chunks of of each dimension. Lower values
+                        will result in lower memory usage at the cost of
+                        speed. Recommended values: 128, 64, 32. Default: None.
+  --cpu-only            CPU only
+  --cpu-offload         Enable CPU offloading
+```
+
+The command will make one prediction for every sequence in the fasta file. Multimers can be predicted and should be entered in the fasta file as a single sequence, with chains seprated by a ":" character.
+
+By default, predictions will be batched together so that shorter sequences are predicted simultaneously. This can be disabled by setting `--max-tokens-per-batch=0`. Batching can significantly improve prediction speed on shorter sequences.
+
+The `--cpu-offload` flag can be useful for making predictions on longer sequences. It will attempt to offload some parameters to the CPU RAM, rather than storing on GPU.
+
+Finally, the ablation experiments for LMs of varying sizes [Lin et al. 2022 table S1](https://www.science.org/doi/abs/10.1126/science.ade2574) are released as `esm.pretrained.esmfold_structure_module_only_*()`. We don't recommend using these models for structure prediction.
+
+
+### Compute embeddings in bulk from FASTA <a name="bulk_fasta"></a>
+
+We provide a command line interface (`esm-extract`) that efficiently extracts embeddings in bulk for a FASTA file from the ESM:
+```
+usage: esm-extract [-h] [--toks_per_batch TOKS_PER_BATCH]
+                   [--repr_layers REPR_LAYERS [REPR_LAYERS ...]] --include
+                   {mean,per_tok,bos,contacts}
+                   [{mean,per_tok,bos,contacts} ...]
+                   [--truncation_seq_length TRUNCATION_SEQ_LENGTH]
+                   model_location fasta_file output_dir
+
+Extract per-token representations and model outputs for sequences in a FASTA
+file
+
+positional arguments:
+  model_location        PyTorch model file OR name of pretrained model to
+                        download (see README for models)
+  fasta_file            FASTA file on which to extract representations
+  output_dir            output directory for extracted representations
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --toks_per_batch TOKS_PER_BATCH
+                        maximum batch size
+  --repr_layers REPR_LAYERS [REPR_LAYERS ...]
+                        layers indices from which to extract representations
+                        (0 to num_layers, inclusive)
+  --include {mean,per_tok,bos,contacts} [{mean,per_tok,bos,contacts} ...]
+                        specify which representations to return
+  --truncation_seq_length TRUNCATION_SEQ_LENGTH
+                        truncate sequences longer than the given value
+```
+
+The following commands allow the extraction of the final-layer embedding for a FASTA file from the ESM-2 model:
+
+```bash
+esm-extract esm2_t33_650M_UR50D examples/data/some_proteins.fasta \
+  examples/data/some_proteins_emb_esm2 --repr_layers 0 32 33 --include
+```
+```bash
+python scripts/extract.py esm2_t33_650M_UR50D examples/data/some_proteins.fasta \
+  examples/data/some_proteins_emb_esm2 --repr_layers 0 32 33 --include mean per_tok
+```
+
+A cuda device is optional and will be auto-detected.
+
+Directory `some_proteins_emb_esm2/` now contains one `.pt` file per FASTA sequence; use `torch.load()` to load them.
+`scripts/extract.py` has flags that determine what's included in the `.pt` file:
+* `--repr-layers` (default: final only) selects which layers to include embeddings from.
+* `--include` specifies what embeddings to save. You can use the following:
+  * `per_tok` includes the full sequence, with an embedding per amino acid (seq_len x hidden_dim).
+  * `mean` includes the embeddings averaged over the full sequence, per layer.
+  * `bos` includes the embeddings from the beginning-of-sequence token.
+  (NOTE: Don't use with the pre-trained models - we trained without bos-token supervision)
+
+
+### CPU offloading for inference with large models <a name="fsdp"></a>
+If you want to load very large models like 15B and/or do inference on long sequences on your machine, regular GPU inference may lead to OOM errors.
+We show how to load the model with Fairscale's [Fully Sharded Data Parallel (FSDP)](https://fairscale.readthedocs.io/en/stable/api/nn/fsdp.html) and
+use its CPU offloading feature.
+This allows to do inference of large models on a single GPU.
+Please check out `examples/esm2_infer_fairscale_fsdp_cpu_offloading.py` for more details.
+
+### Zero-shot variant prediction <a name="zs_variant"></a>
+See "[examples/variant-prediction/](examples/variant-prediction/)" for code and pre-trained weights for the ESM-1v models described in
+[Language models enable zero-shot prediction of the effects of mutations on protein function. (Meier et al. 2021)](https://doi.org/10.1101/2021.07.09.450648).
+
+Note that ESM-2 could be used for variant prediction as well, and is expected to have similar performance to ESM-1v.
+
+### Inverse folding <a name="invf"></a>
+See "[examples/inverse_folding/](examples/inverse_folding/)" for detailed user guide. The ESM-IF1 model is described as `GVPTransformer` in [Learning inverse folding from millions of predicted structures. (Hsu et al. 2022)](https://doi.org/10.1101/2022.04.10.487779).
+
+We also provide a colab notebook for the sequence design and sequence scoring functionalities.
+
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/inverse_folding/notebook_multichain.ipynb)
+
+The ESM-IF1 inverse folding model is built for predicting protein sequences
+from their backbone atom coordinates. We provide scripts here 1) to sample sequence
+designs for a given structure and 2) to score sequences for a given structure.
+
+Trained with 12M protein structures predicted by AlphaFold2, the ESM-IF1
+model consists of invariant geometric input processing layers followed by a
+sequence-to-sequence transformer, and achieves 51% native sequence recovery on
+structurally held-out backbones with 72% recovery for buried residues.
+The model is also trained with span masking to tolerate missing backbone
+coordinates and therefore can predict sequences for partially masked structures.
+
+#### Sample sequence designs for a given structure
+The environment setup is described in [this subsection of examples/inverse_folding](examples/inverse_folding#recommended-environment).
+
+To sample sequences for a given structure in PDB or mmCIF format, use the
+`sample_sequences.py` script. The input file can have either `.pdb` or
+`.cif` as suffix.
+
+For example, to sample 3 sequence designs for the golgi casein kinase structure
+(PDB [5YH2](https://www.rcsb.org/structure/5yh2); [PDB Molecule of the Month
+from January 2022](https://pdb101.rcsb.org/motm/265)), we can run the following
+command from the esm root directory:
+```bash
+python examples/inverse_folding/sample_sequences.py examples/inverse_folding/data/5YH2.pdb \
+  --chain C --temperature 1 --num-samples 3 --outpath examples/inverse_folding/output/sampled_sequences.fasta
+```
+
+The sampled sequences will be saved in a fasta format to the specified output file.
+
+The temperature parameter controls the sharpness of the probability
+distribution for sequence sampling. Higher sampling temperatures yield more
+diverse sequences but likely with lower native sequence recovery.
+The default sampling temperature is 1. To optimize for native sequence
+recovery, we recommend sampling with low temperature such as 1e-6.
+
+#### Scoring sequences
+To score the conditional log-likelihoods for sequences conditioned on a given
+structure, use the `score_log_likelihoods.py` script.
+
+For example, to score the sequences in `examples/inverse_folding/data/5YH2_mutated_seqs.fasta`
+according to the structure in `examples/inverse_folding/data/5YH2.pdb`, we can run
+the following command from the esm root directory:
+```
+python examples/inverse_folding/score_log_likelihoods.py examples/inverse_folding/data/5YH2.pdb \
+  examples/inverse_folding/data/5YH2_mutated_seqs.fasta --chain C \
+  --outpath examples/inverse_folding/output/5YH2_mutated_seqs_scores.csv
+```
+
+The conditional log-likelihoods are saved in a csv format in the specified output path.
+The output values are the average log-likelihoods averaged over all amino acids in a sequence.
+
+For more information, see "[./examples/inverse_folding/](examples/inverse_folding/)" for detailed user guide.
+
+## ESM Metagenomic Atlas <a name="atlas"></a>
+
+Please visit the [ESM Metagenomic Atlas](https://esmatlas.com/) website, and
+see our [blog post](https://ai.facebook.com/blog/protein-folding-esmfold-metagenomics/) to learn more.
+
+Bulk download instructions available at a seperate README [here](scripts/atlas/README.md).
+
+The Atlas resources include a page to [fold a sequence using ESMFold](https://esmatlas.com/resources?action=fold),
+searching a subset of the ESM Atlas by [structure](https://esmatlas.com/resources?action=search_structure) or 
+[sequence](https://esmatlas.com/resources?action=search_sequence),
+as well as an [API](https://esmatlas.com/about#api) to access those resources programmatically.
+
+Foldseek provides search against the Atlas without the length limitation [here](https://search.foldseek.com/search).
+
+
+## Notebooks <a name="notebooks"></a>
+
+### Inverse folding - predicting or scoring sequences based on backbone structures
+
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/inverse_folding/notebook.ipynb)
+
+The ESM-IF1 inverse folding model predicts protein sequences from their backbone atom coordinates, trained with 12M protein structures predicted by AlphaFold2.
+This notetook guide you through examples of sampling sequences, calculating conditional log-likelihoods, and extracting encoder output as structure representation.
+
+### Supervised variant prediction - training a classifier on the embeddings
+
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/sup_variant_prediction.ipynb)
+
+
+To help you get started with using the embeddings, this [jupyter notebook tutorial](examples/sup_variant_prediction.ipynb) shows how to train a supervised variant predictor using embeddings from ESM-1.
+You can adopt a similar protocol to train a model for any downstream task, even with limited data.
+First you can obtain the embeddings for ``examples/data/P62593.fasta`` either by [downloading the precomputed](https://dl.fbaipublicfiles.com/fair-esm/examples/P62593_reprs.tar.gz) embeddings
+as instructed in the notebook or by running the following:
+
+```bash
+# Obtain the embeddings
+python scripts/extract.py esm1v_t33_650M_UR90S_1 examples/data/P62593.fasta \
+  examples/data/P62593_emb_esm1v --repr_layers 33 --include mean
+```
+
+Then, follow the remaining instructions in the tutorial. You can also run the tutorial in a [colab notebook](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/sup_variant_prediction.ipynb).
+
+**Note, alternatively use [the newer instructions for zero-shot variant prediction](examples/variant-prediction/),
+which predicts mutational effects without any supervised training.**
+
+
+### Unsupervised contact prediction
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb)
+
+This [jupyter notebook tutorial](examples/contact_prediction.ipynb) demonstrates contact prediction with both the ESM-2 and MSA Transformer (ESM-MSA-1) models.
+Contact prediction is based on a logistic regression over the model's attention maps.
+This methodology is based on our ICLR 2021 paper,
+[Transformer protein language models are unsupervised structure learners. (Rao et al. 2020)](https://doi.org/10.1101/2020.12.15.422761)
+The MSA Transformer (ESM-MSA-1) takes a multiple sequence alignment (MSA) as input, and uses the tied row self-attention maps in the same way.
+See [MSA Transformer. (Rao et al. 2021)](https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1).
+
+To get unsupervised attention-based contacts, call `model.predict_contacts(tokens)` or `model(tokens, return_contacts=True)`.
+
+
+### ESMStructuralSplitDataset and self-attention contact prediction
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/esm_structural_dataset.ipynb)
+
+And this [jupyter notebook tutorial](examples/esm_structural_dataset.ipynb) shows how to load and index the `ESMStructuralSplitDataset`,
+and computes the self-attention map unsupervised contact predictions using ESM-2.
+
+
+## Available Models and Datasets <a name="available"></a>
+
+### Pre-trained Models <a name="available-models"></a>
+
+| Shorthand | `esm.pretrained.`           | #layers | #params | Dataset | Embedding Dim |  Model URL (automatically downloaded to `~/.cache/torch/hub/checkpoints`) |
+|-----------|---------------------|---------|-------------|---------|---------------|-----------------------------------------------------------------------|
+| ESM-2     | `esm2_t48_15B_UR50D`         | 48           | 15B         | UR50/D 2021_04                           | 5120 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t48_15B_UR50D.pt          |
+|           | `esm2_t36_3B_UR50D`          | 36           | 3B          | UR50/D 2021_04                           | 2560 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt           |
+|           | `esm2_t33_650M_UR50D`        | 33           | 650M        | UR50/D 2021_04                           | 1280 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt         |
+|           | `esm2_t30_150M_UR50D`        | 30           | 150M        | UR50/D 2021_04                           | 640  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t30_150M_UR50D.pt         |
+|           | `esm2_t12_35M_UR50D`         | 12           | 35M         | UR50/D 2021_04                           | 480  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t12_35M_UR50D.pt          |
+|           | `esm2_t6_8M_UR50D`           | 6            | 8M          | UR50/D 2021_04                           | 320  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t6_8M_UR50D.pt            |
+| ESMFold   | `esmfold_v1`                 | 48 (+36)     | 690M (+3B)  | UR50/D 2021_04                           | -    |  https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt               |
+|           | `esmfold_v0`                 | 48 (+36)     | 690M (+3B)  | UR50/D 2021_04                           | -    |  https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v0.pt               |
+|           | `esmfold_structure_module_only_*`              | 0 (+various) | various     | UR50/D 2021_04                           | -    |  https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_structure_module_only_*                  |
+| ESM-IF1   | `esm_if1_gvp4_t16_142M_UR50` | 20           | 124M        | CATH 4.3 + predicted structures for UR50 | 512  | https://dl.fbaipublicfiles.com/fair-esm/models/esm_if1_gvp4_t16_142M_UR50.pt   |
+| ESM-1v    | `esm1v_t33_650M_UR90S_[1-5]` | 33           | 650M        | UR90/S 2020_03                           | 1280 | https://dl.fbaipublicfiles.com/fair-esm/models/esm1v_t33_650M_UR90S_1.pt       |
+| ESM-MSA-1b| `esm_msa1b_t12_100M_UR50S`   | 12           | 100M        | UR50/S + MSA 2018_03                     | 768  | https://dl.fbaipublicfiles.com/fair-esm/models/esm_msa1b_t12_100M_UR50S.pt     |
+| ESM-MSA-1 | `esm_msa1_t12_100M_UR50S`    | 12           | 100M        | UR50/S + MSA 2018_03                     | 768  | https://dl.fbaipublicfiles.com/fair-esm/models/esm_msa1_t12_100M_UR50S.pt      |
+| ESM-1b    | `esm1b_t33_650M_UR50S`       | 33           | 650M        | UR50/S 2018_03                           | 1280 | https://dl.fbaipublicfiles.com/fair-esm/models/esm1b_t33_650M_UR50S.pt         |
+| ESM-1     | `esm1_t34_670M_UR50S`        | 34           | 670M        | UR50/S 2018_03                           | 1280 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t34_670M_UR50S.pt         |
+|           | `esm1_t34_670M_UR50D`        | 34           | 670M        | UR50/D 2018_03                           | 1280 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t34_670M_UR50D.pt         |
+|           | `esm1_t34_670M_UR100`        | 34           | 670M        | UR100 2018_03                            | 1280 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t34_670M_UR100.pt         |
+|           | `esm1_t12_85M_UR50S`         | 12           | 85M         | UR50/S 2018_03                           | 768  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t12_85M_UR50S.pt          |
+|           | `esm1_t6_43M_UR50S`          | 6            | 43M         | UR50/S 2018_03                           | 768  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t6_43M_UR50S.pt           |
+
+
+Here is a chronological list of the released models and the paper they were introduced in:
+
+| Shorthand  | Release Notes |
+|------------|---------------|
+| ESM-1      | Released with Rives et al. 2019 (Aug 2020 update). |
+| ESM-1b     | Released with Rives et al. 2019 (Dec 2020 update). See Appendix B. |
+| ESM-MSA-1  | Released with Rao et al. 2021 (Preprint v1). |
+| ESM-MSA-1b | Released with Rao et al. 2021 (ICML'21 version, June 2021). |
+| ESM-1v     | Released with Meier et al. 2021. |
+| ESM-IF1    | Released with Hsu et al. 2022. |
+| ESM-2      | Released with Lin et al. 2022. |
+
+### ESM Structural Split Dataset <a name="available-esmssd"></a>
+This is a five-fold cross validation dataset of protein domain structures that can be used to measure generalization of representations
+across different levels of structural dissimilarity.
+The dataset implements structural holdouts at the family, superfamily, and fold
+level. The SCOPe database is used to classify domains. Independently for each level of structural hold-out,
+the domains are split into 5 equal sets, i.e. five sets of folds, superfamilies, or families. This ensures
+that for each of the five partitions, structures having the same classification do not appear in both the
+train and test sets. For a given classification level each structure appears in a test set once, so that
+in the cross validation experiment each of the structures will be evaluated exactly once.
+
+The dataset provides 3d coordinates, distance maps, and secondary structure labels.
+For further details on the construction of the dataset
+see [Rives et al. 2019](https://doi.org/10.1101/622803) Appendix A.10.
+
+This [jupyter notebook tutorial](examples/esm_structural_dataset.ipynb) shows how to load and index the `ESMStructuralSplitDataset`.
+
+`ESMStructuralSplitDataset`, upon initializing, will download `splits` and `pkl`.
+We also provide `msas` for each of the domains. The data can be directly downloaded below.
+
+| Name   | Description                                                                   | URL                                                                   |
+|--------|-------------------------------------------------------------------------------|-----------------------------------------------------------------------|
+| splits | train/valid splits                                                            | https://dl.fbaipublicfiles.com/fair-esm/structural-data/splits.tar.gz |
+| pkl    | pkl objects containing sequence, SSP labels, distance map, and 3d coordinates | https://dl.fbaipublicfiles.com/fair-esm/structural-data/pkl.tar.gz    |
+| msas   | a3m files containing MSA for each domain                                      | https://dl.fbaipublicfiles.com/fair-esm/structural-data/msas.tar.gz   |
+
+### Pre-training Dataset Split  <a name="available-pretraining-split"></a>
+The split files establishing which UniRef50 clusters were used as held-out evaluation set for pre-training
+in [Rives et al. 2019](https://doi.org/10.1101/622803) and [Rao et al. 2021](https://doi.org/10.1101/2021.02.12.430858) can be found here:
+* [UniRef50 IDs of evaluation set](https://dl.fbaipublicfiles.com/fair-esm/pretraining-data/uniref201803_ur50_valid_headers.txt.gz): 3.016 M clusters
+* [UniRef100 IDs of evaluation set](https://dl.fbaipublicfiles.com/fair-esm/pretraining-data/uniref201803_ur100_valid_headers.txt.gz): 13.745 M proteins, expanding the same UniRef50 clusters.
+
+These files only contain only the UniRef50 IDs and UniRef100 IDs corresponding to the [UniRef database, 2018-03 release](https://ftp.uniprot.org/pub/databases/uniprot/previous_releases/release-2018_03/uniref/)
+which is released by the UniProt Consortium under a [Creative Commons Attribution (CC BY 4.0) License](https://www.uniprot.org/help/license).
+
+
+### Comparison to related works <a name="perf_related"></a>
+<!--
+DO NOT EDIT THIS TABLE! This is the source of truth:
+https://docs.google.com/spreadsheets/d/1RPvWF47rIMEr-Jg-SRCoGElHcwCl5d7RyEeSyPgp59A/edit#gid=0
+exported via https://www.tablesgenerator.com/html_tables
+-->
+
+<table class="tg">
+<thead>
+  <tr>
+    <th class="tg-0thz"><span style="font-weight:bold">Task</span></th>
+    <th class="tg-j6zm" colspan="3"><span style="font-weight:bold">Unsupervised contact prediction</span></th>
+    <th class="tg-j6zm" colspan="2"><span style="font-weight:bold">Structure Prediction</span></th>
+  </tr>
+</thead>
+<tbody>
+  <tr>
+    <td class="tg-j6zm"><span style="font-weight:bold">Test set</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">Large valid</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">CASP14</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">CAMEO (Apr-Jun 2022)</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">CASP14</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">CAMEO (Apr-Jun 2022)</span></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">Gremlin (Potts)</td>
+    <td class="tg-7zrl">39.3</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">TAPE</td>
+    <td class="tg-7zrl">11.2</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ProtBert-BFD</td>
+    <td class="tg-7zrl">34.1</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">Prot-T5-XL-BFD</td>
+    <td class="tg-7zrl">35.6</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-2b7s">46.1</td>
+    <td class="tg-2b7s">62.6</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">Prot-T5-XL-Ur50 (3B)</td>
+    <td class="tg-7zrl">47.9</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-2b7s">49.8</td>
+    <td class="tg-2b7s">69.4</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-1</td>
+    <td class="tg-7zrl">33.7</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-1b</td>
+    <td class="tg-7zrl">41.1</td>
+    <td class="tg-7zrl">24.4</td>
+    <td class="tg-7zrl">39</td>
+    <td class="tg-2b7s">41.6</td>
+    <td class="tg-2b7s">64.5</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-1v</td>
+    <td class="tg-7zrl">35.3</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-MSA-1b</td>
+    <td class="tg-7zrl">57.4</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (8M)</td>
+    <td class="tg-7zrl">15.9</td>
+    <td class="tg-7zrl">9.8</td>
+    <td class="tg-7zrl">15.7</td>
+    <td class="tg-2b7s">36.7</td>
+    <td class="tg-2b7s">48.1</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (35M)</td>
+    <td class="tg-7zrl">28.8</td>
+    <td class="tg-7zrl">16.4</td>
+    <td class="tg-7zrl">28.4</td>
+    <td class="tg-2b7s">41.4</td>
+    <td class="tg-2b7s">56.4</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (150M)</td>
+    <td class="tg-7zrl">42.2</td>
+    <td class="tg-7zrl">26.8</td>
+    <td class="tg-7zrl">40.1</td>
+    <td class="tg-2b7s">49.0</td>
+    <td class="tg-2b7s">64.9</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (700M)</td>
+    <td class="tg-7zrl">50.1</td>
+    <td class="tg-7zrl">32.5</td>
+    <td class="tg-7zrl">47.6</td>
+    <td class="tg-2b7s">51.3</td>
+    <td class="tg-2b7s">70.1</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (3B)</td>
+    <td class="tg-7zrl">52.7</td>
+    <td class="tg-7zrl">34.0</td>
+    <td class="tg-7zrl">49.9</td>
+    <td class="tg-2b7s">52.5</td>
+    <td class="tg-2b7s">71.8</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (15B)</td>
+    <td class="tg-7zrl">54.5</td>
+    <td class="tg-7zrl">37.0</td>
+    <td class="tg-7zrl">51.7</td>
+    <td class="tg-2b7s">55.4</td>
+    <td class="tg-2b7s">72.1</td>
+  </tr>
+</tbody>
+</table>
+
+Comparison to related protein language models on structure prediction tasks.
+
+* All contact numbers are the top-L,LR precision metric, where long range means sequence separation of at least 24 residues
+* For unsupervised contact prediction, a sparse linear combination of the attention heads is used to directly predict protein contacts,
+fitted with logistic regression on 20 structures.
+For more details on the method, see [Rao et al. 2020](https://doi.org/10.1101/2020.12.15.422761).
+* For structure prediction, an AlphaFold2 structure module is trained directly from the frozen language model embeddings.
+For more details on the method, see [Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574).
+* Direct coupling analysis methods (Gremlin, mfDCA, Psicov) and ESM-MSA-1 use the [trRosetta MSAs](https://yanglab.nankai.edu.cn/trRosetta/benchmark/), while other methods predict from single sequence.
+
+
+## Citations <a name="citations"></a>
+
+If you find the models useful in your research, we ask that you cite the relevant paper:
+
+```bibtex
+@article{rives2019biological,
+  author={Rives, Alexander and Meier, Joshua and Sercu, Tom and Goyal, Siddharth and Lin, Zeming and Liu, Jason and Guo, Demi and Ott, Myle and Zitnick, C. Lawrence and Ma, Jerry and Fergus, Rob},
+  title={Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences},
+  year={2019},
+  doi={10.1101/622803},
+  url={https://www.biorxiv.org/content/10.1101/622803v4},
+  journal={PNAS}
+}
+```
+
+For the self-attention contact prediction:
+
+```bibtex
+@article{rao2020transformer,
+  author = {Rao, Roshan M and Meier, Joshua and Sercu, Tom and Ovchinnikov, Sergey and Rives, Alexander},
+  title={Transformer protein language models are unsupervised structure learners},
+  year={2020},
+  doi={10.1101/2020.12.15.422761},
+  url={https://www.biorxiv.org/content/10.1101/2020.12.15.422761v1},
+  journal={bioRxiv}
+}
+```
+
+For the MSA Transformer:
+
+```bibtex
+@article{rao2021msa,
+  author = {Rao, Roshan and Liu, Jason and Verkuil, Robert and Meier, Joshua and Canny, John F. and Abbeel, Pieter and Sercu, Tom and Rives, Alexander},
+  title={MSA Transformer},
+  year={2021},
+  doi={10.1101/2021.02.12.430858},
+  url={https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1},
+  journal={bioRxiv}
+}
+```
+
+For variant prediction using ESM-1v:
+
+```bibtex
+@article{meier2021language,
+  author = {Meier, Joshua and Rao, Roshan and Verkuil, Robert and Liu, Jason and Sercu, Tom and Rives, Alexander},
+  title = {Language models enable zero-shot prediction of the effects of mutations on protein function},
+  year={2021},
+  doi={10.1101/2021.07.09.450648},
+  url={https://www.biorxiv.org/content/10.1101/2021.07.09.450648v1},
+  journal={bioRxiv}
+}
+```
+
+For inverse folding using ESM-IF1:
+
+```bibtex
+@article{hsu2022learning,
+	author = {Hsu, Chloe and Verkuil, Robert and Liu, Jason and Lin, Zeming and Hie, Brian and Sercu, Tom and Lerer, Adam and Rives, Alexander},
+	title = {Learning inverse folding from millions of predicted structures},
+	year = {2022},
+	doi = {10.1101/2022.04.10.487779},
+	url = {https://www.biorxiv.org/content/early/2022/04/10/2022.04.10.487779},
+	journal = {ICML}
+}
+```
+
+For the ESM-2 language model and ESMFold:
+
+```bibtex
+@article{lin2022language,
+  title={Language models of protein sequences at the scale of evolution enable accurate structure prediction},
+  author={Lin, Zeming and Akin, Halil and Rao, Roshan and Hie, Brian and Zhu, Zhongkai and Lu, Wenting and Smetanin, Nikita and dos Santos Costa, Allan and Fazel-Zarandi, Maryam and Sercu, Tom and Candido, Sal and others},
+  journal={bioRxiv},
+  year={2022},
+  publisher={Cold Spring Harbor Laboratory}
+}
+```
+
+Much of this code builds on the [fairseq](https://github.com/pytorch/fairseq) sequence modeling framework. We use fairseq internally for our protein language modeling research. We highly recommend trying it out if you'd like to pre-train protein language models from scratch.
+
+Additionally, if you would like to use the variant prediction benchmark from Meier et al. (2021), we provide a bibtex file with citations for all data in [./examples/variant-prediction/mutation_data.bib](./examples/variant-prediction/mutation_data.bib). You can cite each paper individually, or add all citations in bulk using the LaTeX command:
+
+```tex
+\nocite{wrenbeck2017deep,klesmith2015comprehensive,haddox2018mapping,romero2015dissecting,firnberg2014comprehensive,deng2012deep,stiffler2015evolvability,jacquier2013capturing,findlay2018comprehensive,mclaughlin2012spatial,kitzman2015massively,doud2016accurate,pokusaeva2019experimental,mishra2016systematic,kelsic2016rna,melnikov2014comprehensive,brenan2016phenotypic,rockah2015systematic,wu2015functional,aakre2015evolving,qi2014quantitative,matreyek2018multiplex,bandaru2017deconstruction,roscoe2013analyses,roscoe2014systematic,mavor2016determination,chan2017correlation,melamed2013deep,starita2013activity,araya2012fundamental}
+```
+
+## License <a name="license"></a>
+
+This source code is licensed under the MIT license found in the `LICENSE` file
+in the root directory of this source tree.
+
+ESM Metagenomic Atlas (also referred to as “ESM Metagenomic Structure Atlas” or “ESM Atlas”) data is available under a CC BY 4.0 license for academic and commercial use. Copyright (c) Meta Platforms, Inc. All Rights Reserved. Use of the ESM Metagenomic Atlas data is subject to the Meta Open Source [Terms of Use](https://opensource.fb.com/legal/terms/) and [Privacy Policy](https://opensource.fb.com/legal/privacy/).

esm/source/__init__.py

@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+"""
+esm 项目包初始化文件
+"""

esm/source/environment.yml

@@ -0,0 +1,36 @@
+name: esmfold
+channels:
+  - conda-forge
+  - bioconda
+  - pytorch
+dependencies:
+  - conda-forge::python=3.7
+  - conda-forge::setuptools=59.5.0
+  - conda-forge::pip
+  - conda-forge::openmm=7.5.1
+  - conda-forge::pdbfixer
+  - conda-forge::cudatoolkit==11.3.*
+  - conda-forge::einops
+  - conda-forge::fairscale
+  - conda-forge::omegaconf
+  - conda-forge::hydra-core
+  - conda-forge::pandas
+  - conda-forge::pytest
+  - bioconda::hmmer==3.3.2
+  - bioconda::hhsuite==3.3.0
+  - bioconda::kalign2==2.04
+  - pytorch::pytorch=1.12.*
+  - pip:
+      - biopython==1.79
+      - deepspeed==0.5.9
+      - dm-tree==0.1.6
+      - ml-collections==0.1.0
+      - numpy==1.21.2
+      - PyYAML==5.4.1
+      - requests==2.26.0
+      - scipy==1.7.1
+      - tqdm==4.62.2
+      - typing-extensions==3.10.0.2
+      - pytorch_lightning==1.5.10
+      - wandb==0.12.21
+      - git+https://github.com/NVIDIA/dllogger.git

esm/source/esm/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .version import version as __version__  # noqa
+
+from .data import Alphabet, BatchConverter, FastaBatchedDataset  # noqa
+from .model.esm1 import ProteinBertModel  # noqa
+from .model.esm2 import ESM2  # noqa
+from .model.msa_transformer import MSATransformer  #noqa
+from . import pretrained  # noqa

esm/source/esm/axial_attention.py

@@ -0,0 +1,239 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import torch
+import torch.nn as nn
+
+
+class RowSelfAttention(nn.Module):
+    """Compute self-attention over rows of a 2D input."""
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        max_tokens_per_msa: int = 2 ** 16,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.scaling = self.head_dim ** -0.5
+        self.max_tokens_per_msa = max_tokens_per_msa
+        self.attn_shape = "hnij"
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout_module = nn.Dropout(dropout)
+
+    def align_scaling(self, q):
+        num_rows = q.size(0)
+        return self.scaling / math.sqrt(num_rows)
+
+    def _batched_forward(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        max_rows = max(1, self.max_tokens_per_msa // num_cols)
+        attns = 0
+        scaling = self.align_scaling(x)
+        for start in range(0, num_rows, max_rows):
+            attn_weights = self.compute_attention_weights(
+                x[start : start + max_rows],
+                scaling,
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask[:, start : start + max_rows]
+                if self_attn_padding_mask is not None
+                else None,
+            )
+            attns += attn_weights
+        attn_probs = attns.softmax(-1)
+        attn_probs = self.dropout_module(attn_probs)
+
+        outputs = []
+        for start in range(0, num_rows, max_rows):
+            output = self.compute_attention_update(x[start : start + max_rows], attn_probs)
+            outputs.append(output)
+
+        output = torch.cat(outputs, 0)
+        return output, attn_probs
+
+    def compute_attention_weights(
+        self,
+        x,
+        scaling: float,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        q = self.q_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+        k = self.k_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+        q *= scaling
+        if self_attn_padding_mask is not None:
+            # Zero out any padded aligned positions - this is important since
+            # we take a sum across the alignment axis.
+            q *= 1 - self_attn_padding_mask.permute(1, 2, 0).unsqueeze(3).unsqueeze(4).to(q)
+
+        attn_weights = torch.einsum(f"rinhd,rjnhd->{self.attn_shape}", q, k)
+
+        if self_attn_mask is not None:
+            raise NotImplementedError
+            # Mask Size: [B x R x C], Weights Size: [H x B x C x C]
+
+        if self_attn_padding_mask is not None:
+            attn_weights = attn_weights.masked_fill(
+                self_attn_padding_mask[:, 0].unsqueeze(0).unsqueeze(2),
+                -10000,
+            )
+
+        return attn_weights
+
+    def compute_attention_update(
+        self,
+        x,
+        attn_probs,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        v = self.v_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+        context = torch.einsum(f"{self.attn_shape},rjnhd->rinhd", attn_probs, v)
+        context = context.contiguous().view(num_rows, num_cols, batch_size, embed_dim)
+        output = self.out_proj(context)
+        return output
+
+    def forward(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        if (num_rows * num_cols > self.max_tokens_per_msa) and not torch.is_grad_enabled():
+            return self._batched_forward(x, self_attn_mask, self_attn_padding_mask)
+        else:
+            scaling = self.align_scaling(x)
+            attn_weights = self.compute_attention_weights(
+                x, scaling, self_attn_mask, self_attn_padding_mask
+            )
+            attn_probs = attn_weights.softmax(-1)
+            attn_probs = self.dropout_module(attn_probs)
+            output = self.compute_attention_update(x, attn_probs)
+            return output, attn_probs
+
+
+class ColumnSelfAttention(nn.Module):
+    """Compute self-attention over columns of a 2D input."""
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        max_tokens_per_msa: int = 2 ** 16,
+    ):
+        super().__init__()
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.scaling = self.head_dim ** -0.5
+        self.max_tokens_per_msa = max_tokens_per_msa
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout_module = nn.Dropout(dropout)
+
+    def _batched_forward(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        max_cols = max(1, self.max_tokens_per_msa // num_rows)
+        outputs = []
+        attns = []
+        for start in range(0, num_cols, max_cols):
+            output, attn = self(
+                x[:, start : start + max_cols],
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask[:, :, start : start + max_cols]
+                if self_attn_padding_mask is not None
+                else None,
+            )
+            outputs.append(output)
+            attns.append(attn)
+        output = torch.cat(outputs, 1)
+        attns = torch.cat(attns, 1)
+        return output, attns
+
+    def compute_attention_update(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        if num_rows == 1:
+            # if there is only 1 position, this is equivalent and doesn't break with padding
+            attn_probs = torch.ones(
+                self.num_heads,
+                num_cols,
+                batch_size,
+                num_rows,
+                num_rows,
+                device=x.device,
+                dtype=x.dtype,
+            )
+            output = self.out_proj(self.v_proj(x))
+        else:
+            q = self.q_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+            k = self.k_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+            v = self.v_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+            q *= self.scaling
+
+            attn_weights = torch.einsum("icnhd,jcnhd->hcnij", q, k)
+
+            if self_attn_mask is not None:
+                raise NotImplementedError
+            if self_attn_padding_mask is not None:
+                attn_weights = attn_weights.masked_fill(
+                    self_attn_padding_mask.permute(2, 0, 1).unsqueeze(0).unsqueeze(3),
+                    -10000,
+                )
+
+            attn_probs = attn_weights.softmax(-1)
+            attn_probs = self.dropout_module(attn_probs)
+            context = torch.einsum("hcnij,jcnhd->icnhd", attn_probs, v)
+            context = context.contiguous().view(num_rows, num_cols, batch_size, embed_dim)
+            output = self.out_proj(context)
+        return output, attn_probs
+
+    def forward(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        # if False and num_rows * num_cols > 2 ** 14 and not torch.is_grad_enabled():
+        if (num_rows * num_cols) > self.max_tokens_per_msa and not torch.is_grad_enabled():
+            return self._batched_forward(
+                x,
+                self_attn_mask,
+                self_attn_padding_mask,
+            )
+        else:
+            return self.compute_attention_update(x, self_attn_mask, self_attn_padding_mask)

esm/source/esm/constants.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# fmt: off
+proteinseq_toks = {
+    'toks': ['L', 'A', 'G', 'V', 'S', 'E', 'R', 'T', 'I', 'D', 'P', 'K', 'Q', 'N', 'F', 'Y', 'M', 'H', 'W', 'C', 'X', 'B', 'U', 'Z', 'O', '.', '-']
+}
+# fmt: on

esm/source/esm/data.py

@@ -0,0 +1,493 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import itertools
+import os
+from typing import Sequence, Tuple, List, Union
+import pickle
+import re
+import shutil
+import torch
+from pathlib import Path
+from esm.constants import proteinseq_toks
+
+RawMSA = Sequence[Tuple[str, str]]
+
+
+class FastaBatchedDataset(object):
+    def __init__(self, sequence_labels, sequence_strs):
+        self.sequence_labels = list(sequence_labels)
+        self.sequence_strs = list(sequence_strs)
+
+    @classmethod
+    def from_file(cls, fasta_file):
+        sequence_labels, sequence_strs = [], []
+        cur_seq_label = None
+        buf = []
+
+        def _flush_current_seq():
+            nonlocal cur_seq_label, buf
+            if cur_seq_label is None:
+                return
+            sequence_labels.append(cur_seq_label)
+            sequence_strs.append("".join(buf))
+            cur_seq_label = None
+            buf = []
+
+        with open(fasta_file, "r") as infile:
+            for line_idx, line in enumerate(infile):
+                if line.startswith(">"):  # label line
+                    _flush_current_seq()
+                    line = line[1:].strip()
+                    if len(line) > 0:
+                        cur_seq_label = line
+                    else:
+                        cur_seq_label = f"seqnum{line_idx:09d}"
+                else:  # sequence line
+                    buf.append(line.strip())
+
+        _flush_current_seq()
+
+        assert len(set(sequence_labels)) == len(
+            sequence_labels
+        ), "Found duplicate sequence labels"
+
+        return cls(sequence_labels, sequence_strs)
+
+    def __len__(self):
+        return len(self.sequence_labels)
+
+    def __getitem__(self, idx):
+        return self.sequence_labels[idx], self.sequence_strs[idx]
+
+    def get_batch_indices(self, toks_per_batch, extra_toks_per_seq=0):
+        sizes = [(len(s), i) for i, s in enumerate(self.sequence_strs)]
+        sizes.sort()
+        batches = []
+        buf = []
+        max_len = 0
+
+        def _flush_current_buf():
+            nonlocal max_len, buf
+            if len(buf) == 0:
+                return
+            batches.append(buf)
+            buf = []
+            max_len = 0
+
+        for sz, i in sizes:
+            sz += extra_toks_per_seq
+            if max(sz, max_len) * (len(buf) + 1) > toks_per_batch:
+                _flush_current_buf()
+            max_len = max(max_len, sz)
+            buf.append(i)
+
+        _flush_current_buf()
+        return batches
+
+
+class Alphabet(object):
+    def __init__(
+        self,
+        standard_toks: Sequence[str],
+        prepend_toks: Sequence[str] = ("<null_0>", "<pad>", "<eos>", "<unk>"),
+        append_toks: Sequence[str] = ("<cls>", "<mask>", "<sep>"),
+        prepend_bos: bool = True,
+        append_eos: bool = False,
+        use_msa: bool = False,
+    ):
+        self.standard_toks = list(standard_toks)
+        self.prepend_toks = list(prepend_toks)
+        self.append_toks = list(append_toks)
+        self.prepend_bos = prepend_bos
+        self.append_eos = append_eos
+        self.use_msa = use_msa
+
+        self.all_toks = list(self.prepend_toks)
+        self.all_toks.extend(self.standard_toks)
+        for i in range((8 - (len(self.all_toks) % 8)) % 8):
+            self.all_toks.append(f"<null_{i  + 1}>")
+        self.all_toks.extend(self.append_toks)
+
+        self.tok_to_idx = {tok: i for i, tok in enumerate(self.all_toks)}
+
+        self.unk_idx = self.tok_to_idx["<unk>"]
+        self.padding_idx = self.get_idx("<pad>")
+        self.cls_idx = self.get_idx("<cls>")
+        self.mask_idx = self.get_idx("<mask>")
+        self.eos_idx = self.get_idx("<eos>")
+        self.all_special_tokens = ['<eos>', '<unk>', '<pad>', '<cls>', '<mask>']
+        self.unique_no_split_tokens = self.all_toks
+
+    def __len__(self):
+        return len(self.all_toks)
+
+    def get_idx(self, tok):
+        return self.tok_to_idx.get(tok, self.unk_idx)
+
+    def get_tok(self, ind):
+        return self.all_toks[ind]
+
+    def to_dict(self):
+        return self.tok_to_idx.copy()
+
+    def get_batch_converter(self, truncation_seq_length: int = None):
+        if self.use_msa:
+            return MSABatchConverter(self, truncation_seq_length)
+        else:
+            return BatchConverter(self, truncation_seq_length)
+
+    @classmethod
+    def from_architecture(cls, name: str) -> "Alphabet":
+        if name in ("ESM-1", "protein_bert_base"):
+            standard_toks = proteinseq_toks["toks"]
+            prepend_toks: Tuple[str, ...] = ("<null_0>", "<pad>", "<eos>", "<unk>")
+            append_toks: Tuple[str, ...] = ("<cls>", "<mask>", "<sep>")
+            prepend_bos = True
+            append_eos = False
+            use_msa = False
+        elif name in ("ESM-1b", "roberta_large"):
+            standard_toks = proteinseq_toks["toks"]
+            prepend_toks = ("<cls>", "<pad>", "<eos>", "<unk>")
+            append_toks = ("<mask>",)
+            prepend_bos = True
+            append_eos = True
+            use_msa = False
+        elif name in ("MSA Transformer", "msa_transformer"):
+            standard_toks = proteinseq_toks["toks"]
+            prepend_toks = ("<cls>", "<pad>", "<eos>", "<unk>")
+            append_toks = ("<mask>",)
+            prepend_bos = True
+            append_eos = False
+            use_msa = True
+        elif "invariant_gvp" in name.lower():
+            standard_toks = proteinseq_toks["toks"]
+            prepend_toks = ("<null_0>", "<pad>", "<eos>", "<unk>")
+            append_toks = ("<mask>", "<cath>", "<af2>")
+            prepend_bos = True
+            append_eos = False
+            use_msa = False
+        else:
+            raise ValueError("Unknown architecture selected")
+        return cls(standard_toks, prepend_toks, append_toks, prepend_bos, append_eos, use_msa)
+
+    def _tokenize(self, text) -> str:
+        return text.split()
+
+    def tokenize(self, text, **kwargs) -> List[str]:
+        """
+        Inspired by https://github.com/huggingface/transformers/blob/master/src/transformers/tokenization_utils.py
+        Converts a string in a sequence of tokens, using the tokenizer.
+
+        Args:
+            text (:obj:`str`):
+                The sequence to be encoded.
+
+        Returns:
+            :obj:`List[str]`: The list of tokens.
+        """
+
+        def split_on_token(tok, text):
+            result = []
+            split_text = text.split(tok)
+            for i, sub_text in enumerate(split_text):
+                # AddedToken can control whitespace stripping around them.
+                # We use them for GPT2 and Roberta to have different behavior depending on the special token
+                # Cf. https://github.com/huggingface/transformers/pull/2778
+                # and https://github.com/huggingface/transformers/issues/3788
+                # We strip left and right by default
+                if i < len(split_text) - 1:
+                    sub_text = sub_text.rstrip()
+                if i > 0:
+                    sub_text = sub_text.lstrip()
+
+                if i == 0 and not sub_text:
+                    result.append(tok)
+                elif i == len(split_text) - 1:
+                    if sub_text:
+                        result.append(sub_text)
+                    else:
+                        pass
+                else:
+                    if sub_text:
+                        result.append(sub_text)
+                    result.append(tok)
+            return result
+
+        def split_on_tokens(tok_list, text):
+            if not text.strip():
+                return []
+
+            tokenized_text = []
+            text_list = [text]
+            for tok in tok_list:
+                tokenized_text = []
+                for sub_text in text_list:
+                    if sub_text not in self.unique_no_split_tokens:
+                        tokenized_text.extend(split_on_token(tok, sub_text))
+                    else:
+                        tokenized_text.append(sub_text)
+                text_list = tokenized_text
+
+            return list(
+                itertools.chain.from_iterable(
+                    (
+                        self._tokenize(token)
+                        if token not in self.unique_no_split_tokens
+                        else [token]
+                        for token in tokenized_text
+                    )
+                )
+            )
+
+        no_split_token = self.unique_no_split_tokens
+        tokenized_text = split_on_tokens(no_split_token, text)
+        return tokenized_text
+
+    def encode(self, text):
+        return [self.tok_to_idx[tok] for tok in self.tokenize(text)]
+
+
+class BatchConverter(object):
+    """Callable to convert an unprocessed (labels + strings) batch to a
+    processed (labels + tensor) batch.
+    """
+
+    def __init__(self, alphabet, truncation_seq_length: int = None):
+        self.alphabet = alphabet
+        self.truncation_seq_length = truncation_seq_length
+
+    def __call__(self, raw_batch: Sequence[Tuple[str, str]]):
+        # RoBERTa uses an eos token, while ESM-1 does not.
+        batch_size = len(raw_batch)
+        batch_labels, seq_str_list = zip(*raw_batch)
+        seq_encoded_list = [self.alphabet.encode(seq_str) for seq_str in seq_str_list]
+        if self.truncation_seq_length:
+            seq_encoded_list = [seq_str[:self.truncation_seq_length] for seq_str in seq_encoded_list]
+        max_len = max(len(seq_encoded) for seq_encoded in seq_encoded_list)
+        tokens = torch.empty(
+            (
+                batch_size,
+                max_len + int(self.alphabet.prepend_bos) + int(self.alphabet.append_eos),
+            ),
+            dtype=torch.int64,
+        )
+        tokens.fill_(self.alphabet.padding_idx)
+        labels = []
+        strs = []
+
+        for i, (label, seq_str, seq_encoded) in enumerate(
+            zip(batch_labels, seq_str_list, seq_encoded_list)
+        ):
+            labels.append(label)
+            strs.append(seq_str)
+            if self.alphabet.prepend_bos:
+                tokens[i, 0] = self.alphabet.cls_idx
+            seq = torch.tensor(seq_encoded, dtype=torch.int64)
+            tokens[
+                i,
+                int(self.alphabet.prepend_bos) : len(seq_encoded)
+                + int(self.alphabet.prepend_bos),
+            ] = seq
+            if self.alphabet.append_eos:
+                tokens[i, len(seq_encoded) + int(self.alphabet.prepend_bos)] = self.alphabet.eos_idx
+
+        return labels, strs, tokens
+
+
+class MSABatchConverter(BatchConverter):
+    def __call__(self, inputs: Union[Sequence[RawMSA], RawMSA]):
+        if isinstance(inputs[0][0], str):
+            # Input is a single MSA
+            raw_batch: Sequence[RawMSA] = [inputs]  # type: ignore
+        else:
+            raw_batch = inputs  # type: ignore
+
+        batch_size = len(raw_batch)
+        max_alignments = max(len(msa) for msa in raw_batch)
+        max_seqlen = max(len(msa[0][1]) for msa in raw_batch)
+
+        tokens = torch.empty(
+            (
+                batch_size,
+                max_alignments,
+                max_seqlen + int(self.alphabet.prepend_bos) + int(self.alphabet.append_eos),
+            ),
+            dtype=torch.int64,
+        )
+        tokens.fill_(self.alphabet.padding_idx)
+        labels = []
+        strs = []
+
+        for i, msa in enumerate(raw_batch):
+            msa_seqlens = set(len(seq) for _, seq in msa)
+            if not len(msa_seqlens) == 1:
+                raise RuntimeError(
+                    "Received unaligned sequences for input to MSA, all sequence "
+                    "lengths must be equal."
+                )
+            msa_labels, msa_strs, msa_tokens = super().__call__(msa)
+            labels.append(msa_labels)
+            strs.append(msa_strs)
+            tokens[i, : msa_tokens.size(0), : msa_tokens.size(1)] = msa_tokens
+
+        return labels, strs, tokens
+
+
+def read_fasta(
+    path,
+    keep_gaps=True,
+    keep_insertions=True,
+    to_upper=False,
+):
+    with open(path, "r") as f:
+        for result in read_alignment_lines(
+            f, keep_gaps=keep_gaps, keep_insertions=keep_insertions, to_upper=to_upper
+        ):
+            yield result
+
+
+def read_alignment_lines(
+    lines,
+    keep_gaps=True,
+    keep_insertions=True,
+    to_upper=False,
+):
+    seq = desc = None
+
+    def parse(s):
+        if not keep_gaps:
+            s = re.sub("-", "", s)
+        if not keep_insertions:
+            s = re.sub("[a-z]", "", s)
+        return s.upper() if to_upper else s
+
+    for line in lines:
+        # Line may be empty if seq % file_line_width == 0
+        if len(line) > 0 and line[0] == ">":
+            if seq is not None:
+                yield desc, parse(seq)
+            desc = line.strip().lstrip(">")
+            seq = ""
+        else:
+            assert isinstance(seq, str)
+            seq += line.strip()
+    assert isinstance(seq, str) and isinstance(desc, str)
+    yield desc, parse(seq)
+
+
+class ESMStructuralSplitDataset(torch.utils.data.Dataset):
+    """
+    Structural Split Dataset as described in section A.10 of the supplement of our paper.
+    https://doi.org/10.1101/622803
+
+    We use the full version of SCOPe 2.07, clustered at 90% sequence identity,
+    generated on January 23, 2020.
+
+    For each SCOPe domain:
+        - We extract the sequence from the corresponding PDB file
+        - We extract the 3D coordinates of the Carbon beta atoms, aligning them
+          to the sequence. We put NaN where Cb atoms are missing.
+        - From the 3D coordinates, we calculate a pairwise distance map, based
+          on L2 distance
+        - We use DSSP to generate secondary structure labels for the corresponding
+          PDB file. This is also aligned to the sequence. We put - where SSP
+          labels are missing.
+
+    For each SCOPe classification level of family/superfamily/fold (in order of difficulty),
+    we have split the data into 5 partitions for cross validation. These are provided
+    in a downloaded splits folder, in the format:
+            splits/{split_level}/{cv_partition}/{train|valid}.txt
+    where train is the partition and valid is the concatentation of the remaining 4.
+
+    For each SCOPe domain, we provide a pkl dump that contains:
+        - seq    : The domain sequence, stored as an L-length string
+        - ssp    : The secondary structure labels, stored as an L-length string
+        - dist   : The distance map, stored as an LxL numpy array
+        - coords : The 3D coordinates, stored as an Lx3 numpy array
+
+    """
+
+    base_folder = "structural-data"
+    file_list = [
+        #  url  tar filename   filename      MD5 Hash
+        (
+            "https://dl.fbaipublicfiles.com/fair-esm/structural-data/splits.tar.gz",
+            "splits.tar.gz",
+            "splits",
+            "456fe1c7f22c9d3d8dfe9735da52411d",
+        ),
+        (
+            "https://dl.fbaipublicfiles.com/fair-esm/structural-data/pkl.tar.gz",
+            "pkl.tar.gz",
+            "pkl",
+            "644ea91e56066c750cd50101d390f5db",
+        ),
+    ]
+
+    def __init__(
+        self,
+        split_level,
+        cv_partition,
+        split,
+        root_path=os.path.expanduser("~/.cache/torch/data/esm"),
+        download=False,
+    ):
+        super().__init__()
+        assert split in [
+            "train",
+            "valid",
+        ], "train_valid must be 'train' or 'valid'"
+        self.root_path = root_path
+        self.base_path = os.path.join(self.root_path, self.base_folder)
+
+        # check if root path has what you need or else download it
+        if download:
+            self.download()
+
+        self.split_file = os.path.join(
+            self.base_path, "splits", split_level, cv_partition, f"{split}.txt"
+        )
+        self.pkl_dir = os.path.join(self.base_path, "pkl")
+        self.names = []
+        with open(self.split_file) as f:
+            self.names = f.read().splitlines()
+
+    def __len__(self):
+        return len(self.names)
+
+    def _check_exists(self) -> bool:
+        for (_, _, filename, _) in self.file_list:
+            fpath = os.path.join(self.base_path, filename)
+            if not os.path.exists(fpath) or not os.path.isdir(fpath):
+                return False
+        return True
+
+    def download(self):
+
+        if self._check_exists():
+            print("Files already downloaded and verified")
+            return
+
+        from torchvision.datasets.utils import download_url
+
+        for url, tar_filename, filename, md5_hash in self.file_list:
+            download_path = os.path.join(self.base_path, tar_filename)
+            download_url(url=url, root=self.base_path, filename=tar_filename, md5=md5_hash)
+            shutil.unpack_archive(download_path, self.base_path)
+
+    def __getitem__(self, idx):
+        """
+        Returns a dict with the following entires
+         - seq : Str (domain sequence)
+         - ssp : Str (SSP labels)
+         - dist : np.array (distance map)
+         - coords : np.array (3D coordinates)
+        """
+        name = self.names[idx]
+        pkl_fname = os.path.join(self.pkl_dir, name[1:3], f"{name}.pkl")
+        with open(pkl_fname, "rb") as f:
+            obj = pickle.load(f)
+        return obj

esm/source/esm/esmfold/v1/categorical_mixture.py

@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+
+
+class CategoricalMixture:
+    def __init__(self, param, bins=50, start=0, end=1):
+        # All tensors are of shape ..., bins.
+        self.logits = param
+        bins = torch.linspace(
+            start, end, bins + 1, device=self.logits.device, dtype=self.logits.dtype
+        )
+        self.v_bins = (bins[:-1] + bins[1:]) / 2
+
+    def log_prob(self, true):
+        # Shapes are:
+        #     self.probs: ... x bins
+        #     true      : ...
+        true_index = (
+            (
+                true.unsqueeze(-1)
+                - self.v_bins[
+                    [
+                        None,
+                    ]
+                    * true.ndim
+                ]
+            )
+            .abs()
+            .argmin(-1)
+        )
+        nll = self.logits.log_softmax(-1)
+        return torch.take_along_dim(nll, true_index.unsqueeze(-1), dim=-1).squeeze(-1)
+
+    def mean(self):
+        return (self.logits.softmax(-1) @ self.v_bins.unsqueeze(1)).squeeze(-1)
+
+
+def categorical_lddt(logits, bins=50):
+    # Logits are ..., 37, bins.
+    return CategoricalMixture(logits, bins=bins).mean()

esm/source/esm/esmfold/v1/esmfold.py

@@ -0,0 +1,364 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import typing as T
+from dataclasses import dataclass
+from functools import partial
+
+import torch
+import torch.nn as nn
+from torch import nn
+from torch.nn import LayerNorm
+
+import esm
+from esm import Alphabet
+from esm.esmfold.v1.categorical_mixture import categorical_lddt
+from esm.esmfold.v1.misc import (
+    batch_encode_sequences,
+    collate_dense_tensors,
+    output_to_pdb,
+)
+from esm.esmfold.v1.trunk import FoldingTrunk, FoldingTrunkConfig
+from openfold.data.data_transforms import make_atom14_masks
+from openfold.np import residue_constants
+from openfold.utils.loss import compute_predicted_aligned_error, compute_tm
+
+
+@dataclass
+class ESMFoldConfig:
+    trunk: T.Any = FoldingTrunkConfig()
+    lddt_head_hid_dim: int = 128
+
+
+load_fn = esm.pretrained.load_model_and_alphabet
+esm_registry = {
+    "esm2_8M": partial(load_fn, "esm2_t6_8M_UR50D_500K"),
+    "esm2_8M_270K": esm.pretrained.esm2_t6_8M_UR50D,
+    "esm2_35M": partial(load_fn, "esm2_t12_35M_UR50D_500K"),
+    "esm2_35M_270K": esm.pretrained.esm2_t12_35M_UR50D,
+    "esm2_150M": partial(load_fn, "esm2_t30_150M_UR50D_500K"),
+    "esm2_150M_270K": partial(load_fn, "esm2_t30_150M_UR50D_270K"),
+    "esm2_650M": esm.pretrained.esm2_t33_650M_UR50D,
+    "esm2_650M_270K": partial(load_fn, "esm2_t33_650M_270K_UR50D"),
+    "esm2_3B": esm.pretrained.esm2_t36_3B_UR50D,
+    "esm2_3B_270K": partial(load_fn, "esm2_t36_3B_UR50D_500K"),
+    "esm2_15B": esm.pretrained.esm2_t48_15B_UR50D,
+}
+
+
+class ESMFold(nn.Module):
+    def __init__(self, esmfold_config=None, **kwargs):
+        super().__init__()
+
+        self.cfg = esmfold_config if esmfold_config else ESMFoldConfig(**kwargs)
+        cfg = self.cfg
+
+        self.distogram_bins = 64
+
+        self.esm, self.esm_dict = esm_registry.get(cfg.esm_type)()
+
+        self.esm.requires_grad_(False)
+        self.esm.half()
+
+        self.esm_feats = self.esm.embed_dim
+        self.esm_attns = self.esm.num_layers * self.esm.attention_heads
+        self.register_buffer("af2_to_esm", ESMFold._af2_to_esm(self.esm_dict))
+        self.esm_s_combine = nn.Parameter(torch.zeros(self.esm.num_layers + 1))
+
+        c_s = cfg.trunk.sequence_state_dim
+        c_z = cfg.trunk.pairwise_state_dim
+
+        self.esm_s_mlp = nn.Sequential(
+            LayerNorm(self.esm_feats),
+            nn.Linear(self.esm_feats, c_s),
+            nn.ReLU(),
+            nn.Linear(c_s, c_s),
+        )
+        if cfg.use_esm_attn_map:
+            self.esm_z_mlp = nn.Sequential(
+                LayerNorm(self.esm_attns),
+                nn.Linear(self.esm_attns, c_z),
+                nn.ReLU(),
+                nn.Linear(c_z, c_z),
+            )
+
+        # 0 is padding, N is unknown residues, N + 1 is mask.
+        self.n_tokens_embed = residue_constants.restype_num + 3
+        self.pad_idx = 0
+        self.unk_idx = self.n_tokens_embed - 2
+        self.mask_idx = self.n_tokens_embed - 1
+        self.embedding = nn.Embedding(self.n_tokens_embed, c_s, padding_idx=0)
+
+        self.trunk = FoldingTrunk(**cfg.trunk)
+
+        self.distogram_head = nn.Linear(c_z, self.distogram_bins)
+        self.ptm_head = nn.Linear(c_z, self.distogram_bins)
+        self.lm_head = nn.Linear(c_s, self.n_tokens_embed)
+        self.lddt_bins = 50
+        self.lddt_head = nn.Sequential(
+            nn.LayerNorm(cfg.trunk.structure_module.c_s),
+            nn.Linear(cfg.trunk.structure_module.c_s, cfg.lddt_head_hid_dim),
+            nn.Linear(cfg.lddt_head_hid_dim, cfg.lddt_head_hid_dim),
+            nn.Linear(cfg.lddt_head_hid_dim, 37 * self.lddt_bins),
+        )
+
+    @staticmethod
+    def _af2_to_esm(d: Alphabet):
+        # Remember that t is shifted from residue_constants by 1 (0 is padding).
+        esm_reorder = [d.padding_idx] + [
+            d.get_idx(v) for v in residue_constants.restypes_with_x
+        ]
+        return torch.tensor(esm_reorder)
+
+    def _af2_idx_to_esm_idx(self, aa, mask):
+        aa = (aa + 1).masked_fill(mask != 1, 0)
+        return self.af2_to_esm[aa]
+
+    def _compute_language_model_representations(
+        self, esmaa: torch.Tensor
+    ) -> torch.Tensor:
+        """Adds bos/eos tokens for the language model, since the structure module doesn't use these."""
+        batch_size = esmaa.size(0)
+
+        bosi, eosi = self.esm_dict.cls_idx, self.esm_dict.eos_idx
+        bos = esmaa.new_full((batch_size, 1), bosi)
+        eos = esmaa.new_full((batch_size, 1), self.esm_dict.padding_idx)
+        esmaa = torch.cat([bos, esmaa, eos], dim=1)
+        # Use the first padding index as eos during inference.
+        esmaa[range(batch_size), (esmaa != 1).sum(1)] = eosi
+
+        res = self.esm(
+            esmaa,
+            repr_layers=range(self.esm.num_layers + 1),
+            need_head_weights=self.cfg.use_esm_attn_map,
+        )
+        esm_s = torch.stack(
+            [v for _, v in sorted(res["representations"].items())], dim=2
+        )
+        esm_s = esm_s[:, 1:-1]  # B, L, nLayers, C
+        esm_z = (
+            res["attentions"].permute(0, 4, 3, 1, 2).flatten(3, 4)[:, 1:-1, 1:-1, :]
+            if self.cfg.use_esm_attn_map
+            else None
+        )
+        return esm_s, esm_z
+
+    def _mask_inputs_to_esm(self, esmaa, pattern):
+        new_esmaa = esmaa.clone()
+        new_esmaa[pattern == 1] = self.esm_dict.mask_idx
+        return new_esmaa
+
+    def forward(
+        self,
+        aa: torch.Tensor,
+        mask: T.Optional[torch.Tensor] = None,
+        residx: T.Optional[torch.Tensor] = None,
+        masking_pattern: T.Optional[torch.Tensor] = None,
+        num_recycles: T.Optional[int] = None,
+    ):
+        """Runs a forward pass given input tokens. Use `model.infer` to
+        run inference from a sequence.
+
+        Args:
+            aa (torch.Tensor): Tensor containing indices corresponding to amino acids. Indices match
+                openfold.np.residue_constants.restype_order_with_x.
+            mask (torch.Tensor): Binary tensor with 1 meaning position is unmasked and 0 meaning position is masked.
+            residx (torch.Tensor): Residue indices of amino acids. Will assume contiguous if not provided.
+            masking_pattern (torch.Tensor): Optional masking to pass to the input. Binary tensor of the same size
+                as `aa`. Positions with 1 will be masked. ESMFold sometimes produces different samples when
+                different masks are provided.
+            num_recycles (int): How many recycle iterations to perform. If None, defaults to training max
+                recycles, which is 3.
+        """
+
+        if mask is None:
+            mask = torch.ones_like(aa)
+
+        B = aa.shape[0]
+        L = aa.shape[1]
+        device = aa.device
+
+        if residx is None:
+            residx = torch.arange(L, device=device).expand_as(aa)
+
+        # === ESM ===
+        esmaa = self._af2_idx_to_esm_idx(aa, mask)
+
+        if masking_pattern is not None:
+            esmaa = self._mask_inputs_to_esm(esmaa, masking_pattern)
+
+        esm_s, esm_z = self._compute_language_model_representations(esmaa)
+
+        # Convert esm_s to the precision used by the trunk and
+        # the structure module. These tensors may be a lower precision if, for example,
+        # we're running the language model in fp16 precision.
+        esm_s = esm_s.to(self.esm_s_combine.dtype)
+        esm_s = esm_s.detach()
+
+        # === preprocessing ===
+        esm_s = (self.esm_s_combine.softmax(0).unsqueeze(0) @ esm_s).squeeze(2)
+
+        s_s_0 = self.esm_s_mlp(esm_s)
+        if self.cfg.use_esm_attn_map:
+            esm_z = esm_z.to(self.esm_s_combine.dtype)
+            esm_z = esm_z.detach()
+            s_z_0 = self.esm_z_mlp(esm_z)
+        else:
+            s_z_0 = s_s_0.new_zeros(B, L, L, self.cfg.trunk.pairwise_state_dim)
+
+        s_s_0 += self.embedding(aa)
+
+        structure: dict = self.trunk(
+            s_s_0, s_z_0, aa, residx, mask, no_recycles=num_recycles
+        )
+        # Documenting what we expect:
+        structure = {
+            k: v
+            for k, v in structure.items()
+            if k
+            in [
+                "s_z",
+                "s_s",
+                "frames",
+                "sidechain_frames",
+                "unnormalized_angles",
+                "angles",
+                "positions",
+                "states",
+            ]
+        }
+
+        disto_logits = self.distogram_head(structure["s_z"])
+        disto_logits = (disto_logits + disto_logits.transpose(1, 2)) / 2
+        structure["distogram_logits"] = disto_logits
+
+        lm_logits = self.lm_head(structure["s_s"])
+        structure["lm_logits"] = lm_logits
+
+        structure["aatype"] = aa
+        make_atom14_masks(structure)
+
+        for k in [
+            "atom14_atom_exists",
+            "atom37_atom_exists",
+        ]:
+            structure[k] *= mask.unsqueeze(-1)
+        structure["residue_index"] = residx
+
+        lddt_head = self.lddt_head(structure["states"]).reshape(
+            structure["states"].shape[0], B, L, -1, self.lddt_bins
+        )
+        structure["lddt_head"] = lddt_head
+        plddt = categorical_lddt(lddt_head[-1], bins=self.lddt_bins)
+        structure["plddt"] = (
+            100 * plddt
+        )  # we predict plDDT between 0 and 1, scale to be between 0 and 100.
+
+        ptm_logits = self.ptm_head(structure["s_z"])
+
+        seqlen = mask.type(torch.int64).sum(1)
+        structure["ptm_logits"] = ptm_logits
+        structure["ptm"] = torch.stack(
+            [
+                compute_tm(
+                    batch_ptm_logits[None, :sl, :sl],
+                    max_bins=31,
+                    no_bins=self.distogram_bins,
+                )
+                for batch_ptm_logits, sl in zip(ptm_logits, seqlen)
+            ]
+        )
+        structure.update(
+            compute_predicted_aligned_error(
+                ptm_logits, max_bin=31, no_bins=self.distogram_bins
+            )
+        )
+
+        return structure
+
+    @torch.no_grad()
+    def infer(
+        self,
+        sequences: T.Union[str, T.List[str]],
+        residx=None,
+        masking_pattern: T.Optional[torch.Tensor] = None,
+        num_recycles: T.Optional[int] = None,
+        residue_index_offset: T.Optional[int] = 512,
+        chain_linker: T.Optional[str] = "G" * 25,
+    ):
+        """Runs a forward pass given input sequences.
+
+        Args:
+            sequences (Union[str, List[str]]): A list of sequences to make predictions for. Multimers can also be passed in,
+                each chain should be separated by a ':' token (e.g. "<chain1>:<chain2>:<chain3>").
+            residx (torch.Tensor): Residue indices of amino acids. Will assume contiguous if not provided.
+            masking_pattern (torch.Tensor): Optional masking to pass to the input. Binary tensor of the same size
+                as `aa`. Positions with 1 will be masked. ESMFold sometimes produces different samples when
+                different masks are provided.
+            num_recycles (int): How many recycle iterations to perform. If None, defaults to training max
+                recycles (cfg.trunk.max_recycles), which is 4.
+            residue_index_offset (int): Residue index separation between chains if predicting a multimer. Has no effect on
+                single chain predictions. Default: 512.
+            chain_linker (str): Linker to use between chains if predicting a multimer. Has no effect on single chain
+                predictions. Default: length-25 poly-G ("G" * 25).
+        """
+        if isinstance(sequences, str):
+            sequences = [sequences]
+
+        aatype, mask, _residx, linker_mask, chain_index = batch_encode_sequences(
+            sequences, residue_index_offset, chain_linker
+        )
+
+        if residx is None:
+            residx = _residx
+        elif not isinstance(residx, torch.Tensor):
+            residx = collate_dense_tensors(residx)
+
+        aatype, mask, residx, linker_mask = map(
+            lambda x: x.to(self.device), (aatype, mask, residx, linker_mask)
+        )
+
+        output = self.forward(
+            aatype,
+            mask=mask,
+            residx=residx,
+            masking_pattern=masking_pattern,
+            num_recycles=num_recycles,
+        )
+
+        output["atom37_atom_exists"] = output[
+            "atom37_atom_exists"
+        ] * linker_mask.unsqueeze(2)
+
+        output["mean_plddt"] = (output["plddt"] * output["atom37_atom_exists"]).sum(
+            dim=(1, 2)
+        ) / output["atom37_atom_exists"].sum(dim=(1, 2))
+        output["chain_index"] = chain_index
+
+        return output
+
+    def output_to_pdb(self, output: T.Dict) -> T.List[str]:
+        """Returns the pbd (file) string from the model given the model output."""
+        return output_to_pdb(output)
+
+    def infer_pdbs(self, seqs: T.List[str], *args, **kwargs) -> T.List[str]:
+        """Returns list of pdb (files) strings from the model given a list of input sequences."""
+        output = self.infer(seqs, *args, **kwargs)
+        return self.output_to_pdb(output)
+
+    def infer_pdb(self, sequence: str, *args, **kwargs) -> str:
+        """Returns the pdb (file) string from the model given an input sequence."""
+        return self.infer_pdbs([sequence], *args, **kwargs)[0]
+
+    def set_chunk_size(self, chunk_size: T.Optional[int]):
+        # This parameter means the axial attention will be computed
+        # in a chunked manner. This should make the memory used more or less O(L) instead of O(L^2).
+        # It's equivalent to running a for loop over chunks of the dimension we're iterative over,
+        # where the chunk_size is the size of the chunks, so 128 would mean to parse 128-lengthed chunks.
+        # Setting the value to None will return to default behavior, disable chunking.
+        self.trunk.set_chunk_size(chunk_size)
+
+    @property
+    def device(self):
+        return self.esm_s_combine.device

esm/source/esm/esmfold/v1/misc.py

@@ -0,0 +1,309 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import typing as T
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import nn
+from openfold.np import residue_constants
+from openfold.np.protein import Protein as OFProtein
+from openfold.np.protein import to_pdb
+from openfold.utils.feats import atom14_to_atom37
+
+
+def encode_sequence(
+    seq: str,
+    residue_index_offset: T.Optional[int] = 512,
+    chain_linker: T.Optional[str] = "G" * 25,
+) -> T.Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if chain_linker is None:
+        chain_linker = ""
+    if residue_index_offset is None:
+        residue_index_offset = 0
+
+    chains = seq.split(":")
+    seq = chain_linker.join(chains)
+
+    unk_idx = residue_constants.restype_order_with_x["X"]
+    encoded = torch.tensor(
+        [residue_constants.restype_order_with_x.get(aa, unk_idx) for aa in seq]
+    )
+    residx = torch.arange(len(encoded))
+
+    if residue_index_offset > 0:
+        start = 0
+        for i, chain in enumerate(chains):
+            residx[start : start + len(chain) + len(chain_linker)] += (
+                i * residue_index_offset
+            )
+            start += len(chain) + len(chain_linker)
+
+    linker_mask = torch.ones_like(encoded, dtype=torch.float32)
+    chain_index = []
+    offset = 0
+    for i, chain in enumerate(chains):
+        if i > 0:
+            chain_index.extend([i - 1] * len(chain_linker))
+        chain_index.extend([i] * len(chain))
+        offset += len(chain)
+        linker_mask[offset : offset + len(chain_linker)] = 0
+        offset += len(chain_linker)
+
+    chain_index = torch.tensor(chain_index, dtype=torch.int64)
+
+    return encoded, residx, linker_mask, chain_index
+
+
+def batch_encode_sequences(
+    sequences: T.Sequence[str],
+    residue_index_offset: T.Optional[int] = 512,
+    chain_linker: T.Optional[str] = "G" * 25,
+) -> T.Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    aatype_list = []
+    residx_list = []
+    linker_mask_list = []
+    chain_index_list = []
+    for seq in sequences:
+        aatype_seq, residx_seq, linker_mask_seq, chain_index_seq = encode_sequence(
+            seq,
+            residue_index_offset=residue_index_offset,
+            chain_linker=chain_linker,
+        )
+        aatype_list.append(aatype_seq)
+        residx_list.append(residx_seq)
+        linker_mask_list.append(linker_mask_seq)
+        chain_index_list.append(chain_index_seq)
+
+    aatype = collate_dense_tensors(aatype_list)
+    mask = collate_dense_tensors(
+        [aatype.new_ones(len(aatype_seq)) for aatype_seq in aatype_list]
+    )
+    residx = collate_dense_tensors(residx_list)
+    linker_mask = collate_dense_tensors(linker_mask_list)
+    chain_index_list = collate_dense_tensors(chain_index_list, -1)
+
+    return aatype, mask, residx, linker_mask, chain_index_list
+
+
+def output_to_pdb(output: T.Dict) -> T.List[str]:
+    """Returns the pbd (file) string from the model given the model output."""
+    # atom14_to_atom37 must be called first, as it fails on latest numpy if the
+    # input is a numpy array. It will work if the input is a torch tensor.
+    final_atom_positions = atom14_to_atom37(output["positions"][-1], output)
+    output = {k: v.to("cpu").numpy() for k, v in output.items()}
+    final_atom_positions = final_atom_positions.cpu().numpy()
+    final_atom_mask = output["atom37_atom_exists"]
+    pdbs = []
+    for i in range(output["aatype"].shape[0]):
+        aa = output["aatype"][i]
+        pred_pos = final_atom_positions[i]
+        mask = final_atom_mask[i]
+        resid = output["residue_index"][i] + 1
+        pred = OFProtein(
+            aatype=aa,
+            atom_positions=pred_pos,
+            atom_mask=mask,
+            residue_index=resid,
+            b_factors=output["plddt"][i],
+            chain_index=output["chain_index"][i] if "chain_index" in output else None,
+        )
+        pdbs.append(to_pdb(pred))
+    return pdbs
+
+
+def collate_dense_tensors(
+    samples: T.List[torch.Tensor], pad_v: float = 0
+) -> torch.Tensor:
+    """
+    Takes a list of tensors with the following dimensions:
+        [(d_11,       ...,           d_1K),
+         (d_21,       ...,           d_2K),
+         ...,
+         (d_N1,       ...,           d_NK)]
+    and stack + pads them into a single tensor of:
+    (N, max_i=1,N { d_i1 }, ..., max_i=1,N {diK})
+    """
+    if len(samples) == 0:
+        return torch.Tensor()
+    if len(set(x.dim() for x in samples)) != 1:
+        raise RuntimeError(
+            f"Samples has varying dimensions: {[x.dim() for x in samples]}"
+        )
+    (device,) = tuple(set(x.device for x in samples))  # assumes all on same device
+    max_shape = [max(lst) for lst in zip(*[x.shape for x in samples])]
+    result = torch.empty(
+        len(samples), *max_shape, dtype=samples[0].dtype, device=device
+    )
+    result.fill_(pad_v)
+    for i in range(len(samples)):
+        result_i = result[i]
+        t = samples[i]
+        result_i[tuple(slice(0, k) for k in t.shape)] = t
+    return result
+
+
+class Attention(nn.Module):
+    def __init__(self, embed_dim, num_heads, head_width, gated=False):
+        super().__init__()
+        assert embed_dim == num_heads * head_width
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_width = head_width
+
+        self.proj = nn.Linear(embed_dim, embed_dim * 3, bias=False)
+        self.o_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.gated = gated
+        if gated:
+            self.g_proj = nn.Linear(embed_dim, embed_dim)
+            torch.nn.init.zeros_(self.g_proj.weight)
+            torch.nn.init.ones_(self.g_proj.bias)
+
+        self.rescale_factor = self.head_width**-0.5
+
+        torch.nn.init.zeros_(self.o_proj.bias)
+
+    def forward(self, x, mask=None, bias=None, indices=None):
+        """
+        Basic self attention with optional mask and external pairwise bias.
+        To handle sequences of different lengths, use mask.
+
+        Inputs:
+          x: batch of input sequneces (.. x L x C)
+          mask: batch of boolean masks where 1=valid, 0=padding position (.. x L_k). optional.
+          bias: batch of scalar pairwise attention biases (.. x Lq x Lk x num_heads). optional.
+
+        Outputs:
+          sequence projection (B x L x embed_dim), attention maps (B x L x L x num_heads)
+        """
+
+        t = rearrange(self.proj(x), "... l (h c) -> ... h l c", h=self.num_heads)
+        q, k, v = t.chunk(3, dim=-1)
+
+        q = self.rescale_factor * q
+        a = torch.einsum("...qc,...kc->...qk", q, k)
+
+        # Add external attention bias.
+        if bias is not None:
+            a = a + rearrange(bias, "... lq lk h -> ... h lq lk")
+
+        # Do not attend to padding tokens.
+        if mask is not None:
+            mask = repeat(
+                mask, "... lk -> ... h lq lk", h=self.num_heads, lq=q.shape[-2]
+            )
+            a = a.masked_fill(mask == False, -np.inf)
+
+        a = F.softmax(a, dim=-1)
+
+        y = torch.einsum("...hqk,...hkc->...qhc", a, v)
+        y = rearrange(y, "... h c -> ... (h c)", h=self.num_heads)
+
+        if self.gated:
+            y = self.g_proj(x).sigmoid() * y
+        y = self.o_proj(y)
+
+        return y, rearrange(a, "... lq lk h -> ... h lq lk")
+
+
+class Dropout(nn.Module):
+    """
+    Implementation of dropout with the ability to share the dropout mask
+    along a particular dimension.
+    """
+
+    def __init__(self, r: float, batch_dim: T.Union[int, T.List[int]]):
+        super(Dropout, self).__init__()
+
+        self.r = r
+        if type(batch_dim) == int:
+            batch_dim = [batch_dim]
+        self.batch_dim = batch_dim
+        self.dropout = nn.Dropout(self.r)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = list(x.shape)
+        if self.batch_dim is not None:
+            for bd in self.batch_dim:
+                shape[bd] = 1
+        return x * self.dropout(x.new_ones(shape))
+
+
+class SequenceToPair(nn.Module):
+    def __init__(self, sequence_state_dim, inner_dim, pairwise_state_dim):
+        super().__init__()
+
+        self.layernorm = nn.LayerNorm(sequence_state_dim)
+        self.proj = nn.Linear(sequence_state_dim, inner_dim * 2, bias=True)
+        self.o_proj = nn.Linear(2 * inner_dim, pairwise_state_dim, bias=True)
+
+        torch.nn.init.zeros_(self.proj.bias)
+        torch.nn.init.zeros_(self.o_proj.bias)
+
+    def forward(self, sequence_state):
+        """
+        Inputs:
+          sequence_state: B x L x sequence_state_dim
+
+        Output:
+          pairwise_state: B x L x L x pairwise_state_dim
+
+        Intermediate state:
+          B x L x L x 2*inner_dim
+        """
+
+        assert len(sequence_state.shape) == 3
+
+        s = self.layernorm(sequence_state)
+        s = self.proj(s)
+        q, k = s.chunk(2, dim=-1)
+
+        prod = q[:, None, :, :] * k[:, :, None, :]
+        diff = q[:, None, :, :] - k[:, :, None, :]
+
+        x = torch.cat([prod, diff], dim=-1)
+        x = self.o_proj(x)
+
+        return x
+
+
+class PairToSequence(nn.Module):
+    def __init__(self, pairwise_state_dim, num_heads):
+        super().__init__()
+
+        self.layernorm = nn.LayerNorm(pairwise_state_dim)
+        self.linear = nn.Linear(pairwise_state_dim, num_heads, bias=False)
+
+    def forward(self, pairwise_state):
+        """
+        Inputs:
+          pairwise_state: B x L x L x pairwise_state_dim
+
+        Output:
+          pairwise_bias: B x L x L x num_heads
+        """
+        assert len(pairwise_state.shape) == 4
+        z = self.layernorm(pairwise_state)
+        pairwise_bias = self.linear(z)
+        return pairwise_bias
+
+
+class ResidueMLP(nn.Module):
+    def __init__(self, embed_dim, inner_dim, norm=nn.LayerNorm, dropout=0):
+        super().__init__()
+
+        self.mlp = nn.Sequential(
+            norm(embed_dim),
+            nn.Linear(embed_dim, inner_dim),
+            nn.ReLU(),
+            nn.Linear(inner_dim, embed_dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return x + self.mlp(x)

esm/source/esm/esmfold/v1/pretrained.py

@@ -0,0 +1,181 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from pathlib import Path
+
+import torch
+
+from esm.esmfold.v1.esmfold import ESMFold
+
+
+def _load_model(model_name):
+    if model_name.endswith(".pt"):  # local, treat as filepath
+        model_path = Path(model_name)
+        model_data = torch.load(str(model_path), map_location="cpu")
+    else:  # load from hub
+        url = f"https://dl.fbaipublicfiles.com/fair-esm/models/{model_name}.pt"
+        model_data = torch.hub.load_state_dict_from_url(url, progress=False, map_location="cpu")
+
+    cfg = model_data["cfg"]["model"]
+    model_state = model_data["model"]
+    model = ESMFold(esmfold_config=cfg)
+
+    expected_keys = set(model.state_dict().keys())
+    found_keys = set(model_state.keys())
+
+    missing_essential_keys = []
+    for missing_key in expected_keys - found_keys:
+        if not missing_key.startswith("esm."):
+            missing_essential_keys.append(missing_key)
+
+    if missing_essential_keys:
+        raise RuntimeError(f"Keys '{', '.join(missing_essential_keys)}' are missing.")
+
+    model.load_state_dict(model_state, strict=False)
+
+    return model
+
+
+def esmfold_v0():
+    """
+    ESMFold v0 model with 3B ESM-2, 48 folding blocks.
+    This version was used for the paper (Lin et al, 2022). It was trained
+    on all PDB chains until 2020-05, to ensure temporal holdout with CASP14
+    and the CAMEO validation and test set reported there.
+    """
+    return _load_model("esmfold_3B_v0")
+
+
+def esmfold_v1():
+    """
+    ESMFold v1 model using 3B ESM-2, 48 folding blocks.
+    ESMFold provides fast high accuracy atomic level structure prediction
+    directly from the individual sequence of a protein. ESMFold uses the ESM2
+    protein language model to extract meaningful representations from the
+    protein sequence.
+    """
+    return _load_model("esmfold_3B_v1")
+
+
+def esmfold_structure_module_only_8M():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_8M")
+
+
+def esmfold_structure_module_only_8M_270K():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_8M_270K")
+
+
+def esmfold_structure_module_only_35M():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_35M")
+
+
+def esmfold_structure_module_only_35M_270K():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_35M_270K")
+
+
+def esmfold_structure_module_only_150M():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_150M")
+
+
+def esmfold_structure_module_only_150M_270K():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_150M_270K")
+
+
+def esmfold_structure_module_only_650M():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_650M")
+
+
+def esmfold_structure_module_only_650M_270K():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_650M_270K")
+
+
+def esmfold_structure_module_only_3B():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_3B")
+
+
+def esmfold_structure_module_only_3B_270K():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_3B_270K")
+
+
+def esmfold_structure_module_only_15B():
+    """
+    ESMFold baseline model using 15B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    The 15B parameter ESM-2 was not trained out to 500K updates
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_15B")

esm/source/esm/esmfold/v1/tri_self_attn_block.py

@@ -0,0 +1,160 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+from openfold.model.triangular_attention import (
+    TriangleAttentionEndingNode,
+    TriangleAttentionStartingNode,
+)
+from openfold.model.triangular_multiplicative_update import (
+    TriangleMultiplicationIncoming,
+    TriangleMultiplicationOutgoing,
+)
+from torch import nn
+
+from esm.esmfold.v1.misc import (
+    Attention,
+    Dropout,
+    PairToSequence,
+    ResidueMLP,
+    SequenceToPair,
+)
+
+
+class TriangularSelfAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        sequence_state_dim,
+        pairwise_state_dim,
+        sequence_head_width,
+        pairwise_head_width,
+        dropout=0,
+        **__kwargs,
+    ):
+        super().__init__()
+
+        assert sequence_state_dim % sequence_head_width == 0
+        assert pairwise_state_dim % pairwise_head_width == 0
+        sequence_num_heads = sequence_state_dim // sequence_head_width
+        pairwise_num_heads = pairwise_state_dim // pairwise_head_width
+        assert sequence_state_dim == sequence_num_heads * sequence_head_width
+        assert pairwise_state_dim == pairwise_num_heads * pairwise_head_width
+        assert pairwise_state_dim % 2 == 0
+
+        self.sequence_state_dim = sequence_state_dim
+        self.pairwise_state_dim = pairwise_state_dim
+
+        self.layernorm_1 = nn.LayerNorm(sequence_state_dim)
+
+        self.sequence_to_pair = SequenceToPair(
+            sequence_state_dim, pairwise_state_dim // 2, pairwise_state_dim
+        )
+        self.pair_to_sequence = PairToSequence(pairwise_state_dim, sequence_num_heads)
+
+        self.seq_attention = Attention(
+            sequence_state_dim, sequence_num_heads, sequence_head_width, gated=True
+        )
+        self.tri_mul_out = TriangleMultiplicationOutgoing(
+            pairwise_state_dim,
+            pairwise_state_dim,
+        )
+        self.tri_mul_in = TriangleMultiplicationIncoming(
+            pairwise_state_dim,
+            pairwise_state_dim,
+        )
+        self.tri_att_start = TriangleAttentionStartingNode(
+            pairwise_state_dim,
+            pairwise_head_width,
+            pairwise_num_heads,
+            inf=1e9,
+        )  # type: ignore
+        self.tri_att_end = TriangleAttentionEndingNode(
+            pairwise_state_dim,
+            pairwise_head_width,
+            pairwise_num_heads,
+            inf=1e9,
+        )  # type: ignore
+
+        self.mlp_seq = ResidueMLP(sequence_state_dim, 4 * sequence_state_dim, dropout=dropout)
+        self.mlp_pair = ResidueMLP(pairwise_state_dim, 4 * pairwise_state_dim, dropout=dropout)
+
+        assert dropout < 0.4
+        self.drop = nn.Dropout(dropout)
+        self.row_drop = Dropout(dropout * 2, 2)
+        self.col_drop = Dropout(dropout * 2, 1)
+
+        torch.nn.init.zeros_(self.tri_mul_in.linear_z.weight)
+        torch.nn.init.zeros_(self.tri_mul_in.linear_z.bias)
+        torch.nn.init.zeros_(self.tri_mul_out.linear_z.weight)
+        torch.nn.init.zeros_(self.tri_mul_out.linear_z.bias)
+        torch.nn.init.zeros_(self.tri_att_start.mha.linear_o.weight)
+        torch.nn.init.zeros_(self.tri_att_start.mha.linear_o.bias)
+        torch.nn.init.zeros_(self.tri_att_end.mha.linear_o.weight)
+        torch.nn.init.zeros_(self.tri_att_end.mha.linear_o.bias)
+
+        torch.nn.init.zeros_(self.sequence_to_pair.o_proj.weight)
+        torch.nn.init.zeros_(self.sequence_to_pair.o_proj.bias)
+        torch.nn.init.zeros_(self.pair_to_sequence.linear.weight)
+        torch.nn.init.zeros_(self.seq_attention.o_proj.weight)
+        torch.nn.init.zeros_(self.seq_attention.o_proj.bias)
+        torch.nn.init.zeros_(self.mlp_seq.mlp[-2].weight)
+        torch.nn.init.zeros_(self.mlp_seq.mlp[-2].bias)
+        torch.nn.init.zeros_(self.mlp_pair.mlp[-2].weight)
+        torch.nn.init.zeros_(self.mlp_pair.mlp[-2].bias)
+
+    def forward(self, sequence_state, pairwise_state, mask=None, chunk_size=None, **__kwargs):
+        """
+        Inputs:
+          sequence_state: B x L x sequence_state_dim
+          pairwise_state: B x L x L x pairwise_state_dim
+          mask: B x L boolean tensor of valid positions
+
+        Output:
+          sequence_state: B x L x sequence_state_dim
+          pairwise_state: B x L x L x pairwise_state_dim
+        """
+        assert len(sequence_state.shape) == 3
+        assert len(pairwise_state.shape) == 4
+        if mask is not None:
+            assert len(mask.shape) == 2
+
+        batch_dim, seq_dim, sequence_state_dim = sequence_state.shape
+        pairwise_state_dim = pairwise_state.shape[3]
+        assert sequence_state_dim == self.sequence_state_dim
+        assert pairwise_state_dim == self.pairwise_state_dim
+        assert batch_dim == pairwise_state.shape[0]
+        assert seq_dim == pairwise_state.shape[1]
+        assert seq_dim == pairwise_state.shape[2]
+
+        # Update sequence state
+        bias = self.pair_to_sequence(pairwise_state)
+
+        # Self attention with bias + mlp.
+        y = self.layernorm_1(sequence_state)
+        y, _ = self.seq_attention(y, mask=mask, bias=bias)
+        sequence_state = sequence_state + self.drop(y)
+        sequence_state = self.mlp_seq(sequence_state)
+
+        # Update pairwise state
+        pairwise_state = pairwise_state + self.sequence_to_pair(sequence_state)
+
+        # Axial attention with triangular bias.
+        tri_mask = mask.unsqueeze(2) * mask.unsqueeze(1) if mask is not None else None
+        pairwise_state = pairwise_state + self.row_drop(
+            self.tri_mul_out(pairwise_state, mask=tri_mask)
+        )
+        pairwise_state = pairwise_state + self.col_drop(
+            self.tri_mul_in(pairwise_state, mask=tri_mask)
+        )
+        pairwise_state = pairwise_state + self.row_drop(
+            self.tri_att_start(pairwise_state, mask=tri_mask, chunk_size=chunk_size)
+        )
+        pairwise_state = pairwise_state + self.col_drop(
+            self.tri_att_end(pairwise_state, mask=tri_mask, chunk_size=chunk_size)
+        )
+
+        # MLP over pairs.
+        pairwise_state = self.mlp_pair(pairwise_state)
+
+        return sequence_state, pairwise_state

esm/source/esm/esmfold/v1/trunk.py

@@ -0,0 +1,243 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import typing as T
+from contextlib import ExitStack
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from openfold.model.structure_module import StructureModule
+
+from esm.esmfold.v1.tri_self_attn_block import TriangularSelfAttentionBlock
+
+
+@dataclass
+class StructureModuleConfig:
+    c_s: int = 384
+    c_z: int = 128
+    c_ipa: int = 16
+    c_resnet: int = 128
+    no_heads_ipa: int = 12
+    no_qk_points: int = 4
+    no_v_points: int = 8
+    dropout_rate: float = 0.1
+    no_blocks: int = 8
+    no_transition_layers: int = 1
+    no_resnet_blocks: int = 2
+    no_angles: int = 7
+    trans_scale_factor: int = 10
+    epsilon: float = 1e-8
+    inf: float = 1e5
+
+
+@dataclass
+class FoldingTrunkConfig:
+    _name: str = "FoldingTrunkConfig"
+    num_blocks: int = 48
+    sequence_state_dim: int = 1024
+    pairwise_state_dim: int = 128
+    sequence_head_width: int = 32
+    pairwise_head_width: int = 32
+    position_bins: int = 32
+    dropout: float = 0
+    layer_drop: float = 0
+    cpu_grad_checkpoint: bool = False
+
+    max_recycles: int = 4
+    chunk_size: T.Optional[int] = None
+
+    structure_module: StructureModuleConfig = StructureModuleConfig()
+
+
+def get_axial_mask(mask):
+    """
+    Helper to convert B x L mask of valid positions to axial mask used
+    in row column attentions.
+
+    Input:
+      mask: B x L tensor of booleans
+
+    Output:
+      mask: B x L x L tensor of booleans
+    """
+
+    if mask is None:
+        return None
+    assert len(mask.shape) == 2
+    batch_dim, seq_dim = mask.shape
+    m = mask.unsqueeze(1).expand(batch_dim, seq_dim, seq_dim)
+    m = m.reshape(batch_dim * seq_dim, seq_dim)
+    return m
+
+
+class RelativePosition(nn.Module):
+    def __init__(self, bins, pairwise_state_dim):
+        super().__init__()
+        self.bins = bins
+
+        # Note an additional offset is used so that the 0th position
+        # is reserved for masked pairs.
+        self.embedding = torch.nn.Embedding(2 * bins + 2, pairwise_state_dim)
+
+    def forward(self, residue_index, mask=None):
+        """
+        Input:
+          residue_index: B x L tensor of indices (dytpe=torch.long)
+          mask: B x L tensor of booleans
+
+        Output:
+          pairwise_state: B x L x L x pairwise_state_dim tensor of embeddings
+        """
+
+        assert residue_index.dtype == torch.long
+        if mask is not None:
+            assert residue_index.shape == mask.shape
+
+        diff = residue_index[:, None, :] - residue_index[:, :, None]
+        diff = diff.clamp(-self.bins, self.bins)
+        diff = diff + self.bins + 1  # Add 1 to adjust for padding index.
+
+        if mask is not None:
+            mask = mask[:, None, :] * mask[:, :, None]
+            diff[mask == False] = 0
+
+        output = self.embedding(diff)
+        return output
+
+
+class FoldingTrunk(nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.cfg = FoldingTrunkConfig(**kwargs)
+        assert self.cfg.max_recycles > 0
+
+        c_s = self.cfg.sequence_state_dim
+        c_z = self.cfg.pairwise_state_dim
+
+        assert c_s % self.cfg.sequence_head_width == 0
+        assert c_z % self.cfg.pairwise_head_width == 0
+        block = TriangularSelfAttentionBlock
+
+        self.pairwise_positional_embedding = RelativePosition(self.cfg.position_bins, c_z)
+
+        self.blocks = nn.ModuleList(
+            [
+                block(
+                    sequence_state_dim=c_s,
+                    pairwise_state_dim=c_z,
+                    sequence_head_width=self.cfg.sequence_head_width,
+                    pairwise_head_width=self.cfg.pairwise_head_width,
+                    dropout=self.cfg.dropout,
+                )
+                for i in range(self.cfg.num_blocks)
+            ]
+        )
+
+        self.recycle_bins = 15
+        self.recycle_s_norm = nn.LayerNorm(c_s)
+        self.recycle_z_norm = nn.LayerNorm(c_z)
+        self.recycle_disto = nn.Embedding(self.recycle_bins, c_z)
+        self.recycle_disto.weight[0].detach().zero_()
+
+        self.structure_module = StructureModule(**self.cfg.structure_module)  # type: ignore
+        self.trunk2sm_s = nn.Linear(c_s, self.structure_module.c_s)
+        self.trunk2sm_z = nn.Linear(c_z, self.structure_module.c_z)
+
+        self.chunk_size = self.cfg.chunk_size
+
+    def set_chunk_size(self, chunk_size):
+        # This parameter means the axial attention will be computed
+        # in a chunked manner. This should make the memory used more or less O(L) instead of O(L^2).
+        # It's equivalent to running a for loop over chunks of the dimension we're iterative over,
+        # where the chunk_size is the size of the chunks, so 128 would mean to parse 128-lengthed chunks.
+        self.chunk_size = chunk_size
+
+    def forward(self, seq_feats, pair_feats, true_aa, residx, mask, no_recycles: T.Optional[int] = None):
+        """
+        Inputs:
+          seq_feats:     B x L x C            tensor of sequence features
+          pair_feats:    B x L x L x C        tensor of pair features
+          residx:        B x L                long tensor giving the position in the sequence
+          mask:          B x L                boolean tensor indicating valid residues
+
+        Output:
+          predicted_structure: B x L x (num_atoms_per_residue * 3) tensor wrapped in a Coordinates object
+        """
+
+        device = seq_feats.device
+        s_s_0 = seq_feats
+        s_z_0 = pair_feats
+
+        if no_recycles is None:
+            no_recycles = self.cfg.max_recycles
+        else:
+            assert no_recycles >= 0, "Number of recycles must not be negative."
+            no_recycles += 1  # First 'recycle' is just the standard forward pass through the model.
+
+        def trunk_iter(s, z, residx, mask):
+            z = z + self.pairwise_positional_embedding(residx, mask=mask)
+
+            for block in self.blocks:
+                s, z = block(s, z, mask=mask, residue_index=residx, chunk_size=self.chunk_size)
+            return s, z
+
+        s_s = s_s_0
+        s_z = s_z_0
+        recycle_s = torch.zeros_like(s_s)
+        recycle_z = torch.zeros_like(s_z)
+        recycle_bins = torch.zeros(*s_z.shape[:-1], device=device, dtype=torch.int64)
+
+        assert no_recycles > 0
+        for recycle_idx in range(no_recycles):
+            with ExitStack() if recycle_idx == no_recycles - 1 else torch.no_grad():
+                # === Recycling ===
+                recycle_s = self.recycle_s_norm(recycle_s.detach())
+                recycle_z = self.recycle_z_norm(recycle_z.detach())
+                recycle_z += self.recycle_disto(recycle_bins.detach())
+
+                s_s, s_z = trunk_iter(s_s_0 + recycle_s, s_z_0 + recycle_z, residx, mask)
+
+                # === Structure module ===
+                structure = self.structure_module(
+                    {"single": self.trunk2sm_s(s_s), "pair": self.trunk2sm_z(s_z)},
+                    true_aa,
+                    mask.float(),
+                )
+
+                recycle_s = s_s
+                recycle_z = s_z
+                # Distogram needs the N, CA, C coordinates, and bin constants same as alphafold.
+                recycle_bins = FoldingTrunk.distogram(
+                    structure["positions"][-1][:, :, :3],
+                    3.375,
+                    21.375,
+                    self.recycle_bins,
+                )
+
+        assert isinstance(structure, dict)  # type: ignore
+        structure["s_s"] = s_s
+        structure["s_z"] = s_z
+
+        return structure
+
+    @staticmethod
+    def distogram(coords, min_bin, max_bin, num_bins):
+        # Coords are [... L x 3 x 3], where it's [N, CA, C] x 3 coordinates.
+        boundaries = torch.linspace(
+            min_bin,
+            max_bin,
+            num_bins - 1,
+            device=coords.device,
+        )
+        boundaries = boundaries**2
+        N, CA, C = [x.squeeze(-2) for x in coords.chunk(3, dim=-2)]
+        # Infer CB coordinates.
+        b = CA - N
+        c = C - CA
+        a = b.cross(c, dim=-1)
+        CB = -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + CA
+        dists = (CB[..., None, :, :] - CB[..., :, None, :]).pow(2).sum(dim=-1, keepdims=True)
+        bins = torch.sum(dists > boundaries, dim=-1)  # [..., L, L]
+        return bins

esm/source/esm/inverse_folding/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from . import gvp_transformer
+from . import util
+from . import multichain_util

esm/source/esm/inverse_folding/features.py

@@ -0,0 +1,352 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# 
+# Portions of this file were adapted from the open source code for the following
+# two papers:
+#
+#   Ingraham, J., Garg, V., Barzilay, R., & Jaakkola, T. (2019). Generative
+#   models for graph-based protein design. Advances in Neural Information
+#   Processing Systems, 32.
+#
+#   Jing, B., Eismann, S., Suriana, P., Townshend, R. J. L., & Dror, R. (2020).
+#   Learning from Protein Structure with Geometric Vector Perceptrons. In
+#   International Conference on Learning Representations.
+#
+# MIT License
+# 
+# Copyright (c) 2020 Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael Townshend, Ron Dror
+# 
+# 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.
+# 
+# ================================================================
+# The below license applies to the portions of the code (parts of 
+# src/datasets.py and src/models.py) adapted from Ingraham, et al.
+# ================================================================
+# 
+# MIT License
+# 
+# Copyright (c) 2019 John Ingraham, Vikas Garg, Regina Barzilay, Tommi Jaakkola
+# 
+# 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.
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .gvp_utils import flatten_graph
+from .gvp_modules import GVP, LayerNorm
+from .util import normalize, norm, nan_to_num, rbf
+
+
+class GVPInputFeaturizer(nn.Module):
+
+    @staticmethod
+    def get_node_features(coords, coord_mask, with_coord_mask=True):
+        # scalar features
+        node_scalar_features = GVPInputFeaturizer._dihedrals(coords)
+        if with_coord_mask:
+            node_scalar_features = torch.cat([
+                node_scalar_features,
+                coord_mask.float().unsqueeze(-1)
+            ], dim=-1) 
+        # vector features
+        X_ca = coords[:, :, 1]
+        orientations = GVPInputFeaturizer._orientations(X_ca)
+        sidechains = GVPInputFeaturizer._sidechains(coords)
+        node_vector_features = torch.cat([orientations, sidechains.unsqueeze(-2)], dim=-2)
+        return node_scalar_features, node_vector_features
+
+    @staticmethod
+    def _orientations(X):
+        forward = normalize(X[:, 1:] - X[:, :-1])
+        backward = normalize(X[:, :-1] - X[:, 1:])
+        forward = F.pad(forward, [0, 0, 0, 1])
+        backward = F.pad(backward, [0, 0, 1, 0])
+        return torch.cat([forward.unsqueeze(-2), backward.unsqueeze(-2)], -2)
+    
+    @staticmethod
+    def _sidechains(X):
+        n, origin, c = X[:, :, 0], X[:, :, 1], X[:, :, 2]
+        c, n = normalize(c - origin), normalize(n - origin)
+        bisector = normalize(c + n)
+        perp = normalize(torch.cross(c, n, dim=-1))
+        vec = -bisector * math.sqrt(1 / 3) - perp * math.sqrt(2 / 3)
+        return vec 
+
+    @staticmethod
+    def _dihedrals(X, eps=1e-7):
+        X = torch.flatten(X[:, :, :3], 1, 2)
+        bsz = X.shape[0]
+        dX = X[:, 1:] - X[:, :-1]
+        U = normalize(dX, dim=-1)
+        u_2 = U[:, :-2]
+        u_1 = U[:, 1:-1]
+        u_0 = U[:, 2:]
+    
+        # Backbone normals
+        n_2 = normalize(torch.cross(u_2, u_1, dim=-1), dim=-1)
+        n_1 = normalize(torch.cross(u_1, u_0, dim=-1), dim=-1)
+    
+        # Angle between normals
+        cosD = torch.sum(n_2 * n_1, -1)
+        cosD = torch.clamp(cosD, -1 + eps, 1 - eps)
+        D = torch.sign(torch.sum(u_2 * n_1, -1)) * torch.acos(cosD)
+    
+        # This scheme will remove phi[0], psi[-1], omega[-1]
+        D = F.pad(D, [1, 2]) 
+        D = torch.reshape(D, [bsz, -1, 3])
+        # Lift angle representations to the circle
+        D_features = torch.cat([torch.cos(D), torch.sin(D)], -1)
+        return D_features
+
+    @staticmethod
+    def _positional_embeddings(edge_index, 
+                               num_embeddings=None,
+                               num_positional_embeddings=16,
+                               period_range=[2, 1000]):
+        # From https://github.com/jingraham/neurips19-graph-protein-design
+        num_embeddings = num_embeddings or num_positional_embeddings
+        d = edge_index[0] - edge_index[1]
+     
+        frequency = torch.exp(
+            torch.arange(0, num_embeddings, 2, dtype=torch.float32,
+                device=edge_index.device)
+            * -(np.log(10000.0) / num_embeddings)
+        )
+        angles = d.unsqueeze(-1) * frequency
+        E = torch.cat((torch.cos(angles), torch.sin(angles)), -1)
+        return E
+
+    @staticmethod
+    def _dist(X, coord_mask, padding_mask, top_k_neighbors, eps=1e-8):
+        """ Pairwise euclidean distances """
+        bsz, maxlen = X.size(0), X.size(1)
+        coord_mask_2D = torch.unsqueeze(coord_mask,1) * torch.unsqueeze(coord_mask,2)
+        residue_mask = ~padding_mask
+        residue_mask_2D = torch.unsqueeze(residue_mask,1) * torch.unsqueeze(residue_mask,2)
+        dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)
+        D = coord_mask_2D * norm(dX, dim=-1)
+    
+        # sorting preference: first those with coords, then among the residues that
+        # exist but are masked use distance in sequence as tie breaker, and then the
+        # residues that came from padding are last
+        seqpos = torch.arange(maxlen, device=X.device)
+        Dseq = torch.abs(seqpos.unsqueeze(1) - seqpos.unsqueeze(0)).repeat(bsz, 1, 1)
+        D_adjust = nan_to_num(D) + (~coord_mask_2D) * (1e8 + Dseq*1e6) + (
+            ~residue_mask_2D) * (1e10)
+    
+        if top_k_neighbors == -1:
+            D_neighbors = D_adjust
+            E_idx = seqpos.repeat(
+                    *D_neighbors.shape[:-1], 1)
+        else:
+            # Identify k nearest neighbors (including self)
+            k = min(top_k_neighbors, X.size(1))
+            D_neighbors, E_idx = torch.topk(D_adjust, k, dim=-1, largest=False)
+    
+        coord_mask_neighbors = (D_neighbors < 5e7)
+        residue_mask_neighbors = (D_neighbors < 5e9)
+        return D_neighbors, E_idx, coord_mask_neighbors, residue_mask_neighbors
+
+
+class Normalize(nn.Module):
+    def __init__(self, features, epsilon=1e-6):
+        super(Normalize, self).__init__()
+        self.gain = nn.Parameter(torch.ones(features))
+        self.bias = nn.Parameter(torch.zeros(features))
+        self.epsilon = epsilon
+
+    def forward(self, x, dim=-1):
+        mu = x.mean(dim, keepdim=True)
+        sigma = torch.sqrt(x.var(dim, keepdim=True) + self.epsilon)
+        gain = self.gain
+        bias = self.bias
+        # Reshape
+        if dim != -1:
+            shape = [1] * len(mu.size())
+            shape[dim] = self.gain.size()[0]
+            gain = gain.view(shape)
+            bias = bias.view(shape)
+        return gain * (x - mu) / (sigma + self.epsilon) + bias
+
+
+class DihedralFeatures(nn.Module):
+    def __init__(self, node_embed_dim):
+        """ Embed dihedral angle features. """
+        super(DihedralFeatures, self).__init__()
+        # 3 dihedral angles; sin and cos of each angle
+        node_in = 6
+        # Normalization and embedding
+        self.node_embedding = nn.Linear(node_in,  node_embed_dim, bias=True)
+        self.norm_nodes = Normalize(node_embed_dim)
+
+    def forward(self, X):
+        """ Featurize coordinates as an attributed graph """
+        V = self._dihedrals(X)
+        V = self.node_embedding(V)
+        V = self.norm_nodes(V)
+        return V
+
+    @staticmethod
+    def _dihedrals(X, eps=1e-7, return_angles=False):
+        # First 3 coordinates are N, CA, C
+        X = X[:,:,:3,:].reshape(X.shape[0], 3*X.shape[1], 3)
+
+        # Shifted slices of unit vectors
+        dX = X[:,1:,:] - X[:,:-1,:]
+        U = F.normalize(dX, dim=-1)
+        u_2 = U[:,:-2,:]
+        u_1 = U[:,1:-1,:]
+        u_0 = U[:,2:,:]
+        # Backbone normals
+        n_2 = F.normalize(torch.cross(u_2, u_1, dim=-1), dim=-1)
+        n_1 = F.normalize(torch.cross(u_1, u_0, dim=-1), dim=-1)
+
+        # Angle between normals
+        cosD = (n_2 * n_1).sum(-1)
+        cosD = torch.clamp(cosD, -1+eps, 1-eps)
+        D = torch.sign((u_2 * n_1).sum(-1)) * torch.acos(cosD)
+
+        # This scheme will remove phi[0], psi[-1], omega[-1]
+        D = F.pad(D, (1,2), 'constant', 0)
+        D = D.view((D.size(0), int(D.size(1)/3), 3))
+        phi, psi, omega = torch.unbind(D,-1)
+
+        if return_angles:
+            return phi, psi, omega
+
+        # Lift angle representations to the circle
+        D_features = torch.cat((torch.cos(D), torch.sin(D)), 2)
+        return D_features
+
+
+class GVPGraphEmbedding(GVPInputFeaturizer):
+
+    def __init__(self, args):
+        super().__init__()
+        self.top_k_neighbors = args.top_k_neighbors
+        self.num_positional_embeddings = 16
+        self.remove_edges_without_coords = True
+        node_input_dim = (7, 3)
+        edge_input_dim = (34, 1)
+        node_hidden_dim = (args.node_hidden_dim_scalar,
+                args.node_hidden_dim_vector)
+        edge_hidden_dim = (args.edge_hidden_dim_scalar,
+                args.edge_hidden_dim_vector)
+        self.embed_node = nn.Sequential(
+            GVP(node_input_dim, node_hidden_dim, activations=(None, None)),
+            LayerNorm(node_hidden_dim, eps=1e-4)
+        )
+        self.embed_edge = nn.Sequential(
+            GVP(edge_input_dim, edge_hidden_dim, activations=(None, None)),
+            LayerNorm(edge_hidden_dim, eps=1e-4)
+        )
+        self.embed_confidence = nn.Linear(16, args.node_hidden_dim_scalar)
+
+    def forward(self, coords, coord_mask, padding_mask, confidence):
+        with torch.no_grad():
+            node_features = self.get_node_features(coords, coord_mask)
+            edge_features, edge_index = self.get_edge_features(
+                coords, coord_mask, padding_mask)
+        node_embeddings_scalar, node_embeddings_vector = self.embed_node(node_features)
+        edge_embeddings = self.embed_edge(edge_features)
+
+        rbf_rep = rbf(confidence, 0., 1.)
+        node_embeddings = (
+            node_embeddings_scalar + self.embed_confidence(rbf_rep),
+            node_embeddings_vector
+        )
+
+        node_embeddings, edge_embeddings, edge_index = flatten_graph(
+            node_embeddings, edge_embeddings, edge_index)
+        return node_embeddings, edge_embeddings, edge_index
+
+    def get_edge_features(self, coords, coord_mask, padding_mask):
+        X_ca = coords[:, :, 1]
+        # Get distances to the top k neighbors
+        E_dist, E_idx, E_coord_mask, E_residue_mask = GVPInputFeaturizer._dist(
+                X_ca, coord_mask, padding_mask, self.top_k_neighbors)
+        # Flatten the graph to be batch size 1 for torch_geometric package 
+        dest = E_idx
+        B, L, k = E_idx.shape[:3]
+        src = torch.arange(L, device=E_idx.device).view([1, L, 1]).expand(B, L, k)
+        # After flattening, [2, B, E]
+        edge_index = torch.stack([src, dest], dim=0).flatten(2, 3)
+        # After flattening, [B, E]
+        E_dist = E_dist.flatten(1, 2)
+        E_coord_mask = E_coord_mask.flatten(1, 2).unsqueeze(-1)
+        E_residue_mask = E_residue_mask.flatten(1, 2)
+        # Calculate relative positional embeddings and distance RBF 
+        pos_embeddings = GVPInputFeaturizer._positional_embeddings(
+            edge_index,
+            num_positional_embeddings=self.num_positional_embeddings,
+        )
+        D_rbf = rbf(E_dist, 0., 20.)
+        # Calculate relative orientation 
+        X_src = X_ca.unsqueeze(2).expand(-1, -1, k, -1).flatten(1, 2)
+        X_dest = torch.gather(
+            X_ca,
+            1,
+            edge_index[1, :, :].unsqueeze(-1).expand([B, L*k, 3])
+        )
+        coord_mask_src = coord_mask.unsqueeze(2).expand(-1, -1, k).flatten(1, 2)
+        coord_mask_dest = torch.gather(
+            coord_mask,
+            1,
+            edge_index[1, :, :].expand([B, L*k])
+        )
+        E_vectors = X_src - X_dest
+        # For the ones without coordinates, substitute in the average vector
+        E_vector_mean = torch.sum(E_vectors * E_coord_mask, dim=1,
+                keepdims=True) / torch.sum(E_coord_mask, dim=1, keepdims=True)
+        E_vectors = E_vectors * E_coord_mask + E_vector_mean * ~(E_coord_mask)
+        # Normalize and remove nans 
+        edge_s = torch.cat([D_rbf, pos_embeddings], dim=-1)
+        edge_v = normalize(E_vectors).unsqueeze(-2)
+        edge_s, edge_v = map(nan_to_num, (edge_s, edge_v))
+        # Also add indications of whether the coordinates are present 
+        edge_s = torch.cat([
+            edge_s,
+            (~coord_mask_src).float().unsqueeze(-1),
+            (~coord_mask_dest).float().unsqueeze(-1),
+        ], dim=-1)
+        edge_index[:, ~E_residue_mask] = -1
+        if self.remove_edges_without_coords:
+            edge_index[:, ~E_coord_mask.squeeze(-1)] = -1
+        return (edge_s, edge_v), edge_index.transpose(0, 1) 

esm/source/esm/inverse_folding/gvp_encoder.py

@@ -0,0 +1,56 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from argparse import Namespace
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .features import GVPGraphEmbedding
+from .gvp_modules import GVPConvLayer, LayerNorm 
+from .gvp_utils import unflatten_graph
+
+
+
+class GVPEncoder(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.embed_graph = GVPGraphEmbedding(args)
+
+        node_hidden_dim = (args.node_hidden_dim_scalar,
+                args.node_hidden_dim_vector)
+        edge_hidden_dim = (args.edge_hidden_dim_scalar,
+                args.edge_hidden_dim_vector)
+        
+        conv_activations = (F.relu, torch.sigmoid)
+        self.encoder_layers = nn.ModuleList(
+                GVPConvLayer(
+                    node_hidden_dim,
+                    edge_hidden_dim,
+                    drop_rate=args.dropout,
+                    vector_gate=True,
+                    attention_heads=0,
+                    n_message=3,
+                    conv_activations=conv_activations,
+                    n_edge_gvps=0,
+                    eps=1e-4,
+                    layernorm=True,
+                ) 
+            for i in range(args.num_encoder_layers)
+        )
+
+    def forward(self, coords, coord_mask, padding_mask, confidence):
+        node_embeddings, edge_embeddings, edge_index = self.embed_graph(
+                coords, coord_mask, padding_mask, confidence)
+        
+        for i, layer in enumerate(self.encoder_layers):
+            node_embeddings, edge_embeddings = layer(node_embeddings,
+                    edge_index, edge_embeddings)
+
+        node_embeddings = unflatten_graph(node_embeddings, coords.shape[0])
+        return node_embeddings

esm/source/esm/inverse_folding/gvp_modules.py

@@ -0,0 +1,475 @@
+# Contents of this file are from the open source code for
+#
+#   Jing, B., Eismann, S., Suriana, P., Townshend, R. J. L., & Dror, R. (2020).
+#   Learning from Protein Structure with Geometric Vector Perceptrons. In
+#   International Conference on Learning Representations.
+#
+# MIT License
+# 
+# Copyright (c) 2020 Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael Townshend, Ron Dror
+# 
+# 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.
+
+import typing as T
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch_geometric.nn import MessagePassing
+
+def tuple_size(tp):
+    return tuple([0 if a is None else a.size() for a in tp])
+
+def tuple_sum(tp1, tp2):
+    s1, v1 = tp1
+    s2, v2 = tp2
+    if v2 is None and v2 is None:
+        return (s1 + s2, None)
+    return (s1 + s2, v1 + v2)
+
+def tuple_cat(*args, dim=-1):
+    '''
+    Concatenates any number of tuples (s, V) elementwise.
+    
+    :param dim: dimension along which to concatenate when viewed
+                as the `dim` index for the scalar-channel tensors.
+                This means that `dim=-1` will be applied as
+                `dim=-2` for the vector-channel tensors.
+    '''
+    dim %= len(args[0][0].shape)
+    s_args, v_args = list(zip(*args))
+    return torch.cat(s_args, dim=dim), torch.cat(v_args, dim=dim)
+
+def tuple_index(x, idx):
+    '''
+    Indexes into a tuple (s, V) along the first dimension.
+    
+    :param idx: any object which can be used to index into a `torch.Tensor`
+    '''
+    return x[0][idx], x[1][idx]
+
+def randn(n, dims, device="cpu"):
+    '''
+    Returns random tuples (s, V) drawn elementwise from a normal distribution.
+    
+    :param n: number of data points
+    :param dims: tuple of dimensions (n_scalar, n_vector)
+    
+    :return: (s, V) with s.shape = (n, n_scalar) and
+             V.shape = (n, n_vector, 3)
+    '''
+    return torch.randn(n, dims[0], device=device), \
+            torch.randn(n, dims[1], 3, device=device)
+
+def _norm_no_nan(x, axis=-1, keepdims=False, eps=1e-8, sqrt=True):
+    '''
+    L2 norm of tensor clamped above a minimum value `eps`.
+    
+    :param sqrt: if `False`, returns the square of the L2 norm
+    '''
+    # clamp is slow
+    # out = torch.clamp(torch.sum(torch.square(x), axis, keepdims), min=eps)
+    out = torch.sum(torch.square(x), axis, keepdims) + eps
+    return torch.sqrt(out) if sqrt else out
+
+def _split(x, nv):
+    '''
+    Splits a merged representation of (s, V) back into a tuple. 
+    Should be used only with `_merge(s, V)` and only if the tuple 
+    representation cannot be used.
+    
+    :param x: the `torch.Tensor` returned from `_merge`
+    :param nv: the number of vector channels in the input to `_merge`
+    '''
+    v = torch.reshape(x[..., -3*nv:], x.shape[:-1] + (nv, 3))
+    s = x[..., :-3*nv]
+    return s, v
+
+def _merge(s, v):
+    '''
+    Merges a tuple (s, V) into a single `torch.Tensor`, where the
+    vector channels are flattened and appended to the scalar channels.
+    Should be used only if the tuple representation cannot be used.
+    Use `_split(x, nv)` to reverse.
+    '''
+    v = torch.reshape(v, v.shape[:-2] + (3*v.shape[-2],))
+    return torch.cat([s, v], -1)
+
+class GVP(nn.Module):
+    '''
+    Geometric Vector Perceptron. See manuscript and README.md
+    for more details.
+    
+    :param in_dims: tuple (n_scalar, n_vector)
+    :param out_dims: tuple (n_scalar, n_vector)
+    :param h_dim: intermediate number of vector channels, optional
+    :param activations: tuple of functions (scalar_act, vector_act)
+    :param tuple_io: whether to keep accepting tuple inputs and outputs when vi
+    or vo = 0
+    '''
+    def __init__(self, in_dims, out_dims, h_dim=None, vector_gate=False,
+                 activations=(F.relu, torch.sigmoid), tuple_io=True,
+                 eps=1e-8):
+        super(GVP, self).__init__()
+        self.si, self.vi = in_dims
+        self.so, self.vo = out_dims
+        self.tuple_io = tuple_io
+        if self.vi: 
+            self.h_dim = h_dim or max(self.vi, self.vo) 
+            self.wh = nn.Linear(self.vi, self.h_dim, bias=False)
+            self.ws = nn.Linear(self.h_dim + self.si, self.so)
+            if self.vo:
+                self.wv = nn.Linear(self.h_dim, self.vo, bias=False)
+                if vector_gate:
+                    self.wg = nn.Linear(self.so, self.vo)
+        else:
+            self.ws = nn.Linear(self.si, self.so)
+        
+        self.vector_gate = vector_gate
+        self.scalar_act, self.vector_act = activations
+        self.eps = eps
+        
+    def forward(self, x):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`, 
+                  or (if vectors_in is 0), a single `torch.Tensor`
+        :return: tuple (s, V) of `torch.Tensor`,
+                 or (if vectors_out is 0), a single `torch.Tensor`
+        '''
+        if self.vi:
+            s, v = x
+            v = torch.transpose(v, -1, -2)
+            vh = self.wh(v)    
+            vn = _norm_no_nan(vh, axis=-2, eps=self.eps)
+            s = self.ws(torch.cat([s, vn], -1))
+            if self.scalar_act:
+                s = self.scalar_act(s)
+            if self.vo: 
+                v = self.wv(vh) 
+                v = torch.transpose(v, -1, -2)
+                if self.vector_gate:
+                    g = self.wg(s).unsqueeze(-1)
+                else:
+                    g = _norm_no_nan(v, axis=-1, keepdims=True, eps=self.eps)
+                if self.vector_act:
+                    g = self.vector_act(g)
+                    v = v * g
+        else:
+            if self.tuple_io:
+                assert x[1] is None
+                x = x[0]
+            s = self.ws(x)
+            if self.scalar_act:
+                s = self.scalar_act(s)
+            if self.vo:
+                v = torch.zeros(list(s.shape)[:-1] + [self.vo, 3],
+                        device=s.device)
+        
+        if self.vo:
+            return (s, v)
+        elif self.tuple_io:
+            return (s, None)
+        else:
+            return s
+
+
+class _VDropout(nn.Module):
+    '''
+    Vector channel dropout where the elements of each
+    vector channel are dropped together.
+    '''
+    def __init__(self, drop_rate):
+        super(_VDropout, self).__init__()
+        self.drop_rate = drop_rate
+
+    def forward(self, x):
+        '''
+        :param x: `torch.Tensor` corresponding to vector channels
+        '''
+        if x is None:
+            return None
+        device = x.device
+        if not self.training:
+            return x
+        mask = torch.bernoulli(
+            (1 - self.drop_rate) * torch.ones(x.shape[:-1], device=device)
+        ).unsqueeze(-1)
+        x = mask * x / (1 - self.drop_rate)
+        return x
+
+class Dropout(nn.Module):
+    '''
+    Combined dropout for tuples (s, V).
+    Takes tuples (s, V) as input and as output.
+    '''
+    def __init__(self, drop_rate):
+        super(Dropout, self).__init__()
+        self.sdropout = nn.Dropout(drop_rate)
+        self.vdropout = _VDropout(drop_rate)
+
+    def forward(self, x):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`,
+                  or single `torch.Tensor` 
+                  (will be assumed to be scalar channels)
+        '''
+        if type(x) is torch.Tensor:
+            return self.sdropout(x)
+        s, v = x
+        return self.sdropout(s), self.vdropout(v)
+
+class LayerNorm(nn.Module):
+    '''
+    Combined LayerNorm for tuples (s, V).
+    Takes tuples (s, V) as input and as output.
+    '''
+    def __init__(self, dims, tuple_io=True, eps=1e-8):
+        super(LayerNorm, self).__init__()
+        self.tuple_io = tuple_io
+        self.s, self.v = dims
+        self.scalar_norm = nn.LayerNorm(self.s)
+        self.eps = eps
+        
+    def forward(self, x):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`,
+                  or single `torch.Tensor` 
+                  (will be assumed to be scalar channels)
+        '''
+        if not self.v:
+            if self.tuple_io:
+                return self.scalar_norm(x[0]), None
+            return self.scalar_norm(x)
+        s, v = x
+        vn = _norm_no_nan(v, axis=-1, keepdims=True, sqrt=False, eps=self.eps)
+        nonzero_mask = (vn > 2 * self.eps)
+        vn = torch.sum(vn * nonzero_mask, dim=-2, keepdim=True
+            ) / (self.eps + torch.sum(nonzero_mask, dim=-2, keepdim=True))
+        vn = torch.sqrt(vn + self.eps)
+        v = nonzero_mask * (v / vn)
+        return self.scalar_norm(s), v
+
+class GVPConv(MessagePassing):
+    '''
+    Graph convolution / message passing with Geometric Vector Perceptrons.
+    Takes in a graph with node and edge embeddings,
+    and returns new node embeddings.
+    
+    This does NOT do residual updates and pointwise feedforward layers
+    ---see `GVPConvLayer`.
+    
+    :param in_dims: input node embedding dimensions (n_scalar, n_vector)
+    :param out_dims: output node embedding dimensions (n_scalar, n_vector)
+    :param edge_dims: input edge embedding dimensions (n_scalar, n_vector)
+    :param n_layers: number of GVPs in the message function
+    :param module_list: preconstructed message function, overrides n_layers
+    :param aggr: should be "add" if some incoming edges are masked, as in
+                 a masked autoregressive decoder architecture
+    '''
+    def __init__(self, in_dims, out_dims, edge_dims, n_layers=3,
+            vector_gate=False, module_list=None, aggr="mean", eps=1e-8,
+            activations=(F.relu, torch.sigmoid)):
+        super(GVPConv, self).__init__(aggr=aggr)
+        self.eps = eps
+        self.si, self.vi = in_dims
+        self.so, self.vo = out_dims
+        self.se, self.ve = edge_dims
+        
+        module_list = module_list or []
+        if not module_list:
+            if n_layers == 1:
+                module_list.append(
+                    GVP((2*self.si + self.se, 2*self.vi + self.ve), 
+                        (self.so, self.vo), activations=(None, None)))
+            else:
+                module_list.append(
+                    GVP((2*self.si + self.se, 2*self.vi + self.ve), out_dims,
+                        vector_gate=vector_gate, activations=activations)
+                )
+                for i in range(n_layers - 2):
+                    module_list.append(GVP(out_dims, out_dims,
+                        vector_gate=vector_gate))
+                module_list.append(GVP(out_dims, out_dims,
+                                       activations=(None, None)))
+        self.message_func = nn.Sequential(*module_list)
+
+    def forward(self, x, edge_index, edge_attr):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`
+        :param edge_index: array of shape [2, n_edges]
+        :param edge_attr: tuple (s, V) of `torch.Tensor`
+        '''
+        x_s, x_v = x
+        message = self.propagate(edge_index, 
+                    s=x_s, v=x_v.reshape(x_v.shape[0], 3*x_v.shape[1]),
+                    edge_attr=edge_attr)
+        return _split(message, self.vo) 
+
+    def message(self, s_i, v_i, s_j, v_j, edge_attr):
+        v_j = v_j.view(v_j.shape[0], v_j.shape[1]//3, 3)
+        v_i = v_i.view(v_i.shape[0], v_i.shape[1]//3, 3)
+        message = tuple_cat((s_j, v_j), edge_attr, (s_i, v_i))
+        message = self.message_func(message)
+        return _merge(*message)
+
+
+class GVPConvLayer(nn.Module):
+    '''
+    Full graph convolution / message passing layer with 
+    Geometric Vector Perceptrons. Residually updates node embeddings with
+    aggregated incoming messages, applies a pointwise feedforward 
+    network to node embeddings, and returns updated node embeddings.
+    
+    To only compute the aggregated messages, see `GVPConv`.
+    
+    :param node_dims: node embedding dimensions (n_scalar, n_vector)
+    :param edge_dims: input edge embedding dimensions (n_scalar, n_vector)
+    :param n_message: number of GVPs to use in message function
+    :param n_feedforward: number of GVPs to use in feedforward function
+    :param drop_rate: drop probability in all dropout layers
+    :param autoregressive: if `True`, this `GVPConvLayer` will be used
+           with a different set of input node embeddings for messages
+           where src >= dst
+    '''
+    def __init__(self, node_dims, edge_dims, vector_gate=False,
+                 n_message=3, n_feedforward=2, drop_rate=.1,
+                 autoregressive=False, attention_heads=0,
+                 conv_activations=(F.relu, torch.sigmoid),
+                 n_edge_gvps=0, layernorm=True, eps=1e-8):
+        
+        super(GVPConvLayer, self).__init__()
+        if attention_heads == 0:
+            self.conv = GVPConv(
+                    node_dims, node_dims, edge_dims, n_layers=n_message,
+                    vector_gate=vector_gate,
+                    aggr="add" if autoregressive else "mean",
+                    activations=conv_activations, 
+                    eps=eps,
+            )
+        else:
+            raise NotImplementedError
+        if layernorm:
+            self.norm = nn.ModuleList([LayerNorm(node_dims, eps=eps) for _ in range(2)])
+        else:
+            self.norm = nn.ModuleList([nn.Identity() for _ in range(2)])
+        self.dropout = nn.ModuleList([Dropout(drop_rate) for _ in range(2)])
+
+        ff_func = []
+        if n_feedforward == 1:
+            ff_func.append(GVP(node_dims, node_dims, activations=(None, None)))
+        else:
+            hid_dims = 4*node_dims[0], 2*node_dims[1]
+            ff_func.append(GVP(node_dims, hid_dims, vector_gate=vector_gate))
+            for i in range(n_feedforward-2):
+                ff_func.append(GVP(hid_dims, hid_dims, vector_gate=vector_gate))
+            ff_func.append(GVP(hid_dims, node_dims, activations=(None, None)))
+        self.ff_func = nn.Sequential(*ff_func)
+
+        self.edge_message_func = None
+        if n_edge_gvps > 0:
+            si, vi = node_dims
+            se, ve = edge_dims
+            module_list = [
+                GVP((2*si + se, 2*vi + ve), edge_dims, vector_gate=vector_gate)
+            ]
+            for i in range(n_edge_gvps - 2):
+                module_list.append(GVP(edge_dims, edge_dims,
+                    vector_gate=vector_gate))
+            if n_edge_gvps > 1:
+                module_list.append(GVP(edge_dims, edge_dims,
+                    activations=(None, None)))
+            self.edge_message_func = nn.Sequential(*module_list)
+            if layernorm:
+                self.edge_norm = LayerNorm(edge_dims, eps=eps)
+            else:
+                self.edge_norm = nn.Identity()
+            self.edge_dropout = Dropout(drop_rate)
+
+    def forward(self, x, edge_index, edge_attr,
+                autoregressive_x=None, node_mask=None):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`
+        :param edge_index: array of shape [2, n_edges]
+        :param edge_attr: tuple (s, V) of `torch.Tensor`
+        :param autoregressive_x: tuple (s, V) of `torch.Tensor`. 
+                If not `None`, will be used as srcqq node embeddings
+                for forming messages where src >= dst. The corrent node 
+                embeddings `x` will still be the base of the update and the 
+                pointwise feedforward.
+        :param node_mask: array of type `bool` to index into the first
+                dim of node embeddings (s, V). If not `None`, only
+                these nodes will be updated.
+        '''
+        if self.edge_message_func:
+            src, dst = edge_index
+            if autoregressive_x is None:
+                x_src = x[0][src], x[1][src]
+            else: 
+                mask = (src < dst).unsqueeze(-1)
+                x_src = (
+                    torch.where(mask, x[0][src], autoregressive_x[0][src]),
+                    torch.where(mask.unsqueeze(-1), x[1][src],
+                        autoregressive_x[1][src])
+                )
+            x_dst = x[0][dst], x[1][dst]
+            x_edge = (
+                torch.cat([x_src[0], edge_attr[0], x_dst[0]], dim=-1),
+                torch.cat([x_src[1], edge_attr[1], x_dst[1]], dim=-2)
+            )
+            edge_attr_dh = self.edge_message_func(x_edge)
+            edge_attr = self.edge_norm(tuple_sum(edge_attr,
+                self.edge_dropout(edge_attr_dh)))
+        
+        if autoregressive_x is not None:
+            # Guarding this import here to remove the dependency on torch_scatter, since this isn't used
+            # in ESM-IF1
+            from torch_scatter import scatter_add
+            src, dst = edge_index
+            mask = src < dst
+            edge_index_forward = edge_index[:, mask]
+            edge_index_backward = edge_index[:, ~mask]
+            edge_attr_forward = tuple_index(edge_attr, mask)
+            edge_attr_backward = tuple_index(edge_attr, ~mask)
+            
+            dh = tuple_sum(
+                self.conv(x, edge_index_forward, edge_attr_forward),
+                self.conv(autoregressive_x, edge_index_backward, edge_attr_backward)
+            )
+                        
+            count = scatter_add(torch.ones_like(dst), dst,
+                        dim_size=dh[0].size(0)).clamp(min=1).unsqueeze(-1)
+            
+            dh = dh[0] / count, dh[1] / count.unsqueeze(-1)
+
+        else:
+            dh = self.conv(x, edge_index, edge_attr)
+        
+        if node_mask is not None:
+            x_ = x
+            x, dh = tuple_index(x, node_mask), tuple_index(dh, node_mask)
+            
+        x = self.norm[0](tuple_sum(x, self.dropout[0](dh)))
+        
+        dh = self.ff_func(x)
+        x = self.norm[1](tuple_sum(x, self.dropout[1](dh)))
+        
+        if node_mask is not None:
+            x_[0][node_mask], x_[1][node_mask] = x[0], x[1]
+            x = x_
+
+        return x, edge_attr

esm/source/esm/inverse_folding/gvp_transformer.py

@@ -0,0 +1,140 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+from typing import Any, Dict, List, Optional, Tuple, NamedTuple
+import torch
+from torch import nn
+from torch import Tensor
+import torch.nn.functional as F
+from scipy.spatial import transform
+
+from esm.data import Alphabet
+
+from .features import DihedralFeatures
+from .gvp_encoder import GVPEncoder
+from .gvp_utils import unflatten_graph
+from .gvp_transformer_encoder import GVPTransformerEncoder
+from .transformer_decoder import TransformerDecoder
+from .util import rotate, CoordBatchConverter 
+
+
+class GVPTransformerModel(nn.Module):
+    """
+    GVP-Transformer inverse folding model.
+
+    Architecture: Geometric GVP-GNN as initial layers, followed by
+    sequence-to-sequence Transformer encoder and decoder.
+    """
+
+    def __init__(self, args, alphabet):
+        super().__init__()
+        encoder_embed_tokens = self.build_embedding(
+            args, alphabet, args.encoder_embed_dim,
+        )
+        decoder_embed_tokens = self.build_embedding(
+            args, alphabet, args.decoder_embed_dim, 
+        )
+        encoder = self.build_encoder(args, alphabet, encoder_embed_tokens)
+        decoder = self.build_decoder(args, alphabet, decoder_embed_tokens)
+        self.args = args
+        self.encoder = encoder
+        self.decoder = decoder
+
+    @classmethod
+    def build_encoder(cls, args, src_dict, embed_tokens):
+        encoder = GVPTransformerEncoder(args, src_dict, embed_tokens)
+        return encoder
+
+    @classmethod
+    def build_decoder(cls, args, tgt_dict, embed_tokens):
+        decoder = TransformerDecoder(
+            args,
+            tgt_dict,
+            embed_tokens,
+        )
+        return decoder
+
+    @classmethod
+    def build_embedding(cls, args, dictionary, embed_dim):
+        num_embeddings = len(dictionary)
+        padding_idx = dictionary.padding_idx
+        emb = nn.Embedding(num_embeddings, embed_dim, padding_idx)
+        nn.init.normal_(emb.weight, mean=0, std=embed_dim ** -0.5)
+        nn.init.constant_(emb.weight[padding_idx], 0)
+        return emb
+
+    def forward(
+        self,
+        coords,
+        padding_mask,
+        confidence,
+        prev_output_tokens,
+        return_all_hiddens: bool = False,
+        features_only: bool = False,
+    ):
+        encoder_out = self.encoder(coords, padding_mask, confidence,
+            return_all_hiddens=return_all_hiddens)
+        logits, extra = self.decoder(
+            prev_output_tokens,
+            encoder_out=encoder_out,
+            features_only=features_only,
+            return_all_hiddens=return_all_hiddens,
+        )
+        return logits, extra
+    
+    def sample(self, coords, partial_seq=None, temperature=1.0, confidence=None, device=None):
+        """
+        Samples sequences based on multinomial sampling (no beam search).
+
+        Args:
+            coords: L x 3 x 3 list representing one backbone
+            partial_seq: Optional, partial sequence with mask tokens if part of
+                the sequence is known
+            temperature: sampling temperature, use low temperature for higher
+                sequence recovery and high temperature for higher diversity
+            confidence: optional length L list of confidence scores for coordinates
+        """
+        L = len(coords)
+        # Convert to batch format
+        batch_converter = CoordBatchConverter(self.decoder.dictionary)
+        batch_coords, confidence, _, _, padding_mask = (
+            batch_converter([(coords, confidence, None)], device=device)
+        )
+        
+        # Start with prepend token
+        mask_idx = self.decoder.dictionary.get_idx('<mask>')
+        sampled_tokens = torch.full((1, 1+L), mask_idx, dtype=int)
+        sampled_tokens[0, 0] = self.decoder.dictionary.get_idx('<cath>')
+        if partial_seq is not None:
+            for i, c in enumerate(partial_seq):
+                sampled_tokens[0, i+1] = self.decoder.dictionary.get_idx(c)
+            
+        # Save incremental states for faster sampling
+        incremental_state = dict()
+        
+        # Run encoder only once
+        encoder_out = self.encoder(batch_coords, padding_mask, confidence)
+        
+        # Make sure all tensors are on the same device if a GPU is present
+        if device:
+            sampled_tokens = sampled_tokens.to(device)
+        
+        # Decode one token at a time
+        for i in range(1, L+1):
+            logits, _ = self.decoder(
+                sampled_tokens[:, :i], 
+                encoder_out,
+                incremental_state=incremental_state,
+            )
+            logits = logits[0].transpose(0, 1)
+            logits /= temperature
+            probs = F.softmax(logits, dim=-1)
+            if sampled_tokens[0, i] == mask_idx:
+                sampled_tokens[:, i] = torch.multinomial(probs, 1).squeeze(-1)
+        sampled_seq = sampled_tokens[0, 1:]
+        
+        # Convert back to string via lookup
+        return ''.join([self.decoder.dictionary.get_tok(a) for a in sampled_seq])

esm/source/esm/inverse_folding/gvp_transformer_encoder.py

@@ -0,0 +1,184 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Contents of this file were adapted from the open source fairseq repository.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import math
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from esm.modules import SinusoidalPositionalEmbedding
+from .features import GVPInputFeaturizer, DihedralFeatures
+from .gvp_encoder import GVPEncoder
+from .transformer_layer import TransformerEncoderLayer
+from .util import nan_to_num, get_rotation_frames, rotate, rbf
+
+
+class GVPTransformerEncoder(nn.Module):
+    """
+    Transformer encoder consisting of *args.encoder.layers* layers. Each layer
+    is a :class:`TransformerEncoderLayer`.
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+        dictionary (~fairseq.data.Dictionary): encoding dictionary
+        embed_tokens (torch.nn.Embedding): input embedding
+    """
+
+    def __init__(self, args, dictionary, embed_tokens):
+        super().__init__()
+        self.args = args
+        self.dictionary = dictionary
+
+        self.dropout_module = nn.Dropout(args.dropout)
+
+        embed_dim = embed_tokens.embedding_dim
+        self.padding_idx = embed_tokens.padding_idx
+
+        self.embed_tokens = embed_tokens
+        self.embed_scale = math.sqrt(embed_dim)
+        self.embed_positions = SinusoidalPositionalEmbedding(
+            embed_dim,
+            self.padding_idx,
+        )
+        self.embed_gvp_input_features = nn.Linear(15, embed_dim)
+        self.embed_confidence = nn.Linear(16, embed_dim)
+        self.embed_dihedrals = DihedralFeatures(embed_dim)
+
+        gvp_args = argparse.Namespace()
+        for k, v in vars(args).items():
+            if k.startswith("gvp_"):
+                setattr(gvp_args, k[4:], v)
+        self.gvp_encoder = GVPEncoder(gvp_args)
+        gvp_out_dim = gvp_args.node_hidden_dim_scalar + (3 *
+                gvp_args.node_hidden_dim_vector)
+        self.embed_gvp_output = nn.Linear(gvp_out_dim, embed_dim)
+
+        self.layers = nn.ModuleList([])
+        self.layers.extend(
+            [self.build_encoder_layer(args) for i in range(args.encoder_layers)]
+        )
+        self.num_layers = len(self.layers)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+    def build_encoder_layer(self, args):
+        return TransformerEncoderLayer(args)
+
+    def forward_embedding(self, coords, padding_mask, confidence):
+        """
+        Args:
+            coords: N, CA, C backbone coordinates in shape length x 3 (atoms) x 3 
+            padding_mask: boolean Tensor (true for padding) of shape length
+            confidence: confidence scores between 0 and 1 of shape length
+        """
+        components = dict()
+        coord_mask = torch.all(torch.all(torch.isfinite(coords), dim=-1), dim=-1)
+        coords = nan_to_num(coords)
+        mask_tokens = (
+            padding_mask * self.dictionary.padding_idx + 
+            ~padding_mask * self.dictionary.get_idx("<mask>")
+        )
+        components["tokens"] = self.embed_tokens(mask_tokens) * self.embed_scale
+        components["diherals"] = self.embed_dihedrals(coords)
+
+        # GVP encoder
+        gvp_out_scalars, gvp_out_vectors = self.gvp_encoder(coords,
+                coord_mask, padding_mask, confidence)
+        R = get_rotation_frames(coords)
+        # Rotate to local rotation frame for rotation-invariance
+        gvp_out_features = torch.cat([
+            gvp_out_scalars,
+            rotate(gvp_out_vectors, R.transpose(-2, -1)).flatten(-2, -1),
+        ], dim=-1)
+        components["gvp_out"] = self.embed_gvp_output(gvp_out_features)
+
+        components["confidence"] = self.embed_confidence(
+             rbf(confidence, 0., 1.))
+
+        # In addition to GVP encoder outputs, also directly embed GVP input node
+        # features to the Transformer
+        scalar_features, vector_features = GVPInputFeaturizer.get_node_features(
+            coords, coord_mask, with_coord_mask=False)
+        features = torch.cat([
+            scalar_features,
+            rotate(vector_features, R.transpose(-2, -1)).flatten(-2, -1),
+        ], dim=-1)
+        components["gvp_input_features"] = self.embed_gvp_input_features(features)
+
+        embed = sum(components.values())
+        # for k, v in components.items():
+        #     print(k, torch.mean(v, dim=(0,1)), torch.std(v, dim=(0,1)))
+
+        x = embed
+        x = x + self.embed_positions(mask_tokens)
+        x = self.dropout_module(x)
+        return x, components 
+
+    def forward(
+        self,
+        coords,
+        encoder_padding_mask,
+        confidence,
+        return_all_hiddens: bool = False,
+    ):
+        """
+        Args:
+            coords (Tensor): backbone coordinates
+                shape batch_size x num_residues x num_atoms (3 for N, CA, C) x 3
+            encoder_padding_mask (ByteTensor): the positions of
+                  padding elements of shape `(batch_size x num_residues)`
+            confidence (Tensor): the confidence score of shape (batch_size x
+                num_residues). The value is between 0. and 1. for each residue
+                coordinate, or -1. if no coordinate is given
+            return_all_hiddens (bool, optional): also return all of the
+                intermediate hidden states (default: False).
+
+        Returns:
+            dict:
+                - **encoder_out** (Tensor): the last encoder layer's output of
+                  shape `(num_residues, batch_size, embed_dim)`
+                - **encoder_padding_mask** (ByteTensor): the positions of
+                  padding elements of shape `(batch_size, num_residues)`
+                - **encoder_embedding** (Tensor): the (scaled) embedding lookup
+                  of shape `(batch_size, num_residues, embed_dim)`
+                - **encoder_states** (List[Tensor]): all intermediate
+                  hidden states of shape `(num_residues, batch_size, embed_dim)`.
+                  Only populated if *return_all_hiddens* is True.
+        """
+        x, encoder_embedding = self.forward_embedding(coords,
+                encoder_padding_mask, confidence)
+        # account for padding while computing the representation
+        x = x * (1 - encoder_padding_mask.unsqueeze(-1).type_as(x))
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        encoder_states = []
+
+        if return_all_hiddens:
+            encoder_states.append(x)
+
+        # encoder layers
+        for layer in self.layers:
+            x = layer(
+                x, encoder_padding_mask=encoder_padding_mask
+            )
+            if return_all_hiddens:
+                assert encoder_states is not None
+                encoder_states.append(x)
+
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        return {
+            "encoder_out": [x],  # T x B x C
+            "encoder_padding_mask": [encoder_padding_mask],  # B x T
+            "encoder_embedding": [encoder_embedding],  # dictionary
+            "encoder_states": encoder_states,  # List[T x B x C]
+        }

esm/source/esm/inverse_folding/gvp_utils.py

@@ -0,0 +1,68 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+
+def flatten_graph(node_embeddings, edge_embeddings, edge_index):
+    """
+    Flattens the graph into a batch size one (with disconnected subgraphs for
+    each example) to be compatible with pytorch-geometric package.
+    Args:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch size x nodes x node_embed_dim
+                - vector: shape batch size x nodes x node_embed_dim x 3
+        edge_embeddings: edge embeddings of in tuple form (scalar, vector)
+                - scalar: shape batch size x edges x edge_embed_dim
+                - vector: shape batch size x edges x edge_embed_dim x 3
+        edge_index: shape batch_size x 2 (source node and target node) x edges
+    Returns:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch total_nodes x node_embed_dim
+                - vector: shape batch total_nodes x node_embed_dim x 3
+        edge_embeddings: edge embeddings of in tuple form (scalar, vector)
+                - scalar: shape batch total_edges x edge_embed_dim
+                - vector: shape batch total_edges x edge_embed_dim x 3
+        edge_index: shape 2 x total_edges
+    """
+    x_s, x_v = node_embeddings
+    e_s, e_v = edge_embeddings
+    batch_size, N = x_s.shape[0], x_s.shape[1]
+    node_embeddings = (torch.flatten(x_s, 0, 1), torch.flatten(x_v, 0, 1))
+    edge_embeddings = (torch.flatten(e_s, 0, 1), torch.flatten(e_v, 0, 1))
+
+    edge_mask = torch.any(edge_index != -1, dim=1)
+    # Re-number the nodes by adding batch_idx * N to each batch
+    edge_index = edge_index + (torch.arange(batch_size, device=edge_index.device) *
+            N).unsqueeze(-1).unsqueeze(-1)
+    edge_index = edge_index.permute(1, 0, 2).flatten(1, 2)
+    edge_mask = edge_mask.flatten()
+    edge_index = edge_index[:, edge_mask] 
+    edge_embeddings = (
+        edge_embeddings[0][edge_mask, :],
+        edge_embeddings[1][edge_mask, :]
+    )
+    return node_embeddings, edge_embeddings, edge_index 
+
+
+def unflatten_graph(node_embeddings, batch_size):
+    """
+    Unflattens node embeddings.
+    Args:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch total_nodes x node_embed_dim
+                - vector: shape batch total_nodes x node_embed_dim x 3
+        batch_size: int
+    Returns:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch size x nodes x node_embed_dim
+                - vector: shape batch size x nodes x node_embed_dim x 3
+    """
+    x_s, x_v = node_embeddings
+    x_s = x_s.reshape(batch_size, -1, x_s.shape[1])
+    x_v = x_v.reshape(batch_size, -1, x_v.shape[1], x_v.shape[2])
+    return (x_s, x_v)
+
+

esm/source/esm/inverse_folding/multichain_util.py

@@ -0,0 +1,152 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import biotite.structure
+import numpy as np
+import torch
+from typing import Sequence, Tuple, List
+
+from esm.inverse_folding.util import (
+    load_structure,
+    extract_coords_from_structure,
+    load_coords,
+    get_sequence_loss,
+    get_encoder_output,
+)
+
+
+def extract_coords_from_complex(structure: biotite.structure.AtomArray):
+    """
+    Args:
+        structure: biotite AtomArray
+    Returns:
+        Tuple (coords_list, seq_list)
+        - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+          coordinates representing the backbone of each chain
+        - seqs: Dictionary mapping chain ids to native sequences of each chain
+    """
+    coords = {}
+    seqs = {}
+    all_chains = biotite.structure.get_chains(structure)
+    for chain_id in all_chains:
+        chain = structure[structure.chain_id == chain_id]
+        coords[chain_id], seqs[chain_id] = extract_coords_from_structure(chain)
+    return coords, seqs
+
+
+def load_complex_coords(fpath, chains):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chains: the chain ids (the order matters for autoregressive model)
+    Returns:
+        Tuple (coords_list, seq_list)
+        - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+          coordinates representing the backbone of each chain
+        - seqs: Dictionary mapping chain ids to native sequences of each chain
+    """
+    structure = load_structure(fpath, chains)
+    return extract_coords_from_complex(structure)
+
+
+def _concatenate_coords(coords, target_chain_id, padding_length=10):
+    """
+    Args:
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        padding_length: Length of padding between concatenated chains
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates, a
+              concatenation of the chains with padding in between
+            - seq is the extracted sequence, with padding tokens inserted
+              between the concatenated chains
+    """
+    pad_coords = np.full((padding_length, 3, 3), np.nan, dtype=np.float32)
+    # For best performance, put the target chain first in concatenation.
+    coords_list = [coords[target_chain_id]]
+    for chain_id in coords:
+        if chain_id == target_chain_id:
+            continue
+        coords_list.append(pad_coords)
+        coords_list.append(coords[chain_id])
+    coords_concatenated = np.concatenate(coords_list, axis=0)
+    return coords_concatenated
+
+
+def sample_sequence_in_complex(model, coords, target_chain_id, temperature=1.,
+        padding_length=10):
+    """
+    Samples sequence for one chain in a complex.
+    Args:
+        model: An instance of the GVPTransformer model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        padding_length: padding length in between chains
+    Returns:
+        Sampled sequence for the target chain
+    """
+    target_chain_len = coords[target_chain_id].shape[0]
+    all_coords = _concatenate_coords(coords, target_chain_id)
+    device = next(model.parameters()).device
+
+    # Supply padding tokens for other chains to avoid unused sampling for speed
+    padding_pattern = ['<pad>'] * all_coords.shape[0]
+    for i in range(target_chain_len):
+        padding_pattern[i] = '<mask>'
+    sampled = model.sample(all_coords, partial_seq=padding_pattern,
+            temperature=temperature, device=device)
+    sampled = sampled[:target_chain_len]
+    return sampled
+
+
+def score_sequence_in_complex(model, alphabet, coords, target_chain_id,
+        target_seq, padding_length=10):
+    """
+    Scores sequence for one chain in a complex.
+    Args:
+        model: An instance of the GVPTransformer model
+        alphabet: Alphabet for the model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        target_seq: Target sequence for the target chain for scoring.
+        padding_length: padding length in between chains
+    Returns:
+        Tuple (ll_fullseq, ll_withcoord)
+        - ll_fullseq: Average log-likelihood over the full target chain
+        - ll_withcoord: Average log-likelihood in target chain excluding those
+            residues without coordinates
+    """
+    all_coords = _concatenate_coords(coords, target_chain_id)
+
+    loss, target_padding_mask = get_sequence_loss(model, alphabet, all_coords,
+            target_seq)
+    ll_fullseq = -np.sum(loss * ~target_padding_mask) / np.sum(
+            ~target_padding_mask)
+
+    # Also calculate average when excluding masked portions
+    coord_mask = np.all(np.isfinite(coords[target_chain_id]), axis=(-1, -2))
+    ll_withcoord = -np.sum(loss * coord_mask) / np.sum(coord_mask)
+    return ll_fullseq, ll_withcoord
+
+
+def get_encoder_output_for_complex(model, alphabet, coords, target_chain_id):
+    """
+    Args:
+        model: An instance of the GVPTransformer model
+        alphabet: Alphabet for the model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+    Returns:
+        Dictionary mapping chain id to encoder output for each chain
+    """
+    all_coords = _concatenate_coords(coords, target_chain_id)
+    all_rep = get_encoder_output(model, alphabet, all_coords)
+    target_chain_len = coords[target_chain_id].shape[0]
+    return all_rep[:target_chain_len]

esm/source/esm/inverse_folding/transformer_decoder.py

@@ -0,0 +1,228 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Contents of this file were adapted from the open source fairseq repository.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Any, Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from esm.modules import SinusoidalPositionalEmbedding
+from .transformer_layer import TransformerDecoderLayer
+
+
+def fill_with_neg_inf(t):
+    """FP16-compatible function that fills a tensor with -inf."""
+    return t.float().fill_(float("-inf")).type_as(t)
+
+
+class TransformerDecoder(nn.Module):
+    """
+    Transformer decoder consisting of *args.decoder.layers* layers. Each layer
+    is a :class:`TransformerDecoderLayer`.
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+        dictionary (~fairseq.data.Dictionary): decoding dictionary
+        embed_tokens (torch.nn.Embedding): output embedding
+        no_encoder_attn (bool, optional): whether to attend to encoder outputs
+            (default: False).
+    """
+
+    def __init__(
+        self,
+        args,
+        dictionary,
+        embed_tokens,
+    ):
+        super().__init__()
+        self.args = args
+        self.dictionary = dictionary
+        self._future_mask = torch.empty(0)
+
+        self.dropout_module = nn.Dropout(args.dropout)
+
+        input_embed_dim = embed_tokens.embedding_dim
+        embed_dim = args.decoder_embed_dim
+        self.embed_dim = embed_dim
+
+        self.padding_idx = embed_tokens.padding_idx
+
+        self.embed_tokens = embed_tokens
+        self.embed_scale = math.sqrt(embed_dim)
+
+        self.project_in_dim = (
+            nn.Linear(input_embed_dim, embed_dim, bias=False)
+            if embed_dim != input_embed_dim
+            else None
+        )
+        self.embed_positions = SinusoidalPositionalEmbedding(
+            embed_dim,
+            self.padding_idx,
+        )
+
+        self.layers = nn.ModuleList([])
+        self.layers.extend(
+            [
+                self.build_decoder_layer(args)
+                for _ in range(args.decoder_layers)
+            ]
+        )
+        self.num_layers = len(self.layers)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+        self.build_output_projection(args, dictionary)
+
+    def build_output_projection(self, args, dictionary):
+        self.output_projection = nn.Linear(
+            args.decoder_embed_dim, len(dictionary), bias=False
+        )
+        nn.init.normal_(
+            self.output_projection.weight, mean=0, std=args.decoder_embed_dim ** -0.5
+        )
+
+    def build_decoder_layer(self, args):
+        return TransformerDecoderLayer(args)
+
+    def forward(
+        self,
+        prev_output_tokens,
+        encoder_out: Optional[Dict[str, List[Tensor]]] = None,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+        features_only: bool = False,
+        return_all_hiddens: bool = False,
+    ):
+        """
+        Args:
+            prev_output_tokens (LongTensor): previous decoder outputs of shape
+                `(batch, tgt_len)`, for teacher forcing
+            encoder_out (optional): output from the encoder, used for
+                encoder-side attention, should be of size T x B x C
+            incremental_state (dict): dictionary used for storing state during
+                :ref:`Incremental decoding`
+            features_only (bool, optional): only return features without
+                applying output layer (default: False).
+
+        Returns:
+            tuple:
+                - the decoder's output of shape `(batch, tgt_len, vocab)`
+                - a dictionary with any model-specific outputs
+        """
+
+        x, extra = self.extract_features(
+            prev_output_tokens,
+            encoder_out=encoder_out,
+            incremental_state=incremental_state,
+        )
+
+        if not features_only:
+            x = self.output_layer(x)
+        x = x.transpose(1, 2) # B x T x C -> B x C x T
+        return x, extra
+
+    def extract_features(
+        self,
+        prev_output_tokens,
+        encoder_out: Optional[Dict[str, List[Tensor]]],
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+    ):
+        """
+        Similar to *forward* but only return features.
+
+        Includes several features from "Jointly Learning to Align and
+        Translate with Transformer Models" (Garg et al., EMNLP 2019).
+
+        Returns:
+            tuple:
+                - the decoder's features of shape `(batch, tgt_len, embed_dim)`
+                - a dictionary with any model-specific outputs
+        """
+        bs, slen = prev_output_tokens.size()
+
+        enc: Optional[Tensor] = None
+        padding_mask: Optional[Tensor] = None
+        if encoder_out is not None and len(encoder_out["encoder_out"]) > 0:
+            enc = encoder_out["encoder_out"][0]
+            assert (
+                enc.size()[1] == bs
+            ), f"Expected enc.shape == (t, {bs}, c) got {enc.shape}"
+        if encoder_out is not None and len(encoder_out["encoder_padding_mask"]) > 0:
+            padding_mask = encoder_out["encoder_padding_mask"][0]
+
+        # embed positions
+        positions = self.embed_positions(
+            prev_output_tokens
+        )
+
+        if incremental_state is not None:
+            prev_output_tokens = prev_output_tokens[:, -1:]
+            positions = positions[:, -1:]
+
+        # embed tokens and positions
+        x = self.embed_scale * self.embed_tokens(prev_output_tokens)
+
+        if self.project_in_dim is not None:
+            x = self.project_in_dim(x)
+
+        x += positions
+
+        x = self.dropout_module(x)
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        self_attn_padding_mask: Optional[Tensor] = None
+        if prev_output_tokens.eq(self.padding_idx).any():
+            self_attn_padding_mask = prev_output_tokens.eq(self.padding_idx)
+
+        # decoder layers
+        attn: Optional[Tensor] = None
+        inner_states: List[Optional[Tensor]] = [x]
+        for idx, layer in enumerate(self.layers):
+            if incremental_state is None:
+                self_attn_mask = self.buffered_future_mask(x)
+            else:
+                self_attn_mask = None
+
+            x, layer_attn, _ = layer(
+                x,
+                enc,
+                padding_mask,
+                incremental_state,
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask,
+                need_attn=False,
+                need_head_weights=False,
+            )
+            inner_states.append(x)
+
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        # T x B x C -> B x C x T
+        x = x.transpose(0, 1)
+
+        return x, {"inner_states": inner_states}
+
+    def output_layer(self, features):
+        """Project features to the vocabulary size."""
+        return self.output_projection(features)
+
+    def buffered_future_mask(self, tensor):
+        dim = tensor.size(0)
+        # self._future_mask.device != tensor.device is not working in TorchScript. This is a workaround.
+        if (
+            self._future_mask.size(0) == 0
+            or (not self._future_mask.device == tensor.device)
+            or self._future_mask.size(0) < dim
+        ):
+            self._future_mask = torch.triu(
+                fill_with_neg_inf(torch.zeros([dim, dim])), 1
+            )
+        self._future_mask = self._future_mask.to(tensor)
+        return self._future_mask[:dim, :dim]