Dies ist ein Versuch, eine Adaption (Low-Rang-Anpassung) für das Proteinsprachenmodell ESM-2 für eine Token-Klassifizierungsaufgabe auszubilden. Insbesondere versuchen wir, einen RNA -Bindungsstellenprädiktor auszubilden. Es gibt noch einige Probleme zu trainieren und Feedback oder Ratschläge wären sehr geschätzt. Dieser Code gilt für ein kleines Modell, sodass er in angemessener Zeit in fast jeder GPU Wandb -Sweep für die Hyperparameter -Suche durchführen sollte. Sie können sich jedoch problemlos gegen größere Modelle austauschen, wenn Sie möchten.
Das Modell selbst
"AmelieSchreiber/esm2_t6_8M_UR50D_lora_rna_binding_sites"kann hier auf dem umarmen Gesicht gefunden werden.
Klonen Sie das Repo, um die Umgebung von The Conda einzurichten und zu rennen:
conda env create -f environment.yml
Dann rennen:
conda activate lora_esm_2
Um den Modelllauf zu trainieren:
from lora_esm2_script import train_protein_model
train_protein_model ()Versuchen Sie zum Laufen:
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 )