Updating README.md to include example usage.

README.md

@@ -16,35 +16,23 @@ Our pre-training dataset contains RefinedWeb, deduplicated PILE, a subset of Red
 
 ## Usage
 
-Below we provide an example of loading the model via [HuggingFace Hub](https://huggingface.co/docs/hub/) as:
-
-```python
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-# obtain access to "meta-llama/Llama-2-7b-hf", then see https://huggingface.co/docs/hub/security-tokens to get a token 
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", token="hf_xxxx")
-
-model_path = "apple/OpenELM-450M"
-
-model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
-model = model.cuda().eval()
-prompt = "Once upon a time there was"
-tokenized_prompt = tokenizer(prompt)
-prompt_tensor = torch.tensor(tokenized_prompt["input_ids"], device="cuda").unsqueeze(0)
-output_ids = model.generate(prompt_tensor, max_new_tokens=256, repetition_penalty=1.2, pad_token_id=0)
-output_ids = output_ids[0].tolist()
-output_text = tokenizer.decode(output_ids, skip_special_tokens=True)
-print(f'{model_path=}, {prompt=}\n')
-print(output_text)
-
-# below is the output:
-"""
-model_path='apple/OpenELM-450M', prompt='Once upon a time there was'
-
-Once upon a time there was a little girl who lived in the woods. She had a big heart and she loved to play with her friends. One day, she decided to go for a walk in the woods. As she walked, she saw a beautiful tree. It was so tall that it looked like a mountain. The tree was covered with leaves and flowers.
-The little girl thought that this tree was very pretty. She wanted to climb up to the tree and see what was inside. So, she went up to the tree and climbed up to the top. She was very excited when she saw that the tree was full of beautiful flowers. She also
-"""
+We have provided an example function to generate output from OpenELM models loaded via [HuggingFace Hub](https://huggingface.co/docs/hub/) in `generate_openelm.py`.
+
+You can try the model by running the following command:
+```
+python generate_openelm.py --checkpoint apple/OpenELM-3B --hf_security_token [HF_SECURITY_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs no_repeat_ngram_size=10
+```
+
+Additional arguments to the huggingface generate function can be passed via `generate_kwargs`. As an example, to speedup the inference, you can try [lookup token speculative generation](https://huggingface.co/docs/transformers/generation_strategies)  by passing the `prompt_lookup_num_tokens` argument as follows:
+```
+python generate_openelm.py --checkpoint apple/OpenELM-3B --hf_security_token [HF_SECURITY_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs no_repeat_ngram_size=10 prompt_lookup_num_tokens=10
 ```
+Alternatively, model-wise speculative generation can be also tried by passing a smaller model checkpoint through the `speculative_model_ckpt` argument, for example:
+```
+python generate_openelm.py --checkpoint apple/OpenELM-3B --hf_security_token [HF_SECURITY_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs no_repeat_ngram_size=10 --speculative_model_ckpt apple/OpenELM-270M
+```
+
+
 
 
 ## Main Results