Update README.md

README.md

@@ -12,25 +12,53 @@ The model has ~102m parameters and was trained for 150000 iterations with a batc
 This model is useful for generating embeddings from SMILES strings.
 
 ```python
-from transformers import RobertaTokenizerFast, RobertaForMaskedLM, DataCollatorWithPadding
+from transformers import AutoModel, AutoTokenizer
 
-tokenizer = RobertaTokenizerFast.from_pretrained("entropy/roberta_zinc_480m", max_len=128)
-model = RobertaForMaskedLM.from_pretrained('entropy/roberta_zinc_480m')
-collator = DataCollatorWithPadding(tokenizer, padding=True, return_tensors='pt')
+tokenizer = AutoTokenizer.from_pretrained("entropy/roberta_zinc_480m")
+roberta_zinc = AutoModel.from_pretrained("entropy/roberta_zinc_480m", 
+                                         add_pooling_layer=False) # model was not trained with a pooler
 
-smiles = ['Brc1cc2c(NCc3ccccc3)ncnc2s1',
- 'Brc1cc2c(NCc3ccccn3)ncnc2s1',
- 'Brc1cc2c(NCc3cccs3)ncnc2s1',
- 'Brc1cc2c(NCc3ccncc3)ncnc2s1',
- 'Brc1cc2c(Nc3ccccc3)ncnc2s1']
+smiles = [
+    "Brc1cc2c(NCc3ccccc3)ncnc2s1",
+    "Brc1cc2c(NCc3ccccn3)ncnc2s1",
+    "Brc1cc2c(NCc3cccs3)ncnc2s1",
+    "Brc1cc2c(NCc3ccncc3)ncnc2s1",
+    "Brc1cc2c(Nc3ccccc3)ncnc2s1"
+]
 
-inputs = collator(tokenizer(smiles))
-outputs = model(**inputs, output_hidden_states=True)
+batch = tokenizer(smiles, return_tensors='pt', padding=True, pad_to_multiple_of=8)
+
+# mean pooling
+outputs = roberta_zinc(**batch, output_hidden_states=True)
 full_embeddings = outputs[1][-1]
-mask = inputs['attention_mask']
+mask = batch['attention_mask']
 embeddings = ((full_embeddings * mask.unsqueeze(-1)).sum(1) / mask.sum(-1).unsqueeze(-1))
 ```
 
+```python
+from sentence_transformers import models, SentenceTransformer
+
+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")
+
+model = SentenceTransformer(modules=[transformer, pooling])
+
+smiles = [
+    "Brc1cc2c(NCc3ccccc3)ncnc2s1",
+    "Brc1cc2c(NCc3ccccn3)ncnc2s1",
+    "Brc1cc2c(NCc3cccs3)ncnc2s1",
+    "Brc1cc2c(NCc3ccncc3)ncnc2s1",
+    "Brc1cc2c(Nc3ccccc3)ncnc2s1"
+]
+
+embeddings2 = model.encode(smiles, convert_to_tensor=True)
+```
+
+
 ## Decoder
 
 There is also a [decoder model](https://huggingface.co/entropy/roberta_zinc_decoder) trained to reconstruct