tabbyAPI/model.py

import gc, time, pathlib
import torch
from datetime import datetime
from exllamav2 import(
    ExLlamaV2,
    ExLlamaV2Config,
    ExLlamaV2Cache,
    ExLlamaV2Cache_8bit,
    ExLlamaV2Tokenizer,
)
from exllamav2.generator import(
    ExLlamaV2StreamingGenerator,
    ExLlamaV2Sampler
)
from typing import List, Optional, Union

# 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

    cache_fp8: bool = False
    draft_enabled: bool = False
    gpu_split_auto: bool = True
    gpu_split: list or None = None

    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
                '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)
                '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_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.
                '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)
        """

        self.quiet = quiet

        self.cache_fp8 = "cache_mode" in kwargs and kwargs["cache_mode"] == "FP8"
        self.gpu_split = kwargs.get("gpu_split", None)
        self.gpu_split_auto = kwargs.get("gpu_split_auto", True)

        self.config = ExLlamaV2Config()
        self.config.model_dir = str(model_directory.resolve())
        self.config.prepare()

        if "max_seq_len" in kwargs: self.config.max_seq_len = kwargs["max_seq_len"]
        if "rope_scale" in kwargs: self.config.scale_pos_emb = kwargs["rope_scale"]
        if "rope_alpha" in kwargs: self.config.scale_alpha_value = kwargs["rope_alpha"]
        if "no_flash_attn" in kwargs: self.config.no_flash_attn = kwargs["no_flash_attn"]

        if "low_mem" in kwargs and kwargs["low_mem"]:
            self.config.set_low_mem()

        chunk_size = min(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_config =  kwargs.get("draft") or {}
        draft_model_name = draft_config.get("draft_model_name")
        enable_draft = bool(draft_config) and draft_model_name is not None

        if bool(draft_config) 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.")
            self.draft_enabled = False
        else:
            self.draft_enabled = enable_draft

        if self.draft_enabled:

            self.draft_config = ExLlamaV2Config()
            draft_model_path = pathlib.Path(kwargs.get("draft_model_dir") or "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.max_seq_len = self.config.max_seq_len

            if "draft_rope_alpha" in kwargs:
                self.draft_config.scale_alpha_value = kwargs.get("draft_rope_alpha") or 1
            else:
                ratio = self.config.max_seq_len / self.draft_config.max_seq_len
                alpha = -0.13436 + 0.80541 * ratio + 0.28833 * ratio ** 2
                self.draft_config.scale_alpha_value = alpha

            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 get_model_path(self):
        model_path = pathlib.Path(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_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 self.draft_enabled:

            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):
        """
        Free all VRAM resources used by this model
        """

        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 = kwargs.get("add_bos_token", True),
                encode_special_tokens = kwargs.get("encode_special_tokens", True)
            )
        if ids:
            # Assume token decoding
            ids = torch.tensor([ids])
            return self.tokenizer.decode(ids, decode_special_tokens = kwargs.get("decode_special_tokens", True))[0]


    def generate(self, prompt: str, **kwargs):
        gen = self.generate_gen(prompt, **kwargs)
        reponse = "".join(gen)
        return reponse

    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)
                '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)
                '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 = kwargs.get("token_healing", False)
        max_tokens = kwargs.get("max_tokens", 150)
        stream_interval = kwargs.get("stream_interval", 0)
        generate_window = min(kwargs.get("generate_window", 512), max_tokens)

        # Sampler settings

        gen_settings = ExLlamaV2Sampler.Settings()
        gen_settings.temperature = kwargs.get("temperature", 1.0)
        gen_settings.top_k = kwargs.get("top_k", 1)
        gen_settings.top_p = kwargs.get("top_p", 1.0)
        gen_settings.min_p = kwargs.get("min_p", 0.0)
        gen_settings.tfs = kwargs.get("tfs", 0.0)
        gen_settings.typical = kwargs.get("typical", 0.0)
        gen_settings.mirostat = kwargs.get("mirostat", False)

        # Default tau and eta fallbacks don't matter if mirostat is off
        gen_settings.mirostat_tau = kwargs.get("mirostat_tau", 1.5)
        gen_settings.mirostat_eta = kwargs.get("mirostat_eta", 0.1)
        gen_settings.token_repetition_penalty = kwargs.get("repetition_penalty", 1.0)
        gen_settings.token_repetition_range = kwargs.get("repetition_range", self.config.max_seq_len)
        gen_settings.token_repetition_decay = kwargs.get("repetition_decay", gen_settings.token_repetition_range)

        stop_conditions: List[Union[str, int]] = kwargs.get("stop", [])
        ban_eos_token = kwargs.get("ban_eos_token", False)

        # Ban the EOS token if specified. If not, append to stop conditions as well.

        if ban_eos_token:
            gen_settings.disallow_tokens(self.tokenizer, [self.tokenizer.eos_token_id])
        else:
            stop_conditions.append(self.tokenizer.eos_token_id)

        # 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

        # Stop conditions           

        self.generator.set_stop_conditions(stop_conditions)

        # Tokenized context

        ids = self.tokenizer.encode(
            prompt,
            add_bos=kwargs.get("add_bos_token", True),
            encode_special_tokens = True
        )

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

            # 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
                full_response += chunk_buffer
                chunk_buffer = ""
                last_chunk_time = now

            if eos or generated_tokens == max_tokens: break

        elapsed_time = last_chunk_time - start_time
        print(f"Response generated in {round(elapsed_time, 2)} seconds ({round(generated_tokens / elapsed_time, 2)} T/s)")