Upload 94 files

modules/AutoGPTQ_loader.py

@@ -0,0 +1,74 @@
+from pathlib import Path
+
+from accelerate.utils import is_xpu_available
+from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
+
+import modules.shared as shared
+from modules.logging_colors import logger
+from modules.models import get_max_memory_dict
+
+
+def load_quantized(model_name):
+    path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
+    pt_path = None
+
+    # Find the model checkpoint
+    if shared.args.checkpoint:
+        pt_path = Path(shared.args.checkpoint)
+    else:
+        for ext in ['.safetensors', '.pt', '.bin']:
+            found = list(path_to_model.glob(f"*{ext}"))
+            if len(found) > 0:
+                if len(found) > 1:
+                    logger.warning(f'More than one {ext} model has been found. The last one will be selected. It could be wrong.')
+
+                pt_path = found[-1]
+                break
+
+    if pt_path is None:
+        logger.error("The model could not be loaded because its checkpoint file in .bin/.pt/.safetensors format could not be located.")
+        return
+
+    use_safetensors = pt_path.suffix == '.safetensors'
+    if not (path_to_model / "quantize_config.json").exists():
+        quantize_config = BaseQuantizeConfig(
+            bits=bits if (bits := shared.args.wbits) > 0 else 4,
+            group_size=gs if (gs := shared.args.groupsize) > 0 else -1,
+            desc_act=shared.args.desc_act
+        )
+    else:
+        quantize_config = None
+
+    # Define the params for AutoGPTQForCausalLM.from_quantized
+    params = {
+        'model_basename': pt_path.stem,
+        'device': "xpu:0" if is_xpu_available() else "cuda:0" if not shared.args.cpu else "cpu",
+        'use_triton': shared.args.triton,
+        'inject_fused_attention': not shared.args.no_inject_fused_attention,
+        'inject_fused_mlp': not shared.args.no_inject_fused_mlp,
+        'use_safetensors': use_safetensors,
+        'trust_remote_code': shared.args.trust_remote_code,
+        'max_memory': get_max_memory_dict(),
+        'quantize_config': quantize_config,
+        'use_cuda_fp16': not shared.args.no_use_cuda_fp16,
+        'disable_exllama': shared.args.disable_exllama,
+        'disable_exllamav2': shared.args.disable_exllamav2,
+    }
+
+    logger.info(f"The AutoGPTQ params are: {params}")
+    model = AutoGPTQForCausalLM.from_quantized(path_to_model, **params)
+
+    # These lines fix the multimodal extension when used with AutoGPTQ
+    if hasattr(model, 'model'):
+        if not hasattr(model, 'dtype'):
+            if hasattr(model.model, 'dtype'):
+                model.dtype = model.model.dtype
+
+        if hasattr(model.model, 'model') and hasattr(model.model.model, 'embed_tokens'):
+            if not hasattr(model, 'embed_tokens'):
+                model.embed_tokens = model.model.model.embed_tokens
+
+            if not hasattr(model.model, 'embed_tokens'):
+                model.model.embed_tokens = model.model.model.embed_tokens
+
+    return model

modules/GPTQ_loader.py

@@ -0,0 +1,171 @@
+import inspect
+import re
+from pathlib import Path
+
+import accelerate
+import torch
+import transformers
+from accelerate.utils import is_xpu_available
+from gptq_for_llama import llama_inference_offload
+from gptq_for_llama.modelutils import find_layers
+from gptq_for_llama.quant import make_quant
+from transformers import AutoConfig, AutoModelForCausalLM
+
+import modules.shared as shared
+from modules.logging_colors import logger
+
+
+# This function is a replacement for the load_quant function in the
+# GPTQ-for_LLaMa repository. It supports more models and branches.
+def _load_quant(model, checkpoint, wbits, groupsize=-1, faster_kernel=False, exclude_layers=None, kernel_switch_threshold=128, eval=True):
+    exclude_layers = exclude_layers or ['lm_head']
+
+    def noop(*args, **kwargs):
+        pass
+
+    config = AutoConfig.from_pretrained(model, trust_remote_code=shared.args.trust_remote_code)
+    torch.nn.init.kaiming_uniform_ = noop
+    torch.nn.init.uniform_ = noop
+    torch.nn.init.normal_ = noop
+
+    torch.set_default_dtype(torch.half)
+    transformers.modeling_utils._init_weights = False
+    torch.set_default_dtype(torch.half)
+    model = AutoModelForCausalLM.from_config(config, trust_remote_code=shared.args.trust_remote_code)
+    torch.set_default_dtype(torch.float)
+    if eval:
+        model = model.eval()
+
+    layers = find_layers(model)
+    for name in exclude_layers:
+        if name in layers:
+            del layers[name]
+
+    gptq_args = inspect.getfullargspec(make_quant).args
+
+    make_quant_kwargs = {
+        'module': model,
+        'names': layers,
+        'bits': wbits,
+    }
+    if 'groupsize' in gptq_args:
+        make_quant_kwargs['groupsize'] = groupsize
+    if 'faster' in gptq_args:
+        make_quant_kwargs['faster'] = faster_kernel
+    if 'kernel_switch_threshold' in gptq_args:
+        make_quant_kwargs['kernel_switch_threshold'] = kernel_switch_threshold
+
+    make_quant(**make_quant_kwargs)
+
+    del layers
+    if checkpoint.endswith('.safetensors'):
+        from safetensors.torch import load_file as safe_load
+        model.load_state_dict(safe_load(checkpoint), strict=False)
+    else:
+        model.load_state_dict(torch.load(checkpoint, weights_only=True), strict=False)
+
+    model.seqlen = 2048
+    return model
+
+
+# Used to locate the .pt/.safetensors quantized file
+def find_quantized_model_file(model_name):
+    if shared.args.checkpoint:
+        return Path(shared.args.checkpoint)
+
+    path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
+    pt_path = None
+    priority_name_list = [
+        Path(f'{shared.args.model_dir}/{model_name}{hyphen}{shared.args.wbits}bit{group}{ext}')
+        for group in ([f'-{shared.args.groupsize}g', ''] if shared.args.groupsize > 0 else [''])
+        for ext in ['.safetensors', '.pt']
+        for hyphen in ['-', f'/{model_name}-', '/']
+    ]
+
+    for path in priority_name_list:
+        if path.exists():
+            pt_path = path
+            break
+
+    # If the model hasn't been found with a well-behaved name, pick the last .pt
+    # or the last .safetensors found in its folder as a last resort
+    if not pt_path:
+        for ext in ['.pt', '.safetensors']:
+            found = list(path_to_model.glob(f"*{ext}"))
+            if len(found) > 0:
+                if len(found) > 1:
+                    logger.warning(f'More than one {ext} model has been found. The last one will be selected. It could be wrong.')
+
+                pt_path = found[-1]
+                break
+
+    return pt_path
+
+
+# The function that loads the model in modules/models.py
+def load_quantized(model_name):
+    if shared.args.model_type is None:
+        logger.error("The model could not be loaded because its type could not be inferred from its name.")
+        logger.error("Please specify the type manually using the --model_type argument.")
+        return None
+
+    # Select the appropriate load_quant function
+    model_type = shared.args.model_type.lower()
+    if shared.args.pre_layer and model_type == 'llama':
+        load_quant = llama_inference_offload.load_quant
+    elif model_type in ('llama', 'opt', 'gptj'):
+        if shared.args.pre_layer:
+            logger.warning("Ignoring --pre_layer because it only works for llama model type.")
+
+        load_quant = _load_quant
+    else:
+        logger.error("Unknown pre-quantized model type specified. Only 'llama', 'opt' and 'gptj' are supported")
+        exit()
+
+    # Find the quantized model weights file (.pt/.safetensors)
+    path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
+    pt_path = find_quantized_model_file(model_name)
+    if not pt_path:
+        logger.error("Could not find the quantized model in .pt or .safetensors format. Exiting.")
+        exit()
+    else:
+        logger.info(f"Found the following quantized model: {pt_path}")
+
+    # qwopqwop200's offload
+    if model_type == 'llama' and shared.args.pre_layer:
+        if len(shared.args.pre_layer) == 1:
+            pre_layer = shared.args.pre_layer[0]
+        else:
+            pre_layer = shared.args.pre_layer
+
+        model = load_quant(str(path_to_model), str(pt_path), shared.args.wbits, shared.args.groupsize, pre_layer)
+    else:
+        threshold = False if model_type == 'gptj' else 128
+        model = load_quant(str(path_to_model), str(pt_path), shared.args.wbits, shared.args.groupsize, kernel_switch_threshold=threshold)
+
+        # accelerate offload (doesn't work properly)
+        if shared.args.gpu_memory or torch.cuda.device_count() > 1 or (is_xpu_available() and torch.xpu.device_count() > 1):
+            if shared.args.gpu_memory:
+                memory_map = list(map(lambda x: x.strip(), shared.args.gpu_memory))
+                max_cpu_memory = shared.args.cpu_memory.strip() if shared.args.cpu_memory is not None else '99GiB'
+                max_memory = {}
+                for i in range(len(memory_map)):
+                    max_memory[i] = f'{memory_map[i]}GiB' if not re.match('.*ib$', memory_map[i].lower()) else memory_map[i]
+
+                max_memory['cpu'] = f'{max_cpu_memory}GiB' if not re.match('.*ib$', max_cpu_memory.lower()) else max_cpu_memory
+            else:
+                max_memory = accelerate.utils.get_balanced_memory(model)
+
+            device_map = accelerate.infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=["LlamaDecoderLayer"])
+            logger.info("Using the following device map for the quantized model:", device_map)
+            # https://huggingface.co/docs/accelerate/package_reference/big_modeling#accelerate.dispatch_model
+            model = accelerate.dispatch_model(model, device_map=device_map, offload_buffers=True)
+
+        # No offload
+        elif not shared.args.cpu:
+            if is_xpu_available():
+                model = model.to(torch.device("xpu:0"))
+            else:
+                model = model.to(torch.device('cuda:0'))
+
+    return model

modules/LoRA.py

