# Copyright 2022 MosaicML LLM Foundry authors # SPDX-License-Identifier: Apache-2.0 import time import warnings from argparse import ArgumentParser, ArgumentTypeError, Namespace from contextlib import nullcontext from typing import Any, Dict, List, Optional, Union import torch from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase, StoppingCriteria, StoppingCriteriaList, TextStreamer) class ChatFormatter: """A class for formatting the chat history. Args: system: The system prompt. If None, a default ChatML-formatted prompt is used. user: The user prompt. If None, a default ChatML value is used. assistant: The assistant prompt. If None, a default ChatML value is used. Attributes: system: The system prompt. user: The user prompt. assistant: The assistant prompt. response_prefix: The response prefix (anything before {} in the assistant format string) """ def __init__(self, system: str, user: str, assistant: str) -> None: self.system = system if system else '<|im_start|>system\nA conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.<|im_end|>\n' self.user = user if user else '<|im_start|>user\n{}<|im_end|>\n' self.assistant = assistant if assistant else '<|im_start|>assistant\n{}<|im_end|>\n' self.response_prefix = self.assistant.split('{}')[0] class Conversation: """A class for interacting with a chat-tuned LLM. Args: model: The model to use for inference. tokenizer: The tokenizer to use for inference. chat_format: The chat format to use for the conversation. generate_kwargs: The keyword arguments to pass to `model.generate`. stop_tokens: The tokens to stop generation on. Attributes: model: The model to use for inference. tokenizer: The tokenizer to use for inference. chat_format: The chat format to use for the conversation. streamer: The streamer to use for inference. generate_kwargs: The keyword arguments to pass to `model.generate`. history: The conversation history. cli_instructions: The instructions to display to the user. """ def __init__(self, model: PreTrainedModel, tokenizer: PreTrainedTokenizerBase, chat_format: ChatFormatter, generate_kwargs: Dict[str, Any], stop_tokens: Optional[List[str]] = None) -> None: if stop_tokens is None: stop_tokens = ['<|endoftext|>', '<|im_end|>'] self.model = model self.tokenizer = tokenizer self.chat_format = chat_format stop_token_ids = self.tokenizer.convert_tokens_to_ids(stop_tokens) if len(stop_token_ids) != len(stop_tokens): warnings.warn( f'Not all stop tokens were found in the tokenizer vocabulary: {stop_tokens}\n' + 'Generation may stop or continue unexpectedly.') class StopOnTokens(StoppingCriteria): def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs: Any) -> bool: del kwargs # unused for stop_id in stop_token_ids: if input_ids[0][-1] == stop_id: return True return False self.streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) self.generate_kwargs = { **generate_kwargs, 'stopping_criteria': StoppingCriteriaList([StopOnTokens()]), 'streamer': self.streamer, } self.history = [] self.cli_instructions = ( 'Enter your message below.\n- Hit return twice to send input to the model\n' + "- Type 'clear' to restart the conversation\n- Type 'history' to see the conversation\n" + "- Type 'quit' to end\n- Type 'system' to change the system prompt\n" ) def _history_as_formatted_str(self) -> str: text = self.chat_format.system + ''.join([ '\n'.join([ self.chat_format.user.format(item[0]), self.chat_format.assistant.format(item[1]), ]) for item in self.history[:-1] ]) text += self.chat_format.user.format(self.history[-1][0]) text += self.chat_format.response_prefix return text def turn(self, user_inp: str) -> None: self.history.append([user_inp, '']) conversation = self._history_as_formatted_str() input_ids = self.tokenizer(conversation, return_tensors='pt').input_ids input_ids = input_ids.to(self.model.device) # also stream to stdout maybe_synchronize() start = time.time() print('Assistant:') gkwargs = {**self.generate_kwargs, 'input_ids': input_ids} # this will stream to stdout, but we need to keep track of the output_ids for saving history output_ids = self.model.generate(**gkwargs) maybe_synchronize() end = time.time() print(f'took {end - start:.2f} seconds') new_tokens = output_ids[0, len(input_ids[0]):] assistant_response = self.tokenizer.decode(new_tokens, skip_special_tokens=True) self.history[-1][-1] = assistant_response def __call__(self) -> None: print(self.cli_instructions) while True: print('User:') user_inp_lines = [] while True: line = input() if line.strip() == '': break user_inp_lines.append(line) user_inp = '\n'.join(user_inp_lines) if user_inp.lower() == 'quit': break elif user_inp.lower() == 'clear': self.history = [] continue elif user_inp == 'history': print(f'history: {self.history}') continue elif user_inp == 'history_fmt': print(f'history: {self._history_as_formatted_str()}') continue elif user_inp == 'system': print('Enter a new system prompt:') new_system = input() sys = f'<|im_start|>system\n{new_system.strip()}.<|im_end|>\n' self.chat_format.system = sys continue self.turn(user_inp) def get_dtype(dtype: str): if dtype == 'fp32': return torch.float32 elif dtype == 'fp16': return torch.float16 elif dtype == 'bf16': return torch.bfloat16 else: raise NotImplementedError( f'dtype {dtype} is not supported. ' + 'We only support fp32, fp16, and bf16 currently') def str2bool(v: Union[str, bool]): if isinstance(v, bool): return v if v.lower() in ('yes', 'true', 't', 'y', '1'): return True elif v.lower() in ('no', 'false', 'f', 'n', '0'): return False else: raise ArgumentTypeError('Boolean value expected.') def str_or_bool(v: Union[str, bool]): if isinstance(v, bool): return v if v.lower() in ('yes', 'true', 't', 'y', '1'): return True elif v.lower() in ('no', 'false', 'f', 'n', '0'): return False else: return v def parse_args() -> Namespace: """Parse commandline arguments.""" parser = ArgumentParser( description='Load a HF CausalLM Model and use it to generate text.') parser.add_argument('-n', '--name_or_path', type=str, required=True) parser.add_argument('--max_new_tokens', type=int, default=512) parser.add_argument('--max_seq_len', type=int, default=None) parser.add_argument('--temperature', type=float, default=1.0) parser.add_argument('--top_k', type=int, default=50) parser.add_argument('--top_p', type=float, default=1.0) parser.add_argument('--do_sample', type=str2bool, nargs='?', const=True, default=True) parser.add_argument('--use_cache', type=str2bool, nargs='?', const=True, default=True) parser.add_argument('--eos_token_id', type=str, default=None) parser.add_argument('--pad_token_id', type=str, default=None) parser.add_argument('--model_dtype', type=str, choices=['fp32', 'fp16', 'bf16'], default=None) parser.add_argument('--autocast_dtype', type=str, choices=['fp32', 'fp16', 'bf16'], default=None) parser.add_argument('--warmup', type=str2bool, nargs='?', const=True, default=True) parser.add_argument('--trust_remote_code', type=str2bool, nargs='?', const=True, default=True) parser.add_argument('--use_auth_token', type=str_or_bool, nargs='?', const=True, default=None) parser.add_argument('--revision', type=str, default=None) parser.add_argument('--device', type=str, default=None) parser.add_argument('--device_map', type=str, default=None) parser.add_argument('--attn_impl', type=str, default=None) parser.add_argument('--seed', type=int, default=42) parser.add_argument('--system_prompt', type=str, default=None) parser.add_argument('--user_msg_fmt', type=str, default=None) parser.add_argument('--assistant_msg_fmt', type=str, default=None) parser.add_argument( '--stop_tokens', type=str, default='<|endoftext|> <|im_end|>', help='A string of tokens to stop generation on; will be split on spaces.' ) return parser.parse_args() def maybe_synchronize(): if torch.cuda.is_available(): torch.cuda.synchronize() def main(args: Namespace) -> None: # Set device or device_map if args.device and args.device_map: raise ValueError('You can only set one of `device` and `device_map`.') if args.device is not None: device = args.device device_map = None else: device = None device_map = args.device_map or 'auto' print(f'Using {device=} and {device_map=}') # Set model_dtype if args.model_dtype is not None: model_dtype = get_dtype(args.model_dtype) else: model_dtype = torch.float32 print(f'Using {model_dtype=}') # Grab config first print(f'Loading HF Config...') from_pretrained_kwargs = { 'use_auth_token': args.use_auth_token, 'trust_remote_code': args.trust_remote_code, 'revision': args.revision, } try: config = AutoConfig.from_pretrained(args.name_or_path, **from_pretrained_kwargs) if args.attn_impl is not None and hasattr(config, 'attn_config'): config.attn_config['attn_impl'] = args.attn_impl if hasattr(config, 'init_device') and device is not None: config.init_device = device if args.max_seq_len is not None and hasattr(config, 'max_seq_len'): config.max_seq_len = args.max_seq_len except Exception as e: raise RuntimeError( 'If you are having auth problems, try logging in via `huggingface-cli login` ' + 'or by setting the environment variable `export HUGGING_FACE_HUB_TOKEN=... ' + 'using your access token from https://huggingface.co/settings/tokens.' ) from e # Load HF Model print(f'Loading HF model with dtype={model_dtype}...') try: model = AutoModelForCausalLM.from_pretrained(args.name_or_path, config=config, torch_dtype=model_dtype, device_map=device_map, **from_pretrained_kwargs) model.eval() print(f'n_params={sum(p.numel() for p in model.parameters())}') if device is not None: print(f'Placing model on {device=}...') model.to(device) except Exception as e: raise RuntimeError( 'Unable to load HF model. ' + 'If you are having auth problems, try logging in via `huggingface-cli login` ' + 'or by setting the environment variable `export HUGGING_FACE_HUB_TOKEN=... ' + 'using your access token from https://huggingface.co/settings/tokens.' ) from e print('\nLoading HF tokenizer...') tokenizer = AutoTokenizer.from_pretrained(args.name_or_path, **from_pretrained_kwargs) if tokenizer.pad_token_id is None: warnings.warn( 'pad_token_id is not set for the tokenizer. Using eos_token_id as pad_token_id.' ) tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = 'left' generate_kwargs = { 'max_new_tokens': args.max_new_tokens, 'temperature': args.temperature, 'top_p': args.top_p, 'top_k': args.top_k, 'use_cache': args.use_cache, 'do_sample': args.do_sample, 'eos_token_id': args.eos_token_id or tokenizer.eos_token_id, 'pad_token_id': args.pad_token_id or tokenizer.eos_token_id, } # Autocast if args.autocast_dtype is not None: autocast_dtype = get_dtype(args.autocast_dtype) autocast_context = torch.autocast(model.device.type, autocast_dtype) print(f'Using autocast with dtype={autocast_dtype}...') else: autocast_context = nullcontext() print('NOT using autocast...') chat_format = ChatFormatter(system=args.system_prompt, user=args.user_msg_fmt, assistant=args.assistant_msg_fmt) conversation = Conversation(model=model, tokenizer=tokenizer, chat_format=chat_format, generate_kwargs=generate_kwargs, stop_tokens=args.stop_tokens.split()) # Warmup if args.warmup: print('Warming up...') with autocast_context: conversation.turn('Write a welcome message to the user.') conversation.history = [] print('Starting conversation...') with autocast_context: conversation() if __name__ == '__main__': main(parse_args())