esm2_loras

ESM2_LORAS

Il s'agit d'une tentative de formation d'une adaptation de bas rang (LORA) pour le modèle de langue protéique ESM-2 pour une tâche de classification de jeton. En particulier, nous essayons de former un prédicteur de site de liaison d'ARN. Il y a encore des problèmes à résoudre et tout commentaire ou conseil serait très apprécié. Ce code est destiné à un petit modèle, il devrait donc effectuer des balayages WANDB pour la recherche d'hyperparamètre dans un temps raisonnable sur presque tous les GPU. Vous pouvez facilement échanger contre des modèles plus grands si vous le souhaitez.

Poids du modèle et configuration

Le modèle lui-même

 "AmelieSchreiber/esm2_t6_8M_UR50D_lora_rna_binding_sites"

peut être trouvé sur le visage étreint ici.

Configuration de ce repo

Pour configurer l'environnement Conda, clonez le repo et exécutez:

 conda env create -f environment.yml

Puis courez:

 conda activate lora_esm_2

Pour former le modèle Run:

 from lora_esm2_script import train_protein_model

train_protein_model ()

Pour utiliser, essayez de courir:

 from transformers import AutoModelForTokenClassification , AutoTokenizer
from peft import PeftModel
import torch
import numpy as np
import random



# Path to the saved LoRA model
model_path = "esm2_t6_8M-finetuned-lora_2023-08-03_18-32-25"
# ESM2 base model
base_model_path = "facebook/esm2_t6_8M_UR50D"

# Load the model
base_model = AutoModelForTokenClassification . from_pretrained ( base_model_path )
loaded_model = PeftModel . from_pretrained ( base_model , model_path )

# Load the tokenizer
loaded_tokenizer = AutoTokenizer . from_pretrained ( model_path )

# New unseen protein sequence
new_protein_sequence = "FDLNDFLEQKVLVRMEAIINSMTMKERAKPEIIKGSRKRRIAAGSGMQVQDVNRLLKQFDDMQRMMKKM"

# Tokenize the new sequence
inputs = loaded_tokenizer ( new_protein_sequence , truncation = True , padding = 'max_length' , max_length = 512 , return_tensors = "pt" )

# Make predictions
with torch . no_grad ():
    outputs = loaded_model ( ** inputs )
    logits = outputs . logits
    predictions = torch . argmax ( logits , dim = 2 )

# Print logits for debugging
print ( "Logits:" , logits )

# Convert predictions to a list
predicted_labels = predictions . squeeze (). tolist ()

# Get input IDs to identify padding and special tokens
input_ids = inputs [ 'input_ids' ]. squeeze (). tolist ()

# Define a set of token IDs that correspond to special tokens
special_tokens_ids = { loaded_tokenizer . cls_token_id , loaded_tokenizer . pad_token_id , loaded_tokenizer . eos_token_id }

# Filter the predicted labels using the special_tokens_ids to remove predictions for special tokens
binding_sites = [ label for label , token_id in zip ( predicted_labels , input_ids ) if token_id not in special_tokens_ids ]

print ( "Predicted binding sites:" , binding_sites )