Emaad commited on
Commit
14ec7d1
1 Parent(s): 0233cd1

Upload 3 files

Browse files
Files changed (4) hide show
  1. .gitattributes +1 -0
  2. README.md +87 -0
  3. images/architecture.png +3 -0
  4. images/huanglogo.jpeg +0 -0
.gitattributes CHANGED
@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
32
  *.zip filter=lfs diff=lfs merge=lfs -text
33
  *.zst filter=lfs diff=lfs merge=lfs -text
34
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
32
  *.zip filter=lfs diff=lfs merge=lfs -text
33
  *.zst filter=lfs diff=lfs merge=lfs -text
34
  *tfevents* filter=lfs diff=lfs merge=lfs -text
35
+ images/architecture.png filter=lfs diff=lfs merge=lfs -text
README.md CHANGED
@@ -1,3 +1,90 @@
1
  ---
2
  license: mit
 
 
 
 
 
 
 
 
3
  ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
  license: mit
3
+ library_name: pytorch
4
+ tags:
5
+ - biology
6
+ - microscopy
7
+ - text-to-image
8
+ - transformers
9
+ metrics:
10
+ - accuracy
11
  ---
12
+ [![Huang Lab](images/huanglogo.jpeg)](huanglab.ucsf.edu)
13
+
14
+ # CELL-E 2
15
+
16
+
17
+ ## Model description
18
+ [![CELL-E_2](images/architecture.png)](https://github.com/BoHuangLab/CELL-E_2)
19
+
20
+ CELL-E 2 is the second iteration of the original [CELL-E](https://www.biorxiv.org/content/10.1101/2022.05.27.493774v1) model which utilizes an amino acid sequence and nucleus image to make predictions of subcellular protein localization with respect to the nucleus.
21
+
22
+ CELL-E 2 is novel bidirectional transformer that can generate images depicting protein subcellular localization from the amino acid sequences (and *vice versa*).
23
+ 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.
24
+ We trained on the [Human Protein Atlas](https://www.proteinatlas.org) (HPA) and the [OpenCell](https://opencell.czbiohub.org) datasets.
25
+
26
+ CELL-E 2 utilizes pretrained amino acid embeddings from [ESM-2](https://github.com/facebookresearch/esm).
27
+ Localization is predicted as a binary image atop the provided nucleus. The logit values are weighted against these binary images to produce a heatmap of expected localization.
28
+
29
+ ## Model variations
30
+
31
+ 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.
32
+ We annotate the most useful models under Notes, however other models can be used if memory constraints are present.
33
+ Since these models share similarities with BERT, the embeddings from any of these models may be benefical for downstream tasks.
34
+
35
+ **HPA Models**:
36
+ HPA models are trained on the HPA dataset. They are best for general purpose predictions as they include a variety of cell types.
37
+
38
+ | Model | Size | Notes
39
+ |------------------------|--------------------------------|-------|
40
+ | [`HPA_480`](https://huggingface.co/HuangLab/CELL-E_2_HPA_480) | 4.73 GB | **Best for Image Prediction** |
41
+ | [`HPA_640`](https://huggingface.co/HuangLab/CELL-E_2_HPA_640) | 6.31 GB | |
42
+ | [`HPA_1280`](https://huggingface.co/HuangLab/CELL-E_2_HPA_1280) | 10.8 GB | |
43
+ | [`HPA_2560`](https://huggingface.co/HuangLab/CELL-E_2_HPA_2560) | 17.5 GB | **Best for Sequence Prediction** |
44
+
45
+ **OpenCell Models**:
46
+ 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.
47
+
48
+ | Model | Size | Notes
49
+ |------------------------|--------------------------------|-------|
50
+ | [`HPA_480`](https://huggingface.co/HuangLab/CELL-E_2_OpenCell_480) | 4.73 GB | |
51
+ | [`HPA_640`](https://huggingface.co/HuangLab/CELL-E_2_OpenCell_640) | 6.31 GB | |
52
+ | [`HPA_1280`](https://huggingface.co/HuangLab/CELL-E_2_OpenCel_1280) | 10.8 GB | |
53
+ | [`HPA_2560`](https://huggingface.co/HuangLab/CELL-E_2_OpenCell_2560) | 17.5 GB | **Best for Sequence Prediction** |
54
+
55
+ **Finetuned HPA Models**:
56
+ 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.
57
+
58
+ | Model | Size | Notes
59
+ |------------------------|--------------------------------|-------|
60
+ | [`HPA_480`](https://huggingface.co/HuangLab/CELL-E_2_HPA_Finetuned_480) | 4.73 GB | **Best for Image Prediction** |
61
+ | [`HPA_640`](https://huggingface.co/HuangLab/CELL-E_2_HPA_Finetuned_640) | 6.31 GB | |
62
+ | [`HPA_1280`](https://huggingface.co/HuangLab/CELL-E_2_HPA_Finetuned_1280) | 10.8 GB | |
63
+ | [`HPA_2560`](https://huggingface.co/HuangLab/CELL-E_2_HPA_Finetuned_2560) | 17.5 GB | |
64
+
65
+
66
+ ### How to use
67
+
68
+ The full codebase is available on [GitHub](https://github.com/BoHuangLab/CELL-E_2).
69
+ Download the model and make sure ```nuclues_vqgan.yaml```, ```threshold_vqgan.yaml```, ```config.yaml```, and ```model.ckpt``` are present.
70
+ ```
71
+ Here is how to use this model to do sequence prediction:
72
+
73
+ ```python
74
+ configs = OmegaConf.load(configs/config.yaml);
75
+ model = instantiate_from_config(configs.model).to(device);
76
+ model.sample(text=sequence, condition=nucleus)
77
+ ```
78
+
79
+
80
+ ### BibTeX entry and citation info
81
+
82
+ ```bibtex
83
+ @article{,
84
+ author = {Emaad Khwaja and
85
+ Yun S Song and
86
+ Aaron Agarunov and
87
+ Bo Huang},
88
+ title = {{CELL-E 2:} Translating Proteins to Pictures and Back with a Bidirectional Text-to-Image Transforme},
89
+ }
90
+ ```
images/architecture.png ADDED

Git LFS Details

  • SHA256: 95a3488ee3ee4f4deb17d9c6f3573e41690d589bb64a61a4aae2ab6cf471f0a0
  • Pointer size: 132 Bytes
  • Size of remote file: 2.77 MB
images/huanglogo.jpeg ADDED