rixty
/
KoboldAI-Client
mirror of https://github.com/KoboldAI/KoboldAI-Client.git
1
1
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #26 from VE-FORBRYDERNE/sp-patch 

More softprompting bug fixes
This commit is contained in:
henk717 2021-11-11 17:05:32 +01:00 committed by GitHub
parent e6599db78f 1fadcbe1e3
commit 8ad3863854
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
aiserver.py
View File

@ -591,6 +591,12 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
return stopping_criteria
transformers.generation_utils.GenerationMixin._get_stopping_criteria = new_get_stopping_criteria
def get_hidden_size_from_model(model):
try:
return int(model.transformer.hidden_size)
except:
return int(model.transformer.embed_dim)
# If custom GPT Neo model was chosen
if(vars.model == "NeoCustom"):
model_config = open(vars.custmodpth + "/config.json", "r")
@ -632,17 +638,20 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
if(vars.hascuda):
if(vars.usegpu):
model = AutoModelForCausalLM.from_pretrained(vars.model, device=0)
vars.modeldim = int(model.transformer.hidden_size)
vars.modeldim = get_hidden_size_from_model(model)
model = model.to(0)
generator = model.generate
elif(vars.breakmodel): # Use both RAM and VRAM (breakmodel)
model = AutoModelForCausalLM.from_pretrained(vars.model)
vars.modeldim = get_hidden_size_from_model(model)
device_config(model)
else:
model = AutoModelForCausalLM.from_pretrained(vars.model)
vars.modeldim = get_hidden_size_from_model(model)
generator = model.generate
else:
model = AutoModelForCausalLM.from_pretrained(vars.model)
vars.modeldim = get_hidden_size_from_model(model)
generator = model.generate
# Suppress Author's Note by flagging square brackets (Old implementation)
@ -718,6 +727,8 @@ def do_connect():
emit('from_server', {'cmd': 'connected', 'smandelete': vars.smandelete, 'smanrename': vars.smanrename})
if(vars.remote):
emit('from_server', {'cmd': 'runs_remotely'})
if(vars.allowsp):
emit('from_server', {'cmd': 'allowsp', 'data': vars.allowsp})
if(not vars.gamestarted):
setStartState()
@ -961,8 +972,6 @@ def get_message(msg):
#==================================================================#
def setStartState():
txt = "<span>Welcome to <span class=\"color_cyan\">KoboldAI</span>! You are running <span class=\"color_green\">"+getmodelname()+"</span>.<br/>"
if(vars.allowsp):
emit('from_server', {'cmd': 'allowsp', 'data': vars.allowsp}, broadcast=True)
if(not vars.noai):
txt = txt + "Please load a game or enter a prompt below to begin!</span>"
else: