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