mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-10-30 06:00:15 +01:00
138 lines
4.5 KiB
Python
138 lines
4.5 KiB
Python
import re
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
from modules import loaders, shared, ui
|
|
|
|
|
|
def get_model_settings_from_yamls(model):
|
|
settings = shared.model_config
|
|
model_settings = {}
|
|
for pat in settings:
|
|
if re.match(pat.lower(), model.lower()):
|
|
for k in settings[pat]:
|
|
model_settings[k] = settings[pat][k]
|
|
|
|
return model_settings
|
|
|
|
|
|
def infer_loader(model_name):
|
|
path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
|
|
model_settings = get_model_settings_from_yamls(model_name)
|
|
if not path_to_model.exists():
|
|
loader = None
|
|
elif Path(f'{shared.args.model_dir}/{model_name}/quantize_config.json').exists() or ('wbits' in model_settings and type(model_settings['wbits']) is int and model_settings['wbits'] > 0):
|
|
loader = 'AutoGPTQ'
|
|
elif len(list(path_to_model.glob('*ggml*.bin'))) > 0:
|
|
loader = 'llama.cpp'
|
|
elif re.match(r'.*ggml.*\.bin', model_name.lower()):
|
|
loader = 'llama.cpp'
|
|
elif re.match(r'.*rwkv.*\.pth', model_name.lower()):
|
|
loader = 'RWKV'
|
|
else:
|
|
loader = 'Transformers'
|
|
|
|
return loader
|
|
|
|
|
|
# UI: update the command-line arguments based on the interface values
|
|
def update_model_parameters(state, initial=False):
|
|
elements = ui.list_model_elements() # the names of the parameters
|
|
gpu_memories = []
|
|
|
|
for i, element in enumerate(elements):
|
|
if element not in state:
|
|
continue
|
|
|
|
value = state[element]
|
|
if element.startswith('gpu_memory'):
|
|
gpu_memories.append(value)
|
|
continue
|
|
|
|
if initial and vars(shared.args)[element] != vars(shared.args_defaults)[element]:
|
|
continue
|
|
|
|
# Setting null defaults
|
|
if element in ['wbits', 'groupsize', 'model_type'] and value == 'None':
|
|
value = vars(shared.args_defaults)[element]
|
|
elif element in ['cpu_memory'] and value == 0:
|
|
value = vars(shared.args_defaults)[element]
|
|
|
|
# Making some simple conversions
|
|
if element in ['wbits', 'groupsize', 'pre_layer']:
|
|
value = int(value)
|
|
elif element == 'cpu_memory' and value is not None:
|
|
value = f"{value}MiB"
|
|
|
|
if element in ['pre_layer']:
|
|
value = [value] if value > 0 else None
|
|
|
|
setattr(shared.args, element, value)
|
|
|
|
found_positive = False
|
|
for i in gpu_memories:
|
|
if i > 0:
|
|
found_positive = True
|
|
break
|
|
|
|
if not (initial and vars(shared.args)['gpu_memory'] != vars(shared.args_defaults)['gpu_memory']):
|
|
if found_positive:
|
|
shared.args.gpu_memory = [f"{i}MiB" for i in gpu_memories]
|
|
else:
|
|
shared.args.gpu_memory = None
|
|
|
|
|
|
# UI: update the state variable with the model settings
|
|
def apply_model_settings_to_state(model, state):
|
|
model_settings = get_model_settings_from_yamls(model)
|
|
if 'loader' not in model_settings:
|
|
loader = infer_loader(model)
|
|
if 'wbits' in model_settings and type(model_settings['wbits']) is int and model_settings['wbits'] > 0:
|
|
loader = 'AutoGPTQ'
|
|
|
|
# If the user is using an alternative GPTQ loader, let them keep using it
|
|
if not (loader == 'AutoGPTQ' and state['loader'] in ['GPTQ-for-LLaMa', 'ExLlama', 'ExLlama_HF']):
|
|
state['loader'] = loader
|
|
|
|
for k in model_settings:
|
|
if k in state:
|
|
if k in ['wbits', 'groupsize']:
|
|
state[k] = str(model_settings[k])
|
|
else:
|
|
state[k] = model_settings[k]
|
|
|
|
return state
|
|
|
|
|
|
# Save the settings for this model to models/config-user.yaml
|
|
def save_model_settings(model, state):
|
|
if model == 'None':
|
|
yield ("Not saving the settings because no model is loaded.")
|
|
return
|
|
|
|
with Path(f'{shared.args.model_dir}/config-user.yaml') as p:
|
|
if p.exists():
|
|
user_config = yaml.safe_load(open(p, 'r').read())
|
|
else:
|
|
user_config = {}
|
|
|
|
model_regex = model + '$' # For exact matches
|
|
for _dict in [user_config, shared.model_config]:
|
|
if model_regex not in _dict:
|
|
_dict[model_regex] = {}
|
|
|
|
if model_regex not in user_config:
|
|
user_config[model_regex] = {}
|
|
|
|
for k in ui.list_model_elements():
|
|
if k == 'loader' or k in loaders.loaders_and_params[state['loader']]:
|
|
user_config[model_regex][k] = state[k]
|
|
shared.model_config[model_regex][k] = state[k]
|
|
|
|
output = yaml.dump(user_config, sort_keys=False)
|
|
with open(p, 'w') as f:
|
|
f.write(output)
|
|
|
|
yield (f"Settings for {model} saved to {p}")
|