Update README.md

README.md

@@ -24,21 +24,42 @@ goal-directed molecule generation tasks.
 ## How to load
 
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
-tokenizer = AutoTokenizer.from_pretrained("chandar-lab/NovoMolGen_32M_SMILES_BPE", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("chandar-lab/NovoMolGen_32M_SMILES_BPE", trust_remote_code=True)
+>>> from transformers import AutoTokenizer, AutoModelForCausalLM
+>>> tokenizer = AutoTokenizer.from_pretrained("chandar-lab/NovoMolGen_32M_SMILES_BPE", trust_remote_code=True)
+>>> model = AutoModelForCausalLM.from_pretrained("chandar-lab/NovoMolGen_32M_SMILES_BPE", trust_remote_code=True)
 ```
 
 ## Quick-start (FlashAttention + bf16)
 
 ```python
-from accelerate import Accelerator
+>>> from accelerate import Accelerator
 
-acc = Accelerator(mixed_precision='bf16')
-model = acc.prepare(model)
+>>> acc = Accelerator(mixed_precision='bf16')
+>>> model = acc.prepare(model)
 
-outputs = model.sample(tokenizer=tokenizer, batch_size=4)
-print(outputs['SMILES'])
+>>> outputs = model.sample(tokenizer=tokenizer, batch_size=4)
+>>> print(outputs['SMILES'])
+```
+
+## Transformers-native HF checkpoint (`revision="hf-checkpoint"`)
+
+We also publish a Transformers-native checkpoint on the `hf-checkpoint` revision. This version loads directly with `AutoModelForCausalLM` and works out-of-the-box with `.generate(...)`.
+
+```python 
+>>> import torch
+>>> from transformers import AutoTokenizer, AutoModelForCausalLM
+
+>>> model = AutoModelForCausalLM.from_pretrained("chandar-lab/NovoMolGen_32M_SMILES_BPE", revision='hf-checkpoint', device_map='auto')
+>>> tokenizer = AutoTokenizer.from_pretrained("chandar-lab/NovoMolGen_32M_SMILES_BPE", revision='hf-checkpoint')
+
+>>> input_ids = torch.tensor([[tokenizer.bos_token_id]]).expand(4, -1).contiguous().to(model.device)
+>>> outs = model.generate(input_ids=input_ids, temperature=1.0, max_length=64, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+
+>>> molecules = [t.replace(" ", "") for t in tokenizer.batch_decode(outs, skip_special_tokens=True)]
+['CCO[C@H](CNC(=O)N(CC(=O)OC(C)(C)C)c1cccc(Br)n1)C(F)(F)F',
+'CCn1nnnc1CNc1ncnc(N[C@H]2CCO[C@@H](C)C2)c1C',
+'CC(C)(O)CNC(=O)CC[C@H]1C[C@@H](NC(=O)COCC(F)F)C1',
+'Cc1ncc(C(=O)N2C[C@H]3[C@H](CNC(=O)c4cnn[nH]4)CCC[C@H]3C2)n1C']
 ```
 
 ## Citation