add OpenELM

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,47 @@
+Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+Disclaimer: IMPORTANT:  This Apple software is supplied to you by Apple
+Inc. ("Apple") in consideration of your agreement to the following
+terms, and your use, installation, modification or redistribution of
+this Apple software constitutes acceptance of these terms.  If you do
+not agree with these terms, please do not use, install, modify or
+redistribute this Apple software.
+
+In consideration of your agreement to abide by the following terms, and
+subject to these terms, Apple grants you a personal, non-exclusive
+license, under Apple's copyrights in this original Apple software (the
+"Apple Software"), to use, reproduce, modify and redistribute the Apple
+Software, with or without modifications, in source and/or binary forms;
+provided that if you redistribute the Apple Software in its entirety and
+without modifications, you must retain this notice and the following
+text and disclaimers in all such redistributions of the Apple Software.
+Neither the name, trademarks, service marks or logos of Apple Inc. may
+be used to endorse or promote products derived from the Apple Software
+without specific prior written permission from Apple.  Except as
+expressly stated in this notice, no other rights or licenses, express or
+implied, are granted by Apple herein, including but not limited to any
+patent rights that may be infringed by your derivative works or by other
+works in which the Apple Software may be incorporated.
+
+The Apple Software is provided by Apple on an "AS IS" basis.  APPLE
+MAKES NO WARRANTIES, EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION
+THE IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY AND FITNESS
+FOR A PARTICULAR PURPOSE, REGARDING THE APPLE SOFTWARE OR ITS USE AND
+OPERATION ALONE OR IN COMBINATION WITH YOUR PRODUCTS.
+
+IN NO EVENT SHALL APPLE BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL
+OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+INTERRUPTION) ARISING IN ANY WAY OUT OF THE USE, REPRODUCTION,
+MODIFICATION AND/OR DISTRIBUTION OF THE APPLE SOFTWARE, HOWEVER CAUSED
+AND WHETHER UNDER THEORY OF CONTRACT, TORT (INCLUDING NEGLIGENCE),
+STRICT LIABILITY OR OTHERWISE, EVEN IF APPLE HAS BEEN ADVISED OF THE
+POSSIBILITY OF SUCH DAMAGE.
+
+
+-------------------------------------------------------------------------------
+SOFTWARE DISTRIBUTED IN THIS REPOSITORY:
+
+This software includes a number of subcomponents with separate
+copyright notices and license terms - please see the file ACKNOWLEDGEMENTS.
+-------------------------------------------------------------------------------

README.md

