a14abfe21c
These are commonly seen in huggingface provided chat templates and aren't that difficult to add in. For feature parity, honor the add_bos_token and ban_eos_token parameters when constructing the prompt. Signed-off-by: kingbri <bdashore3@proton.me>
567 lines
24 KiB
Python
567 lines
24 KiB
Python
import gc
|
|
import pathlib
|
|
import time
|
|
import torch
|
|
from exllamav2 import(
|
|
ExLlamaV2,
|
|
ExLlamaV2Config,
|
|
ExLlamaV2Cache,
|
|
ExLlamaV2Cache_8bit,
|
|
ExLlamaV2Tokenizer,
|
|
ExLlamaV2Lora
|
|
)
|
|
from exllamav2.generator import(
|
|
ExLlamaV2StreamingGenerator,
|
|
ExLlamaV2Sampler
|
|
)
|
|
|
|
from gen_logging import log_generation_params, log_prompt, log_response
|
|
from typing import List, Optional, Union
|
|
from templating import PromptTemplate, find_template_from_model, get_template_from_model_json, get_template_from_file
|
|
from utils import coalesce, unwrap
|
|
|
|
# Bytes to reserve on first device when loading with auto split
|
|
auto_split_reserve_bytes = 96 * 1024**2
|
|
|
|
class ModelContainer:
|
|
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
|
|
|
|
cache_fp8: bool = False
|
|
gpu_split_auto: bool = True
|
|
gpu_split: Optional[list] = None
|
|
|
|
active_loras: List[ExLlamaV2Lora] = []
|
|
|
|
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)
|
|
"""
|
|
|
|
self.quiet = quiet
|
|
|
|
self.cache_fp8 = "cache_mode" in kwargs and kwargs["cache_mode"] == "FP8"
|
|
self.gpu_split = kwargs.get("gpu_split")
|
|
self.gpu_split_auto = unwrap(kwargs.get("gpu_split_auto"), True)
|
|
|
|
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"), 1.0)
|
|
|
|
# Automatically calculate rope alpha
|
|
self.config.scale_alpha_value = unwrap(
|
|
kwargs.get("rope_alpha"),
|
|
self.calculate_rope_alpha(base_seq_len)
|
|
)
|
|
|
|
# Turn off flash attention?
|
|
self.config.no_flash_attn = 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()
|
|
"""
|
|
|
|
# Set prompt template override if provided
|
|
prompt_template_name = kwargs.get("prompt_template")
|
|
if prompt_template_name:
|
|
print(f"Attempting to load prompt template with name {prompt_template_name}")
|
|
# Read the template
|
|
self.prompt_template = get_template_from_file(prompt_template_name)
|
|
else:
|
|
# Then try finding the template from the tokenizer_config.json
|
|
self.prompt_template = get_template_from_model_json(
|
|
pathlib.Path(self.config.model_dir) / "tokenizer_config.json",
|
|
"chat_template",
|
|
"from_tokenizer_config"
|
|
)
|
|
|
|
# Try finding the chat template from the model's config.json
|
|
# TODO: This may not even be used with huggingface models, mark for removal.
|
|
if self.prompt_template is None:
|
|
self.prompt_template = get_template_from_model_json(
|
|
pathlib.Path(self.config.model_config),
|
|
"chat_template",
|
|
"from_model_config"
|
|
)
|
|
|
|
# If that fails, attempt fetching from model name
|
|
if self.prompt_template is None:
|
|
template_match = find_template_from_model(model_directory)
|
|
if template_match:
|
|
self.prompt_template = get_template_from_file(template_match)
|
|
|
|
# Catch all for template lookup errors
|
|
if self.prompt_template:
|
|
print(f"Using template {self.prompt_template.name} for chat completions.")
|
|
else:
|
|
print(
|
|
"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:
|
|
if hasattr(self.config, "num_experts_per_token"):
|
|
self.config.num_experts_per_token = num_experts_override
|
|
else:
|
|
print(" !! Warning: Currently installed ExLlamaV2 does not support overriding MoE experts")
|
|
|
|
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:
|
|
print("A draft config was found but a model name was not given. Please check your config.yml! Skipping draft load.")
|
|
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 calculate_rope_alpha(self, base_seq_len):
|
|
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
|
|
alpha = 1 if ratio <= 1.0 else -0.13436 + 0.80541 * ratio + 0.28833 * ratio ** 2
|
|
return alpha
|
|
|
|
def get_model_path(self, is_draft: bool = False):
|
|
model_path = pathlib.Path(self.draft_config.model_dir if is_draft else self.config.model_dir)
|
|
return model_path
|
|
|
|
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:
|
|
print("One of your loras does not have a name. Please check your config.yml! Skipping lora load.")
|
|
failure.append(lora_name)
|
|
continue
|
|
|
|
print(f"Loading lora: {lora_name} at scaling {lora_scaling}")
|
|
lora_path = lora_directory / lora_name
|
|
self.active_loras.append(ExLlamaV2Lora.from_directory(self.model, lora_path, lora_scaling))
|
|
print("Lora successfully loaded.")
|
|
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)
|
|
"""
|
|
|
|
# Load tokenizer
|
|
self.tokenizer = ExLlamaV2Tokenizer(self.config)
|
|
|
|
# Load draft model if a config is present
|
|
if self.draft_config:
|
|
self.draft_model = ExLlamaV2(self.draft_config)
|
|
if not self.quiet:
|
|
print("Loading draft model: " + self.draft_config.model_dir)
|
|
|
|
self.draft_cache = ExLlamaV2Cache(self.draft_model, lazy = True)
|
|
reserve = [auto_split_reserve_bytes] + [0] * 16
|
|
yield from self.draft_model.load_autosplit_gen(self.draft_cache, reserve_vram = 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)
|
|
|
|
# Load model
|
|
self.model = ExLlamaV2(self.config)
|
|
if not self.quiet:
|
|
print("Loading model: " + self.config.model_dir)
|
|
|
|
if not self.gpu_split_auto:
|
|
for value in self.model.load_gen(self.gpu_split, callback_gen = progress_callback):
|
|
if isinstance(value, str):
|
|
yield value
|
|
|
|
if self.cache_fp8:
|
|
self.cache = ExLlamaV2Cache_8bit(self.model, lazy = self.gpu_split_auto)
|
|
else:
|
|
self.cache = ExLlamaV2Cache(self.model, lazy = self.gpu_split_auto)
|
|
|
|
if self.gpu_split_auto:
|
|
reserve = [auto_split_reserve_bytes] + [0] * 16
|
|
yield from self.model.load_autosplit_gen(self.cache, reserve_vram = 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.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)
|
|
|
|
print("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
|
|
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
# Common function for token operations
|
|
def get_tokens(self, text: Optional[str], ids: Optional[List[int]], **kwargs):
|
|
if text:
|
|
# Assume token encoding
|
|
return self.tokenizer.encode(
|
|
text,
|
|
add_bos = unwrap(kwargs.get("add_bos_token"), True),
|
|
encode_special_tokens = unwrap(kwargs.get("encode_special_tokens"), True)
|
|
)
|
|
if ids:
|
|
# Assume token decoding
|
|
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 generate(self, prompt: str, **kwargs):
|
|
generation = list(self.generate_gen(prompt, **kwargs))
|
|
if generation:
|
|
response = "".join(map(lambda chunk: chunk[0], generation))
|
|
return response, generation[-1][1], generation[-1][2]
|
|
else:
|
|
return "", 0, 0
|
|
|
|
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)
|
|
'repetition_penalty' (float): Token repetition/presence penalty (default: 1.15)
|
|
'repetition_range' (int): Repetition 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 past
|
|
the 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 = min(unwrap(kwargs.get("generate_window"), 512), max_tokens)
|
|
|
|
# Sampler settings
|
|
gen_settings = ExLlamaV2Sampler.Settings()
|
|
|
|
# Warn of unsupported settings if the setting is enabled
|
|
if (unwrap(kwargs.get("mirostat"), False)) and not hasattr(gen_settings, "mirostat"):
|
|
print(" !! Warning: Currently installed ExLlamaV2 does not support Mirostat sampling")
|
|
|
|
if (unwrap(kwargs.get("min_p"), 0.0)) not in [0.0, 1.0] and not hasattr(gen_settings, "min_p"):
|
|
print(" !! Warning: Currently installed ExLlamaV2 does not support min-P sampling")
|
|
|
|
if (unwrap(kwargs.get("tfs"), 0.0)) not in [0.0, 1.0] and not hasattr(gen_settings, "tfs"):
|
|
print(" !! Warning: Currently installed ExLlamaV2 does not support tail-free sampling (TFS)")
|
|
|
|
if (unwrap(kwargs.get("temperature_last"), False)) and not hasattr(gen_settings, "temperature_last"):
|
|
print(" !! Warning: Currently installed ExLlamaV2 does not support temperature_last")
|
|
|
|
# Apply settings
|
|
gen_settings.temperature = unwrap(kwargs.get("temperature"), 1.0)
|
|
gen_settings.temperature_last = unwrap(kwargs.get("temperature_last"), False)
|
|
gen_settings.top_k = unwrap(kwargs.get("top_k"), 0)
|
|
gen_settings.top_p = unwrap(kwargs.get("top_p"), 1.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)
|
|
|
|
# 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)
|
|
gen_settings.token_repetition_penalty = unwrap(kwargs.get("repetition_penalty"), 1.0)
|
|
gen_settings.token_repetition_range = unwrap(kwargs.get("repetition_range"), self.config.max_seq_len)
|
|
|
|
# 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
|
|
fallback_decay = 0 if gen_settings.token_repetition_range <= 0 else 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")
|
|
|
|
# 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,
|
|
add_bos_token = add_bos_token,
|
|
ban_eos_token = ban_eos_token,
|
|
stop_conditions = stop_conditions,
|
|
logit_bias = logit_bias
|
|
)
|
|
|
|
# Log prompt to console
|
|
log_prompt(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, bias in logit_bias.items():
|
|
gen_settings.token_bias[token] = bias
|
|
|
|
# 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 = self.tokenizer.encode(
|
|
prompt,
|
|
add_bos = add_bos_token,
|
|
encode_special_tokens = True
|
|
)
|
|
context_len = len(ids[0])
|
|
|
|
if context_len > self.config.max_seq_len:
|
|
print(
|
|
f"WARNING: The context length {context_len} is greater than the max_seq_len {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((1, 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((1, 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]
|
|
|
|
self.generator.begin_stream(active_ids, gen_settings, token_healing = token_healing, loras = self.active_loras)
|
|
|
|
# Generate
|
|
chunk, eos, tokens = self.generator.stream()
|
|
|
|
if token_healing:
|
|
ids[:, -1] = self.generator.sequence_ids[:, -2] # Extract healed token
|
|
token_healing = False
|
|
|
|
save_tokens = torch.cat((save_tokens, tokens), 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):
|
|
yield chunk_buffer, prompt_tokens, generated_tokens
|
|
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 {round(elapsed_time, 2)} seconds"
|
|
itemization = []
|
|
extra_parts = []
|
|
|
|
# Add tokens per second
|
|
itemization.append(f"{'Indeterminate' if elapsed_time == 0 else round(generated_tokens / elapsed_time, 2)} 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
|
|
print(initial_response + " (" + ", ".join(itemization) + ") " + " ".join(extra_parts))
|