@@ -0,0 +1,153 @@
+from pathlib import Path
+
+import torch
+from peft import PeftModel
+from transformers import is_torch_xpu_available
+
+import modules.shared as shared
+from modules.logging_colors import logger
+from modules.models import reload_model
+
+
+def add_lora_to_model(lora_names):
+    if 'GPTQForCausalLM' in shared.model.__class__.__name__ or shared.args.loader == 'AutoGPTQ':
+        add_lora_autogptq(lora_names)
+    elif shared.model.__class__.__name__ in ['Exllamav2Model', 'Exllamav2HF'] or shared.args.loader in ['ExLlamav2', 'ExLlamav2_HF']:
+        add_lora_exllamav2(lora_names)
+    else:
+        add_lora_transformers(lora_names)
+
+
+def get_lora_path(lora_name):
+    p = Path(lora_name)
+    if p.exists():
+        lora_name = p.parts[-1]
+
+    return Path(f"{shared.args.lora_dir}/{lora_name}")
+
+
+def add_lora_exllamav2(lora_names):
+
+    from exllamav2 import ExLlamaV2Lora
+
+    if isinstance(shared.model.loras, list):
+        for lora in shared.model.loras:
+            lora.unload()
+
+    if len(lora_names) > 0:
+        logger.info("Applying the following LoRAs to {}: {}".format(shared.model_name, ', '.join(lora_names)))
+        shared.model.loras = []
+        for lora_name in lora_names:
+            lora_path = get_lora_path(lora_name)
+            if shared.model.__class__.__name__ == 'Exllamav2Model':
+                lora = ExLlamaV2Lora.from_directory(shared.model.model, str(lora_path))
+            else:
+                lora = ExLlamaV2Lora.from_directory(shared.model.ex_model, str(lora_path))
+
+            shared.model.loras.append(lora)
+
+        shared.lora_names = lora_names
+    else:
+        shared.lora_names = []
+        shared.model.loras = None
+
+
+def add_lora_autogptq(lora_names):
+    '''
+    Adapted from https://github.com/Ph0rk0z/text-generation-webui-testing
+    '''
+
+    try:
+        from auto_gptq import get_gptq_peft_model
+        from auto_gptq.utils.peft_utils import GPTQLoraConfig
+    except:
+        logger.error("This version of AutoGPTQ does not support LoRA. You need to install from source or wait for a new release.")
+        return
+
+    if len(lora_names) == 0:
+        reload_model()
+
+        shared.lora_names = []
+        return
+    else:
+        if len(lora_names) > 1:
+            logger.warning('AutoGPTQ can only work with 1 LoRA at the moment. Only the first one in the list will be loaded.')
+        if not shared.args.no_inject_fused_attention:
+            logger.warning('Fused Atttention + AutoGPTQ may break Lora loading. Disable it.')
+
+        peft_config = GPTQLoraConfig(
+            inference_mode=True,
+        )
+
+        lora_path = get_lora_path(lora_names[0])
+        logger.info("Applying the following LoRAs to {}: {}".format(shared.model_name, ', '.join([lora_names[0]])))
+        shared.model = get_gptq_peft_model(shared.model, peft_config, lora_path)
+        shared.lora_names = [lora_names[0]]
+        return
+
+
+def add_lora_transformers(lora_names):
+    prior_set = set(shared.lora_names)
+    added_set = set(lora_names) - prior_set
+    removed_set = prior_set - set(lora_names)
+
+    # If no LoRA needs to be added or removed, exit
+    if len(added_set) == 0 and len(removed_set) == 0:
+        return
+
+    # Add a LoRA when another LoRA is already present
+    if len(removed_set) == 0 and len(prior_set) > 0 and "__merged" not in shared.model.peft_config.keys():
+        logger.info(f"Adding the LoRA(s) named {added_set} to the model")
+        for lora in added_set:
+            shared.model.load_adapter(get_lora_path(lora), lora)
+
+        if len(lora_names) > 1:
+            merge_loras()
+
+        shared.lora_names = lora_names
+        return
+
+    # If any LoRA needs to be removed, start over
+    if len(removed_set) > 0:
+        shared.model = shared.model.unload()
+
+    if len(lora_names) > 0:
+        params = {}
+        if not shared.args.cpu:
+            if shared.args.load_in_4bit or shared.args.load_in_8bit:
+                params['peft_type'] = shared.model.dtype
+            else:
+                params['dtype'] = shared.model.dtype
+                if hasattr(shared.model, "hf_device_map"):
+                    params['device_map'] = {"base_model.model." + k: v for k, v in shared.model.hf_device_map.items()}
+
+        logger.info("Applying the following LoRAs to {}: {}".format(shared.model_name, ', '.join(lora_names)))
+        shared.model = PeftModel.from_pretrained(shared.model, get_lora_path(lora_names[0]), adapter_name=lora_names[0], **params)
+        for lora in lora_names[1:]:
+            shared.model.load_adapter(get_lora_path(lora), lora)
+
+        if len(lora_names) > 1:
+            merge_loras()
+
+        if not shared.args.load_in_8bit and not shared.args.cpu:
+            shared.model.half()
+            if not hasattr(shared.model, "hf_device_map"):
+                if torch.backends.mps.is_available():
+                    device = torch.device('mps')
+                    shared.model = shared.model.to(device)
+                elif is_torch_xpu_available():
+                    device = torch.device("xpu:0")
+                    shared.model = shared.model.to(device)
+                else:
+                    shared.model = shared.model.cuda()
+
+    shared.lora_names = lora_names
+
+
+def merge_loras():
+    if len(list({shared.model.peft_config[adapter].r for adapter in shared.model.peft_config.keys()})) > 1:
+        logger.warning("The loaded LoRAs cannot be merged, as they have dissimilar ranks. Only the first one will be active.")
+        return
+
+    shared.model.add_weighted_adapter(shared.lora_names, [1] * len(shared.lora_names), "__merged")
+    shared.model.set_adapter("__merged")

modules/RoPE.py

@@ -0,0 +1,18 @@
+def get_alpha_value(alpha, base):
+    '''
+    Gets alpha_value from alpha_value and rope_freq_base
+    '''
+    if base > 0:
+        return (base / 10000.) ** (63 / 64.)
+    else:
+        return alpha
+
+
+def get_rope_freq_base(alpha, base):
+    '''
+    Gets rope_freq_base from alpha_value and rope_freq_base
+    '''
+    if base > 0:
+        return base
+    else:
+        return 10000 * alpha ** (64 / 63.)

modules/block_requests.py

@@ -0,0 +1,47 @@
+import builtins
+import io
+
+import requests
+
+from modules.logging_colors import logger
+
+original_open = open
+original_get = requests.get
+
+
+class RequestBlocker:
+
+    def __enter__(self):
+        requests.get = my_get
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        requests.get = original_get
+
+
+class OpenMonkeyPatch:
+
+    def __enter__(self):
+        builtins.open = my_open
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        builtins.open = original_open
+
+
+def my_get(url, **kwargs):
+    logger.info('Unwanted HTTP request redirected to localhost :)')
+    kwargs.setdefault('allow_redirects', True)
+    return requests.api.request('get', 'http://127.0.0.1/', **kwargs)
+
+
+# Kindly provided by our friend WizardLM-30B
+def my_open(*args, **kwargs):
+    filename = str(args[0])
+    if filename.endswith('index.html'):
+        with original_open(*args, **kwargs) as f:
+            file_contents = f.read()
+
+        file_contents = file_contents.replace(b'\t\t<script\n\t\t\tsrc="https://cdnjs.cloudflare.com/ajax/libs/iframe-resizer/4.3.7/iframeResizer.contentWindow.min.js"\n\t\t\tasync\n\t\t></script>', b'')
+        file_contents = file_contents.replace(b'cdnjs.cloudflare.com', b'127.0.0.1')
+        return io.BytesIO(file_contents)
+    else:
+        return original_open(*args, **kwargs)

modules/callbacks.py

@@ -0,0 +1,103 @@
+import gc
+import traceback
+from queue import Queue
+from threading import Thread
+
+import torch
+import transformers
+from transformers import is_torch_xpu_available
+
+import modules.shared as shared
+
+
+class StopNowException(Exception):
+    pass
+
+
+class _StopEverythingStoppingCriteria(transformers.StoppingCriteria):
+    def __init__(self):
+        transformers.StoppingCriteria.__init__(self)
+
+    def __call__(self, input_ids: torch.LongTensor, _scores: torch.FloatTensor) -> bool:
+        return shared.stop_everything
+
+
+class Stream(transformers.StoppingCriteria):
+    def __init__(self, callback_func=None):
+        self.callback_func = callback_func
+
+    def __call__(self, input_ids, scores) -> bool:
+        if self.callback_func is not None:
+            self.callback_func(input_ids[0])
+
+        return False
+
+
+class Iteratorize:
+
+    """
+    Transforms a function that takes a callback
+    into a lazy iterator (generator).
+
+    Adapted from: https://stackoverflow.com/a/9969000
+    """
+
+    def __init__(self, func, args=None, kwargs=None, callback=None):
+        self.mfunc = func
+        self.c_callback = callback
+        self.q = Queue()
+        self.sentinel = object()
+        self.args = args or []
+        self.kwargs = kwargs or {}
+        self.stop_now = False
+
+        def _callback(val):
+            if self.stop_now or shared.stop_everything:
+                raise StopNowException
+            self.q.put(val)
+
+        def gentask():
+            try:
+                ret = self.mfunc(callback=_callback, *args, **self.kwargs)
+            except StopNowException:
+                pass
+            except:
+                traceback.print_exc()
+                pass
+
+            clear_torch_cache()
+            self.q.put(self.sentinel)
+            if self.c_callback:
+                self.c_callback(ret)
+
+        self.thread = Thread(target=gentask)
+        self.thread.start()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        obj = self.q.get(True, None)
+        if obj is self.sentinel:
+            raise StopIteration
+        else:
+            return obj
+
+    def __del__(self):
+        clear_torch_cache()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop_now = True
+        clear_torch_cache()
+
+
+def clear_torch_cache():
+    gc.collect()
+    if not shared.args.cpu:
+        if is_torch_xpu_available():
+            torch.xpu.empty_cache()
+        else:
+            torch.cuda.empty_cache()

modules/chat.py

