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