Add readme and sample inference code.

README.md

@@ -8,9 +8,9 @@ license_link: LICENSE
 
 *Sachin Mehta, Mohammad Hossein Sekhavat, Qingqing Cao, Maxwell Horton, Yanzi Jin, Chenfan Sun, Iman Mirzadeh, Mahyar Najibi, Dmitry Belenko, Peter Zatloukal, Mohammad Rastegari*
 
-We introduce **OpenELM**, a family of **Open** **E**fficient **L**anguage **M**odels. OpenELM uses a layer-wise scaling strategy to efficiently allocate parameters within each layer of the transformer model, leading to enhanced accuracy. We pretrained OpenELM models using the [CoreNet](https://github.com/apple/corenet) library. We release both pretrained and instruction tuned models with 270M, 450M, 1.1B and 3B parameters. We release the complete framework, encompassing data preparation, training, fine-tuning, and evaluation procedures, alongside multiple pre-trained checkpoints and training logs, to facilitate open research.
+We introduce **OpenELM**, a family of **Open**-source **E**fficient **L**anguage **M**odels. We release both pretrained and instruction tuned models with 270M, 450M, 1.1B and 3B parameters.  
 
-Our pre-training dataset contains RefinedWeb, deduplicated PILE, a subset of RedPajama, and a subset of Dolma v1.6, totaling approximately 1.8 trillion tokens. Please check license agreements and terms of these datasets before using them.
+Our pre-training dataset contains RefinedWeb, deduplicated PILE, a subset of RedPajama, and a subset of Dolma v1.6, totaling approximately 1.8 trillion tokens. 
 
 
 
@@ -28,11 +28,14 @@ Additional arguments to the hugging face generate function can be passed via `ge
 ```
 python generate_openelm.py --model apple/OpenELM-3B --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, try model-wise speculative generation with an [assistive model](https://huggingface.co/blog/assisted-generation) by passing a smaller model through the `assistant_model` argument, for example:
+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 apple/OpenELM-3B --hf_access_token [HF_ACCESS_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs repetition_penalty=1.2 --assistant_model [SMALLER_MODEL]
+python generate_openelm.py --model apple/OpenELM-3B --hf_access_token [HF_ACCESS_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs repetition_penalty=1.2 --assistant_model apple/OpenELM-270M
 ```
 
+
+
+
 ## Main Results
 
 ### Zero-Shot
@@ -105,11 +108,10 @@ pip install tokenizers>=0.15.2 transformers>=4.38.2 sentencepiece>=0.2.0
 
 ```bash
 
-# OpenELM-3B
-hf_model=apple/OpenELM-3B
+# OpenELM-270M
+hf_model=OpenELM-270M
 
-# this flag is needed because lm-eval-harness set add_bos_token to False by default, but OpenELM uses LLaMA tokenizer which requires add_bos_token to be True
-tokenizer=meta-llama/Llama-2-7b-hf
+# this flag is needed because lm-eval-harness set add_bos_token to False by default, but OpenELM uses LLaMa tokenizer which requires add_bos_token to be True
 add_bos_token=True
 batch_size=1
 
@@ -118,7 +120,7 @@ mkdir lm_eval_output
 shot=0
 task=arc_challenge,arc_easy,boolq,hellaswag,piqa,race,winogrande,sciq,truthfulqa_mc2
 lm_eval --model hf \
-        --model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token},tokenizer=${tokenizer} \
+        --model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token} \
         --tasks ${task} \
         --device cuda:0 \
         --num_fewshot ${shot} \
@@ -128,7 +130,7 @@ lm_eval --model hf \
 shot=5
 task=mmlu,winogrande
 lm_eval --model hf \
-        --model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token},tokenizer=${tokenizer} \
+        --model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token} \
         --tasks ${task} \
         --device cuda:0 \
         --num_fewshot ${shot} \
@@ -138,7 +140,7 @@ lm_eval --model hf \
 shot=25
 task=arc_challenge,crows_pairs_english
 lm_eval --model hf \