@@ -0,0 +1,927 @@
+import base64
+import copy
+import functools
+import html
+import json
+import re
+from datetime import datetime
+from functools import partial
+from pathlib import Path
+
+import gradio as gr
+import yaml
+from jinja2.sandbox import ImmutableSandboxedEnvironment
+from PIL import Image
+
+import modules.shared as shared
+from modules import utils
+from modules.extensions import apply_extensions
+from modules.html_generator import chat_html_wrapper, make_thumbnail
+from modules.logging_colors import logger
+from modules.text_generation import (
+    generate_reply,
+    get_encoded_length,
+    get_max_prompt_length
+)
+from modules.utils import delete_file, get_available_characters, save_file
+
+# Copied from the Transformers library
+jinja_env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True)
+
+
+def str_presenter(dumper, data):
+    """
+    Copied from https://github.com/yaml/pyyaml/issues/240
+    Makes pyyaml output prettier multiline strings.
+    """
+
+    if data.count('\n') > 0:
+        return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
+
+    return dumper.represent_scalar('tag:yaml.org,2002:str', data)
+
+
+yaml.add_representer(str, str_presenter)
+yaml.representer.SafeRepresenter.add_representer(str, str_presenter)
+
+
+def get_generation_prompt(renderer, impersonate=False, strip_trailing_spaces=True):
+    '''
+    Given a Jinja template, reverse-engineers the prefix and the suffix for
+    an assistant message (if impersonate=False) or an user message
+    (if impersonate=True)
+    '''
+
+    if impersonate:
+        messages = [
+            {"role": "user", "content": "<<|user-message-1|>>"},
+            {"role": "user", "content": "<<|user-message-2|>>"},
+        ]
+    else:
+        messages = [
+            {"role": "assistant", "content": "<<|user-message-1|>>"},
+            {"role": "assistant", "content": "<<|user-message-2|>>"},
+        ]
+
+    prompt = renderer(messages=messages)
+
+    suffix_plus_prefix = prompt.split("<<|user-message-1|>>")[1].split("<<|user-message-2|>>")[0]
+    suffix = prompt.split("<<|user-message-2|>>")[1]
+    prefix = suffix_plus_prefix[len(suffix):]
+
+    if strip_trailing_spaces:
+        prefix = prefix.rstrip(' ')
+
+    return prefix, suffix
+
+
+def generate_chat_prompt(user_input, state, **kwargs):
+    impersonate = kwargs.get('impersonate', False)
+    _continue = kwargs.get('_continue', False)
+    also_return_rows = kwargs.get('also_return_rows', False)
+    history = kwargs.get('history', state['history'])['internal']
+
+    # Templates
+    chat_template = jinja_env.from_string(state['chat_template_str'])
+    instruction_template = jinja_env.from_string(state['instruction_template_str'])
+    chat_renderer = partial(chat_template.render, add_generation_prompt=False, name1=state['name1'], name2=state['name2'])
+    instruct_renderer = partial(instruction_template.render, add_generation_prompt=False)
+
+    messages = []
+
+    if state['mode'] == 'instruct':
+        renderer = instruct_renderer
+        if state['custom_system_message'].strip() != '':
+            messages.append({"role": "system", "content": state['custom_system_message']})
+    else:
+        renderer = chat_renderer
+        if state['context'].strip() != '':
+            context = replace_character_names(state['context'], state['name1'], state['name2'])
+            messages.append({"role": "system", "content": context})
+
+    insert_pos = len(messages)
+    for user_msg, assistant_msg in reversed(history):
+        user_msg = user_msg.strip()
+        assistant_msg = assistant_msg.strip()
+
+        if assistant_msg:
+            messages.insert(insert_pos, {"role": "assistant", "content": assistant_msg})
+
+        if user_msg not in ['', '<|BEGIN-VISIBLE-CHAT|>']:
+            messages.insert(insert_pos, {"role": "user", "content": user_msg})
+
+    user_input = user_input.strip()
+    if user_input and not impersonate and not _continue:
+        messages.append({"role": "user", "content": user_input})
+
+    def remove_extra_bos(prompt):
+        for bos_token in ['<s>', '<|startoftext|>']:
+            while prompt.startswith(bos_token):
+                prompt = prompt[len(bos_token):]
+
+        return prompt
+
+    def make_prompt(messages):
+        if state['mode'] == 'chat-instruct' and _continue:
+            prompt = renderer(messages=messages[:-1])
+        else:
+            prompt = renderer(messages=messages)
+
+        if state['mode'] == 'chat-instruct':
+            outer_messages = []
+            if state['custom_system_message'].strip() != '':
+                outer_messages.append({"role": "system", "content": state['custom_system_message']})
+
+            prompt = remove_extra_bos(prompt)
+            command = state['chat-instruct_command']
+            command = command.replace('<|character|>', state['name2'] if not impersonate else state['name1'])
+            command = command.replace('<|prompt|>', prompt)
+
+            if _continue:
+                prefix = get_generation_prompt(renderer, impersonate=impersonate, strip_trailing_spaces=False)[0]
+                prefix += messages[-1]["content"]
+            else:
+                prefix = get_generation_prompt(renderer, impersonate=impersonate)[0]
+                if not impersonate:
+                    prefix = apply_extensions('bot_prefix', prefix, state)
+
+            outer_messages.append({"role": "user", "content": command})
+            outer_messages.append({"role": "assistant", "content": prefix})
+
+            prompt = instruction_template.render(messages=outer_messages)
+            suffix = get_generation_prompt(instruct_renderer, impersonate=False)[1]
+            prompt = prompt[:-len(suffix)]
+
+        else:
+            if _continue:
+                suffix = get_generation_prompt(renderer, impersonate=impersonate)[1]
+                prompt = prompt[:-len(suffix)]
+            else:
+                prefix = get_generation_prompt(renderer, impersonate=impersonate)[0]
+                if state['mode'] == 'chat' and not impersonate:
+                    prefix = apply_extensions('bot_prefix', prefix, state)
+
+                prompt += prefix
+
+        prompt = remove_extra_bos(prompt)
+        return prompt
+
+    # Handle truncation
+    max_length = get_max_prompt_length(state)
+    prompt = make_prompt(messages)
+    encoded_length = get_encoded_length(prompt)
+
+    while len(messages) > 0 and encoded_length > max_length:
+
+        # Remove old message, save system message
+        if len(messages) > 2 and messages[0]['role'] == 'system':
+            messages.pop(1)
+
+        # Remove old message when no system message is present
+        elif len(messages) > 1 and messages[0]['role'] != 'system':
+            messages.pop(0)
+
+        # Resort to truncating the user input
+        else:
+
+            user_message = messages[-1]['content']
+
+            # Bisect the truncation point
+            left, right = 0, len(user_message) - 1
+
+            while right - left > 1:
+                mid = (left + right) // 2
+
+                messages[-1]['content'] = user_message[mid:]
+                prompt = make_prompt(messages)
+                encoded_length = get_encoded_length(prompt)
+
+                if encoded_length <= max_length:
+                    right = mid
+                else:
+                    left = mid
+
+            messages[-1]['content'] = user_message[right:]
+            prompt = make_prompt(messages)
+            encoded_length = get_encoded_length(prompt)
+            if encoded_length > max_length:
+                logger.error(f"Failed to build the chat prompt. The input is too long for the available context length.\n\nTruncation length: {state['truncation_length']}\nmax_new_tokens: {state['max_new_tokens']} (is it too high?)\nAvailable context length: {max_length}\n")
+                raise ValueError
+            else:
+                logger.warning(f"The input has been truncated. Context length: {state['truncation_length']}, max_new_tokens: {state['max_new_tokens']}, available context length: {max_length}.")
+                break
+
+        prompt = make_prompt(messages)
+        encoded_length = get_encoded_length(prompt)
+
+    if also_return_rows:
+        return prompt, [message['content'] for message in messages]
+    else:
+        return prompt
+
+
+def get_stopping_strings(state):
+    stopping_strings = []
+    renderers = []
+
+    if state['mode'] in ['instruct', 'chat-instruct']:
+        template = jinja_env.from_string(state['instruction_template_str'])
+        renderer = partial(template.render, add_generation_prompt=False)
+        renderers.append(renderer)
+
+    if state['mode'] in ['chat', 'chat-instruct']:
+        template = jinja_env.from_string(state['chat_template_str'])
+        renderer = partial(template.render, add_generation_prompt=False, name1=state['name1'], name2=state['name2'])
+        renderers.append(renderer)
+
+    for renderer in renderers:
+        prefix_bot, suffix_bot = get_generation_prompt(renderer, impersonate=False)
+        prefix_user, suffix_user = get_generation_prompt(renderer, impersonate=True)
+
+        stopping_strings += [
+            suffix_user + prefix_bot,
+            suffix_user + prefix_user,
+            suffix_bot + prefix_bot,
+            suffix_bot + prefix_user,
+        ]
+
+    if 'stopping_strings' in state and isinstance(state['stopping_strings'], list):
+        stopping_strings += state.pop('stopping_strings')
+
+    return list(set(stopping_strings))
+
+
+def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_message=True, for_ui=False):
+    history = state['history']
+    output = copy.deepcopy(history)
+    output = apply_extensions('history', output)
+    state = apply_extensions('state', state)
+
+    visible_text = None
+    stopping_strings = get_stopping_strings(state)
+    is_stream = state['stream']
+
+    # Prepare the input
+    if not (regenerate or _continue):
+        visible_text = html.escape(text)
+
+        # Apply extensions
+        text, visible_text = apply_extensions('chat_input', text, visible_text, state)
+        text = apply_extensions('input', text, state, is_chat=True)
+
+        output['internal'].append([text, ''])
+        output['visible'].append([visible_text, ''])
+
+        # *Is typing...*
+        if loading_message:
+            yield {
+                'visible': output['visible'][:-1] + [[output['visible'][-1][0], shared.processing_message]],
+                'internal': output['internal']
+            }
+    else:
+        text, visible_text = output['internal'][-1][0], output['visible'][-1][0]
+        if regenerate:
+            if loading_message:
+                yield {
+                    'visible': output['visible'][:-1] + [[visible_text, shared.processing_message]],
+                    'internal': output['internal'][:-1] + [[text, '']]
+                }
+        elif _continue:
+            last_reply = [output['internal'][-1][1], output['visible'][-1][1]]
+            if loading_message:
+                yield {
+                    'visible': output['visible'][:-1] + [[visible_text, last_reply[1] + '...']],
+                    'internal': output['internal']
+                }
+
+    if shared.model_name == 'None' or shared.model is None:
+        raise ValueError("No model is loaded! Select one in the Model tab.")
+
+    # Generate the prompt
+    kwargs = {
+        '_continue': _continue,
+        'history': output if _continue else {k: v[:-1] for k, v in output.items()}
+    }
+    prompt = apply_extensions('custom_generate_chat_prompt', text, state, **kwargs)
+    if prompt is None:
+        prompt = generate_chat_prompt(text, state, **kwargs)
+
+    # Generate
+    reply = None
+    for j, reply in enumerate(generate_reply(prompt, state, stopping_strings=stopping_strings, is_chat=True, for_ui=for_ui)):
+
+        # Extract the reply
+        visible_reply = reply
+        if state['mode'] in ['chat', 'chat-instruct']:
+            visible_reply = re.sub("(<USER>|<user>|{{user}})", state['name1'], reply)
+
+        visible_reply = html.escape(visible_reply)
+
+        if shared.stop_everything:
+            output['visible'][-1][1] = apply_extensions('output', output['visible'][-1][1], state, is_chat=True)
+            yield output
+            return
+
+        if _continue:
+            output['internal'][-1] = [text, last_reply[0] + reply]
+            output['visible'][-1] = [visible_text, last_reply[1] + visible_reply]
+            if is_stream:
+                yield output
+        elif not (j == 0 and visible_reply.strip() == ''):
+            output['internal'][-1] = [text, reply.lstrip(' ')]
+            output['visible'][-1] = [visible_text, visible_reply.lstrip(' ')]
+            if is_stream:
+                yield output
+
+    output['visible'][-1][1] = apply_extensions('output', output['visible'][-1][1], state, is_chat=True)
+    yield output
+
+
+def impersonate_wrapper(text, state):
+
+    static_output = chat_html_wrapper(state['history'], state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu'])
+
+    if shared.model_name == 'None' or shared.model is None:
+        logger.error("No model is loaded! Select one in the Model tab.")
+        yield '', static_output
+        return
+
+    prompt = generate_chat_prompt('', state, impersonate=True)
+    stopping_strings = get_stopping_strings(state)
+
+    yield text + '...', static_output
+    reply = None
+    for reply in generate_reply(prompt + text, state, stopping_strings=stopping_strings, is_chat=True):
+        yield (text + reply).lstrip(' '), static_output
+        if shared.stop_everything:
+            return
+
+
+def generate_chat_reply(text, state, regenerate=False, _continue=False, loading_message=True, for_ui=False):
+    history = state['history']
+    if regenerate or _continue:
+        text = ''
+        if (len(history['visible']) == 1 and not history['visible'][0][0]) or len(history['internal']) == 0:
+            yield history
+            return
+
+    for history in chatbot_wrapper(text, state, regenerate=regenerate, _continue=_continue, loading_message=loading_message, for_ui=for_ui):
+        yield history
+
+
+def character_is_loaded(state, raise_exception=False):
+    if state['mode'] in ['chat', 'chat-instruct'] and state['name2'] == '':
+        logger.error('It looks like no character is loaded. Please load one under Parameters > Character.')
+        if raise_exception:
+            raise ValueError
+
+        return False
+    else:
+        return True
+
+
+def generate_chat_reply_wrapper(text, state, regenerate=False, _continue=False):
+    '''
+    Same as above but returns HTML for the UI
+    '''
+
+    if not character_is_loaded(state):
+        return
+
+    if state['start_with'] != '' and not _continue:
+        if regenerate:
+            text, state['history'] = remove_last_message(state['history'])
+            regenerate = False
+
+        _continue = True
+        send_dummy_message(text, state)
+        send_dummy_reply(state['start_with'], state)
+
+    for i, history in enumerate(generate_chat_reply(text, state, regenerate, _continue, loading_message=True, for_ui=True)):
+        yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu']), history
+
+
+def remove_last_message(history):
+    if len(history['visible']) > 0 and history['internal'][-1][0] != '<|BEGIN-VISIBLE-CHAT|>':
+        last = history['visible'].pop()
+        history['internal'].pop()
+    else:
+        last = ['', '']
+
+    return html.unescape(last[0]), history
+
+
+def send_last_reply_to_input(history):
+    if len(history['visible']) > 0:
+        return html.unescape(history['visible'][-1][1])
+    else:
+        return ''
+
+
+def replace_last_reply(text, state):
+    history = state['history']
+
+    if len(text.strip()) == 0:
+        return history
+    elif len(history['visible']) > 0:
+        history['visible'][-1][1] = html.escape(text)
+        history['internal'][-1][1] = apply_extensions('input', text, state, is_chat=True)
+
+    return history
+
+
+def send_dummy_message(text, state):
+    history = state['history']
+    history['visible'].append([html.escape(text), ''])
+    history['internal'].append([apply_extensions('input', text, state, is_chat=True), ''])
+    return history
+
+
+def send_dummy_reply(text, state):
+    history = state['history']
+    if len(history['visible']) > 0 and not history['visible'][-1][1] == '':
+        history['visible'].append(['', ''])
+        history['internal'].append(['', ''])
+
+    history['visible'][-1][1] = html.escape(text)
+    history['internal'][-1][1] = apply_extensions('input', text, state, is_chat=True)
+    return history
+
+
+def redraw_html(history, name1, name2, mode, style, character, reset_cache=False):
+    return chat_html_wrapper(history, name1, name2, mode, style, character, reset_cache=reset_cache)
+
+
+def start_new_chat(state):
+    mode = state['mode']
+    history = {'internal': [], 'visible': []}
+
+    if mode != 'instruct':
+        greeting = replace_character_names(state['greeting'], state['name1'], state['name2'])
+        if greeting != '':
+            history['internal'] += [['<|BEGIN-VISIBLE-CHAT|>', greeting]]
+            history['visible'] += [['', apply_extensions('output', greeting, state, is_chat=True)]]
+
+    unique_id = datetime.now().strftime('%Y%m%d-%H-%M-%S')
+    save_history(history, unique_id, state['character_menu'], state['mode'])
+
+    return history
+
+
+def get_history_file_path(unique_id, character, mode):
+    if mode == 'instruct':
+        p = Path(f'logs/instruct/{unique_id}.json')
+    else:
+        p = Path(f'logs/chat/{character}/{unique_id}.json')
+
+    return p
+
+
+def save_history(history, unique_id, character, mode):
+    if shared.args.multi_user:
+        return
+
+    p = get_history_file_path(unique_id, character, mode)
+    if not p.parent.is_dir():
+        p.parent.mkdir(parents=True)
+
+    with open(p, 'w', encoding='utf-8') as f:
+        f.write(json.dumps(history, indent=4))
+
+
+def rename_history(old_id, new_id, character, mode):
+    if shared.args.multi_user:
+        return
+
+    old_p = get_history_file_path(old_id, character, mode)
+    new_p = get_history_file_path(new_id, character, mode)
+    if new_p.parent != old_p.parent:
+        logger.error(f"The following path is not allowed: {new_p}.")
+    elif new_p == old_p:
+        logger.info("The provided path is identical to the old one.")
+    else:
+        logger.info(f"Renaming {old_p} to {new_p}")
+        old_p.rename(new_p)
+
+
+def find_all_histories(state):
+    if shared.args.multi_user:
+        return ['']
+
+    if state['mode'] == 'instruct':
+        paths = Path('logs/instruct').glob('*.json')
+    else:
+        character = state['character_menu']
+
+        # Handle obsolete filenames and paths
+        old_p = Path(f'logs/{character}_persistent.json')
+        new_p = Path(f'logs/persistent_{character}.json')
+        if old_p.exists():
+            logger.warning(f"Renaming {old_p} to {new_p}")
+            old_p.rename(new_p)
+        if new_p.exists():
+            unique_id = datetime.now().strftime('%Y%m%d-%H-%M-%S')
+            p = get_history_file_path(unique_id, character, state['mode'])
+            logger.warning(f"Moving {new_p} to {p}")
+            p.parent.mkdir(exist_ok=True)
+            new_p.rename(p)
+
+        paths = Path(f'logs/chat/{character}').glob('*.json')
+
+    histories = sorted(paths, key=lambda x: x.stat().st_mtime, reverse=True)
+    histories = [path.stem for path in histories]
+
+    return histories
+
+
+def load_latest_history(state):
+    '''
+    Loads the latest history for the given character in chat or chat-instruct
+    mode, or the latest instruct history for instruct mode.
+    '''
+
+    if shared.args.multi_user:
+        return start_new_chat(state)
+
+    histories = find_all_histories(state)
+
+    if len(histories) > 0:
+        history = load_history(histories[0], state['character_menu'], state['mode'])
+    else:
+        history = start_new_chat(state)
+
+    return history
+
+
+def load_history_after_deletion(state, idx):
+    '''
+    Loads the latest history for the given character in chat or chat-instruct
+    mode, or the latest instruct history for instruct mode.
+    '''
+
+    if shared.args.multi_user:
+        return start_new_chat(state)
+
+    histories = find_all_histories(state)
+    idx = min(int(idx), len(histories) - 1)
+    idx = max(0, idx)
+
+    if len(histories) > 0:
+        history = load_history(histories[idx], state['character_menu'], state['mode'])
+    else:
+        history = start_new_chat(state)
+        histories = find_all_histories(state)
+
+    return history, gr.update(choices=histories, value=histories[idx])
+
+
+def update_character_menu_after_deletion(idx):
+    characters = utils.get_available_characters()
+    idx = min(int(idx), len(characters) - 1)
+    idx = max(0, idx)
+    return gr.update(choices=characters, value=characters[idx])
+
+
+def load_history(unique_id, character, mode):
+    p = get_history_file_path(unique_id, character, mode)
+
+    f = json.loads(open(p, 'rb').read())
+    if 'internal' in f and 'visible' in f:
+        history = f
+    else:
+        history = {
+            'internal': f['data'],
+            'visible': f['data_visible']
+        }
+
+    return history
+
+
+def load_history_json(file, history):
+    try:
+        file = file.decode('utf-8')
+        f = json.loads(file)
+        if 'internal' in f and 'visible' in f:
+            history = f
+        else:
+            history = {
+                'internal': f['data'],
+                'visible': f['data_visible']
+            }
+
+        return history
+    except:
+        return history
+
+
+def delete_history(unique_id, character, mode):
+    p = get_history_file_path(unique_id, character, mode)
+    delete_file(p)
+
+
+def replace_character_names(text, name1, name2):
+    text = text.replace('{{user}}', name1).replace('{{char}}', name2)
+    return text.replace('<USER>', name1).replace('<BOT>', name2)
+
+
+def generate_pfp_cache(character):
+    cache_folder = Path(shared.args.disk_cache_dir)
+    if not cache_folder.exists():
+        cache_folder.mkdir()
+
+    for path in [Path(f"characters/{character}.{extension}") for extension in ['png', 'jpg', 'jpeg']]:
+        if path.exists():
+            original_img = Image.open(path)
+            original_img.save(Path(f'{cache_folder}/pfp_character.png'), format='PNG')
+
+            thumb = make_thumbnail(original_img)
+            thumb.save(Path(f'{cache_folder}/pfp_character_thumb.png'), format='PNG')
+
+            return thumb
+
+    return None
+
+
+def load_character(character, name1, name2):
+    context = greeting = ""
+    greeting_field = 'greeting'
+    picture = None
+
+    filepath = None
+    for extension in ["yml", "yaml", "json"]:
+        filepath = Path(f'characters/{character}.{extension}')
+        if filepath.exists():
+            break
+
+    if filepath is None or not filepath.exists():
+        logger.error(f"Could not find the character \"{character}\" inside characters/. No character has been loaded.")
+        raise ValueError
+
+    file_contents = open(filepath, 'r', encoding='utf-8').read()
+    data = json.loads(file_contents) if extension == "json" else yaml.safe_load(file_contents)
+    cache_folder = Path(shared.args.disk_cache_dir)
+
+    for path in [Path(f"{cache_folder}/pfp_character.png"), Path(f"{cache_folder}/pfp_character_thumb.png")]:
+        if path.exists():
+            path.unlink()
+
+    picture = generate_pfp_cache(character)
+
+    # Finding the bot's name
+    for k in ['name', 'bot', '<|bot|>', 'char_name']:
+        if k in data and data[k] != '':
+            name2 = data[k]
+            break
+
+    # Find the user name (if any)
+    for k in ['your_name', 'user', '<|user|>']:
+        if k in data and data[k] != '':
+            name1 = data[k]
+            break
+
+    if 'context' in data:
+        context = data['context'].strip()
+    elif "char_persona" in data:
+        context = build_pygmalion_style_context(data)
+        greeting_field = 'char_greeting'
+
+    greeting = data.get(greeting_field, greeting)
+    return name1, name2, picture, greeting, context
+
+
+def load_instruction_template(template):
+    for filepath in [Path(f'instruction-templates/{template}.yaml'), Path('instruction-templates/Alpaca.yaml')]:
+        if filepath.exists():
+            break
+    else:
+        return ''
+
+    file_contents = open(filepath, 'r', encoding='utf-8').read()
+    data = yaml.safe_load(file_contents)
+    if 'instruction_template' in data:
+        return data['instruction_template']
+    else:
+        return jinja_template_from_old_format(data)
+
+
+@functools.cache
+def load_character_memoized(character, name1, name2):
+    return load_character(character, name1, name2)
+
+
+@functools.cache
+def load_instruction_template_memoized(template):
+    return load_instruction_template(template)
+
+
+def upload_character(file, img, tavern=False):
+    decoded_file = file if isinstance(file, str) else file.decode('utf-8')
+    try:
+        data = json.loads(decoded_file)
+    except:
+        data = yaml.safe_load(decoded_file)
+
+    if 'char_name' in data:
+        name = data['char_name']
+        greeting = data['char_greeting']
+        context = build_pygmalion_style_context(data)
+        yaml_data = generate_character_yaml(name, greeting, context)
+    else:
+        name = data['name']
+        yaml_data = generate_character_yaml(data['name'], data['greeting'], data['context'])
+
+    outfile_name = name
+    i = 1
+    while Path(f'characters/{outfile_name}.yaml').exists():
+        outfile_name = f'{name}_{i:03d}'
+        i += 1
+
+    with open(Path(f'characters/{outfile_name}.yaml'), 'w', encoding='utf-8') as f:
+        f.write(yaml_data)
+
+    if img is not None:
+        img.save(Path(f'characters/{outfile_name}.png'))
+
+    logger.info(f'New character saved to "characters/{outfile_name}.yaml".')
+    return gr.update(value=outfile_name, choices=get_available_characters())
+
+
+def build_pygmalion_style_context(data):
+    context = ""
+    if 'char_persona' in data and data['char_persona'] != '':
+        context += f"{data['char_name']}'s Persona: {data['char_persona']}\n"
+
+    if 'world_scenario' in data and data['world_scenario'] != '':
+        context += f"Scenario: {data['world_scenario']}\n"
+
+    if 'example_dialogue' in data and data['example_dialogue'] != '':
+        context += f"{data['example_dialogue'].strip()}\n"
+
+    context = f"{context.strip()}\n"
+    return context
+
+
+def upload_tavern_character(img, _json):
+    _json = {'char_name': _json['name'], 'char_persona': _json['description'], 'char_greeting': _json['first_mes'], 'example_dialogue': _json['mes_example'], 'world_scenario': _json['scenario']}
+    return upload_character(json.dumps(_json), img, tavern=True)
+
+
+def check_tavern_character(img):
+    if "chara" not in img.info:
+        return "Not a TavernAI card", None, None, gr.update(interactive=False)
+
+    decoded_string = base64.b64decode(img.info['chara']).replace(b'\\r\\n', b'\\n')
+    _json = json.loads(decoded_string)
+    if "data" in _json:
+        _json = _json["data"]
+
+    return _json['name'], _json['description'], _json, gr.update(interactive=True)
+
+
+def upload_your_profile_picture(img):
+    cache_folder = Path(shared.args.disk_cache_dir)
+    if not cache_folder.exists():
+        cache_folder.mkdir()
+
+    if img is None:
+        if Path(f"{cache_folder}/pfp_me.png").exists():
+            Path(f"{cache_folder}/pfp_me.png").unlink()
+    else:
+        img = make_thumbnail(img)
+        img.save(Path(f'{cache_folder}/pfp_me.png'))
+        logger.info(f'Profile picture saved to "{cache_folder}/pfp_me.png"')
+
+
+def generate_character_yaml(name, greeting, context):
+    data = {
+        'name': name,
+        'greeting': greeting,
+        'context': context,
+    }
+
+    data = {k: v for k, v in data.items() if v}  # Strip falsy
+    return yaml.dump(data, sort_keys=False, width=float("inf"))
+
+
+def generate_instruction_template_yaml(instruction_template):
+    data = {
+        'instruction_template': instruction_template
+    }
+
+    return my_yaml_output(data)
+
+
+def save_character(name, greeting, context, picture, filename):
+    if filename == "":
+        logger.error("The filename is empty, so the character will not be saved.")
+        return
+
+    data = generate_character_yaml(name, greeting, context)
+    filepath = Path(f'characters/{filename}.yaml')
+    save_file(filepath, data)
+    path_to_img = Path(f'characters/{filename}.png')
+    if picture is not None:
+        picture.save(path_to_img)
+        logger.info(f'Saved {path_to_img}.')
+
+
+def delete_character(name, instruct=False):
+    for extension in ["yml", "yaml", "json"]:
+        delete_file(Path(f'characters/{name}.{extension}'))
+
+    delete_file(Path(f'characters/{name}.png'))
+
+
+def jinja_template_from_old_format(params, verbose=False):
+    MASTER_TEMPLATE = """
+{%- set ns = namespace(found=false) -%}
+{%- for message in messages -%}
+    {%- if message['role'] == 'system' -%}
+        {%- set ns.found = true -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- if not ns.found -%}
+    {{- '<|PRE-SYSTEM|>' + '<|SYSTEM-MESSAGE|>' + '<|POST-SYSTEM|>' -}}
+{%- endif %}
+{%- for message in messages %}
+    {%- if message['role'] == 'system' -%}
+        {{- '<|PRE-SYSTEM|>' + message['content'] + '<|POST-SYSTEM|>' -}}
+    {%- else -%}
+        {%- if message['role'] == 'user' -%}
+            {{-'<|PRE-USER|>' + message['content'] + '<|POST-USER|>'-}}
+        {%- else -%}
+            {{-'<|PRE-ASSISTANT|>' + message['content'] + '<|POST-ASSISTANT|>' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{-'<|PRE-ASSISTANT-GENERATE|>'-}}
+{%- endif -%}
+"""
+
+    if 'context' in params and '<|system-message|>' in params['context']:
+        pre_system = params['context'].split('<|system-message|>')[0]
+        post_system = params['context'].split('<|system-message|>')[1]
+    else:
+        pre_system = ''
+        post_system = ''
+
+    pre_user = params['turn_template'].split('<|user-message|>')[0].replace('<|user|>', params['user'])
+    post_user = params['turn_template'].split('<|user-message|>')[1].split('<|bot|>')[0]
+
+    pre_assistant = '<|bot|>' + params['turn_template'].split('<|bot-message|>')[0].split('<|bot|>')[1]
+    pre_assistant = pre_assistant.replace('<|bot|>', params['bot'])
+    post_assistant = params['turn_template'].split('<|bot-message|>')[1]
+
+    def preprocess(string):
+        return string.replace('\n', '\\n').replace('\'', '\\\'')
+
+    pre_system = preprocess(pre_system)
+    post_system = preprocess(post_system)
+    pre_user = preprocess(pre_user)
+    post_user = preprocess(post_user)
+    pre_assistant = preprocess(pre_assistant)
+    post_assistant = preprocess(post_assistant)
+
+    if verbose:
+        print(
+            '\n',
+            repr(pre_system) + '\n',
+            repr(post_system) + '\n',
+            repr(pre_user) + '\n',
+            repr(post_user) + '\n',
+            repr(pre_assistant) + '\n',
+            repr(post_assistant) + '\n',
+        )
+
+    result = MASTER_TEMPLATE
+    if 'system_message' in params:
+        result = result.replace('<|SYSTEM-MESSAGE|>', preprocess(params['system_message']))
+    else:
+        result = result.replace('<|SYSTEM-MESSAGE|>', '')
+
+    result = result.replace('<|PRE-SYSTEM|>', pre_system)
+    result = result.replace('<|POST-SYSTEM|>', post_system)
+    result = result.replace('<|PRE-USER|>', pre_user)
+    result = result.replace('<|POST-USER|>', post_user)
+    result = result.replace('<|PRE-ASSISTANT|>', pre_assistant)
+    result = result.replace('<|PRE-ASSISTANT-GENERATE|>', pre_assistant.rstrip(' '))
+    result = result.replace('<|POST-ASSISTANT|>', post_assistant)
+
+    result = result.strip()
+
+    return result
+
+
+def my_yaml_output(data):
+    '''
+    pyyaml is very inconsistent with multiline strings.
+    for simple instruction template outputs, this is enough.
+    '''
+    result = ""
+    for k in data:
+        result += k + ": |-\n"
+        for line in data[k].splitlines():
+            result += "  " + line.rstrip(' ') + "\n"
+
+    return result

modules/ctransformers_model.py

@@ -0,0 +1,79 @@
+from ctransformers import AutoConfig, AutoModelForCausalLM
+
+from modules import shared
+from modules.callbacks import Iteratorize
+from modules.logging_colors import logger
+
+
+class CtransformersModel:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def from_pretrained(cls, path):
+        result = cls()
+
+        config = AutoConfig.from_pretrained(
+            str(path),
+            threads=shared.args.threads if shared.args.threads != 0 else -1,
+            gpu_layers=shared.args.n_gpu_layers,
+            batch_size=shared.args.n_batch,
+            context_length=shared.args.n_ctx,
+            stream=True,
+            mmap=not shared.args.no_mmap,
+            mlock=shared.args.mlock
+        )
+
+        result.model = AutoModelForCausalLM.from_pretrained(
+            str(result.model_dir(path) if result.model_type_is_auto() else path),
+            model_type=(None if result.model_type_is_auto() else shared.args.model_type),
+            config=config
+        )
+
+        logger.info(f'Using ctransformers model_type: {result.model.model_type} for {result.model.model_path}')
+        return result, result
+
+    def model_type_is_auto(self):
+        return shared.args.model_type is None or shared.args.model_type == "Auto" or shared.args.model_type == "None"
+
+    def model_dir(self, path):
+        if path.is_file():
+            return path.parent
+
+        return path
+
+    def encode(self, string, **kwargs):
+        return self.model.tokenize(string)
+
+    def decode(self, ids):
+        return self.model.detokenize(ids)
+
+    def generate(self, prompt, state, callback=None):
+        prompt = prompt if type(prompt) is str else prompt.decode()
+        # ctransformers uses -1 for random seed
+        generator = self.model(
+            prompt=prompt,
+            max_new_tokens=state['max_new_tokens'],
+            temperature=state['temperature'],
+            top_p=state['top_p'],
+            top_k=state['top_k'],
+            repetition_penalty=state['repetition_penalty'],
+            last_n_tokens=state['repetition_penalty_range'],
+            seed=int(state['seed'])
+        )
+
+        output = ""
+        for token in generator:
+            if callback:
+                callback(token)
+
+            output += token
+
+        return output
+
+    def generate_with_streaming(self, *args, **kwargs):
+        with Iteratorize(self.generate, args, kwargs, callback=None) as generator:
+            reply = ''
+            for token in generator:
+                reply += token
+                yield reply

modules/deepspeed_parameters.py

@@ -0,0 +1,74 @@
+def generate_ds_config(ds_bf16, train_batch_size, nvme_offload_dir):
+    '''
+    DeepSpeed configuration
+    https://huggingface.co/docs/transformers/main_classes/deepspeed
+    '''
+
+    if nvme_offload_dir:
+        ds_config = {
+            "fp16": {
+                "enabled": not ds_bf16,
+            },
+            "bf16": {
+                "enabled": ds_bf16,
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "offload_param": {
+                    "device": "nvme",
+                    "nvme_path": nvme_offload_dir,
+                    "pin_memory": True,
+                    "buffer_count": 5,
+                    "buffer_size": 1e9,
+                    "max_in_cpu": 1e9
+                },
+                "overlap_comm": True,
+                "reduce_bucket_size": "auto",
+                "contiguous_gradients": True,
+                "sub_group_size": 1e8,
+                "stage3_prefetch_bucket_size": "auto",
+                "stage3_param_persistence_threshold": "auto",
+                "stage3_max_live_parameters": "auto",
+                "stage3_max_reuse_distance": "auto",
+            },
+            "aio": {
+                "block_size": 262144,
+                "queue_depth": 32,
+                "thread_count": 1,
+                "single_submit": False,
+                "overlap_events": True
+            },
+            "steps_per_print": 2000,
+            "train_batch_size": train_batch_size,
+            "train_micro_batch_size_per_gpu": 1,
+            "wall_clock_breakdown": False
+        }
+    else:
+        ds_config = {
+            "fp16": {
+                "enabled": not ds_bf16,
+            },
+            "bf16": {
+                "enabled": ds_bf16,
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "offload_param": {
+                    "device": "cpu",
+                    "pin_memory": True
+                },
+                "overlap_comm": True,
+                "contiguous_gradients": True,
+                "reduce_bucket_size": "auto",
+                "stage3_prefetch_bucket_size": "auto",
+                "stage3_param_persistence_threshold": "auto",
+                "stage3_max_live_parameters": "auto",
+                "stage3_max_reuse_distance": "auto",
+            },
+            "steps_per_print": 2000,
+            "train_batch_size": train_batch_size,
+            "train_micro_batch_size_per_gpu": 1,
+            "wall_clock_breakdown": False
+        }
+
+    return ds_config

modules/evaluate.py

@@ -0,0 +1,153 @@
+import datetime
+from pathlib import Path
+
+import pandas as pd
+import torch
+from datasets import load_dataset
+from tqdm import tqdm
+
+from modules import shared
+from modules.logging_colors import logger
+from modules.models import clear_torch_cache, load_model, unload_model
+from modules.models_settings import get_model_metadata, update_model_parameters
+from modules.text_generation import encode
+
+
+def load_past_evaluations():
+    if Path('logs/evaluations.csv').exists():
+        df = pd.read_csv(Path('logs/evaluations.csv'), dtype=str)
+        df['Perplexity'] = pd.to_numeric(df['Perplexity'])
+        return df
+    else:
+        return pd.DataFrame(columns=['Model', 'LoRAs', 'Dataset', 'Perplexity', 'stride', 'max_length', 'Date', 'Comment'])
+
+
+past_evaluations = load_past_evaluations()
+
+
+def save_past_evaluations(df):
+    global past_evaluations
+    past_evaluations = df
+    filepath = Path('logs/evaluations.csv')
+    filepath.parent.mkdir(parents=True, exist_ok=True)
+    df.to_csv(filepath, index=False)
+
+
+def calculate_perplexity(models, input_dataset, stride, _max_length):
+    '''
+    Based on:
+    https://huggingface.co/docs/transformers/perplexity#calculating-ppl-with-fixedlength-models
+    '''
+
+    if not shared.args.no_use_fast:
+        logger.warning("--no_use_fast is not being used. If tokenizing the input dataset takes a long time, consider loading the model with that option checked.")
+
+    global past_evaluations
+    cumulative_log = ''
+    cumulative_log += "Loading the input dataset...\n\n"
+    yield cumulative_log
+
+    # Copied from https://github.com/qwopqwop200/GPTQ-for-LLaMa/blob/triton/utils/datautils.py
+    if input_dataset == 'wikitext':
+        data = load_dataset('wikitext', 'wikitext-2-raw-v1', split='test')
+        text = "\n\n".join(data['text'])
+    elif input_dataset == 'ptb':
+        data = load_dataset('ptb_text_only', 'penn_treebank', split='validation')
+        text = "\n\n".join(data['sentence'])
+    elif input_dataset == 'ptb_new':
+        data = load_dataset('ptb_text_only', 'penn_treebank', split='test')
+        text = " ".join(data['sentence'])
+    else:
+        with open(Path(f'training/datasets/{input_dataset}.txt'), 'r', encoding='utf-8') as f:
+            text = f.read()
+
+    for model in models:
+        if is_in_past_evaluations(model, input_dataset, stride, _max_length):
+            cumulative_log += f"`{model}` has already been tested. Ignoring.\n\n"
+            yield cumulative_log
+            continue
+
+        if model != 'current model':
+            try:
+                yield cumulative_log + f"Loading `{model}`...\n\n"
+                model_settings = get_model_metadata(model)
+                shared.settings.update({k: v for k, v in model_settings.items() if k in shared.settings})  # hijacking the interface defaults
+                update_model_parameters(model_settings)  # hijacking the command-line arguments
+                unload_model()
+                shared.model, shared.tokenizer = load_model(model)
+            except:
+                cumulative_log += f"Failed to load `{model}`. Moving on.\n\n"
+                yield cumulative_log
+                continue
+
+        cumulative_log += f"Processing `{shared.model_name}`...\n\n"
+        yield cumulative_log + "Tokenizing the input dataset...\n\n"
+        encodings = encode(text, add_special_tokens=False)
+        seq_len = encodings.shape[1]
+        if _max_length:
+            max_length = _max_length
+        elif hasattr(shared.model.config, 'max_position_embeddings'):
+            max_length = shared.model.config.max_position_embeddings
+        else:
+            max_length = 2048
+
+        nlls = []
+        prev_end_loc = 0
+        for begin_loc in tqdm(range(0, seq_len, stride)):
+            yield cumulative_log + f"Evaluating... {100*begin_loc/seq_len:.2f}%"
+            end_loc = min(begin_loc + max_length, seq_len)
+            trg_len = end_loc - prev_end_loc  # may be different from stride on last loop
+            input_ids = encodings[:, begin_loc:end_loc]
+            target_ids = input_ids.clone()
+            target_ids[:, :-trg_len] = -100
+            clear_torch_cache()
+            with torch.no_grad():
+                outputs = shared.model(input_ids=input_ids, labels=target_ids)
+
+                # loss is calculated using CrossEntropyLoss which averages over valid labels
+                # N.B. the model only calculates loss over trg_len - 1 labels, because it internally shifts the labels
+                # to the left by 1.
+                neg_log_likelihood = outputs.loss
+
+            nlls.append(neg_log_likelihood)
+            prev_end_loc = end_loc
+            if end_loc == seq_len:
+                break
+
+        ppl = torch.exp(torch.stack(nlls).mean())
+        add_entry_to_past_evaluations(float(ppl), shared.model_name, input_dataset, stride, _max_length)
+        save_past_evaluations(past_evaluations)
+        cumulative_log += f"The perplexity for `{shared.model_name}` is: {float(ppl)}\n\n"
+        yield cumulative_log
+
+
+def add_entry_to_past_evaluations(perplexity, model, dataset, stride, max_length):
+    global past_evaluations
+    entry = {
+        'Model': model,
+        'LoRAs': ', '.join(shared.lora_names) or '-',
+        'Dataset': dataset,
+        'Perplexity': perplexity,
+        'stride': str(stride),
+        'max_length': str(max_length),
+        'Date': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
+        'Comment': ''
+    }
+    past_evaluations = pd.concat([past_evaluations, pd.DataFrame([entry])], ignore_index=True)
+
+
+def is_in_past_evaluations(model, dataset, stride, max_length):
+    entries = past_evaluations[(past_evaluations['Model'] == model) &
+                               (past_evaluations['Dataset'] == dataset) &
+                               (past_evaluations['max_length'] == str(max_length)) &
+                               (past_evaluations['stride'] == str(stride))]
+
+    if entries.shape[0] > 0:
+        return True
+    else:
+        return False
+
+
+def generate_markdown_table():
+    sorted_df = past_evaluations.sort_values(by=['Dataset', 'stride', 'Perplexity', 'Date'])
+    return sorted_df

modules/exllamav2.py

@@ -0,0 +1,149 @@
+import traceback
+from pathlib import Path
+
+import torch
+from exllamav2 import (
+    ExLlamaV2,
+    ExLlamaV2Cache,
+    ExLlamaV2Cache_8bit,
+    ExLlamaV2Config,
+    ExLlamaV2Tokenizer
+)
+from exllamav2.generator import ExLlamaV2Sampler, ExLlamaV2StreamingGenerator
+
+from modules import shared
+from modules.logging_colors import logger
+from modules.text_generation import get_max_prompt_length
+
+try:
+    import flash_attn
+except ModuleNotFoundError:
+    logger.warning(
+        'You are running ExLlamaV2 without flash-attention. This will cause the VRAM usage '
+        'to be a lot higher than it could be.\n'
+        'Try installing flash-attention following the instructions here: '
+        'https://github.com/Dao-AILab/flash-attention#installation-and-features'
+    )
+    pass
+except Exception:
+    logger.warning('Failed to load flash-attention due to the following error:\n')
+    traceback.print_exc()
+
+
+class Exllamav2Model:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def from_pretrained(self, path_to_model):
+
+        path_to_model = Path(f'{shared.args.model_dir}') / Path(path_to_model)
+
+        config = ExLlamaV2Config()
+        config.model_dir = str(path_to_model)
+        config.prepare()
+
+        config.max_seq_len = shared.args.max_seq_len
+        config.scale_pos_emb = shared.args.compress_pos_emb
+        config.scale_alpha_value = shared.args.alpha_value
+        config.no_flash_attn = shared.args.no_flash_attn
+        config.num_experts_per_token = int(shared.args.num_experts_per_token)
+
+        model = ExLlamaV2(config)
+
+        split = None
+        if shared.args.gpu_split:
+            split = [float(alloc) for alloc in shared.args.gpu_split.split(",")]
+
+        model.load(split)
+
+        tokenizer = ExLlamaV2Tokenizer(config)
+        if shared.args.cache_8bit:
+            cache = ExLlamaV2Cache_8bit(model)
+        else:
+            cache = ExLlamaV2Cache(model)
+
+        generator = ExLlamaV2StreamingGenerator(model, cache, tokenizer)
+
+        result = self()
+        result.model = model
+        result.cache = cache
+        result.tokenizer = tokenizer
+        result.generator = generator
+        result.loras = None
+        return result, result
+
+    def encode(self, string, **kwargs):
+        return self.tokenizer.encode(string, add_bos=True, encode_special_tokens=True)
+
+    def decode(self, ids, **kwargs):
+        if isinstance(ids, list):
+            ids = torch.tensor([ids])
+        elif isinstance(ids, torch.Tensor) and ids.numel() == 1:
+            ids = ids.view(1, -1)
+
+        return self.tokenizer.decode(ids, decode_special_tokens=True)[0]
+
+    def get_logits(self, token_ids, **kwargs):
+        self.cache.current_seq_len = 0
+        if token_ids.shape[-1] > 1:
+            self.model.forward(token_ids[:, :-1], self.cache, input_mask=None, preprocess_only=True, loras=self.loras)
+
+        return self.model.forward(token_ids[:, -1:], self.cache, input_mask=None, loras=self.loras, **kwargs).float().cpu()
+
+    def generate_with_streaming(self, prompt, state):
+        settings = ExLlamaV2Sampler.Settings()
+
+        settings.token_repetition_penalty = state['repetition_penalty']
+        settings.token_repetition_range = -1 if state['repetition_penalty_range'] <= 0 else state['repetition_penalty_range']
+
+        settings.token_frequency_penalty = state['frequency_penalty']
+        settings.token_presence_penalty = state['presence_penalty']
+
+        settings.temperature = state['temperature']
+        settings.top_k = state['top_k']
+        settings.top_p = state['top_p']
+        settings.top_a = state['top_a']
+        settings.min_p = state['min_p']
+        settings.tfs = state['tfs']
+        settings.typical = state['typical_p']
+
+        settings.temperature_last = state['temperature_last']
+
+        settings.mirostat = state['mirostat_mode'] == 2
+        settings.mirostat_tau = state['mirostat_tau']
+        settings.mirostat_eta = state['mirostat_eta']
+
+        if state['ban_eos_token']:
+            settings.disallow_tokens(self.tokenizer, [self.tokenizer.eos_token_id])
+
+        if state['custom_token_bans']:
+            to_ban = [int(x) for x in state['custom_token_bans'].split(',')]
+            if len(to_ban) > 0:
+                settings.disallow_tokens(self.tokenizer, to_ban)
+
+        ids = self.tokenizer.encode(prompt, add_bos=state['add_bos_token'], encode_special_tokens=True)
+        ids = ids[:, -get_max_prompt_length(state):]
+
+        if state['auto_max_new_tokens']:
+            max_new_tokens = state['truncation_length'] - ids.shape[-1]
+        else:
+            max_new_tokens = state['max_new_tokens']
+
+        self.generator.begin_stream(ids, settings, loras=self.loras)
+
+        decoded_text = ''
+        for i in range(max_new_tokens):
+            chunk, eos, _ = self.generator.stream()
+            if eos or shared.stop_everything:
+                break
+
+            decoded_text += chunk
+            yield decoded_text
+
+    def generate(self, prompt, state):
+        output = ''
+        for output in self.generate_with_streaming(prompt, state):
+            pass
+
+        return output

modules/exllamav2_hf.py

@@ -0,0 +1,170 @@
+import os
+import traceback
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+import torch
+from exllamav2 import (
+    ExLlamaV2,
+    ExLlamaV2Cache,
+    ExLlamaV2Cache_8bit,
+    ExLlamaV2Config
+)
+from torch.nn import CrossEntropyLoss
+from transformers import GenerationConfig, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from modules import shared
+from modules.logging_colors import logger
+
+try:
+    import flash_attn
+except ModuleNotFoundError:
+    logger.warning(
+        'You are running ExLlamaV2 without flash-attention. This will cause the VRAM usage '
+        'to be a lot higher than it could be.\n'
+        'Try installing flash-attention following the instructions here: '
+        'https://github.com/Dao-AILab/flash-attention#installation-and-features'
+    )
+    pass
+except Exception:
+    logger.warning('Failed to load flash-attention due to the following error:\n')
+    traceback.print_exc()
+
+
+class Exllamav2HF(PreTrainedModel):
+    def __init__(self, config: ExLlamaV2Config):
+        super().__init__(PretrainedConfig())
+        self.ex_config = config
+        self.ex_model = ExLlamaV2(config)
+        split = None
+        if shared.args.gpu_split:
+            split = [float(alloc) for alloc in shared.args.gpu_split.split(",")]
+
+        self.ex_model.load(split)
+        self.generation_config = GenerationConfig()
+        self.loras = None
+
+        if shared.args.cache_8bit:
+            self.ex_cache = ExLlamaV2Cache_8bit(self.ex_model)
+        else:
+            self.ex_cache = ExLlamaV2Cache(self.ex_model)
+
+        self.past_seq = None
+        if shared.args.cfg_cache:
+            if shared.args.cache_8bit:
+                self.ex_cache_negative = ExLlamaV2Cache_8bit(self.ex_model)
+            else:
+                self.ex_cache_negative = ExLlamaV2Cache(self.ex_model)
+
+            self.past_seq_negative = None
+
+    def _validate_model_class(self):
+        pass
+
+    def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]):
+        pass
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {'input_ids': input_ids, **kwargs}
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device(0)
+
+    def __call__(self, *args, **kwargs):
+        use_cache = kwargs.get('use_cache', True)
+        labels = kwargs.get('labels', None)
+        past_key_values = kwargs.get('past_key_values', None)
+
+        if len(args) > 0:
+            if not shared.args.cfg_cache:
+                logger.error("Please enable the cfg-cache option to use CFG with ExLlamav2_HF.")
+                return
+
+            input_ids = args[0]
+            is_negative = True
+            past_seq = self.past_seq_negative
+            ex_cache = self.ex_cache_negative
+        else:
+            input_ids = kwargs['input_ids']
+            is_negative = False
+            past_seq = self.past_seq
+            ex_cache = self.ex_cache
+
+        seq = input_ids[0].tolist()
+        if is_negative and past_key_values is not None:
+            seq = past_key_values + seq
+
+        seq_tensor = torch.tensor(seq)
+        reset = True
+
+        # Make the forward call
+        if labels is None:
+            if past_seq is not None:
+                min_length = min(past_seq.shape[0], seq_tensor.shape[0])
+                indices = torch.nonzero(~torch.eq(past_seq[:min_length], seq_tensor[:min_length]))
+                if len(indices) > 0:
+                    longest_prefix = indices[0].item()
+                else:
+                    longest_prefix = min_length
+
+                if longest_prefix > 0:
+                    reset = False
+                    ex_cache.current_seq_len = longest_prefix
+                    if len(seq_tensor) - longest_prefix > 1:
+                        self.ex_model.forward(seq_tensor[longest_prefix:-1].view(1, -1), ex_cache, preprocess_only=True, loras=self.loras)
+                    elif len(seq_tensor) == longest_prefix:
+                        # Very tricky: if the prefix we are reusing *is* the input_ids, then we have to back up the cache pointer by one,
+                        # because we feed input_ids[-1] to forward() below, but that last token is already in the cache!
+                        ex_cache.current_seq_len -= 1
+
+            if reset:
+                ex_cache.current_seq_len = 0
+                if len(seq_tensor) > 1:
+                    self.ex_model.forward(seq_tensor[:-1].view(1, -1), ex_cache, preprocess_only=True, loras=self.loras)
+
+            logits = self.ex_model.forward(seq_tensor[-1:].view(1, -1), ex_cache, loras=self.loras).to(input_ids.device).float()
+        else:
+            ex_cache.current_seq_len = 0
+            logits = self.ex_model.forward(seq_tensor.view(1, -1), ex_cache, last_id_only=False, loras=self.loras).float()
+
+        if is_negative:
+            self.past_seq_negative = seq_tensor
+        else:
+            self.past_seq = seq_tensor
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, logits.shape[-1])
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        return CausalLMOutputWithPast(logits=logits, past_key_values=seq if use_cache else None, loss=loss)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):
+        assert len(model_args) == 0 and len(kwargs) == 0, "extra args is currently not supported"
+        if isinstance(pretrained_model_name_or_path, str):
+            pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
+
+        pretrained_model_name_or_path = Path(f'{shared.args.model_dir}') / Path(pretrained_model_name_or_path)
+
+        config = ExLlamaV2Config()
+        config.model_dir = str(pretrained_model_name_or_path)
+        config.prepare()
+
+        config.max_seq_len = shared.args.max_seq_len
+        config.scale_pos_emb = shared.args.compress_pos_emb
+        config.scale_alpha_value = shared.args.alpha_value
+        config.no_flash_attn = shared.args.no_flash_attn
+        config.num_experts_per_token = int(shared.args.num_experts_per_token)
+
+        return Exllamav2HF(config)

