Update README.md

README.md

@@ -3,22 +3,25 @@ library_name: transformers
 tags: []
 ---
 
-ESM++ is a faithful implementation of [ESMC](https://www.evolutionaryscale.ai/blog/esm-cambrian) that allows for batching and standard Huggingface compatibility without requiring the ESM package.
+# ESM++
+ESM++ is a faithful implementation of [ESMC](https://www.evolutionaryscale.ai/blog/esm-cambrian) ([license](https://www.evolutionaryscale.ai/policies/cambrian-open-license-agreement)that allows for batching and standard Huggingface compatibility without requiring the ESM Python package.
+The small version corresponds to the 300 million parameter version of ESMC.
 
-Use with transformers
+
+## Use with transformers
 ```python
 from transformers import AutoModelForMaskedLM #AutoModel also works
 model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_small', trust_remote_code=True)
 tokenizer = model.tokenizer
 
 sequences = ['MPRTEIN', 'MSEQWENCE']
-tokenized = tokenizer(sequences, return_tensors='pt')
+tokenized = tokenizer(sequences, padding=True, return_tensors='pt')
 
 # tokenized['labels'] = tokenized['input_ids'].clone() # correctly mask input_ids and set unmasked instances of labels to -100 for MLM training
 
 output = model(**tokenized) # get all hidden states with output_hidden_states=True
 print(output.logits.shape) # language modeling logits, (batch_size, seq_len, vocab_size), (2, 11, 64)
-print(output.last_hidden_state) # last hidden state of the model, (batch_size, seq_len, hidden_size), (2, 11, 960)
+print(output.last_hidden_state.shape) # last hidden state of the model, (batch_size, seq_len, hidden_size), (2, 11, 960)
 print(output.loss) # language modeling loss if you passed labels
 #print(output.hidden_states) # all hidden states if you passed output_hidden_states=True (in tuple)
 ```
@@ -29,12 +32,40 @@ ESM++ also supports sequence and token level classification tasks like ESM2. Sim
 from transformers import AutoModelForSequenceClassification, AutoModelForTokenClassification
 
 model = AutoModelForSequenceClassification.from_pretrained('Synthyra/ESMplusplus_small', num_labels=2, trust_remote_code=True)
-logits = model(**tokenized)
+logits = model(**tokenized).logits
 print(logits.shape) # (batch_size, num_labels), (2, 2)
 ```
 
+ESM++ weights are fp32 by default. You can load them in fp16 or bf16 like this:
+```python
+import torch
+model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_small', trust_remote_code=True, torch_dtype=torch.float16) # or torch.bfloat16
+```
+
+### Comparison across floating-point precision and implementations
+We measured the difference of the last hidden states of the fp32 weights vs. fp16 or bf16. We find that the fp16 is closer to the fp32 outputs, so we recommend loading in fp16.
+Please note that Evolutionary Scale loads ESMC in bf16 by default, which has its share of advantages and disadvantages in inference / training - so load whichever you like for half precision.
+Average MSE FP32 vs. FP16: 0.00000003
+Average MSE FP32 vs. BF16: 0.00000140
+
+We also measured the difference between the outputs of ESM++ vs. ESMC (both in bfloat16) on 1000 random sequences to ensure compliance with the ESM package.
+Average MSE of last hidden state: 7.74e-10
+
+You can load the weights from the ESM package instead of transformers by replacing .from_pretrained(...) to .from_pretrained_esm('esmc_300m')
+
+## Model probes
+We employ linear probing techniques on various PLMs and standard datasets, similar our previous [paper](https://www.biorxiv.org/content/10.1101/2024.07.30.605924v1), to access the intrinsic correlation between pooled hidden states and valuable properties. ESMC (and thus ESM++) perform very well.
+
+The plot below showcases performance normalized between the negative control (random vector embeddings) and the best performer. Classification task scores are averaged between MCC and F1 (or F1max for multilabel) and regression tasks are averaged between Spearman rho and R2.
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/88EweoycOIhJBk2RvGCKn.png){ width=50% }
 
+## Inference speeds
+We look at various ESM models and their throughput on H100. Adding efficient batching between ESMC and ESM++ significantly improves the throughput. ESM++ small is even faster than ESM2-35M with long sequences!
+The most gains will be seen with PyTorch > 2.5 on linux machines.
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/62f2bd3bdb7cbd214b658c48/NNmbupdIxuieuyjc8vMuj.png)
 
-Measured difference between this implementation and version loaded with ESM package (1000 random sequences)
+### Citation
+If you use any of this implementation or work please cite it (as well as the ESMC preprint)
 
-Average MSE: 7.742734737803403e-10
+### Note on the name
+The original thought was ESMC++ but anything with C would technically go against the ESM license agreement - so ESM++. Open to suggestions!