generate_openelm.py

"""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)