Update README.md

README.md

@@ -1,3 +1,17 @@
+---
+datasets:
+- andrewdalpino/AmiGO-Boost
+metrics:
+- precision
+- recall
+- f1
+base_model:
+- EvolutionaryScale/esmc-600m-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.
@@ -17,7 +31,7 @@ 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 | int6w | 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 |
 
@@ -34,17 +48,12 @@ Then, we'll load the model weights from HuggingFace Hub and the GO graph using `
 ```python
 import torch
 
-import obonet
-
 from esm.tokenization import EsmSequenceTokenizer
 
 from esmc_function_classifier.model import EsmcGoTermClassifier
 
 
-model_name = "andrewdalpino/ESMC-300M-Protein-Function"
-
-# Visit https://geneontology.org/docs/download-ontology/ to download.
-go_db_path = "./dataset/go-basic.obo"
+model_name = "andrewdalpino/ESMC-600M-Protein-Function"
 
 sequence = "MPPKGHKKTADGDFRPVNSAGNTIQAKQKYSIDDLLYPKSTIKNLAKETLPDDAIISKDALTAIQRAATLFVSYMASHGNASAEAGGRKKIT"
 
@@ -54,10 +63,6 @@ tokenizer = EsmSequenceTokenizer()
 
 model = EsmcGoTermClassifier.from_pretrained(model_name)
 
-graph = obonet.read_obo(go_db_path)
-
-model.load_gene_ontology(graph)
-
 out = tokenizer(sequence, max_length=2048, truncation=True)
 
 input_ids = torch.tensor(out["input_ids"], dtype=torch.int64)
@@ -67,12 +72,28 @@ go_term_probabilities = model.predict_terms(
 )
 ```
 
-You can also out put the gene-ontology (GO) `networkx` subgraph for a given sequence like in the example below.
+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.
 
 ```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
@@ -90,4 +111,4 @@ The training code can be found at [https://github.com/andrewdalpino/ESMC-Functio
 ## 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.
+>- M. Ashburner, et al. Gene Ontology: tool for the unification of biology, 2000.