@@ -0,0 +1,191 @@
+---
+license: other
+license_name: apple-sample-code-license
+license_link: LICENSE
+---
+
+# OpenELM
+
+*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**-source **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.  
+
+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.
+
+See the list below for the details of each model:
+
+- [OpenELM-270M](https://huggingface.co/apple/OpenELM-270M)                   
+- [OpenELM-450M](https://huggingface.co/apple/OpenELM-450M)                   
+- [OpenELM-1_1B](https://huggingface.co/apple/OpenELM-1_1B)                   
+- [OpenELM-3B](https://huggingface.co/apple/OpenELM-3B)                       
+- [OpenELM-270M-Instruct](https://huggingface.co/apple/OpenELM-270M-Instruct) 
+- [OpenELM-450M-Instruct](https://huggingface.co/apple/OpenELM-450M-Instruct) 
+- [OpenELM-1_1B-Instruct](https://huggingface.co/apple/OpenELM-1_1B-Instruct) 
+- [OpenELM-3B-Instruct](https://huggingface.co/apple/OpenELM-3B-Instruct)     
+
+
+```python
+
+from transformers import AutoModelForCausalLM
+
+openelm_270m = AutoModelForCausalLM.from_pretrained("apple/OpenELM-270M", trust_remote_code=True)
+openelm_450m = AutoModelForCausalLM.from_pretrained("apple/OpenELM-450M", trust_remote_code=True)
+openelm_1b = AutoModelForCausalLM.from_pretrained("apple/OpenELM-1_1B", trust_remote_code=True)
+openelm_3b = AutoModelForCausalLM.from_pretrained("apple/OpenELM-3B", trust_remote_code=True)
+
+openelm_270m_instruct = AutoModelForCausalLM.from_pretrained("apple/OpenELM-270M-Instruct", trust_remote_code=True)
+openelm_450m_instruct = AutoModelForCausalLM.from_pretrained("apple/OpenELM-450M-Instruct", trust_remote_code=True)
+openelm_1b_instruct = AutoModelForCausalLM.from_pretrained("apple/OpenELM-1_1B-Instruct", trust_remote_code=True)
+openelm_3b_instruct = AutoModelForCausalLM.from_pretrained("apple/OpenELM-3B-Instruct", trust_remote_code=True)
+
+```
+
+## Usage
+
+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 --model [MODEL_NAME] --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 [MODEL_NAME] --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:
+```
+python generate_openelm.py --model [MODEL_NAME] --hf_access_token [HF_ACCESS_TOKEN] --prompt 'Once upon a time there was' --generate_kwargs repetition_penalty=1.2 --assistant_model [SMALLER_MODEL_NAME]
+```
+
+
+## Main Results
+
+### Zero-Shot
+
+| **Model Size**                                                              | **ARC-c** | **ARC-e** | **BoolQ** | **HellaSwag** | **PIQA**  | **SciQ**  | **WinoGrande** | **Average** |
+|-----------------------------------------------------------------------------|-----------|-----------|-----------|---------------|-----------|-----------|----------------|-------------|
+| [OpenELM-270M](https://huggingface.co/apple/OpenELM-270M)                   | 26.45     | 45.08     | **53.98** | 46.71         | 69.75     | **84.70** | **53.91**      | 54.37       |
+| [OpenELM-270M-Instruct](https://huggingface.co/apple/OpenELM-270M-Instruct) | **30.55** | **46.68** | 48.56     | **52.07**     | **70.78** | 84.40     | 52.72          | **55.11**   |
+| [OpenELM-450M](https://huggingface.co/apple/OpenELM-450M)                   | 27.56     | 48.06     | 55.78     | 53.97         | 72.31     | 87.20     | 58.01          | 57.56       |
+| [OpenELM-450M-Instruct](https://huggingface.co/apple/OpenELM-450M-Instruct) | **30.38** | **50.00** | **60.37** | **59.34**     | **72.63** | **88.00** | **58.96**      | **59.95**   |
+| [OpenELM-1_1B](https://huggingface.co/apple/OpenELM-1_1B)                   | 32.34     | **55.43** | 63.58     | 64.81         | **75.57** | **90.60** | 61.72          | 63.44       |
+| [OpenELM-1_1B-Instruct](https://huggingface.co/apple/OpenELM-1_1B-Instruct) | **37.97** | 52.23     | **70.00** | **71.20**     | 75.03     | 89.30     | **62.75**      | **65.50**   |
+| [OpenELM-3B](https://huggingface.co/apple/OpenELM-3B)                       | 35.58     | 59.89     | 67.40     | 72.44         | 78.24     | **92.70** | 65.51          | 67.39       |
+| [OpenELM-3B-Instruct](https://huggingface.co/apple/OpenELM-3B-Instruct)     | **39.42** | **61.74** | **68.17** | **76.36**     | **79.00** | 92.50     | **66.85**      | **69.15**   |
+
+### LLM360
+
+| **Model Size**                                                              | **ARC-c** | **HellaSwag** | **MMLU**  | **TruthfulQA** | **WinoGrande** | **Average** |
+|-----------------------------------------------------------------------------|-----------|---------------|-----------|----------------|----------------|-------------|
+| [OpenELM-270M](https://huggingface.co/apple/OpenELM-270M)                   | 27.65     | 47.15         | 25.72     | **39.24**      | **53.83**      | 38.72       |
+| [OpenELM-270M-Instruct](https://huggingface.co/apple/OpenELM-270M-Instruct) | **32.51** | **51.58**     | **26.70** | 38.72          | 53.20          | **40.54**   |
+| [OpenELM-450M](https://huggingface.co/apple/OpenELM-450M)                   | 30.20     | 53.86         | **26.01** | 40.18          | 57.22          | 41.50       |
+| [OpenELM-450M-Instruct](https://huggingface.co/apple/OpenELM-450M-Instruct) | **33.53** | **59.31**     | 25.41     | **40.48**      | **58.33**      | **43.41**   |
+| [OpenELM-1_1B](https://huggingface.co/apple/OpenELM-1_1B)                   | 36.69     | 65.71         | **27.05** | 36.98          | 63.22          | 45.93       |
+| [OpenELM-1_1B-Instruct](https://huggingface.co/apple/OpenELM-1_1B-Instruct) | **41.55** | **71.83**     | 25.65     | **45.95**      | **64.72**      | **49.94**   |
+| [OpenELM-3B](https://huggingface.co/apple/OpenELM-3B)                       | 42.24     | 73.28         | **26.76** | 34.98          | 67.25          | 48.90       |
+| [OpenELM-3B-Instruct](https://huggingface.co/apple/OpenELM-3B-Instruct)     | **47.70** | **76.87**     | 24.80     | **38.76**      | **67.96**      | **51.22**   |
+
+
+### OpenLLM Leaderboard
+
+| **Model Size**                                                              | **ARC-c** | **CrowS-Pairs** | **HellaSwag** | **MMLU**  | **PIQA**  | **RACE**  | **TruthfulQA** | **WinoGrande** | **Average** |
+|-----------------------------------------------------------------------------|-----------|-----------------|---------------|-----------|-----------|-----------|----------------|----------------|-------------|
+| [OpenELM-270M](https://huggingface.co/apple/OpenELM-270M)                   | 27.65     | **66.79**       | 47.15         | 25.72     | 69.75     | 30.91     | **39.24**      | **53.83**      | 45.13       |
+| [OpenELM-270M-Instruct](https://huggingface.co/apple/OpenELM-270M-Instruct) | **32.51** | 66.01           | **51.58**     | **26.70** | **70.78** | 33.78     | 38.72          | 53.20          | **46.66**   |
+| [OpenELM-450M](https://huggingface.co/apple/OpenELM-450M)                   | 30.20     | **68.63**       | 53.86         | **26.01** | 72.31     | 33.11     | 40.18          | 57.22          | 47.69       |
+| [OpenELM-450M-Instruct](https://huggingface.co/apple/OpenELM-450M-Instruct) | **33.53** | 67.44           | **59.31**     | 25.41     | **72.63** | **36.84** | **40.48**      | **58.33**      | **49.25**   |
+| [OpenELM-1_1B](https://huggingface.co/apple/OpenELM-1_1B)                   | 36.69     | **71.74**       | 65.71         | **27.05** | **75.57** | 36.46     | 36.98          | 63.22          | 51.68       |
+| [OpenELM-1_1B-Instruct](https://huggingface.co/apple/OpenELM-1_1B-Instruct) | **41.55** | 71.02           | **71.83**     | 25.65     | 75.03     | **39.43** | **45.95**      | **64.72**      | **54.40**   |
+| [OpenELM-3B](https://huggingface.co/apple/OpenELM-3B)                       | 42.24     | **73.29**       | 73.28         | **26.76** | 78.24     | **38.76** | 34.98          | 67.25          | 54.35       |
+| [OpenELM-3B-Instruct](https://huggingface.co/apple/OpenELM-3B-Instruct)     | **47.70** | 72.33           | **76.87**     | 24.80     | **79.00** | 38.47     | **38.76**      | **67.96**      | **55.73**   |
+
+See the technical report for more results and comparison.
+
+## Evaluation
+
+### Setup
+
+Install the following dependencies:
+
+```bash
+
+# install public lm-eval-harness
+
+harness_repo="public-lm-eval-harness"
+git clone https://github.com/EleutherAI/lm-evaluation-harness ${harness_repo}
+cd ${harness_repo}
+# use main branch on 03-15-2024, SHA is dc90fec
+git checkout dc90fec
+pip install -e .
+cd ..
+
+# 66d6242 is the main branch on 2024-04-01 
+pip install datasets@git+https://github.com/huggingface/datasets.git@66d6242
+pip install tokenizers>=0.15.2 transformers>=4.38.2 sentencepiece>=0.2.0
+
+```
+
+### Evaluate OpenELM
+
+```bash
+
+# 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
+add_bos_token=True
+batch_size=1
+
+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} \
+        --tasks ${task} \
+        --device cuda:0 \
+        --num_fewshot ${shot} \
+        --output_path ./lm_eval_output/${hf_model//\//_}_${task//,/_}-${shot}shot \
+        --batch_size ${batch_size} 2>&1 | tee ./lm_eval_output/eval-${hf_model//\//_}_${task//,/_}-${shot}shot.log
+
+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} \
+        --tasks ${task} \
+        --device cuda:0 \
+        --num_fewshot ${shot} \
+        --output_path ./lm_eval_output/${hf_model//\//_}_${task//,/_}-${shot}shot \
+        --batch_size ${batch_size} 2>&1 | tee ./lm_eval_output/eval-${hf_model//\//_}_${task//,/_}-${shot}shot.log
+
+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} \
+        --tasks ${task} \
+        --device cuda:0 \
+        --num_fewshot ${shot} \
+        --output_path ./lm_eval_output/${hf_model//\//_}_${task//,/_}-${shot}shot \
+        --batch_size ${batch_size} 2>&1 | tee ./lm_eval_output/eval-${hf_model//\//_}_${task//,/_}-${shot}shot.log
+
+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} \
+        --tasks ${task} \
+        --device cuda:0 \
+        --num_fewshot ${shot} \
+        --output_path ./lm_eval_output/${hf_model//\//_}_${task//,/_}-${shot}shot \
+        --batch_size ${batch_size} 2>&1 | tee ./lm_eval_output/eval-${hf_model//\//_}_${task//,/_}-${shot}shot.log
+
+```
+
+
+## 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.

generate_openelm.py

@@ -0,0 +1,240 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (C) 2024 Apple Inc. All Rights Reserved.
+#
+
+"""Module to generate OpenELM output given a model and an input prompt."""
+import os
+import logging
+import time
+import argparse
+from typing import Optional, Union
+import torch
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+def generate(
+    prompt: str,
+    model: Union[str, AutoModelForCausalLM],
+    hf_access_token: str = None,
+    tokenizer: Union[str, AutoTokenizer] = 'meta-llama/Llama-2-7b-hf',
+    device: Optional[str] = None,
+    max_length: int = 1024,
+    assistant_model: Optional[Union[str, AutoModelForCausalLM]] = None,
+    generate_kwargs: Optional[dict] = None,
+) -> str:
+    """ Generates output given a prompt.
+
+    Args:
+        prompt: The string prompt.
+        model: The LLM Model. If a string is passed, it should be the path to
+            the hf converted checkpoint.
+        hf_access_token: Hugging face access token.
+        tokenizer: Tokenizer instance. If model is set as a string path,
+            the tokenizer will be loaded from the checkpoint.
+        device: String representation of device to run the model on. If None
+            and cuda available it would be set to cuda:0 else cpu.
+        max_length: Maximum length of tokens, input prompt + generated tokens.
+        assistant_model: If set, this model will be used for
+            speculative generation. If a string is passed, it should be the
+            path to the hf converted checkpoint.
+        generate_kwargs: Extra kwargs passed to the hf generate function.
+
+    Returns:
+        output_text: output generated as a string.
+        generation_time: generation time in seconds.
+
+    Raises:
+        ValueError: If device is set to CUDA but no CUDA device is detected.
+        ValueError: If tokenizer is not set.
+        ValueError: If hf_access_token is not specified.
+    """
+    if not device:
+        if torch.cuda.is_available() and torch.cuda.device_count():
+            device = "cuda:0"
+            logging.warning(
+                'inference device is not set, using cuda:0, %s',
+                torch.cuda.get_device_name(0)
+            )
+        else:
+            device = 'cpu'
+            logging.warning(
+                (
+                    'No CUDA device detected, using cpu, '
+                    'expect slower speeds.'
+                )
+            )
+
+    if 'cuda' in device and not torch.cuda.is_available():
+        raise ValueError('CUDA device requested but no CUDA device detected.')
+
+    if not tokenizer:
+        raise ValueError('Tokenizer is not set in the generate function.')
+
+    if not hf_access_token:
+        raise ValueError((
+            'Hugging face access token needs to be specified. '
+            'Please refer to https://huggingface.co/docs/hub/security-tokens'
+            ' to obtain one.'
+            )
+        )
+
+    if isinstance(model, str):
+        checkpoint_path = model
+        model = AutoModelForCausalLM.from_pretrained(
+            checkpoint_path,
+            trust_remote_code=True
+        )
+    model.to(device).eval()
+    if isinstance(tokenizer, str):
+        tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer,
+            token=hf_access_token,
+        )
+
+    # Speculative mode
+    draft_model = None
+    if assistant_model:
+        draft_model = assistant_model
+        if isinstance(assistant_model, str):
+            draft_model = AutoModelForCausalLM.from_pretrained(
+                assistant_model,
+                trust_remote_code=True
+            )
+        draft_model.to(device).eval()
+
+    # Prepare the prompt
+    tokenized_prompt = tokenizer(prompt)
+    tokenized_prompt = torch.tensor(
+        tokenized_prompt['input_ids'],
+        device=device
+    )
+
+    tokenized_prompt = tokenized_prompt.unsqueeze(0)
+
+    # Generate
+    stime = time.time()
+    output_ids = model.generate(
+        tokenized_prompt,
+        max_length=max_length,
+        pad_token_id=0,
+        assistant_model=draft_model,
+        **(generate_kwargs if generate_kwargs else {}),
+    )
+    generation_time = time.time() - stime
+
+    output_text = tokenizer.decode(
+        output_ids[0].tolist(),
+        skip_special_tokens=True
+    )
+
+    return output_text, generation_time
+
+
+def openelm_generate_parser():
+    """Argument Parser"""
+
+    class KwargsParser(argparse.Action):
+        """Parser action class to parse kwargs of form key=value"""
+        def __call__(self, parser, namespace, values, option_string=None):
+            setattr(namespace, self.dest, dict())
+            for val in values:
+                if '=' not in val:
+                    raise ValueError(
+                        (
+                            'Argument parsing error, kwargs are expected in'
+                            ' the form of key=value.'
+                        )
+                    )
+                kwarg_k, kwarg_v = val.split('=')
+                try:
+                    converted_v = int(kwarg_v)
+                except ValueError:
+                    try:
+                        converted_v = float(kwarg_v)
+                    except ValueError:
+                        converted_v = kwarg_v            
+                getattr(namespace, self.dest)[kwarg_k] = converted_v
+
+    parser = argparse.ArgumentParser('OpenELM Generate Module')
+    parser.add_argument(
+        '--model',
+        dest='model',
+        help='Path to the hf converted model.',
+        required=True,
+        type=str,
+    )
+    parser.add_argument(
+        '--hf_access_token',
+        dest='hf_access_token',
+        help='Hugging face access token, starting with "hf_".',
+        type=str,
+    )
+    parser.add_argument(
+      '--prompt',
+      dest='prompt',
+      help='Prompt for LLM call.',
+      default='',
+      type=str,
+    )
+    parser.add_argument(
+        '--device',
+        dest='device',
+        help='Device used for inference.',
+        type=str,
+    )
+    parser.add_argument(
+        '--max_length',
+        dest='max_length',
+        help='Maximum length of tokens.',
+        default=256,
+        type=int,
+    )
+    parser.add_argument(
+        '--assistant_model',
+        dest='assistant_model',
+        help=(
+            (
+                'If set, this is used as a draft model '
+                'for assisted speculative generation.'
+            )
+        ),
+        type=str,
+    )
+    parser.add_argument(
+        '--generate_kwargs',
+        dest='generate_kwargs',
+        help='Additional kwargs passed to the HF generate function.',
+        type=str,
+        nargs='*',
+        action=KwargsParser,
+    )
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = openelm_generate_parser()
+    prompt = args.prompt
+
+    output_text, genertaion_time = generate(
+        prompt=prompt,
+        model=args.model,
+        device=args.device,
+        max_length=args.max_length,
+        assistant_model=args.assistant_model,
+        generate_kwargs=args.generate_kwargs,
+        hf_access_token=args.hf_access_token,
+    )
+
+    print_txt = (
+        f'\r\n{"=" * os.get_terminal_size().columns}\r\n'
+        '\033[1m Prompt + Generated Output\033[0m\r\n'
+        f'{"-" * os.get_terminal_size().columns}\r\n'
+        f'{output_text}\r\n'
+        f'{"-" * os.get_terminal_size().columns}\r\n'
+        '\r\nGeneration took'
+        f'\033[1m\033[92m {round(genertaion_time, 2)} \033[0m'
+        'seconds.\r\n'
+    )
+    print(print_txt)