Update README.md

README.md

@@ -18,7 +18,7 @@ CELL-E 2 is the second iteration of the original [CELL-E](https://www.biorxiv.or
 
 CELL-E 2 is novel bidirectional transformer that can generate images depicting protein subcellular localization from the amino acid sequences (and *vice versa*). 
 CELL-E 2 not only captures the spatial complexity of protein localization and produce probability estimates of localization atop a nucleus image, but also being able to generate sequences from images, enabling *de novo* protein design. 
-We trained on the [Human Protein Atlas](https://www.proteinatlas.org) and the [OpenCell](https://opencell.czbiohub.org) datasets.
+We trained on the [Human Protein Atlas](https://www.proteinatlas.org) (HPA) and the [OpenCell](https://opencell.czbiohub.org) datasets.
 
 CELL-E 2 utilizes pretrained amino acid embeddings from [ESM-2](https://github.com/facebookresearch/esm).
 
@@ -26,25 +26,40 @@ CELL-E 2 utilizes pretrained amino acid embeddings from [ESM-2](https://github.c
 ## Model variations
 
 We have made several versions of CELL-E 2 available. The naming scheme follows the structure ```training set_hidden size``` where the hidden size is set to the embedding dimension of the pretrained ESM-2 model.
+We annotate the most useful models under Notes, however other models can be used if memory constraints are present.
+Since these models share similarities with BERT, the embeddings from any of these models may be benefical for downstream tasks.
 
-### HPA Models:
+**HPA Models**:
+HPA models are trained on the HPA dataset. They are best for general purpose predictions as they include a variety of cell types.
 
-| Model | #params | Language |
+| Model | Size | Notes
 |------------------------|--------------------------------|-------|
-| [`bert-base-uncased`](https://huggingface.co/bert-base-uncased) | 110M   | English |
-| [`bert-large-uncased`](https://huggingface.co/bert-large-uncased)              | 340M    | English | sub 
-| [`bert-base-cased`](https://huggingface.co/bert-base-cased)        | 110M    | English |
-| [`bert-large-cased`](https://huggingface.co/bert-large-cased) | 340M    |  English |
+| [`HPA_480`](https://huggingface.co/HuangLab/CELL-E_2_HPA_480) | 4.73 GB | **Best for Image Prediction** |
+| [`HPA_640`](https://huggingface.co/HuangLab/CELL-E_2_HPA_640) | 6.31 GB | |
+| [`HPA_1280`](https://huggingface.co/HuangLab/CELL-E_2_HPA_1280) | 10.8 GB | |
+| [`HPA_2560`](https://huggingface.co/HuangLab/CELL-E_2_HPA_2560) | 17.5 GB | **Best for Sequence Prediction** |
 
-## Intended uses & limitations
+**OpenCell Models**:
+OpenCell models are trained on the OpenCell dataset. These only contain HEK cells and should ideally only be used for predictions on HEK cells. They perform well on image prediction but the generate heatmaps contain little information. 
+
+| Model | Size | Notes
+|------------------------|--------------------------------|-------|
+| [`HPA_480`](https://huggingface.co/HuangLab/CELL-E_2_OpenCell_480) | 4.73 GB | |
+| [`HPA_640`](https://huggingface.co/HuangLab/CELL-E_2_OpenCell_640) | 6.31 GB | |
+| [`HPA_1280`](https://huggingface.co/HuangLab/CELL-E_2_OpenCel_1280) | 10.8 GB | |
+| [`HPA_2560`](https://huggingface.co/HuangLab/CELL-E_2_OpenCell_2560) | 17.5 GB | **Best for Sequence Prediction** |
+
+**Finetuned HPA Models**:
+These models were used the HPA models as checkpoints, but then were finetuned on the OpenCell dataset. We found that they improve image generation capabilities, but did not necessary see an improvement in sequence prediction. 
+
+| Model | Size | Notes
+|------------------------|--------------------------------|-------|
+| [`HPA_480`](https://huggingface.co/HuangLab/CELL-E_2_HPA_Finetuned_480) | 4.73 GB | **Best for Image Prediction** |
+| [`HPA_640`](https://huggingface.co/HuangLab/CELL-E_2_HPA_Finetuned_640) | 6.31 GB | |
+| [`HPA_1280`](https://huggingface.co/HuangLab/CELL-E_2_HPA_Finetuned_1280) | 10.8 GB | |
+| [`HPA_2560`](https://huggingface.co/HuangLab/CELL-E_2_HPA_Finetuned_2560) | 17.5 GB | |
 
-You can use the raw model for either masked language modeling or next sentence prediction, but it's mostly intended to
-be fine-tuned on a downstream task. See the [model hub](https://huggingface.co/models?filter=bert) to look for
-fine-tuned versions of a task that interests you.
 
-Note that this model is primarily aimed at being fine-tuned on tasks that use the whole sentence (potentially masked)
-to make decisions, such as sequence classification, token classification or question answering. For tasks such as text
-generation you should look at model like GPT2.
 
 ### How to use