La arquitectura BERT (aplicado al análisis de secuencias de proteinas)

La arquitectura BERT (aplicado al análisis de secuencias de proteinas)#

La arquitectura BERT (Bidirectional Encoder Representations from Transformers) sólo usa el encoder del Transformer, eliminando la pila de decodificación.

El modelo BERT original se lanzó poco después del Generative Pre-trained Transformer (GPT) de OpenAI, y ambos se basaron en el trabajo de la arquitectura Transformer propuesto el año anterior. Mientras que GPT se centró en la generación de lenguaje natural (NLG), BERT priorizó la comprensión del lenguaje natural (NLU).

Arquitecturas de sólo codificador (encoder)#

El Transformer en 2017 inició una carrera para producir nuevos modelos que se basaran en su diseño innovador. OpenAI en junio de 2018, crea GPT: un modelo de solo decodificador que sobresalió en NLG, y finalmente pasó a impulsar ChatGPT en iteraciones posteriores. Google lanza BERT cuatro meses después: un modelo de solo codificador diseñado para NLU. Ambas arquitecturas pueden producir modelos muy capaces, pero las tareas que pueden realizar son ligeramente diferentes. A continuación se ofrece una descripción general de cada arquitectura.

  • Objetivo : hacer predicciones sobre palabras dentro de una secuencia de entrada.

  • Visión general : aceptar una secuencia de entrada y crear representaciones vectoriales numéricas enriquecidas para cada token. No tienen el decodificador que generan nuevas palabras como un chatbot, usan pilas de codificadores o encoders. Están orientados a reconocimiento de entidades con nombre (NER) o análisis de sentimientos. Las representaciones vectoriales enriquecidas creadas por los bloques del codificador son las que le dan a BERT una comprensión profunda del texto de entrada.

Nota : Es técnicamente posible generar texto con BERT, pero como veremos, esto no es para lo que se pretendía la arquitectura, y los resultados no rivalizan de ninguna manera con los modelos de solo decodificador

Arquitecturas de sólo decodificador (decoder)#

  • Objetivo : Predecir una nueva secuencia de salida en respuesta a una secuencia de entrada.

  • Visión general : omitien el bloque de codificador por completo y apilan varios decodificadores o decoders juntos en un solo modelo. Estos modelos aceptan mensajes como entradas y generan respuestas mediante la predicción de la siguiente token, uno a la vez, en una tarea conocida como Next Token Prediction (NTP).

Nota : los más conocidos por el público en general debido al uso generalizado de ChatGPT, que funciona con modelos de solo decodificador (GPT-3.5 y GPT-4)

En entrenamiento previo#

GPT también popularizó el entrenamiento previo del modelo. El entrenamiento previo implica entrenar un modelo grande para adquirir una comprensión amplia del lenguaje (que abarca aspectos como el uso de palabras y los patrones gramaticales) con el fin de producir un modelo fundamental independiente de la tarea.

El ajuste fino implica modificar sólo las capas finales que aparecen como lineales (en color morado en el gráfico siguiente), dejando el resto de pesos sin alterar.

Dependiendo de la tarea, el cabezal de clasificación (capa lineal en morado) se puede cambiar para que contenga un número diferente de neuronas de salida. Para una tarea de clasificación multiclase con 10 clases, la cabeza se puede cambiar para que tenga 10 neuronas en la capa de salida.

El contexto bidireccional#

BERT predice la probabilidad de observar ciertas palabras dado que se han observado palabras anteriores. Esto es algo genérico en todos los modelos de lenguaje. GPT está entrenado para predecir la siguiente palabra más probable en una secuencia.

La bidireccionalidad es quizás la característica más significativa de BERT, esto es algo que deriva de usar sólo el codificador (que por naturaleza es bidireccional) y descartar el decodificador (que por naturaleza es unidireccional).

Bidireccional denota que cada palabra en la secuencia de entrada puede obtener contexto de las palabras anteriores y posteriores (denominadas contexto izquierdo y contexto derecho, respectivamente). En términos técnicos, decimos que el mecanismo de atención puede atender a los tokens anteriores y posteriores de cada palabra.

Un

hombre

estaba

[MASK]

en

la

orilla

