Update README.md

README.md

@@ -25,6 +25,25 @@ These checkpoints use the same architecture and starting checkpoint as the offic
 - stage1-step20000-tokens42B
 - stage1-step30000-tokens63B
 
+## Inference
+You can access these checkpoints using the standard Hugging Face Transformers library:
+
+```
+from transformers import AutoModelForCausalLM, AutoTokenizer
+olmo_early_training = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-0425-1B-early-training")
+tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-2-0425-1B-early-training")
+message = ["The capital of the United States is "]
+inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
+
+response = olmo_early_training.generate(**inputs, max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
+print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])
+```
+
+To access a specific checkpoint, you can specify the revision:
+```
+olmo_early_training = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-0425-1B-early-training", revision="stage1-step20000-tokens42B")
+```
+
 ## Model Description
 - Developed by: Allen Institute for AI (Ai2)
 - Model type: a Transformer style autoregressive language model.