Duplicate from AmelieSchreiber/cafa_5_protein_function_prediction

Co-authored-by: Amelie Schreiber <AmelieSchreiber@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,145 @@
+---
+license: mit
+language:
+- en
+library_name: transformers
+tags:
+- ems
+- esm2
+- biology
+- protein
+- protein language model
+- cafa 5
+- protein function prediction
+datasets:
+- AmelieSchreiber/cafa_5
+metrics:
+- f1
+- recall
+- precision
+---
+# ESM-2 for Protein Function Prediction
+
+Please also see the more recent fine-tuned model [AmelieSchreiber/esm2_t6_8M_finetuned_cafa5](https://huggingface.co/AmelieSchreiber/esm2_t6_8M_finetuned_cafa5). 
+
+This model is not intended for protein function prediction, but rather as a checkpoint for further fine-tuning, especially
+with Low Rank Adaptation (LoRA). This is an experimental model fine-tuned from the 
+[esm2_t6_8M_UR50D](https://huggingface.co/facebook/esm2_t6_8M_UR50D) model 
+for multi-label classification. In particular, the model is fine-tuned on the CAFA-5 protein sequence dataset available 
+[here](https://huggingface.co/datasets/AmelieSchreiber/cafa_5). More precisely, the `train_sequences.fasta` file is the 
+list of protein sequences that were trained on, and the 
+`train_terms.tsv` file contains the gene ontology protein function labels for each protein sequence. For more details on using 
+ESM-2 models for multi-label sequence classification, [see here](https://huggingface.co/docs/transformers/model_doc/esm). 
+Due to the potentially complicated class weighting necessary for the hierarchical ontology, further fine-tuning will be necessary. 
+
+## Fine-Tuning
+
+The model was fine-tuned for 7 epochs at a learning rate of `5e-5`, and achieves the following metrics:
+```
+Validation Loss: 0.0027,
+Validation Micro F1: 0.3672,
+Validation Macro F1: 0.9967,
+Validation Micro Precision: 0.6052,
+Validation Macro Precision: 0.9996,
+Validation Micro Recall: 0.2626,
+Validation Macro Recall: 0.9966
+```
+
+## Using the model
+
+First, download the `train_sequences.fasta` file and the `train_terms.tsv` file, and provide the local paths in the code below:
+
+```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 (Replace with your local paths)
+fasta_file = "/Users/amelieschreiber/.cursor-tutor/projects/python/cafa5/cafa-5-protein-function-prediction/Train/train_sequences.fasta"
+tsv_file = "/Users/amelieschreiber/.cursor-tutor/projects/python/cafa5/cafa-5-protein-function-prediction/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:]
+
+# tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
+seq_length = 1022
+# tokenized_data = tokenizer(list(fasta_data.values()), padding=True, truncation=True, return_tensors="pt", max_length=seq_length)
+
+unique_terms = list(set(term for terms in tsv_data.values() for term in terms))
+```
+
+
+Second, downlowd the file `go-basic.obo` [from here](https://huggingface.co/datasets/AmelieSchreiber/cafa_5)
+and store the file locally, then provide the local path in the the code below:
+
+```python
+import torch
+from transformers import AutoTokenizer, EsmForSequenceClassification
+from sklearn.metrics import precision_recall_fscore_support
+
+# 1. Parsing the go-basic.obo file
+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
+
+parsed_terms = parse_obo_file("go-basic.obo")  # Replace `go-basic.obo` with your path
+
+# 2. Load the saved model and tokenizer
+model_path = "AmelieSchreiber/cafa_5_protein_function_prediction"
+loaded_model = EsmForSequenceClassification.from_pretrained(model_path)
+loaded_tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+# 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 an example sequence
+example_sequence = "MAYLGSLVQRRLELASGDRLEASLGVGSELDVRGDRVKAVGSLDLEEGRLEQAGVSMA"  # Replace with your protein sequence
+predicted_functions = predict_protein_function(example_sequence, loaded_model, loaded_tokenizer, parsed_terms)
+print(predicted_functions)
+```

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b979922aa035be9cf44e205a6c018b31a2cd73f38cf09a8bdd957ec8f63d8adf
+size 34583185

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "mask_token": "<mask>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

tokenizer_config.json

@@ -0,0 +1,5 @@
+{
+  "clean_up_tokenization_spaces": true,
+  "model_max_length": 1024,
+  "tokenizer_class": "EsmTokenizer"
+}

vocab.txt

@@ -0,0 +1,33 @@
+<cls>
+<pad>
+<eos>
+<unk>
+L
+A
+G
+V
+S
+E
+R
+T
+I
+D
+P
+K
+Q
+N
+F
+Y
+M
+H
+W
+C
+X
+B
+U
+Z
+O
+.
+-
+<null_1>
+<mask>