de

un

río

\(\leftarrow\)

\(\leftarrow\)

\(\leftarrow\)

\(\rightarrow\)

\(\rightarrow\)

\(\rightarrow\)

\(\rightarrow\)

\(\rightarrow\)

\(\rightarrow\)

Para BERT, la tarea aquí es predecir la palabra enmascarada indicada por [MASK], dado que esta palabra tiene palabras tanto a la izquierda como a la derecha, las palabras de ambos lados se pueden usar para proporcionar contexto.

Como modelo bidireccional, BERT adolece de dos grandes inconvenientes:

  • Aumento del tiempo de entrenamiento.

  • Bajo rendimiento en la generación de idiomas: es más adecuado para tareas como análisis de sentimientos.

Pre-entrenamiento de un modelo BERT#

El entrenamiento de un modelo bidireccional requiere tareas que permitan usar tanto el contexto izquierdo como el derecho para realizar predicciones. A tal fin diseñaron dos tareas:

  • La tarea de Modelo de Lenguaje Enmascarado (MLM, Masked Language Modeling): El modelo debe entrenarse para usar el contexto izquierdo y el contexto derecho de una secuencia de entrada para realizar una predicción. Esto se logra enmascarando aleatoriamente el 15% de las palabras en los datos de entrenamiento y entrenando a BERT para predecir la palabra que falta. Esta tarea ya existía en linguística y se conocía como tarea de Cloze (L. W. Taylor, “Cloze Procedure”: A New Tool for Measuring Readability (1953), Journalism Quarterly, 30(4), 415–433), aunque por la popularidad de BERT se conoce como tarea MLM.

The cat sat on it because it was a nice rug.
The cat sat on it [MASK] it was a nice rug.

BERT tomará una secuencia de entrada de un máximo de 512 tokens tanto para BERT Base como para BERT Large. Si se encuentran menos del número máximo de tokens en la secuencia, se agregará relleno mediante tokens para alcanzar el recuento máximo de 512.

El número de tokens de salida también será exactamente igual al número de tokens de entrada. Si existe un token enmascarado en la posición \(i\) en la secuencia de entrada, la predicción de BERT se encontrará en la posición \(i\) en la secuencia de salida.

El error se calcula utilizando una función de pérdida, que suele ser la función de pérdida de entropía cruzada en las posiciones \(i\) donde se encuentren los token enmascarados

  • La tarea de Predicción de Siguiente Oración (NSP, Next Sentence Predict): el objetivo es clasificar si un segmento (generalmente una oración) sigue lógicamente a otro. Al pre-entrenarse para NSP, BERT es capaz de desarrollar una comprensión del flujo entre oraciones en texto en prosa, una habilidad que es útil para una amplia gama de problemas de NLU, tales como: Pares de oraciones en paráfrasis, Pares hipótesis-premisa en implicación, Pares de preguntas y pasajes en las respuestas a preguntas.

La entrada para NSP consta del primer y segundo segmento (denotados A y B) separados por un token con un segundo token al final. En realidad, BERT espera que al menos un token por secuencia de entrada denote el final de la secuencia, independientemente de si se está realizando NSP o no

Los datos de entrenamiento se pueden generar fácilmente a partir de cualquier corpus monolingüe seleccionando oraciones con su siguiente oración el 50% de las veces y una oración aleatoria para el 50% restante de las oraciones.

Incrustaciones en un modelo BERT#

La incrustación que entra al transformer es suma la incrustación del token, más la de segmento (sentencia, frase, palabra) y la codificación posicional.

La incrustación de token es aprendida para cada token, igual que en el modelo Transformer.

La incrustación por segmento se utilizan para distinguir las oraciones dentro de un par de oraciones, como en tareas de clasificación de relaciones entre frases. BERT acepta una secuencia de hasta dos oraciones, A y B. A cada token se le asigna una etiqueta de segmento, usando valores \(0\) y \(1\):

[CLS]

La

casa

es

grande

[SEP]

El

perro

duerme

[SEP]

\(\downarrow\)

\(\downarrow\)

\(\downarrow\)

\(\downarrow\)

\(\downarrow\)

\(\downarrow\)

\(\downarrow\)

\(\downarrow\)

