"""The model container class for ExLlamaV2 models.""" import gc from itertools import zip_longest import pathlib import time import torch from exllamav2 import ( ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache, ExLlamaV2Cache_8bit, ExLlamaV2Cache_Q4, ExLlamaV2Tokenizer, ExLlamaV2Lora, ) from exllamav2.generator import ExLlamaV2StreamingGenerator, ExLlamaV2Sampler from loguru import logger from typing import List, Optional, Union from backends.exllamav2.grammar import ExLlamaV2Grammar from common.gen_logging import log_generation_params, log_prompt, log_response from common.templating import ( PromptTemplate, find_template_from_model, get_template_from_model_json, get_template_from_file, ) from common.utils import coalesce, unwrap class ExllamaV2Container: """The model container class for ExLlamaV2 models.""" # Exl2 vars config: Optional[ExLlamaV2Config] = None draft_config: Optional[ExLlamaV2Config] = None model: Optional[ExLlamaV2] = None draft_model: Optional[ExLlamaV2] = None cache: Optional[ExLlamaV2Cache] = None draft_cache: Optional[ExLlamaV2Cache] = None tokenizer: Optional[ExLlamaV2Tokenizer] = None generator: Optional[ExLlamaV2StreamingGenerator] = None prompt_template: Optional[PromptTemplate] = None active_loras: List[ExLlamaV2Lora] = [] # Internal config vars cache_mode: str = "FP16" use_cfg: bool = False # GPU split vars gpu_split: Optional[list] = None gpu_split_auto: bool = True autosplit_reserve: List[float] = [96 * 1024**2] # Load state model_is_loading: bool = False model_loaded: bool = False def __init__(self, model_directory: pathlib.Path, quiet=False, **kwargs): """ Create model container Args: model_dir (int): Model directory containing config.json, tokenizer.model etc. quiet (bool): Suppress console output load_progress_callback (function, optional): A function to call for each module loaded. Prototype: def progress(loaded_modules: int, total_modules: int, loading_draft: bool) **kwargs: `cache_mode` (str): Sets cache mode, "FP16" or "FP8" (defaulf: "FP16") 'max_seq_len' (int): Override model's default max sequence length (default: 4096) 'rope_scale' (float): Set RoPE scaling factor for model (default: 1.0) 'rope_alpha' (float): Set RoPE alpha (NTK) factor for model (default: 1.0) 'prompt_template' (str): Manually sets the prompt template for this model (default: None) 'chunk_size' (int): Sets the maximum chunk size for the model (default: 2048) Inferencing in chunks reduces overall VRAM overhead by processing very long sequences in smaller batches. This limits the size of temporary buffers needed for the hidden state and attention weights. 'draft_model_dir' (str): Draft model directory 'draft_rope_scale' (float): Set RoPE scaling factor for draft model (default: 1.0) 'draft_rope_alpha' (float): RoPE alpha (NTK) factor for draft model. By default, the draft model's alpha value is calculated automatically to scale to the size of the full model. 'lora_dir' (str): LoRA directory 'loras' (list[dict]): List of loras to be loaded, consisting of 'name' and 'scaling' 'gpu_split_auto' (bool): Automatically split model across available devices (default: True) 'gpu_split' (list[float]): Allocation for weights and (some) tensors, per device 'no_flash_attn' (bool): Turns off flash attention (increases vram usage) (default: False) 'use_cfg" (bool): Enables CFG support. Disables flash attention (default: False) """ self.quiet = quiet self.cache_mode = unwrap(kwargs.get("cache_mode"), "FP16") # Turn off GPU split if the user is using 1 GPU gpu_count = torch.cuda.device_count() gpu_split_auto = unwrap(kwargs.get("gpu_split_auto"), True) if gpu_count > 1 and gpu_split_auto: # Auto GPU split parameters self.gpu_split_auto = gpu_split_auto autosplit_reserve_megabytes = unwrap(kwargs.get("autosplit_reserve"), [96]) self.autosplit_reserve = list( map(lambda value: value * 1024**2, autosplit_reserve_megabytes) ) elif gpu_count > 1: # Manual GPU split self.gpu_split = kwargs.get("gpu_split") self.gpu_split_auto = False else: # One GPU setup self.gpu_split_auto = False logger.info("Disabling GPU split because one GPU is in use.") self.config = ExLlamaV2Config() self.config.model_dir = str(model_directory.resolve()) # Make the max seq len 4096 before preparing the config # This is a better default than 2038 self.config.max_seq_len = 4096 self.config.prepare() # Then override the base_seq_len if present override_base_seq_len = kwargs.get("override_base_seq_len") if override_base_seq_len: self.config.max_seq_len = override_base_seq_len # Grab the base model's sequence length before overrides for # rope calculations base_seq_len = self.config.max_seq_len # Set the target seq len if present target_max_seq_len = kwargs.get("max_seq_len") if target_max_seq_len: self.config.max_seq_len = target_max_seq_len # Set the rope scale self.config.scale_pos_emb = unwrap( kwargs.get("rope_scale"), self.config.scale_pos_emb ) # Automatically calculate rope alpha self.config.scale_alpha_value = unwrap( kwargs.get("rope_alpha"), self.calculate_rope_alpha(base_seq_len) ) # Enable CFG if present self.use_cfg = unwrap(kwargs.get("use_cfg"), False) # Enable fasttensors loading if present self.config.fasttensors = unwrap(kwargs.get("fasttensors"), False) # Turn off flash attention if CFG is on # Workaround until batched FA2 is fixed in exllamav2 upstream self.config.no_flash_attn = ( True if self.use_cfg else unwrap(kwargs.get("no_flash_attention"), False) ) # low_mem is currently broken in exllamav2. Don't use it until it's # fixed. """ if "low_mem" in kwargs and kwargs["low_mem"]: self.config.set_low_mem() """ # Try to set prompt template self.prompt_template = self.find_prompt_template( kwargs.get("prompt_template"), model_directory ) # Catch all for template lookup errors if self.prompt_template: logger.info( f"Using template {self.prompt_template.name} " "for chat completions." ) else: logger.warning( "Chat completions are disabled because a prompt " "template wasn't provided or auto-detected." ) # Set num of experts per token if provided num_experts_override = kwargs.get("num_experts_per_token") if num_experts_override: self.config.num_experts_per_token = kwargs.get("num_experts_per_token") chunk_size = min( unwrap(kwargs.get("chunk_size"), 2048), self.config.max_seq_len ) self.config.max_input_len = chunk_size self.config.max_attn_size = chunk_size**2 draft_args = unwrap(kwargs.get("draft"), {}) draft_model_name = draft_args.get("draft_model_name") enable_draft = draft_args and draft_model_name # Always disable draft if params are incorrectly configured if draft_args and draft_model_name is None: logger.warning( "Draft model is disabled because a model name " "wasn't provided. Please check your config.yml!" ) enable_draft = False if enable_draft: self.draft_config = ExLlamaV2Config() draft_model_path = pathlib.Path( unwrap(draft_args.get("draft_model_dir"), "models") ) draft_model_path = draft_model_path / draft_model_name self.draft_config.model_dir = str(draft_model_path.resolve()) self.draft_config.prepare() self.draft_config.scale_pos_emb = unwrap( draft_args.get("draft_rope_scale"), 1.0 ) # Automatically calculate draft rope alpha self.draft_config.scale_alpha_value = unwrap( draft_args.get("draft_rope_alpha"), self.calculate_rope_alpha(self.draft_config.max_seq_len), ) self.draft_config.max_seq_len = self.config.max_seq_len if "chunk_size" in kwargs: self.draft_config.max_input_len = kwargs["chunk_size"] self.draft_config.max_attn_size = kwargs["chunk_size"] ** 2 def find_prompt_template(self, prompt_template_name, model_directory): """Tries to find a prompt template using various methods""" logger.info("Attempting to load a prompt template if present.") find_template_functions = [ lambda: get_template_from_model_json( pathlib.Path(self.config.model_dir) / "tokenizer_config.json", "chat_template", "from_tokenizer_config", ), lambda: get_template_from_file(find_template_from_model(model_directory)), ] # Add lookup from prompt template name if provided if prompt_template_name: find_template_functions.insert( 0, lambda: get_template_from_file(prompt_template_name) ) for func in find_template_functions: try: prompt_template = func() if prompt_template is not None: return prompt_template except (FileNotFoundError, LookupError): continue def calculate_rope_alpha(self, base_seq_len): """Calculate the rope alpha value for a given sequence length.""" ratio = self.config.max_seq_len / base_seq_len # Default to a 1 alpha if the sequence length is ever less # than or equal to 1 if ratio <= 1.0: alpha = 1 else: alpha = -0.13436 + 0.80541 * ratio + 0.28833 * ratio**2 return alpha def get_model_path(self, is_draft: bool = False): """Get the path for this model.""" model_path = pathlib.Path( self.draft_config.model_dir if is_draft else self.config.model_dir ) return model_path def get_model_parameters(self): model_params = { "name": self.get_model_path().name, "rope_scale": self.config.scale_pos_emb, "rope_alpha": self.config.scale_alpha_value, "max_seq_len": self.config.max_seq_len, "cache_mode": self.cache_mode, "num_experts_per_token": self.config.num_experts_per_token, "use_cfg": self.use_cfg, "prompt_template": self.prompt_template.name if self.prompt_template else None, } if self.draft_config: draft_model_params = { "name": self.get_model_path(is_draft=True).name, "rope_scale": self.draft_config.scale_pos_emb, "rope_alpha": self.draft_config.scale_alpha_value, "max_seq_len": self.draft_config.max_seq_len, } model_params["draft"] = draft_model_params return model_params def load(self, progress_callback=None): """ Load model Args: progress_callback (function, optional): A function to call for each module loaded. Prototype: def progress(loaded_modules: int, total_modules: int) """ for _ in self.load_gen(progress_callback): pass def load_loras(self, lora_directory: pathlib.Path, **kwargs): """ Load loras """ loras = unwrap(kwargs.get("loras"), []) success: List[str] = [] failure: List[str] = [] for lora in loras: lora_name = lora.get("name") lora_scaling = unwrap(lora.get("scaling"), 1.0) if lora_name is None: logger.warning( "One of your loras does not have a name. Please check your " "config.yml! Skipping lora load." ) failure.append(lora_name) continue logger.info(f"Loading lora: {lora_name} at scaling {lora_scaling}") lora_path = lora_directory / lora_name # FIXME(alpin): Does self.model need to be passed here? self.active_loras.append( ExLlamaV2Lora.from_directory(self.model, lora_path, lora_scaling) ) logger.info(f"Lora successfully loaded: {lora_name}") success.append(lora_name) # Return success and failure names return {"success": success, "failure": failure} def load_gen(self, progress_callback=None): """ Load model, generator function Args: progress_callback (function, optional): A function to call for each module loaded. Prototype: def progress(loaded_modules: int, total_modules: int) """ # Notify that the model is being loaded self.model_is_loading = True # Load tokenizer self.tokenizer = ExLlamaV2Tokenizer(self.config) # Calculate autosplit reserve for all GPUs gpu_count = torch.cuda.device_count() autosplit_reserve = self.autosplit_reserve + [0] * ( gpu_count - len(self.autosplit_reserve) ) # Load draft model if a config is present if self.draft_config: self.draft_model = ExLlamaV2(self.draft_config) if not self.quiet: logger.info("Loading draft model: " + self.draft_config.model_dir) self.draft_cache = ExLlamaV2Cache(self.draft_model, lazy=True) yield from self.draft_model.load_autosplit_gen( self.draft_cache, reserve_vram=autosplit_reserve, last_id_only=True, callback_gen=progress_callback, ) # Test VRAM allocation with a full-length forward pass input_ids = torch.zeros((1, self.config.max_input_len), dtype=torch.long) self.draft_model.forward(input_ids, cache=self.cache, preprocess_only=True) self.model = ExLlamaV2(self.config) if not self.quiet: logger.info("Loading model: " + self.config.model_dir) # Load model with manual split # Entrypoint for single GPU users if not self.gpu_split_auto: logger.info("Loading with a manual GPU split (or a one GPU setup)") for value in self.model.load_gen( self.gpu_split, callback_gen=progress_callback, ): if value: yield value batch_size = 2 if self.use_cfg else 1 if self.cache_mode == "Q4": self.cache = ExLlamaV2Cache_Q4( self.model, lazy=self.gpu_split_auto, batch_size=batch_size ) elif self.cache_mode == "FP8": self.cache = ExLlamaV2Cache_8bit( self.model, lazy=self.gpu_split_auto, batch_size=batch_size ) else: self.cache = ExLlamaV2Cache( self.model, lazy=self.gpu_split_auto, batch_size=batch_size ) # Load model with autosplit if self.gpu_split_auto: logger.info("Loading with autosplit") for value in self.model.load_autosplit_gen( self.cache, reserve_vram=autosplit_reserve, last_id_only=True, callback_gen=progress_callback, ): if value: yield value # Test VRAM allocation with a full-length forward pass input_ids = torch.zeros((1, self.config.max_input_len), dtype=torch.long) self.model.forward(input_ids, cache=self.cache, preprocess_only=True) # Create generator self.generator = ExLlamaV2StreamingGenerator( self.model, self.cache, self.tokenizer, self.draft_model, self.draft_cache, ) # Always return logprobs and logits self.generator.return_probabilities = True self.generator.return_logits = True # Clean up any extra vram usage from torch and cuda # (Helps reduce VRAM bottlenecking on Windows) gc.collect() torch.cuda.empty_cache() # Update model load state self.model_is_loading = False self.model_loaded = True logger.info("Model successfully loaded.") def unload(self, loras_only: bool = False): """ Free all VRAM resources used by this model """ for lora in self.active_loras: lora.unload() self.active_loras = [] # Unload the entire model if not just unloading loras if not loras_only: if self.model: self.model.unload() self.model = None if self.draft_model: self.draft_model.unload() self.draft_model = None self.config = None self.cache = None self.tokenizer = None self.generator = None # Set all model state variables to False self.model_is_loading = False self.model_loaded = False gc.collect() torch.cuda.empty_cache() logger.info("Loras unloaded." if loras_only else "Model unloaded.") def encode_tokens(self, text: str, **kwargs): """Wrapper to encode tokens from a text string""" return ( self.tokenizer.encode( text, add_bos=unwrap(kwargs.get("add_bos_token"), True), encode_special_tokens=unwrap(kwargs.get("encode_special_tokens"), True), ) .flatten() .tolist() ) def decode_tokens(self, ids: List[int], **kwargs): """Wrapper to decode tokens from a list of IDs""" ids = torch.tensor([ids]) return self.tokenizer.decode( ids, decode_special_tokens=unwrap(kwargs.get("decode_special_tokens"), True), )[0] def get_special_tokens(self, add_bos_token: bool, ban_eos_token: bool): return { "bos_token": self.tokenizer.bos_token if add_bos_token else "", "eos_token": self.tokenizer.eos_token if not ban_eos_token else "", "pad_token": self.tokenizer.pad_token, "unk_token": self.tokenizer.unk_token, } def get_logprobs(self, token_ids: torch.Tensor, token_probs: torch.Tensor): top_tokens = list( map( lambda index: self.tokenizer.extended_id_to_piece.get( index, self.tokenizer.id_to_piece[index] ), token_ids.flatten().tolist(), ) ) top_values = torch.log(token_probs).flatten().tolist() # Cannot return -inf in JSON cleaned_values = list( map(lambda value: -1000 if value == float("-inf") else value, top_values) ) return dict(zip_longest(top_tokens, cleaned_values)) def generate(self, prompt: str, **kwargs): """Generate a response to a prompt""" generations = list(self.generate_gen(prompt, **kwargs)) joined_generation = { "text": "", "prompt_tokens": 0, "generation_tokens": 0, "offset": [], "token_probs": {}, "logprobs": [], } if generations: for generation in generations: joined_generation["text"] += unwrap(generation.get("text"), "") joined_generation["offset"].append(unwrap(generation.get("offset"), [])) joined_generation["token_probs"].update( unwrap(generation.get("token_probs"), {}) ) # Include empty logprob dicts for index preservation joined_generation["logprobs"].append( unwrap(generation.get("logprobs"), {}) ) joined_generation["prompt_tokens"] = unwrap( generations[-1].get("prompt_tokens"), 0 ) joined_generation["generation_tokens"] = unwrap( generations[-1].get("generated_tokens"), 0 ) return joined_generation def check_unsupported_settings(self, **kwargs): """Check and warn the user if a sampler is unsupported. Meant for dev wheels!""" pass # pylint: disable=too-many-locals,too-many-branches,too-many-statements def generate_gen(self, prompt: str, **kwargs): """ Create generator function for prompt completion Args: prompt (str): Input prompt **kwargs: 'token_healing' (bool): Use token healing (default: False) 'temperature' (float): Sampling temperature (default: 1.0) 'temperature_last' (bool): Apply temperature after all other samplers (default: False) 'top_k' (int): Sampling top-K (default: 0) 'top_p' (float): Sampling top-P (default: 1.0) 'min_p' (float): Sampling min-P (default: 0.0) 'tfs' (float): Tail-free sampling (default: 0.0) 'typical' (float): Sampling typical (default: 0.0) 'mirostat' (bool): Use Mirostat (default: False) 'mirostat_tau' (float) Mirostat tau parameter (default: 1.5) 'mirostat_eta' (float) Mirostat eta parameter (default: 0.1) 'frequency_penalty' (float): Token frequency penalty (default: 0.0) 'presence_penalty' (float): Token presence penalty (default: 0.0) 'repetition_penalty' (float): Token repetition penalty (default: 1.15) 'penalty_range' (int): Penalty range (default: whole context) 'repetition_decay' (int): Repetition penalty range (default: same as range) 'stop' (List[Union[str, int]]): List of stop strings/tokens to end response (default: [EOS]) 'max_tokens' (int): Max no. tokens in response (default: 150) 'add_bos_token' (bool): Adds the BOS token to the start of the prompt (default: True) 'ban_eos_token' (bool): Bans the EOS token from generation (default: False) 'logit_bias' (Dict[int, float]): Biases specific tokens to either show up more or less (default: None) 'stream_interval' (float): Interval in seconds between each output chunk (default: immediate) 'generate_window' (int): Space to reserve at the end of the model's context when generating. Rolls context window by the same amount if context length is exceeded to allow generating pastthe models max_seq_len. """ token_healing = unwrap(kwargs.get("token_healing"), False) max_tokens = unwrap(kwargs.get("max_tokens"), 150) stream_interval = unwrap(kwargs.get("stream_interval"), 0) generate_window = max( unwrap(kwargs.get("generate_window"), 512), self.config.max_seq_len // 8 ) # Sampler settings gen_settings = ExLlamaV2Sampler.Settings() # Check unsupported settings for dev wheels self.check_unsupported_settings(**kwargs) # Apply settings gen_settings.temperature = unwrap(kwargs.get("temperature"), 1.0) gen_settings.temperature_last = unwrap(kwargs.get("temperature_last"), False) gen_settings.smoothing_factor = unwrap(kwargs.get("smoothing_factor"), 0.0) gen_settings.top_k = unwrap(kwargs.get("top_k"), 0) gen_settings.top_p = unwrap(kwargs.get("top_p"), 1.0) gen_settings.top_a = unwrap(kwargs.get("top_a"), 0.0) gen_settings.min_p = unwrap(kwargs.get("min_p"), 0.0) gen_settings.tfs = unwrap(kwargs.get("tfs"), 1.0) gen_settings.typical = unwrap(kwargs.get("typical"), 1.0) gen_settings.mirostat = unwrap(kwargs.get("mirostat"), False) # DynaTemp settings max_temp = unwrap(kwargs.get("max_temp"), 1.0) min_temp = unwrap(kwargs.get("min_temp"), 1.0) if max_temp > min_temp: gen_settings.max_temp = max_temp gen_settings.min_temp = min_temp gen_settings.temp_exponent = unwrap(kwargs.get("temp_exponent"), 1.0) else: # Force to default values gen_settings.max_temp = 1.0 gen_settings.min_temp = 1.0 gen_settings.temp_exponent = 1.0 # Warn if max/min temp values are > 0 # and if they're less than or equal to each other if max_temp < min_temp or ( 1 not in {min_temp, max_temp} and max_temp == min_temp ): logger.warning( "Max temp is less than or equal to min temp, skipping DynaTemp." ) # Default tau and eta fallbacks don't matter if mirostat is off gen_settings.mirostat_tau = unwrap(kwargs.get("mirostat_tau"), 1.5) gen_settings.mirostat_eta = unwrap(kwargs.get("mirostat_eta"), 0.1) # Set CFG scale and negative prompt cfg_scale = unwrap(kwargs.get("cfg_scale"), 1.0) negative_prompt = None if cfg_scale not in [None, 1.0]: if self.use_cfg: gen_settings.cfg_scale = cfg_scale # If the negative prompt is empty, use the BOS token negative_prompt = unwrap( kwargs.get("negative_prompt"), self.tokenizer.bos_token ) else: logger.warning( "CFG is currently disabled. " "Please reload your model with use_cfg = True.", ) gen_settings.token_repetition_penalty = unwrap( kwargs.get("repetition_penalty"), 1.0 ) gen_settings.token_frequency_penalty = unwrap( kwargs.get("frequency_penalty"), 0.0 ) gen_settings.token_presence_penalty = unwrap( kwargs.get("presence_penalty"), 0.0 ) # Applies for all penalties despite being called token_repetition_range gen_settings.token_repetition_range = unwrap( kwargs.get("penalty_range"), self.config.max_seq_len ) # Dynamically scale penalty range to output tokens # Only do this if freq/pres pen is enabled # and the repetition range is -1 auto_scale_penalty_range = ( gen_settings.token_frequency_penalty != 0 or gen_settings.token_presence_penalty != 0 ) and gen_settings.token_repetition_range == -1 # Always make sure the fallback is 0 if range < 0 # It's technically fine to use -1, but this just validates the passed # fallback # Always default to 0 if something goes wrong if gen_settings.token_repetition_range < 0: fallback_decay = 0 else: fallback_decay = gen_settings.token_repetition_range gen_settings.token_repetition_decay = coalesce( kwargs.get("repetition_decay"), fallback_decay, 0 ) stop_conditions: List[Union[str, int]] = unwrap(kwargs.get("stop"), []) add_bos_token = unwrap(kwargs.get("add_bos_token"), True) ban_eos_token = unwrap(kwargs.get("ban_eos_token"), False) logit_bias = kwargs.get("logit_bias") # Logprobs request_logprobs = unwrap(kwargs.get("logprobs"), 0) self.generator.return_top_tokens = request_logprobs # Override sampler settings for temp = 0 if gen_settings.temperature == 0: gen_settings.temperature = 1.0 gen_settings.top_k = 1 gen_settings.top_p = 0 gen_settings.typical = 0 # Log generation options to console # Some options are too large, so log the args instead log_generation_params( max_tokens=max_tokens, **vars(gen_settings), token_healing=token_healing, auto_scale_penalty_range=auto_scale_penalty_range, generate_window=generate_window, add_bos_token=add_bos_token, ban_eos_token=ban_eos_token, logprobs=request_logprobs, stop_conditions=stop_conditions, logit_bias=logit_bias, ) # Log prompt to console log_prompt(prompt, negative_prompt) # Set logit bias if logit_bias: # Create a vocab tensor if it doesn't exist for token biasing if gen_settings.token_bias is None: padding = -self.tokenizer.config.vocab_size % 32 gen_settings.token_bias = torch.zeros( (self.tokenizer.config.vocab_size + padding,), dtype=torch.float, ) # Map logits to the tensor with their biases for token_id, bias in logit_bias.items(): if 0 <= token_id < len(self.tokenizer.id_to_piece): gen_settings.token_bias[token_id] = bias else: logger.warning( f"Logit bias: Token {token_id} not present " "in the model's vocab. Skipping." ) # Initialize grammar handler grammar_handler = ExLlamaV2Grammar() gen_settings.filters = [] # Add JSON schema filter if it exists json_schema = unwrap(kwargs.get("json_schema")) if json_schema: grammar_handler.add_json_schema_filter( json_schema, gen_settings, self.model, self.tokenizer ) # Add EBNF filter if it exists grammar_string = unwrap(kwargs.get("grammar_string")) if grammar_string: grammar_handler.add_ebnf_filter( grammar_string, gen_settings, self.model, self.tokenizer ) # Ban the EOS token if specified. If not, append to stop conditions # as well. # Set this below logging to avoid polluting the stop strings array if ban_eos_token: gen_settings.disallow_tokens(self.tokenizer, [self.tokenizer.eos_token_id]) else: stop_conditions.append(self.tokenizer.eos_token_id) # Stop conditions self.generator.set_stop_conditions(stop_conditions) # Tokenized context ids, offsets = self.tokenizer.encode( [prompt, negative_prompt] if negative_prompt and gen_settings.cfg_scale not in [None, 1.0] else prompt, add_bos=add_bos_token, encode_special_tokens=True, return_offsets=True, ) mask = ( self.tokenizer.padding_mask(ids) if self.use_cfg and gen_settings.cfg_scale not in [None, 1.0] else None ) context_len = len(ids[0]) if context_len > self.config.max_seq_len: logger.warning( f"Context length {context_len} is greater than max_seq_len " f"{self.config.max_seq_len}. Generation is truncated and " "metrics may not be accurate." ) prompt_tokens = ids.shape[-1] # Begin generated_tokens = 0 full_response = "" start_time = time.time() last_chunk_time = start_time save_tokens = torch.empty((ids.shape[0], 0), dtype=torch.bool) chunk_buffer = "" chunk_tokens = 0 while True: # Ingest prompt if chunk_tokens == 0: ids = torch.cat((ids, save_tokens), dim=-1) save_tokens = torch.empty((ids.shape[0], 0), dtype=torch.bool) overflow = ids.shape[-1] + generate_window - self.config.max_seq_len active_ids = ids[:, max(0, overflow) :] chunk_tokens = self.config.max_seq_len - active_ids.shape[-1] # Split for exllama versions that have CFG if self.use_cfg: self.generator.begin_stream( active_ids, gen_settings, token_healing=token_healing, loras=self.active_loras, input_mask=mask, position_offsets=offsets, ) else: self.generator.begin_stream( active_ids, gen_settings, token_healing=token_healing, loras=self.active_loras, ) # Reset offsets for subsequent passes if the context is truncated offsets = None if auto_scale_penalty_range: gen_settings.token_repetition_range = generated_tokens # Run dict generation # Guarantees return of chunk, eos, and chunk_token_ids raw_generation = self.generator.stream_ex() if token_healing: # Extract healed token ids[:, -1] = self.generator.sequence_ids[:, -2] token_healing = False # Get parameters that will always exist chunk = raw_generation["chunk"] eos = raw_generation["eos"] tokens = raw_generation["chunk_token_ids"] save_tokens = torch.cat( (save_tokens, tokens.expand(save_tokens.shape[0], -1)), dim=-1 ) chunk_buffer += chunk generated_tokens += 1 chunk_tokens -= 1 # Yield output now = time.time() elapsed = now - last_chunk_time if chunk_buffer != "" and ( elapsed > stream_interval or eos or generated_tokens == max_tokens ): generation = { "text": chunk_buffer, "prompt_tokens": prompt_tokens, "generated_tokens": generated_tokens, "offset": len(full_response), } if request_logprobs > 0: # Get top tokens and probs top_tokens = unwrap( raw_generation.get("top_tokens"), torch.empty((1, 0, 1), dtype=torch.long), ) top_probs = unwrap( raw_generation.get("top_probs"), torch.empty((1, 0, 1), dtype=torch.float), ) if top_tokens.numel() > 0 and top_probs.numel() > 0: logprobs = self.get_logprobs(top_tokens, top_probs) generation["logprobs"] = logprobs # The first logprob is the selected token prob generation["token_probs"] = { token: logprobs[token] for token in list(logprobs.keys())[:1] } yield generation full_response += chunk_buffer chunk_buffer = "" last_chunk_time = now if eos or generated_tokens == max_tokens: break # Print response log_response(full_response) elapsed_time = last_chunk_time - start_time initial_response = ( f"Metrics: {generated_tokens} tokens generated in " f"{round(elapsed_time, 2)} seconds" ) itemization = [] extra_parts = [] # Add tokens per second tokens_per_second = ( "Indeterminate" if elapsed_time == 0 else round(generated_tokens / elapsed_time, 2) ) itemization.append(f"{tokens_per_second} T/s") # Add context (original token count) if ids is not None: itemization.append(f"context {context_len} tokens") if context_len > self.config.max_seq_len: extra_parts.append("<-- Not accurate (truncated)") # Print output logger.info( initial_response + " (" + ", ".join(itemization) + ") " + " ".join(extra_parts) )