mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-12-26 14:20:40 +01:00
78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
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(self, path):
|
|
result = self()
|
|
|
|
config = AutoConfig.from_pretrained(
|
|
str(path),
|
|
threads=shared.args.threads,
|
|
gpu_layers=shared.args.n_gpu_layers,
|
|
batch_size=shared.args.n_batch,
|
|
context_length=shared.args.n_ctx,
|
|
stream=True
|
|
)
|
|
|
|
self.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: {self.model.model_type} for {self.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
|