\(\downarrow\)

\(\downarrow\)

0

0

0

0

0

0

1

1

1

1

La incrustación posicional captura la posición del token en la secuencia para permitir que el modelo entienda el orden de las palabras. A diferencia del Transformer, las codificaciones posicionales en BERT son fijas y no son generadas por una función de senos y cosenos. Esto significa que BERT está restringido a 512 tokens en su secuencia de entrada tanto para BERT Base como para BERT Large. Ya se ha comentado que las secuencias en BERT tienen un máximo de 512 tokens.

Token especiales en BERT#

Como se ve en el gráfico anterior BERT tiene algunos token especiales:

  • [PAD] : Token de relleno para completar las secuencias hasta 512, cuando son más cortas.

  • [UNK] : Token desconocido, cuando no está en el vocabulario.

  • [CLS] : Token de clasifiación, se espera al principio de cada secuencia.

  • [SEP] : Token separador, separa dos segmentos de una sóla secuencia.

  • [MASK] : Token de máscara, se usa para entrenar a BERT.

Dimensiones usadas en la arquitectura BERT#

En el Transformer original las dimensiones son \(d_{model}=512\) y el número de cabeceras es \(A=8\), quedando \(d_k=512/8=64\).

Mientras que el BERT base \(d_{model}=768\) y el número de cabeceras es \(A=12\), quedando \(d_k=768/12=64\).

Y el BERT large \(d_{model}=1024\) y el número de cabeceras es \(A=16\), quedando \(d_k=1024/16=64\).

Ambas versiones de BERT también difieren en algunas características en cuanto número de capas y número de neuronas en las capas de paso adelante

Ajuste Fino#

El pre-entrenamiento es el primer paso en el marco de trabajo BERT que tiene 2 sub-etapas:

  • Definir el modelo de arquitectura: número de capas, de cabeceras, dimensiones y otros bloques del modelo.

  • Entrenar el modelo en las tareas MLM y NSP

El segundo paso es el ajuste fino que se puede descomponer en dos etapas:

  • Descargar un modelo elegido con los parámetros BERT preentrenados.

  • Realizar un ajuste fino de los parámetros para algunas tareas tales como Recognizing Textual Entailment (RTE) y Situations With Adversarial Generations (SWAG)

A continuación se muestras los pasos necesarios para un ajuste fino de un modelo pre-entrenado

Se importan librerías#

En el presente cuaderno se realiza un proceso de aplicación primero de un módelo BERT para realizar un ajuste fino sobre secuencias de proteinas que se pretenden clasificar y posteriormente se aplica el modelo T5 para caracterizar, mediante las incrustaciones, familias de proteinas que se representan en una gráfica después de realizar una reducción dimensional.

Los modelos parte de la publicación científica:

[Elnaggar et al., 2021]

Los autores han compartido código en GitHub para saber como manejar el modelo:

Acceder a ProtTrans

Este repositorio se actualiza periódicamente con nuevos modelos pre-entrenados para proteínas como parte del apoyo a la comunidad bioinformática en general y a la investigación de Covid-19 específicamente a través del proyecto Acelerar la investigación del SARS-CoV-2 con aprendizaje por transferencia utilizando modelos de modelado de lenguaje pre-entrenados.

import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from sklearn.model_selection import train_test_split
from transformers import BertTokenizer, BertConfig
from transformers import BertForSequenceClassification, get_linear_schedule_with_warmup
from tqdm import tqdm, trange  #for progress bars
import pandas as pd
import io
import os # accessing directory structure
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from IPython.display import Image #for image rendering
ajusteFino = False  # Poner la variable a True para activar el ajuste fino del modelo

Carga de Datos#

Se utilizará el conjunto de datos PFAM:

La versión 36.0 de Pfam se produjo en el Instituto Europeo de Bioinformática utilizando una base de datos de secuencias llamada Pfamseq, que se basa en UniProt.

Para simplicar hemos accedido a unos datos y se encuentran en moodle o upm-drive bajo la carpeta PFAM/random_split/random_split

data_partitions_dirpath = './data/PFAM/random_split/random_split'
print('Available dataset partitions: ', os.listdir(data_partitions_dirpath))
Available dataset partitions:  ['dev', 'test', 'train']

