File size: 5,089 Bytes
54f047d 5d4253c 54f047d 5d4253c 54f047d 5d4253c 859db5c 5d4253c 859db5c 5d4253c 9c6eb7c 5d4253c 859db5c 5d4253c 9c6eb7c 5d4253c 9c6eb7c 5d4253c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | ---
datasets:
- andrewdalpino/AmiGO-Boost
metrics:
- precision
- recall
- f1
base_model:
- EvolutionaryScale/esmc-300m-2024-12
pipeline_tag: text-classification
tags:
- gene-ontology
---
# ESMC Protein Function Predictor
An Evolutionary-scale Model (ESM) for protein function prediction from amino acid sequences using the Gene Ontology (GO). Based on the ESM Cambrian Transformer architecture, pre-trained on [UniRef](https://www.uniprot.org/help/uniref), [MGnify](https://www.ebi.ac.uk/metagenomics), and the Joint Genome Institute's database and fine-tuned on the [AmiGO Boost](https://huggingface.co/datasets/andrewdalpino/AmiGO-Boost) protein function dataset, this model predicts the GO subgraph for a particular protein sequence - giving you insight into the molecular function, biological process, and location of the activity inside the cell.
## What are GO terms?
> "The Gene Ontology (GO) is a concept hierarchy that describes the biological function of genes and gene products at different levels of abstraction (Ashburner et al., 2000). It is a good model to describe the multi-faceted nature of protein function."
> "GO is a directed acyclic graph. The nodes in this graph are functional descriptors (terms or classes) connected by relational ties between them (is_a, part_of, etc.). For example, terms 'protein binding activity' and 'binding activity' are related by an is_a relationship; however, the edge in the graph is often reversed to point from binding towards protein binding. This graph contains three subgraphs (subontologies): Molecular Function (MF), Biological Process (BP), and Cellular Component (CC), defined by their root nodes. Biologically, each subgraph represent a different aspect of the protein's function: what it does on a molecular level (MF), which biological processes it participates in (BP) and where in the cell it is located (CC)."
From [CAFA 5 Protein Function Prediction](https://www.kaggle.com/competitions/cafa-5-protein-function-prediction/data)
## Pretrained Models
The following pretrained models are available on HuggingFace Hub.
| Name | Embedding Dim. | Attn. Heads | Encoder Layers | Context Length | QAT | Total Parameters |
|---|---|---|---|---|---|---|
| [andrewdalpino/ESMC-300M-Protein-Function](https://huggingface.co/andrewdalpino/ESMC-300M-Protein-Function) | 960 | 15 | 30 | 2048 | None | 361M |
| [andrewdalpino/ESMC-300M-QAT-Protein-Function](https://huggingface.co/andrewdalpino/ESMC-300M-QAT-Protein-Function) | 960 | 15 | 30 | 2048 | int8w | 361M |
| [andrewdalpino/ESMC-600M-Protein-Function](https://huggingface.co/andrewdalpino/ESMC-600M-Protein-Function) | 1152 | 18 | 36 | 2048 | None | 644M |
| [andrewdalpino/ESMC-600M-QAT-Protein-Function](https://huggingface.co/andrewdalpino/ESMC-600M-QAT-Protein-Function) | 1152 | 18 | 36 | 2048 | int8w | 644M |
## Basic Pretrained Example
First, install the `esmc_function_classifier` package using [pip](https://pypi.org/project/pip/).
```sh
pip install esmc_function_classifier
```
Then, we'll load the model weights from HuggingFace Hub, tokenize the amino acid sequence, and infer the GO terms.
```python
import torch
from esm.tokenization import EsmSequenceTokenizer
from esmc_function_classifier.model import EsmcGoTermClassifier
model_name = "andrewdalpino/ESMC-300M-Protein-Function"
sequence = "MPPKGHKKTADGDFRPVNSAGNTIQAKQKYSIDDLLYPKSTIKNLAKETLPDDAIISKDALTAIQRAATLFVSYMASHGNASAEAGGRKKIT"
top_p = 0.5
tokenizer = EsmSequenceTokenizer()
model = EsmcGoTermClassifier.from_pretrained(model_name)
out = tokenizer(sequence, max_length=2048, truncation=True)
input_ids = torch.tensor(out["input_ids"], dtype=torch.int64)
go_term_probabilities = model.predict_terms(
input_ids, top_p=top_p
)
```
You can also output the gene-ontology (GO) `networkx` subgraph for a given sequence like in the example below. You'll need an up-to-date gene ontology database that you can import using the `obonet` package.
```sh
pip install obonet
```
```python
import networkx as nx
import obonet
# Visit https://geneontology.org/docs/download-ontology/ to download.
go_db_path = "./dataset/go-basic.obo"
graph = obonet.read_obo(go_db_path)
model.load_gene_ontology(graph)
subgraph, go_term_probabilities = model.predict_subgraph(
input_ids, top_p=top_p
)
json = nx.node_link_data(subgraph)
print(json)
```
### Quantized Model
To quantize the model weights using int8 call the `quantize_weights()` method. Any model can be quantized, but we recommend one that has been quantization-aware trained (QAT) for the best performance. The `group_size` argument controls the granularity at which quantization scales are computed.
```python
model.quantize_weights(group_size=64)
```
## Code Repository
The training code can be found at [https://github.com/andrewdalpino/ESMC-Function-Classifier](https://github.com/andrewdalpino/ESMC-Function-Classifier).
## References:
>- T. Hayes, et al. Simulating 500 million years of evolution with a language model, 2024.
>- M. Ashburner, et al. Gene Ontology: tool for the unification of biology, 2000. |