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