Para hacer unas pruebas de concepto nos quedamos con uno de los conjuntos reducidos: dev#

def read_all_shards(partition='dev', data_dir=data_partitions_dirpath):
    shards = []
    for fn in os.listdir(os.path.join(data_dir, partition)):
        with open(os.path.join(data_dir, partition, fn)) as f:
            shards.append(pd.read_csv(f, index_col=None))
    return pd.concat(shards)

#test = read_all_shards('test')
dev = read_all_shards('dev')
#train = read_all_shards('train')

#partitions = {'test': test, 'dev': dev, 'train': train}
partitions = {'dev': dev}
for name, df in partitions.items():
    print('Dataset partition "%s" has %d sequences' % (name, len(df)))
Dataset partition "dev" has 126171 sequences
dev.head()
family_id sequence_name family_accession aligned_sequence sequence
0 zf-Tim10_DDP N1QB11_PSEFD/15-76 PF02953.15 ..RMEKKQMKDFMNMYSNLVQRCFNDCV...........TD.F...... RMEKKQMKDFMNMYSNLVQRCFNDCVTDFTSKSLQSKEEGCVMRCV...
1 DNA_primase_S A8XA78_CAEBR/105-345 PF01896.19 FDID..LTDYDNIRNCCKEATVCPKCWKFMVLAVKILDFLLDDMFG... FDIDLTDYDNIRNCCKEATVCPKCWKFMVLAVKILDFLLDDMFGFN...
2 Col_cuticle_N A8XBM5_CAEBR/9-56 PF01484.17 ASAAILSGATIVGCLFFAAQIFNEVNSLYDDVMVDMDAFKVKSNIA... ASAAILSGATIVGCLFFAAQIFNEVNSLYDDVMVDMDAFKVKSNIAWD
3 GST_C_3 W4XBU3_STRPU/120-207 PF14497.6 KD.................................KLKESLPKTVN... KDKLKESLPKTVNPILLKFLEKALEDNPNGNGYFVGQDATMVEFVY...
4 Ada_Zn_binding E8U5K2_DEIML/9-73 PF02805.16 DRWQAVVQRE...AAQ.DG...LFLYAVRTTGIYCRPSCPSRRPR.... DRWQAVVQREAAQDGLFLYAVRTTGIYCRPSCPSRRPRRENVTFFE...

Elegimos aquellas familias de proteinas que presentan una mayor frecuencia en el dataset#

Las familias con un mayor número de secuencias

dev.groupby('family_id').size().sort_values(ascending=False).head(10)
family_id
Methyltransf_25    454
LRR_1              240
Acetyltransf_7     219
His_kinase         192
Bac_transf         190
Lum_binding        187
DNA_binding_1      168
Chromate_transp    157
Lipase_GDSL_2      156
DnaJ_CXXCXGXG      151
dtype: int64
Familias = ['Methyltransf_25','LRR_1','Acetyltransf_7','His_kinase','Bac_transf','Lum_binding','DNA_binding_1','Chromate_transp',
            'Lipase_GDSL_2','DnaJ_CXXCXGXG']
Familias = ['Methyltransf_25','LRR_1']

Se filtran las proteinas de las dos familias más frecuentes para hacer un ajuste fino de un clasificador de 2 clases#

sequences = []
labels = []
for i in range(dev.values.shape[0]): # Versión para generar libro jupyter en HTML
#for i in tqdm(range(dev.values.shape[0])): # Versión para uso online con barra de progreso
    if dev.values[i][0] in Familias:
        sequences.append(dev.values[i][4])
        labels.append(Familias.index(dev.values[i][0]))

len(labels)
694

Depuración de los datos#

Algunos aminoácidos raros (U,Z,O,B) se cambian por (X)

import re
maxLen = 0
for i in range(len(sequences)):
    sequences[i] =  re.sub(r"[UZOB]", "X", sequences[i])
    maxLen = max(maxLen, len(sequences[i]))
"La longitud máxima de la secuencia es " + str(maxLen) 
'La longitud máxima de la secuencia es 108'

Dataset Base resultante#

df = pd.DataFrame({
    "Secuencia": sequences,
    "Clase": labels
})

