README.md

---
library_name: peft
tags:
- esm
- esm2
- ESM-2
- protein language model
- LoRA
- Low Rank Adaptation
- biology
- CAFA-5
- protein function prediction
datasets:
- AmelieSchreiber/cafa_5
license: mit
language:
- en
---

# ESM-2 LoRA for CAFA-5 Protein Function Prediction
This is a Low Rank Adaptation (LoRA) of [cafa_5_protein_function_prediction](https://huggingface.co/AmelieSchreiber/cafa_5_protein_function_prediction), 
which is a fine-tuned (without LoRA) version of `facebook/esm2_t6_8M_UR50D`, for the same task. For more information 
on training a sequence classifier langauge model with LoRA [see here](https://github.com/huggingface/peft/blob/main/examples/sequence_classification/LoRA.ipynb). 
Note, this is for natural language processing and must be adapted to our use case using a protein language model like ESM-2. 

## Training procedure
 Using Hugging Face's Parameter Efficient Fine-Tuning (PEFT) library, a Low Rank Adaptation was trained for 
 3 epochs on the CAFA-5 protein sequences dataset at an 80/20 train/test split. The dataset can be 
 [found here](https://huggingface.co/datasets/AmelieSchreiber/cafa_5). Somewhat naively, the model was trained on
 the `train_sequences.fasta` file of protein sequences, with the `train_terms.tsv` file serving as the labels. 
 The gene ontology used is a hierarchy, and so the labels lower in the hierchay should be weighted more, or the 
 graph structure should be taken into account. The model achieved the following metrics:

```
Epoch: 3,
Validation Loss: 0.0031,
Validation Micro F1: 0.3752,
Validation Macro F1: 0.9968,
Validation Micro Precision: 0.5287,
Validation Macro Precision: 0.9992,
Validation Micro Recall: 0.2911,
Validation Macro Recall: 0.9968
```

Future iterations of this model will likely need to take into account class weighting. 

### Framework versions

- PEFT 0.4.0

## Using the Model

To use the model, try downloading the data [from here](https://huggingface.co/datasets/AmelieSchreiber/cafa_5), 
adjust the paths to the files in the code below to their local paths on your machine, and try running:

```python
import os
import numpy as np
import torch
from transformers import AutoTokenizer, EsmForSequenceClassification, AdamW
from torch.nn.functional import binary_cross_entropy_with_logits
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score, precision_score, recall_score
from accelerate import Accelerator
from Bio import SeqIO

# Step 1: Data Preprocessing
fasta_file = "data/Train/train_sequences.fasta"
tsv_file = "data/Train/train_terms.tsv"

fasta_data = {}
tsv_data = {}

for record in SeqIO.parse(fasta_file, "fasta"):
    fasta_data[record.id] = str(record.seq)

with open(tsv_file, 'r') as f:
    for line in f:
        parts = line.strip().split("\t")
        tsv_data[parts[0]] = parts[1:]

unique_terms = list(set(term for terms in tsv_data.values() for term in terms))

def parse_fasta(file_path):
    """
    Parses a FASTA file and returns a list of sequences.
    """
    with open(file_path, 'r') as f:
        content = f.readlines()

    sequences = []
    current_sequence = ""

    for line in content:
        if line.startswith(">"):
            if current_sequence:
                sequences.append(current_sequence)
                current_sequence = ""
        else:
            current_sequence += line.strip()

    if current_sequence:
        sequences.append(current_sequence)

    return sequences

# Parse the provided FASTA file
fasta_file_path = "data/Test/testsuperset.fasta"
protein_sequences = parse_fasta(fasta_file_path)
# protein_sequences[:3]  # Displaying the first 3 sequences for verification

import torch
from transformers import AutoTokenizer, EsmForSequenceClassification
from sklearn.metrics import precision_recall_fscore_support

# 1. Parsing the go-basic.obo file (Assuming this is still needed)
def parse_obo_file(file_path):
    with open(file_path, 'r') as f:
        data = f.read().split("[Term]")
        
    terms = []
    for entry in data[1:]:
        lines = entry.strip().split("\n")
        term = {}
        for line in lines:
            if line.startswith("id:"):
                term["id"] = line.split("id:")[1].strip()
            elif line.startswith("name:"):
                term["name"] = line.split("name:")[1].strip()
            elif line.startswith("namespace:"):
                term["namespace"] = line.split("namespace:")[1].strip()
            elif line.startswith("def:"):
                term["definition"] = line.split("def:")[1].split('"')[1]
        terms.append(term)
    return terms

# Let's assume the path to go-basic.obo is as follows (please modify if different)
obo_file_path = "data/Train/go-basic.obo"  
parsed_terms = parse_obo_file("data/Train/go-basic.obo")  # Replace with your path

# 2. Load the saved model and tokenizer
# Assuming the model path provided is correct
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel, PeftConfig

# Load the tokenizer and model
model_id = "AmelieSchreiber/esm2_t6_8M_UR50D_cafa5_lora"  # Replace with your Hugging Face hub model name
tokenizer = AutoTokenizer.from_pretrained(model_id)

# First, we load the underlying base model
base_model = AutoModelForSequenceClassification.from_pretrained(model_id)

# Then, we load the model with PEFT
model = PeftModel.from_pretrained(base_model, model_id)
loaded_model = model
loaded_tokenizer = AutoTokenizer.from_pretrained(model_id)

# 3. The predict_protein_function function
def predict_protein_function(sequence, model, tokenizer, go_terms):
    inputs = tokenizer(sequence, return_tensors="pt", padding=True, truncation=True, max_length=1022)
    model.eval()
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.sigmoid(outputs.logits)
        predicted_indices = torch.where(predictions > 0.05)[1].tolist()
    
    functions = []
    for idx in predicted_indices:
        term_id = unique_terms[idx]  # Use the unique_terms list from your training script
        for term in go_terms:
            if term["id"] == term_id:
                functions.append(term["name"])
                break
                
    return functions

# 4. Predicting protein function for the sequences in the FASTA file
protein_functions = {}
for seq in protein_sequences[:20]:  # Using only the first 3 sequences for demonstration
    predicted_functions = predict_protein_function(seq, loaded_model, loaded_tokenizer, parsed_terms)
    protein_functions[seq[:20] + "..."] = predicted_functions  # Using first 20 characters as key

protein_functions
```