modules/extensions.py

@@ -0,0 +1,232 @@
+import traceback
+from functools import partial
+from inspect import signature
+
+import gradio as gr
+
+import extensions
+import modules.shared as shared
+from modules.logging_colors import logger
+
+state = {}
+available_extensions = []
+setup_called = set()
+
+
+def apply_settings(extension, name):
+    if not hasattr(extension, 'params'):
+        return
+
+    for param in extension.params:
+        _id = f"{name}-{param}"
+        shared.default_settings[_id] = extension.params[param]
+        if _id in shared.settings:
+            extension.params[param] = shared.settings[_id]
+
+
+def load_extensions():
+    global state, setup_called
+    state = {}
+    for i, name in enumerate(shared.args.extensions):
+        if name in available_extensions:
+            if name != 'api':
+                logger.info(f'Loading the extension "{name}"')
+            try:
+                try:
+                    exec(f"import extensions.{name}.script")
+                except ModuleNotFoundError:
+                    logger.error(f"Could not import the requirements for '{name}'. Make sure to install the requirements for the extension.\n\nLinux / Mac:\n\npip install -r extensions/{name}/requirements.txt --upgrade\n\nWindows:\n\npip install -r extensions\\{name}\\requirements.txt --upgrade\n\nIf you used the one-click installer, paste the command above in the terminal window opened after launching the cmd script for your OS.")
+                    raise
+
+                extension = getattr(extensions, name).script
+
+                # Only run setup() and apply settings from settings.yaml once
+                if extension not in setup_called:
+                    apply_settings(extension, name)
+                    if hasattr(extension, "setup"):
+                        extension.setup()
+
+                    setup_called.add(extension)
+
+                state[name] = [True, i]
+            except:
+                logger.error(f'Failed to load the extension "{name}".')
+                traceback.print_exc()
+
+
+# This iterator returns the extensions in the order specified in the command-line
+def iterator():
+    for name in sorted(state, key=lambda x: state[x][1]):
+        if state[name][0]:
+            yield getattr(extensions, name).script, name
+
+
+# Extension functions that map string -> string
+def _apply_string_extensions(function_name, text, state, is_chat=False):
+    for extension, _ in iterator():
+        if hasattr(extension, function_name):
+            func = getattr(extension, function_name)
+
+            # Handle old extensions without the 'state' arg or
+            # the 'is_chat' kwarg
+            count = 0
+            has_chat = False
+            for k in signature(func).parameters:
+                if k == 'is_chat':
+                    has_chat = True
+                else:
+                    count += 1
+
+            if count == 2:
+                args = [text, state]
+            else:
+                args = [text]
+
+            if has_chat:
+                kwargs = {'is_chat': is_chat}
+            else:
+                kwargs = {}
+
+            text = func(*args, **kwargs)
+
+    return text
+
+
+# Extension functions that map string -> string
+def _apply_chat_input_extensions(text, visible_text, state):
+    for extension, _ in iterator():
+        if hasattr(extension, 'chat_input_modifier'):
+            text, visible_text = extension.chat_input_modifier(text, visible_text, state)
+
+    return text, visible_text
+
+
+# custom_generate_chat_prompt handling - currently only the first one will work
+def _apply_custom_generate_chat_prompt(text, state, **kwargs):
+    for extension, _ in iterator():
+        if hasattr(extension, 'custom_generate_chat_prompt'):
+            return extension.custom_generate_chat_prompt(text, state, **kwargs)
+
+    return None
+
+
+# Extension that modifies the input parameters before they are used
+def _apply_state_modifier_extensions(state):
+    for extension, _ in iterator():
+        if hasattr(extension, "state_modifier"):
+            state = getattr(extension, "state_modifier")(state)
+
+    return state
+
+
+# Extension that modifies the chat history before it is used
+def _apply_history_modifier_extensions(history):
+    for extension, _ in iterator():
+        if hasattr(extension, "history_modifier"):
+            history = getattr(extension, "history_modifier")(history)
+
+    return history
+
+
+# Extension functions that override the default tokenizer output - The order of execution is not defined
+def _apply_tokenizer_extensions(function_name, state, prompt, input_ids, input_embeds):
+    for extension, _ in iterator():
+        if hasattr(extension, function_name):
+            prompt, input_ids, input_embeds = getattr(extension, function_name)(state, prompt, input_ids, input_embeds)
+
+    return prompt, input_ids, input_embeds
+
+
+# Allow extensions to add their own logits processors to the stack being run.
+# Each extension would call `processor_list.append({their LogitsProcessor}())`.
+def _apply_logits_processor_extensions(function_name, processor_list, input_ids):
+    for extension, _ in iterator():
+        if hasattr(extension, function_name):
+            result = getattr(extension, function_name)(processor_list, input_ids)
+            if type(result) is list:
+                processor_list = result
+
+    return processor_list
+
+
+# Get prompt length in tokens after applying extension functions which override the default tokenizer output
+# currently only the first one will work
+def _apply_custom_tokenized_length(prompt):
+    for extension, _ in iterator():
+        if hasattr(extension, 'custom_tokenized_length'):
+            return getattr(extension, 'custom_tokenized_length')(prompt)
+
+    return None
+
+
+# Custom generate reply handling - currently only the first one will work
+def _apply_custom_generate_reply():
+    for extension, _ in iterator():
+        if hasattr(extension, 'custom_generate_reply'):
+            return getattr(extension, 'custom_generate_reply')
+
+    return None
+
+
+def _apply_custom_css():
+    all_css = ''
+    for extension, _ in iterator():
+        if hasattr(extension, 'custom_css'):
+            all_css += getattr(extension, 'custom_css')()
+
+    return all_css
+
+
+def _apply_custom_js():
+    all_js = ''
+    for extension, _ in iterator():
+        if hasattr(extension, 'custom_js'):
+            all_js += getattr(extension, 'custom_js')()
+
+    return all_js
+
+
+def create_extensions_block():
+    to_display = []
+    for extension, name in iterator():
+        if hasattr(extension, "ui") and not (hasattr(extension, 'params') and extension.params.get('is_tab', False)):
+            to_display.append((extension, name))
+
+    # Creating the extension ui elements
+    if len(to_display) > 0:
+        with gr.Column(elem_id="extensions"):
+            for row in to_display:
+                extension, _ = row
+                extension.ui()
+
+
+def create_extensions_tabs():
+    for extension, name in iterator():
+        if hasattr(extension, "ui") and (hasattr(extension, 'params') and extension.params.get('is_tab', False)):
+            display_name = getattr(extension, 'params', {}).get('display_name', name)
+            with gr.Tab(display_name, elem_classes="extension-tab"):
+                extension.ui()
+
+
+EXTENSION_MAP = {
+    "input": partial(_apply_string_extensions, "input_modifier"),
+    "output": partial(_apply_string_extensions, "output_modifier"),
+    "chat_input": _apply_chat_input_extensions,
+    "state": _apply_state_modifier_extensions,
+    "history": _apply_history_modifier_extensions,
+    "bot_prefix": partial(_apply_string_extensions, "bot_prefix_modifier"),
+    "tokenizer": partial(_apply_tokenizer_extensions, "tokenizer_modifier"),
+    'logits_processor': partial(_apply_logits_processor_extensions, 'logits_processor_modifier'),
+    "custom_generate_chat_prompt": _apply_custom_generate_chat_prompt,
+    "custom_generate_reply": _apply_custom_generate_reply,
+    "tokenized_length": _apply_custom_tokenized_length,
+    "css": _apply_custom_css,
+    "js": _apply_custom_js
+}
+
+
+def apply_extensions(typ, *args, **kwargs):
+    if typ not in EXTENSION_MAP:
+        raise ValueError(f"Invalid extension type {typ}")
+
+    return EXTENSION_MAP[typ](*args, **kwargs)