df.head()
Secuencia Clase
0 VLDLGCANGATSRALADLGARVTGVDVSARLIELARQREAARPRGV... 0
1 ALDAGAGPGVLTSFLMKKNPNLKWMACDISEDMVQYCKLVYPKVDW... 0
2 VLDIACGEGYGTALIGKYAQKAVGVDIDDTCIQWGTQHYAAANNKL... 0
3 NVRILDLSRQKFAVFPKEIWEL 1
4 ILSVGCGSGFVEHCLLQMRPNIQMHCNDSSDSPLRWIAQSLPEEQL... 0

Tokenizar los datos de ajuste fino#

Divide el texto de la revisión en tokens individuales, agrega los tokens especiales y controla el relleno. Es importante seleccionar el tokenizador adecuado para el modelo.

Se usará un tokenizador previamente entrenado del repositorio del modelo ProtTrans: Toward Understanding the Language of Life Through Self-Supervised Learning [Elnaggar et al., 2021]

En el trabajo de los autores se comentan varios modelos. Para este ajuste fino se elige el preentrenamiento con arquitectura BERT:

  • Rostlab/prot_bert .

from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False)
print(tokenizer)
BertTokenizer(name_or_path='Rostlab/prot_bert', vocab_size=30, model_max_length=1000000000000000019884624838656, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'unk_token': '[UNK]', 'sep_token': '[SEP]', 'pad_token': '[PAD]', 'cls_token': '[CLS]', 'mask_token': '[MASK]'}, clean_up_tokenization_spaces=True, added_tokens_decoder={
	0: AddedToken("[PAD]", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	1: AddedToken("[UNK]", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	2: AddedToken("[CLS]", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	3: AddedToken("[SEP]", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	4: AddedToken("[MASK]", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
}
)

Proceso de codificación: conversión de texto en tokens a ID de token#

Se utiliza la función encode_plus que devuelve un diccionario con los siguientes campos:

  • input_ids : IDs de los tokens. Es el mismo valor que devuelve la función: input_ids.

  • token_type_ids : Da un valor 0 para la oración A y un valor 1 para la oración B.

  • attention_mask : Una lista de 0 y 1 donde 0 indica que un token debe ser ignorado durante el proceso de atención y 1 indica que un token no debe ser ignorado.

token_ids = []
attention_masks = []

# Encode each review
for i in tqdm(range(len(sequences))):  
    batch_encoder = tokenizer.encode_plus(
        sequences[i],
        max_length = 128,
        padding = 'max_length',
        truncation = True,
        return_tensors = 'pt')

    token_ids.append(batch_encoder['input_ids'])
    attention_masks.append(batch_encoder['attention_mask'])
Hide code cell output
  0%|                                                                                          | 0/694 [00:00<?, ?it/s]
 20%|███████████████▋                                                              | 140/694 [00:00<00:00, 1396.16it/s]
 40%|███████████████████████████████▍                                              | 280/694 [00:00<00:00, 1183.30it/s]
 58%|█████████████████████████████████████████████                                 | 401/694 [00:00<00:00, 1122.69it/s]
 74%|█████████████████████████████████████████████████████████▉                    | 515/694 [00:00<00:00, 1115.45it/s]
 91%|███████████████████████████████████████████████████████████████████████▎      | 634/694 [00:00<00:00, 1139.88it/s]
100%|██████████████████████████████████████████████████████████████████████████████| 694/694 [00:00<00:00, 1149.93it/s]

# Convert token IDs and attention mask lists to PyTorch tensors
token_ids = torch.cat(token_ids, dim=0)
attention_masks = torch.cat(attention_masks, dim=0)
labels = torch.tensor(labels)
token_ids[0][0:10]
tensor([2, 1, 3, 0, 0, 0, 0, 0, 0, 0])

Dividir los datos en entrenamiento y validación y crear los conjuntos de datos (datasets) y cargadores (dataloaders)#

val_size = 0.25

# Split the token IDs
train_ids, val_ids = train_test_split(
                        token_ids,
                        test_size=val_size,
                        shuffle=False)

# Split the attention masks
train_masks, val_masks = train_test_split(
                            attention_masks,
                            test_size=val_size,
                            shuffle=False)

# Split the labels
train_labels, val_labels = train_test_split(
                                labels,
                                test_size=val_size,
                                shuffle=False)
# Create the DataLoaders
train_data = TensorDataset(train_ids, train_masks, train_labels)
train_dataloader = DataLoader(train_data, shuffle=True, batch_size=16)
val_data = TensorDataset(val_ids, val_masks, val_labels)
val_dataloader = DataLoader(val_data, batch_size=16)

Modelo BERT#

Configuración del modelo BERT#

El siguiente paso es cargar un modelo BERT previamente entrenado para ajustarlo. Se puede importar un modelo desde el repositorio de modelos de Hugging Face de manera similar a como se hizo con el tokenizador. Hugging Face tiene muchas versiones de BERT con cabezales de clasificación ya adjuntos, lo que hace que este proceso sea muy conveniente. Estos son algunos ejemplos de modelos con cabezales de clasificación preconfigurados:

  • BertForMaskedLM

  • BertForNextSentencePrediction

  • BertForSequenceClassification

  • BertForMultipleChoice

  • BertForTokenClassification

  • BertForQuestionAnswering

Por supuesto, es posible importar un modelo BERT sin cabeza y crear su propio cabezal de clasificación desde cero en PyTorch o Tensorflow.

Se utiliza el modelo pre-entrenado y con la salida adapatada a un clasificador (BertForSequenceClassification) indicando en el número de clases a identificar como 2

from transformers import BertForSequenceClassification

model = BertForSequenceClassification.from_pretrained(
    'Rostlab/prot_bert',
    num_labels=2)

configuration = model.config
configuration
Some weights of BertForSequenceClassification were not initialized from the model checkpoint at Rostlab/prot_bert and are newly initialized: ['classifier.bias', 'classifier.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
BertConfig {
  "architectures": [
    "BertForMaskedLM"
  ],
  "attention_probs_dropout_prob": 0.0,
  "classifier_dropout": null,
  "hidden_act": "gelu",
  "hidden_dropout_prob": 0.0,
  "hidden_size": 1024,
  "initializer_range": 0.02,
  "intermediate_size": 4096,
  "layer_norm_eps": 1e-12,
  "max_position_embeddings": 40000,
  "model_type": "bert",
  "num_attention_heads": 16,
  "num_hidden_layers": 30,
  "pad_token_id": 0,
  "position_embedding_type": "absolute",
  "torch_dtype": "float32",
  "transformers_version": "4.53.0",
  "type_vocab_size": 2,
  "use_cache": true,
  "vocab_size": 30
}

Optimizador, función de pérdida y programador#

Se requiere un optimizador para calcular los cambios necesarios en cada peso y sesgo. El habitual es AdamW.

Los modelos de lenguaje suelen utilizar la función de pérdida de entropía cruzada.

Los programadores (scheduler) están diseñados para disminuir gradualmente la tasa de aprendizaje a medida que continúa el proceso de entrenamiento, reduciendo el tamaño de los cambios realizados en cada peso y el sesgo en cada paso del optimizador.

EPOCHS = 1

# Optimizer
#optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=2e-5, eps = 1e-8 )
optimizer = torch.optim.AdamW(model.parameters())

# Loss function
loss_function = nn.CrossEntropyLoss()

# Scheduler
num_training_steps = EPOCHS * len(train_dataloader)
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps)

Bucle con el ajuste fino#

Por la carga del proceso se necesita utilizar GPU o entorno CUDA. En CPU se lanza sólo una ejecución de concepto truncada.

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device
device(type='cpu')
if ajusteFino:
    for epoch in range(0, EPOCHS):
    
        model.train()
        training_loss = 0
    
        for batch in tqdm(train_dataloader, desc="Época=" + str(epoch)):
    
            batch_token_ids = batch[0].to(device)
            batch_attention_mask = batch[1].to(device)
            batch_labels = batch[2].to(device)
    
            model.zero_grad()
    
            loss, logits = model(
                batch_token_ids,
                token_type_ids = None,
                attention_mask=batch_attention_mask,
                labels=batch_labels,
                return_dict=False)
    
            training_loss += loss.item()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
    
        average_train_loss = training_loss / len(train_dataloader)
    filename = "data/proteinSeqClass.pt"
    torch.save(model.state_dict(), filename)
    print("Proceso de ajuste fino finalizado. average_train_loss=", average_train_loss)
else:
    print("El ajuste fino no está activo. Se recupera un modelo previamente entrenado almacenado en proteinSeqClass.pt")
    filename = "data/proteinSeqClass.pt"
    model.load_state_dict(torch.load(filename, map_location=torch.device(device)))
El ajuste fino no está activo. Se recupera un modelo previamente entrenado almacenado en proteinSeqClass.pt

Validación contra el conjunto de validación#

def calculate_accuracy(preds, labels):
    """ Calculate the accuracy of model predictions against true labels.

    Parameters:
        preds (np.array): The predicted label from the model
        labels (np.array): The true label

    Returns:
        accuracy (float): The accuracy as a percentage of the correct
            predictions.
    """
    pred_flat = np.argmax(preds, axis=1).flatten()
    labels_flat = labels.flatten()
    accuracy = np.sum(pred_flat == labels_flat) / len(labels_flat)

    return accuracy
model.eval()
val_loss = 0
val_accuracy = 0

for batch in tqdm(val_dataloader):

    batch_token_ids = batch[0].to(device)
    batch_attention_mask = batch[1].to(device)
    batch_labels = batch[2].to(device)

    with torch.no_grad():
        (loss, logits) = model(
            batch_token_ids,
            attention_mask = batch_attention_mask,
            labels = batch_labels,
            token_type_ids = None,
            return_dict=False)

    logits = logits.detach().cpu().numpy()
    label_ids = batch_labels.to('cpu').numpy()
    val_loss += loss.item()
    val_accuracy += calculate_accuracy(logits, label_ids)


average_val_accuracy = val_accuracy / len(val_dataloader)
Hide code cell output
  0%|                                                                                           | 0/11 [00:00<?, ?it/s]
  9%|███████▌                                                                           | 1/11 [00:32<05:27, 32.72s/it]
 18%|███████████████                                                                    | 2/11 [00:59<04:21, 29.05s/it]
 27%|██████████████████████▋                                                            | 3/11 [01:24<03:40, 27.51s/it]
 36%|██████████████████████████████▏                                                    | 4/11 [01:50<03:07, 26.82s/it]
 45%|█████████████████████████████████████▋                                             | 5/11 [02:15<02:37, 26.23s/it]
 55%|█████████████████████████████████████████████▎                                     | 6/11 [02:41<02:10, 26.15s/it]
 64%|████████████████████████████████████████████████████▊                              | 7/11 [03:06<01:43, 25.82s/it]
 73%|████████████████████████████████████████████████████████████▎                      | 8/11 [03:31<01:16, 25.44s/it]
 82%|███████████████████████████████████████████████████████████████████▉               | 9/11 [03:58<00:51, 25.84s/it]
 91%|██████████████████████████████████████████████████████████████████████████▌       | 10/11 [04:11<00:21, 21.95s/it]
100%|██████████████████████████████████████████████████████████████████████████████████| 11/11 [04:20<00:00, 17.85s/it]
100%|██████████████████████████████████████████████████████████████████████████████████| 11/11 [04:20<00:00, 23.65s/it]

average_val_accuracy
np.float64(0.586038961038961)

Extracción de incrustaciones (“embeddings”) para caracterizar de forma abstracta las proteinas#

Para este ejercicio se amplia el número de familias, adoptando las 10 más frecuentes. Por otro lado se reduce el número de elementos por familia (hasta 20) porque no se va a realizar ningún proceso de ajuste fino, sino una inferencia de las incrustaciones para caracterizar las proteinas.

## Se toman las 10 familias de mayor frecuencia
Familias = ['Methyltransf_25','LRR_1','Acetyltransf_7','His_kinase','Bac_transf','Lum_binding','DNA_binding_1','Chromate_transp',
            'Lipase_GDSL_2','DnaJ_CXXCXGXG']
total = np.zeros(20)
sequences = []
labels = []
for i in range(dev.values.shape[0]): # Para obtener el libro HTML con jupyter-book
#for i in tqdm(range(dev.values.shape[0])):  # Uso online con barra de progreso
    if dev.values[i][0] in Familias:
        ID = Familias.index(dev.values[i][0])
        if total[ID]<20:
            total[ID] += 1 
            sequences.append(dev.values[i][4])
            labels.append(ID)
len(sequences), len(labels)
(200, 200)

Se aplica un modelo de los autores (ProtBert-BFD)#

Para generar las incrustaciones de los aminoacidos de las proteinas se pueden adoptar varios modelos como ProtT5 o ProtBert-BFD, que se basa en un Encoder y que entrenaron los autores en su trabajo:

ProtTrans: Toward Understanding the Language of Life Through Self-Supervised Learning (Elnaggar et al., 2022)

Se puede acceder a ejemplos de código ProtTrans en GitHub

El modelo ProtT5 requiere instalar la librería adicional sentencepiece que desafortunadamente no está disponible en Python 3.13, por lo que se adopta el modelo ProtBert-BFD para hacer el ejemplo de uso

from transformers import AutoTokenizer, AutoModel, pipeline
tokenizer = AutoTokenizer.from_pretrained("Rostlab/prot_bert_bfd", do_lower_case=False )
model = AutoModel.from_pretrained("Rostlab/prot_bert_bfd")

Se concatena el modelo y el tokenizador de ProtBert-BFD#

La función pipeline ejecuta la tokenización y a continuación la incrustación para en conjunto, obtener un extractor de características a partir de la secuencia de palabras de proteinas.

fe = pipeline('feature-extraction', model=model, tokenizer=tokenizer,device=0 )
Device set to use cpu

Depuración de aminoácidos ambiguos#

import re
# replace all rare/ambiguous amino acids by X and introduce white-space between all amino acids
sequence_examples = [" ".join(list(re.sub(r"[UZOB]", "X", sequence))) for sequence in sequences]
sequence_examples[0]
'V L D L G C A N G A T S R A L A D L G A R V T G V D V S A R L I E L A R Q R E A A R P R G V R Y L V G D A A H L P D L A D A S F D R I T A S M V L M D I E N A E G A I R E V A R L L R P G G'

Extracción de características#

features = []
for sequ in sequence_examples: # para obtener HTML con jupyter-book
#for sequ in tqdm(sequence_examples): # Uso online con barra de progreso
    embed = fe(sequ)
    seq_len = len(sequ.replace(" ", ""))
    start_Idx = 1
    end_Idx = seq_len+1
    seq_emb = embed[0][start_Idx:end_Idx]
    seq_featur = []
    for emb in seq_emb:
        seq_featur.append(np.asarray(emb))
    features.append(np.asarray(seq_featur))

Adoptar una incrustación única por proteina#

La incrustación única es el vector media de una secuencia de incrustaciones

lstFullProtein = []
for seq_featur in features:
    fullProtein = np.mean(seq_featur, axis=0) # shape (1024)
    lstFullProtein.append(fullProtein)
arrLabels = np.asarray(labels)
arrFullProtein = np.asarray(lstFullProtein)
arrFullProtein.shape, arrLabels.shape
((200, 1024), (200,))

Realizar una reducción dimensional para hacer una representación 2D#

Se realiza una reducción para representar el resultado con el algoritmo PCA, Análisis de Componentes Principales.

from sklearn.decomposition import PCA
pca = PCA(n_components=2)
X_pca = pca.fit_transform(arrFullProtein)
colors = ['b', 'g', 'r', 'c','m', 'y','k', 'aqua', 'tomato', 'olive']
markers = ['.', 'o', 'v', '^', '<', '>', '8', 's', 'p', '*']
plt.subplots(figsize=(16, 8))
for l, c, m in zip(np.unique(arrLabels), colors, markers):
    plt.scatter(X_pca[arrLabels==l, 0], X_pca[arrLabels==l, 1], c=c, label=Familias[l], marker=m)
plt.xlabel('PC 1')
plt.ylabel('PC 2')
plt.legend(loc='lower left')
plt.show()
_images/8dc48188e667c8340a1d3920286215dba2b13fb257e6a9f4aeac24b52288aeca.png