supporting-blog-content/using-openelm-models/OpenELM/generate_openelm.py

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