Modelo de lenguaje de proteínas de propósito general. Se puede utilizar para predecir directamente a partir de las secuencias individuales: la estructura, la función y otras propiedades.
Los modelos de lenguaje destilan información a partir de bases de datos masivas de secuencias de proteínas.
ESM-2 se carga como una biblioteca de funciones de Python. Para utilizarlo podemos crearnos un entorno virtual de Python e instalar las bibliotecas mediante PIP; alternativamente, podemos usarlo cargando los módulos necesarios sin necesidad de crear el entorno.
module load rama0.3 GCCcore/12.3.0 Python/3.11.3 CUDA/12.2.0
Como introducción al ESM-2 podemos utilizar de ejemplo un código Python y un archivo sbatch para ejecutarlo en la cola de gráficas.
El resultado final serán dos archivos: protein1_representation.npy y protein2_representation.npy. Que son las representaciones de las estructuras de las proteínas en formato NumPy.
Ambos archivos se pueden usar de base y modificarlos para adaptarlos a nuestras necesidades.
Empezamos por el código de Python:
import esm
import torch
import numpy as np
def main():
# Cargar el modelo ESM-2 (podemos elegir otras variantes del modelo ESM-2 que nos resulten más convenientes)
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
model.eval() # Set the model to evaluation mode
# Mover el modelo a la GPU si está disponible
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = torch.device("cuda")
model = model.to(device)
# Definir la secuencia de las proteínas
sequences = [
("protein1", "MKTAYIAKQRQISFVKSHFSRQDILDLI"),
("protein2", "MVHLTPEEKSAVTALWGKVNVDEVGGEALGRLLVVYPWTQRF"),
]
# Preparar los datos para el modelo
batch_converter = alphabet.get_batch_converter()
batch_labels, batch_strs, batch_tokens = batch_converter(sequences)
# Mover datos a la GPU si está disponible
batch_tokens = batch_tokens.to(device)
# Extraer representaciones
with torch.no_grad():
results = model(batch_tokens, repr_layers=[33])
token_representations = results["representations"][33]
# Almacenar representaciones
for i, (label, seq) in enumerate(sequences):
representation = token_representations[i, 1 : len(seq) + 1].cpu().numpy()
output_file = f"{label}_representation.npy"
np.save(output_file, representation)
print(f"Guardadas representaciones para {label} en {output_file}")
if __name__ == "__main__":
main()
Código sbatch en el cual establecemos la cantidad de RAM reservada, el número de nodos, la cola gpu...:
#!/bin/bash
#SBATCH --job-name=esm2_job ## Nombre del trabajo
#SBATCH --output=esm2_job_%j.out ## Nombre del archivo de volcado del reultado de SLURM
#SBATCH --error=esm2_job_%j.err ## Nombre del archivo de volcado de errores de SLURM
#SBATCH --partition=gpu ## Solicitar la cola adecuada (por ahora la cola gpu es la única que tiene gráficas)
##SBATCH --nodes=1 ## Solicitamos un único nodo
#SBATCH --ntasks-per-node=1 ## Especificamos una única tarea por nodo
#SBATCH --mem=16G ## Reservamos 16GB de RAM
#SBATCH --time=02:00:00 ## Reservamos dos horas de tiempo de cómputo
# Carga de módulos
module load rama0.3 GCCcore/12.3.0 Python/3.11.3 CUDA/12.2.0
# Ejecución del script de Python
python run_esm2.py
Los ejemplos se encuentran disponibles aquí:
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 and Privacy Policy.