modules/github.py

@@ -0,0 +1,38 @@
+import subprocess
+from pathlib import Path
+
+new_extensions = set()
+
+
+def clone_or_pull_repository(github_url):
+    global new_extensions
+
+    repository_folder = Path("extensions")
+    repo_name = github_url.rstrip("/").split("/")[-1].split(".")[0]
+
+    # Check if the repository folder exists
+    if not repository_folder.exists():
+        repository_folder.mkdir(parents=True)
+
+    repo_path = repository_folder / repo_name
+
+    # Check if the repository is already cloned
+    if repo_path.exists():
+        yield f"Updating {github_url}..."
+        # Perform a 'git pull' to update the repository
+        try:
+            pull_output = subprocess.check_output(["git", "-C", repo_path, "pull"], stderr=subprocess.STDOUT)
+            yield "Done."
+            return pull_output.decode()
+        except subprocess.CalledProcessError as e:
+            return str(e)
+
+    # Clone the repository
+    try:
+        yield f"Cloning {github_url}..."
+        clone_output = subprocess.check_output(["git", "clone", github_url, repo_path], stderr=subprocess.STDOUT)
+        new_extensions.add(repo_name)
+        yield f"The extension `{repo_name}` has been downloaded.\n\nPlease close the the web UI completely and launch it again to be able to load it."
+        return clone_output.decode()
+    except subprocess.CalledProcessError as e:
+        return str(e)