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