text-generation-webui/modules/llamacpp_model.py

133 lines
4.2 KiB
Python
Raw Normal View History

import re
from functools import partial
import torch
2023-03-19 07:42:10 +01:00
2023-04-01 02:18:05 +02:00
from modules import shared
2023-03-31 19:27:01 +02:00
from modules.callbacks import Iteratorize
from modules.logging_colors import logger
from modules.text_generation import get_max_prompt_length
2023-03-31 19:27:01 +02:00
import llama_cpp
if torch.cuda.is_available() and not torch.version.hip:
try:
import llama_cpp_cuda
except:
llama_cpp_cuda = None
else:
llama_cpp_cuda = None
def llama_cpp_lib():
if shared.args.cpu or llama_cpp_cuda is None:
return llama_cpp
else:
return llama_cpp_cuda
2023-03-19 07:42:10 +01:00
def ban_eos_logits_processor(eos_token, input_ids, logits):
logits[eos_token] = -float('inf')
return logits
2023-03-19 07:42:10 +01:00
class LlamaCppModel:
def __init__(self):
self.initialized = False
def __del__(self):
self.model.__del__()
2023-03-19 07:42:10 +01:00
@classmethod
def from_pretrained(self, path):
Llama = llama_cpp_lib().Llama
LlamaCache = llama_cpp_lib().LlamaCache
2023-03-19 07:42:10 +01:00
result = self()
cache_capacity = 0
if shared.args.cache_capacity is not None:
if 'GiB' in shared.args.cache_capacity:
cache_capacity = int(re.sub('[a-zA-Z]', '', shared.args.cache_capacity)) * 1000 * 1000 * 1000
elif 'MiB' in shared.args.cache_capacity:
cache_capacity = int(re.sub('[a-zA-Z]', '', shared.args.cache_capacity)) * 1000 * 1000
else:
cache_capacity = int(shared.args.cache_capacity)
logger.info("Cache capacity is " + str(cache_capacity) + " bytes")
params = {
'model_path': str(path),
'n_ctx': shared.args.n_ctx,
'seed': int(shared.args.llama_cpp_seed),
'n_threads': shared.args.threads or None,
'n_batch': shared.args.n_batch,
'use_mmap': not shared.args.no_mmap,
'use_mlock': shared.args.mlock,
2023-07-12 16:05:13 +02:00
'low_vram': shared.args.low_vram,
'n_gpu_layers': shared.args.n_gpu_layers,
2023-08-07 02:49:27 +02:00
'rope_freq_base': 10000 * shared.args.alpha_value ** (64 / 63.),
'rope_freq_scale': 1.0 / shared.args.compress_pos_emb,
2023-07-24 21:37:03 +02:00
'n_gqa': shared.args.n_gqa or None,
'rms_norm_eps': shared.args.rms_norm_eps or None,
}
2023-06-06 18:06:05 +02:00
result.model = Llama(**params)
if cache_capacity > 0:
result.model.set_cache(LlamaCache(capacity_bytes=cache_capacity))
# This is ugly, but the model and the tokenizer are the same object in this library.
return result, result
def encode(self, string):
if type(string) is str:
string = string.encode()
2023-06-06 18:06:05 +02:00
return self.model.tokenize(string)
2023-03-19 07:42:10 +01:00
def decode(self, tokens):
return self.model.detokenize(tokens)
2023-06-17 01:35:38 +02:00
def generate(self, prompt, state, callback=None):
LogitsProcessorList = llama_cpp_lib().LogitsProcessorList
2023-06-17 01:35:38 +02:00
prompt = prompt if type(prompt) is str else prompt.decode()
# Handle truncation
prompt = self.encode(prompt)
prompt = prompt[-get_max_prompt_length(state):]
prompt = self.decode(prompt).decode('utf-8')
completion_chunks = self.model.create_completion(
2023-06-17 01:35:38 +02:00
prompt=prompt,
max_tokens=state['max_new_tokens'],
temperature=state['temperature'],
top_p=state['top_p'],
top_k=state['top_k'],
repeat_penalty=state['repetition_penalty'],
2023-06-18 00:08:25 +02:00
tfs_z=state['tfs'],
2023-06-17 01:35:38 +02:00
mirostat_mode=int(state['mirostat_mode']),
mirostat_tau=state['mirostat_tau'],
mirostat_eta=state['mirostat_eta'],
stream=True,
logits_processor=LogitsProcessorList([
partial(ban_eos_logits_processor, self.model.token_eos()),
]) if state['ban_eos_token'] else None,
)
2023-06-06 18:06:05 +02:00
output = ""
for completion_chunk in completion_chunks:
text = completion_chunk['choices'][0]['text']
output += text
if callback:
callback(text)
2023-06-06 18:06:05 +02:00
return output
2023-03-19 07:42:10 +01:00
2023-06-17 01:35:38 +02:00
def generate_with_streaming(self, *args, **kwargs):
with Iteratorize(self.generate, args, kwargs, callback=None) as generator:
2023-03-31 19:27:01 +02:00
reply = ''
2023-03-19 07:42:10 +01:00
for token in generator:
reply += token
yield reply