-        --model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token},tokenizer=${tokenizer} \
+        --model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token} \
         --tasks ${task} \
         --device cuda:0 \
         --num_fewshot ${shot} \
@@ -148,7 +150,7 @@ lm_eval --model hf \
 shot=10
 task=hellaswag
 lm_eval --model hf \
-        --model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token},tokenizer=${tokenizer} \
+        --model_args pretrained=${hf_model},trust_remote_code=True,add_bos_token=${add_bos_token} \
         --tasks ${task} \
         --device cuda:0 \
         --num_fewshot ${shot} \
@@ -160,30 +162,5 @@ lm_eval --model hf \
 
 ## Bias, Risks, and Limitations
 
-The release of OpenELM models aims to empower and enrich the open research community by providing access to state-of-the-art language models. Trained on publicly available datasets, these models are made available without any safety guarantees. Consequently, there exists the possibility of these models producing outputs that are inaccurate, harmful, biased, or objectionable in response to user prompts. Thus, it is imperative for users and developers to undertake thorough safety testing and implement appropriate filtering mechanisms tailored to their specific requirements.
-
-## Citation
-
-If you find our work useful, please cite:
-
-```BibTex 
-@article{mehtaOpenELMEfficientLanguage2024,
-	title = {{OpenELM}: {An} {Efficient} {Language} {Model} {Family} with {Open} {Training} and {Inference} {Framework}},
-	shorttitle = {{OpenELM}},
-	url = {https://arxiv.org/abs/2404.14619v1},
-	language = {en},
-	urldate = {2024-04-24},
-	journal = {arXiv.org},
-	author = {Mehta, Sachin and Sekhavat, Mohammad Hossein and Cao, Qingqing and Horton, Maxwell and Jin, Yanzi and Sun, Chenfan and Mirzadeh, Iman and Najibi, Mahyar and Belenko, Dmitry and Zatloukal, Peter and Rastegari, Mohammad},
-	month = apr,
-	year = {2024},
-}
-
-@inproceedings{mehta2022cvnets, 
-     author = {Mehta, Sachin and Abdolhosseini, Farzad and Rastegari, Mohammad}, 
-     title = {CVNets: High Performance Library for Computer Vision}, 
-     year = {2022}, 
-     booktitle = {Proceedings of the 30th ACM International Conference on Multimedia}, 
-     series = {MM '22} 
-}
-```
+Our OpenELM models are not trained with any safety guarantees, the model outputs can be potentially inaccurate, harmful or contain biased information. produce inaccurate, biased or other objectionable responses to user prompts. Therefore, users and developers should conduct extensive safety testing and filtering suited to their specific needs.
+

modeling_openelm.py

@@ -783,7 +783,7 @@ class OpenELMModel(OpenELMPreTrainedModel):
             )
 
         if self.config._attn_implementation == "sdpa" and attention_mask is not None:
-            # For dynamo, rather use a check on fullgraph=True once this is possible (https://github.com/pytorch/pytorch/pull/120400).
+            # TODO: For dynamo, rather use a check on fullgraph=True once this is possible (https://github.com/pytorch/pytorch/pull/120400).
             is_tracing = (
                 torch.jit.is_tracing()
                 or isinstance(input_tensor, torch.fx.Proxy)
@@ -967,7 +967,7 @@ class OpenELMForCausalLM(OpenELMPreTrainedModel):
             input_ids = input_ids[:, past_length:]
             position_ids = position_ids[:, past_length:]
 
-        # we should only keep a `cache_position` in generate, and do +=1.
+        # TODO @gante we should only keep a `cache_position` in generate, and do +=1.
         # same goes for position ids. Could also help with continued generation.
         cache_position = torch.arange(
             past_length,
@@ -981,7 +981,7 @@ class OpenELMForCausalLM(OpenELMPreTrainedModel):
         else:
             # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
             # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
-            # We could use `next_tokens` directly instead.
+            # TODO: use `next_tokens` directly instead.
             model_inputs = {"input_ids": input_ids.contiguous()}
 
         model_inputs.update(