Update README.md

README.md

@@ -1,199 +1,164 @@
 ---
 library_name: transformers
-tags: []
+tags:
+- chemistry
+- molecule
+license: mit
 ---
 
-# Model Card for Model ID
-
-<!-- Provide a quick summary of what the model is/does. -->
-
-
-
-## Model Details
+# Model Card for Roberta Zinc Enamine Decomposer
 
 ### Model Description
 
-<!-- Provide a longer summary of what this model is. -->
-
-This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+`roberta_zinc_enamine_decomposer` is trained to "decompose" a molecule SMILES embedding into 
+two "building block embeddings" representing Enamine building blocks expected to assemble 
+into the input molecule.
 
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
+The model is trained to convert embeddings from the 
+[roberta_zinc_480m](https://huggingface.co/entropy/roberta_zinc_480m) model, or compressed 
+embeddings from the 
+[roberta_zinc_compression_encoder](https://huggingface.co/entropy/roberta_zinc_compression_encoder) 
+model. The decomposer can map from any input size (32, 64, 128, 256, 512, 768) to any 
+output size (same values). For an input of shape `(batch_size, d_in)`, the output will be 
+of shape `(batch_size, 2, d_out)` (two building block embeddings per input)
 
-### Model Sources [optional]
+- **Developed by:** Karl Heyer
+- **License:** MIT
 
-<!-- Provide the basic links for the model. -->
-
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
-
-## Uses
-
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
 
 ### Direct Use
 
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
-
-[More Information Needed]
-
-### Downstream Use [optional]
-
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
-
-[More Information Needed]
-
-### Out-of-Scope Use
-
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
-
-[More Information Needed]
-
-## Bias, Risks, and Limitations
-
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
-
-[More Information Needed]
-
-### Recommendations
-
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
-
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
-
-## How to Get Started with the Model
-
-Use the code below to get started with the model.
-
-[More Information Needed]
-
-## Training Details
-
-### Training Data
-
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
-
-[More Information Needed]
+Usage examples. Note that input SMILES strings should be canonicalized.
+
+```python
+from sentence_transformers import models, SentenceTransformer
+from transformers import AutoModel
+import torch
+
+transformer = models.Transformer("entropy/roberta_zinc_480m", 
+                                 max_seq_length=256, 
+                                 model_args={"add_pooling_layer": False})
+
+pooling = models.Pooling(transformer.get_word_embedding_dimension(), 
+                         pooling_mode="mean")
+
+roberta_zinc = SentenceTransformer(modules=[transformer, pooling])
+
+decomposer = AutoModel.from_pretrained("entropy/roberta_zinc_enamine_decomposer", 
+                                       trust_remote_code=True)
+
+# smiles should be canonicalized
+smiles = [
+    "COc1ccc(F)cc1S(=O)(=O)OC(=O)C(C)c1ccon1",
+    "C#Cc1cc(C(F)(F)F)ccc1Nc1ccc(OC)c(S(=O)(=O)Cl)c1",
+    "COc1ccc(NC(=O)c2ccccc2Nc2ccc(OC)c(S(=O)(=O)Cl)c2)c(OC)c1",
+    "COc1ccc(OC(=O)c2noc3c2COCC3)cc1S(=O)(=O)Cl",
+    "COc1ccc(N2CCC(C(=O)N3CCCc4ccccc43)CC2)cc1S(=O)(=O)Cl",
+    "COc1ccc(F)cc1S(=O)(=O)OC(=O)C1(C)CCCNS1(=O)=O",
+    "COc1ccc(F)cc1S(=O)(=O)OC(=O)c1cc(-n2c(C)ccc2C)ccc1Cl",
+    "COc1ccc(F)cc1S(=O)(=O)OC(=O)c1cnc2c(c1)OCC(=O)N2"
+]
+
+# embed smiles
+embeddings = roberta_zinc.encode(smiles, convert_to_tensor=True)
+print(embeddings.shape)
+# torch.Size([8, 768])
+
+# decompose from 768 to 512
+# {input_size: [B, input_size]} -> {output_size: [len(output_sizes), B, 2, output_size]}
+output_sizes = [512]
+decomposed_embeddings = decomposer.decompose({embeddings.shape[1]: embeddings}, 
+                                             output_sizes)
+
+for k,v in decomposed_embeddings.items():
+    print(k,v.shape)
+    
+# 512 torch.Size([1, 8, 2, 512])
+
+
+# compress inputs to all sizes
+# [B, input_size] -> {compressed_size: [B, compressed_size]}
+sizes = [32, 64, 128, 256, 512, 768]
+embedding_dict = decomposer.compress(embeddings, sizes)
+for k,v in embedding_dict.items():
+    print(k, v.shape)
+# 32 torch.Size([8, 32])
+# 64 torch.Size([8, 64])
+# 128 torch.Size([8, 128])
+# 256 torch.Size([8, 256])
+# 512 torch.Size([8, 512])
+# 768 torch.Size([8, 768])
+
+# decompose all compressed inputs to all output sizes
+# {input_size: [B, input_size]} -> {output_size: [len(output_sizes), B, 2, output_size]}
+decomposed_embeddings = decomposer.decompose(embedding_dict, sizes)
+for k,v in decomposed_embeddings.items():
+    print(k,v.shape)
+# 32 torch.Size([6, 8, 2, 32])
+# 64 torch.Size([6, 8, 2, 64])
+# 128 torch.Size([6, 8, 2, 128])
+# 256 torch.Size([6, 8, 2, 256])
+# 512 torch.Size([6, 8, 2, 512])
+# 768 torch.Size([6, 8, 2, 768])
+    
+
+# for routing multiple inputs to multiple outputs, 
+# output tensors are stacked in order of `config.comp_sizes` used
+
+input_size = 128
+input_index = decomposer.config.comp_sizes.index(input_size)
+output_size = 512
+
+# outputs at `output_size` that came specificaly from the `input_size` input
+out1 = decomposed_embeddings[output_size][input_index]
+
+# compute only `input_size` to `output_size`, no stacking/routing
+out2 = decomposer.decompose({input_size: embedding_dict[input_size]}, 
+                            [output_size])[output_size]
+
+torch.allclose(out1, out2, atol=5e-6)
+```
 
 ### Training Procedure
 
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
-
-#### Preprocessing [optional]
-
-[More Information Needed]
+#### Preprocessing
 
+A dataset of 50M molecules was created by assembing a set of 80k [Enamine](https://enamine.net/)
+building blocks using in-silico forward synthesis. Product molecules and building blocks were 
+canonicalized and embedded with the 
+[roberta_zinc_480m](https://huggingface.co/entropy/roberta_zinc_480m) model.
 
 #### Training Hyperparameters
 
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
-
-#### Speeds, Sizes, Times [optional]
-
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
-
-[More Information Needed]
-
-## Evaluation
-
-<!-- This section describes the evaluation protocols and provides the results. -->
-
-### Testing Data, Factors & Metrics
-
-#### Testing Data
-
-<!-- This should link to a Dataset Card if possible. -->
+The model was trained for 6 epochs with a batch size of 2048, learning rate of 1e-3, 
+cosine scheduling, weight decay of 0.01 and 10% warmup.
 
-[More Information Needed]
+#### Training Loss
 
-#### Factors
+During training, the model is loaded with frozen, pre-trained embedding compression heads from 
+the [roberta_zinc_compression_encoder](https://huggingface.co/entropy/roberta_zinc_compression_encoder) 
+model and frozen, pre-computed Enamine building block embeddings at all compression sizes.
 
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+The training input is a batch of full size (768) embeddings from the 
+[roberta_zinc_480m](https://huggingface.co/entropy/roberta_zinc_480m) model. The embeddings 
+are first compressed to all compression sizes. The compressed embeddings are then used to 
+predict decomposed embeddings at all compression sizes.
 
-[More Information Needed]
+For the loss, the predicted decomposed embeddings are compared to the ground truth via cosine 
+similarity. We then sample 3072 reference embeddings from the pre-computed Enamine building 
+block embeddings (reference embeddings). For all sizes of outputs and ground truth, we compute 
+the pair-wise cosine similarity between the predicted/targets and the reference embeddings. 
+We then compute the row-wise pearson correlation between the similarity matrices for the 
+predicted and targets.
 
-#### Metrics
+## Model Card Authors
 
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
-
-[More Information Needed]
-
-### Results
-
-[More Information Needed]
-
-#### Summary
-
-
-
-## Model Examination [optional]
-
-<!-- Relevant interpretability work for the model goes here -->
-
-[More Information Needed]
-
-## Environmental Impact
-
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
-
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
-
-- **Hardware Type:** [More Information Needed]
-- **Hours used:** [More Information Needed]
-- **Cloud Provider:** [More Information Needed]
-- **Compute Region:** [More Information Needed]
-- **Carbon Emitted:** [More Information Needed]
-
-## Technical Specifications [optional]
-
-### Model Architecture and Objective
-
-[More Information Needed]
-
-### Compute Infrastructure
-
-[More Information Needed]
-
-#### Hardware
-
-[More Information Needed]
-
-#### Software
-
-[More Information Needed]
-
-## Citation [optional]
-
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
-
-**BibTeX:**
-
-[More Information Needed]
-
-**APA:**
-
-[More Information Needed]
-
-## Glossary [optional]
-
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
-
-[More Information Needed]
-
-## More Information [optional]
-
-[More Information Needed]
-
-## Model Card Authors [optional]
-
-[More Information Needed]
+Karl Heyer
 
 ## Model Card Contact
 
-[More Information Needed]
+[email protected]
+
+---
+license: mit
+---