Update README with usage example.

README.md

@@ -42,34 +42,21 @@ openelm_3b_instruct = AutoModelForCausalLM.from_pretrained("apple/OpenELM-3B-Ins
 
 ## Usage
 
-Below we provide an example of loading the model via [HuggingFace Hub](https://huggingface.co/docs/hub/) as:
+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`.
 
-```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
-"""
+You can try the model by running the following command:
+```
+python generate_openelm.py --model [OPEN_ELM_MODEL] --hf_access_token [HF_ACCESS_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs repetition_penalty=1.2
+```
+Please refer to [this link](https://huggingface.co/docs/hub/security-tokens) to obtain your hugging face access token.
+
+Additional arguments to the hugging face 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 --model [OPEN_ELM_MODEL] --hf_access_token [HF_ACCESS_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs repetition_penalty=1.2 prompt_lookup_num_tokens=10
+```
+Alternatively, model-wise speculative generation with an [assistive model](https://huggingface.co/blog/assisted-generation) can be also tried by passing a smaller model model through the `assistant_model` argument, for example:
+```
+python generate_openelm.py --model [OPEN_ELM_MODEL] --hf_access_token [HF_ACCESS_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs repetition_penalty=1.2 --assistant_model [SMALLER_MODEL]
 ```