rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Fully migrated text to AI calculation to new backend. Fully supports sentence splitting so prompt, actions, etc won't get cut by partial sentences

Browse Source
This commit is contained in:
ebolam
2022-09-22 11:47:59 -04:00
parent faed1444ef
commit d1c2e6e506
2 changed files with 201 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
aiserver.py
View File

@@ -3211,24 +3211,25 @@ def lua_compute_context(submission, entries, folders, kwargs):
while(folders[i] is not None):
allowed_folders.add(int(folders[i]))
i += 1
winfo, mem, anotetxt, _ = calcsubmitbudgetheader(
submission,
allowed_entries=allowed_entries,
allowed_folders=allowed_folders,
force_use_txt=True,
scan_story=kwargs["scan_story"] if kwargs["scan_story"] != None else True,
)
#winfo, mem, anotetxt, _ = calcsubmitbudgetheader(
# submission,
# allowed_entries=allowed_entries,
# allowed_folders=allowed_folders,
# force_use_txt=True,
# scan_story=kwargs["scan_story"] if kwargs["scan_story"] != None else True,
#)
if koboldai_vars.alt_gen:
txt, _, _ = koboldai_vars.calc_ai_text()
print("Using Alt Gen")
else:
txt, _, _ = calcsubmitbudget(
len(actions),
winfo,
mem,
anotetxt,
actions,
)
#txt, _, _ = calcsubmitbudget(
# len(actions),
# winfo,
# mem,
# anotetxt,
# actions,
#)
txt, _, _ = koboldai_vars.calc_ai_text(method=1)
return utils.decodenewlines(tokenizer.decode(txt))
#==================================================================#
@@ -4826,15 +4827,16 @@ def calcsubmit(txt):
anoteadded = False # In case our budget runs out before we hit A.N. depth
actionlen = len(koboldai_vars.actions)
winfo, mem, anotetxt, found_entries = calcsubmitbudgetheader(txt)
#winfo, mem, anotetxt, found_entries = calcsubmitbudgetheader(txt)
# For all transformers models
if(koboldai_vars.model != "InferKit"):
if koboldai_vars.alt_gen:
subtxt, min, max = koboldai_vars.calc_ai_text(submitted_text=txt)
print("Using Alt Gen")
logger.debug("Using Alt Gen")
else:
subtxt, min, max = calcsubmitbudget(actionlen, winfo, mem, anotetxt, koboldai_vars.actions, submission=txt)
#subtxt, min, max = calcsubmitbudget(actionlen, winfo, mem, anotetxt, koboldai_vars.actions, submission=txt)
subtxt, min, max = koboldai_vars.calc_ai_text(submitted_text=txt, method=1)
if(actionlen == 0):
if(not koboldai_vars.use_colab_tpu and koboldai_vars.model not in ["Colab", "API", "CLUSTER", "OAI", "TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX"]):
generate(subtxt, min, max, found_entries=found_entries)
@@ -4976,13 +4978,14 @@ def _generate(txt, minimum, maximum, found_entries):
encoded = []
for i in range(koboldai_vars.numseqs):
txt = utils.decodenewlines(tokenizer.decode(genout[i, -already_generated:]))
winfo, mem, anotetxt, _found_entries = calcsubmitbudgetheader(txt, force_use_txt=True, actions=koboldai_vars.actions)
#winfo, mem, anotetxt, _found_entries = calcsubmitbudgetheader(txt, force_use_txt=True, actions=koboldai_vars.actions)
found_entries[i].update(_found_entries)
if koboldai_vars.alt_gen:
txt, _, _ = koboldai_vars.calc_ai_text(submitted_text=txt)
print("Using Alt Gen")
logger.debug("Using Alt Gen")
else:
txt, _, _ = calcsubmitbudget(len(koboldai_vars.actions), winfo, mem, anotetxt, koboldai_vars.actions, submission=txt)
#txt, _, _ = calcsubmitbudget(len(koboldai_vars.actions), winfo, mem, anotetxt, koboldai_vars.actions, submission=txt)
txt, _, _ = koboldai_vars.calc_ai_text(submitted_text=txt, method=1)
encoded.append(torch.tensor(txt, dtype=torch.long, device=genout.device))
max_length = len(max(encoded, key=len))
encoded = torch.stack(tuple(torch.nn.functional.pad(e, (max_length - len(e), 0), value=model.config.pad_token_id or model.config.eos_token_id) for e in encoded))
@@ -5530,13 +5533,14 @@ def tpumtjgenerate(txt, minimum, maximum, found_entries=None):
encoded = []
for i in range(koboldai_vars.numseqs):
txt = utils.decodenewlines(tokenizer.decode(past[i]))
winfo, mem, anotetxt, _found_entries = calcsubmitbudgetheader(txt, force_use_txt=True, actions=koboldai_vars.actions)
#winfo, mem, anotetxt, _found_entries = calcsubmitbudgetheader(txt, force_use_txt=True, actions=koboldai_vars.actions)
found_entries[i].update(_found_entries)
if koboldai_vars.alt_gen:
txt, _, _ = koboldai_vars.calc_ai_text(submitted_text=txt)
print("Using Alt Gen: {}".format(tokenizer.decode(txt)))
logger.debug("Using Alt Gen")
else:
txt, _, _ = calcsubmitbudget(len(koboldai_vars.actions), winfo, mem, anotetxt, koboldai_vars.actions, submission=txt)
#txt, _, _ = calcsubmitbudget(len(koboldai_vars.actions), winfo, mem, anotetxt, koboldai_vars.actions, submission=txt)
txt, _, _ = koboldai_vars.calc_ai_text(submitted_text=txt, method=1)
encoded.append(np.array(txt, dtype=np.uint32))
max_length = len(max(encoded, key=len))
encoded = np.stack(tuple(np.pad(e, (max_length - len(e), 0), constant_values=tpu_mtj_backend.pad_token_id) for e in encoded))

310
koboldai_settings.py
View File

@@ -118,7 +118,7 @@ class koboldai_vars(object):
if self.tokenizer is None:
used_tokens = 99999999999999999999999
else:
used_tokens = 0 if self.sp_length is None else self.sp_length
used_tokens = 0 if self.sp_length is None else self.sp_length + len(self.tokenizer._koboldai_header)
text = ""
# TODO: We may want to replace the "text" variable with a list-type
@@ -152,24 +152,76 @@ class koboldai_vars(object):
context.append({"type": "world_info", "text": wi_text})
text += wi_text
#we're going to split our actions by sentence for better context. We'll add in which actions the sentence covers. Prompt will be added at a -1 ID
actions = {i: self.actions[i] for i in range(len(self.actions))}
actions[-1] = self.prompt
action_text = self.prompt + str(self.actions)
###########action_text_split = [sentence, actions used in sentence, token length, included in AI context]################
action_text_split = [[x+" ", [], 0 if self.tokenizer is None else len(self.tokenizer.encode(x+" ")), False] for x in re.split("(?<=[.!?])\s+", action_text)]
#The last action shouldn't have the extra space from the sentence splitting, so let's remove it
action_text_split[-1][0] = action_text_split[-1][0][:-1]
action_text_split[-1][2] = 0 if self.tokenizer is None else len(self.tokenizer.encode(action_text_split[-1][0]))
Action_Position = [-1, len(actions[-1])] #First element is the action item, second is how much text is left
Sentence_Position = [0, len(action_text_split[0][0])]
while True:
advance_action = False
advance_sentence = False
if Action_Position[1] <= Sentence_Position[1]:
#We have enough text in the sentence to completely cover the action. Advance it to the next action
advance_action = True
if Sentence_Position[1] <= Action_Position[1]:
advance_sentence = True
if Action_Position[0] not in action_text_split[Sentence_Position[0]][1]:
#Since this action is in the sentence, add it to the list if it's not already there
action_text_split[Sentence_Position[0]][1].append(Action_Position[0])
#Fix the text length leftovers first since they interact with each other
if not advance_action:
Action_Position[1] -= Sentence_Position[1]
if not advance_sentence:
Sentence_Position[1] -= Action_Position[1]
if advance_action:
Action_Position[0] += 1
if Action_Position[0] >= max(actions):
break
Action_Position[1] = len(actions[Action_Position[0]])
if advance_sentence:
Sentence_Position[0] += 1
if Sentence_Position[0] >= len(action_text_split):
break
Sentence_Position[1] = len(action_text_split[Sentence_Position[0]][0])
#OK, action_text_split now contains a list of [sentence including trailing space if needed, [action IDs that sentence includes]]
#Add prompt lenght/text if we're set to always use prompt
if self.useprompt:
self.max_prompt_length if self.prompt_length > self.max_prompt_length else self.prompt_length
prompt_length = 0
prompt_text = ""
for item in action_text_split:
if -1 not in item[1]:
#We've finished going through our prompt. Stop
break
if prompt_length + item[2] < self.max_prompt_length:
prompt_length += item[2]
item[3] = True
prompt_text += item[0]
if prompt_length + used_tokens < token_budget:
used_tokens += self.max_prompt_length if self.prompt_length > self.max_prompt_length else self.prompt_length
used_tokens += prompt_length
#Find World Info entries in prompt
for wi in self.worldinfo_v2:
if wi['uid'] not in used_world_info:
#Check to see if we have the keys/secondary keys in the text so far
match = False
for key in wi['key']:
if key in self.prompt:
if key in prompt_text:
match = True
break
if wi['selective'] and match:
match = False
for key in wi['keysecondary']:
if key in self.prompt:
if key in prompt_text:
match=True
break
if match:
@@ -181,7 +233,123 @@ class koboldai_vars(object):
text += wi_text
self.worldinfo_v2.set_world_info_used(wi['uid'])
prompt_text = self.prompt
prompt_text = prompt_text
if self.tokenizer and self.prompt_length > self.max_prompt_length:
prompt_text = self.tokenizer.decode(self.tokenizer.encode(self.prompt)[-self.max_prompt_length-1:])
#We'll add the prompt text AFTER we go through the game text as the world info needs to come first if we're in method 1 rather than method 2
self.prompt_in_ai = True
else:
self.prompt_in_ai = False
#remove author's notes token length (Add the text later)
used_tokens += self.authornote_length
#Start going through the actions backwards, adding it to the text if it fits and look for world info entries
game_text = ""
game_context = []
authors_note_final = self.authornotetemplate.replace("<|>", self.authornote)
used_all_tokens = False
for action in range(len(self.actions)):
self.actions.set_action_in_ai(action, used=False)
for i in range(len(action_text_split)-1, -1, -1):
if action_text_split[i][3]:
#We've hit an item we've already included. Stop
break;
if len(action_text_split) - i - 1 == self.andepth and self.authornote != "":
game_text = "{}{}".format(authors_note_final, game_text)
game_context.insert(0, {"type": "authors_note", "text": authors_note_final})
length = 0 if self.tokenizer is None else len(self.tokenizer.encode(action_text_split[i][0]))
if length+used_tokens <= token_budget and not used_all_tokens:
used_tokens += length
selected_text = action_text_split[i][0]
game_text = "{}{}".format(selected_text, game_text)
game_context.insert(0, {"type": "action", "text": selected_text})
for action in action_text_split[i][1]:
if action >= 0:
self.actions.set_action_in_ai(action)
#Now we need to check for used world info entries
for wi in self.worldinfo_v2:
if wi['uid'] not in used_world_info:
#Check to see if we have the keys/secondary keys in the text so far
match = False
for key in wi['key']:
if key in selected_text:
match = True
break
if wi['selective'] and match:
match = False
for key in wi['keysecondary']:
if key in selected_text:
match=True
break
if method == 1:
if len(action_text_split) - i > self.widepth:
match = False
if match:
if used_tokens+0 if 'token_length' not in wi or wi['token_length'] is None else wi['token_length'] <= token_budget:
used_tokens+=wi['token_length']
used_world_info.append(wi['uid'])
wi_text = wi["content"]
if method == 1:
text = "{}{}".format(wi_text, game_text)
context.insert(0, {"type": "world_info", "text": wi_text})
else:
game_text = "{}{}".format(wi_text, game_text)
game_context.insert(0, {"type": "world_info", "text": wi_text})
self.worldinfo_v2.set_world_info_used(wi['uid'])
else:
used_all_tokens = True
#if we don't have enough actions to get to author's note depth then we just add it right before the game text
if len(action_text_split) < self.andepth and self.authornote != "":
game_text = "{}{}".format(authors_note_final, game_text)
game_context.insert(0, {"type": "authors_note", "text": authors_note_final})
if self.useprompt:
text += prompt_text
context.append({"type": "prompt", "text": prompt_text})
else self.useprompt:
prompt_length = 0
prompt_text = ""
for item in action_text_split:
if -1 not in item[1]:
#We've finished going through our prompt. Stop
break
if prompt_length + item[2] < self.max_prompt_length:
prompt_length += item[2]
item[3] = True
prompt_text += item[0]
if prompt_length + used_tokens < token_budget:
used_tokens += prompt_length
#Find World Info entries in prompt
for wi in self.worldinfo_v2:
if wi['uid'] not in used_world_info:
#Check to see if we have the keys/secondary keys in the text so far
match = False
for key in wi['key']:
if key in prompt_text:
match = True
break
if wi['selective'] and match:
match = False
for key in wi['keysecondary']:
if key in prompt_text:
match=True
break
if match:
if used_tokens+0 if 'token_length' not in wi else wi['token_length'] <= token_budget:
used_tokens+=wi['token_length']
used_world_info.append(wi['uid'])
wi_text = wi['content']
context.append({"type": "world_info", "text": wi_text})
text += wi_text
self.worldinfo_v2.set_world_info_used(wi['uid'])
prompt_text = prompt_text
if self.tokenizer and self.prompt_length > self.max_prompt_length:
prompt_text = self.tokenizer.decode(self.tokenizer.encode(self.prompt)[-self.max_prompt_length-1:])
@@ -191,136 +359,6 @@ class koboldai_vars(object):
else:
self.prompt_in_ai = False
#remove author's notes token length (Add the text later)
used_tokens += self.authornote_length
#Start going through the actions backwards, adding it to the text if it fits and look for world info entries
game_text = ""
game_context = []
authors_note_final = self.authornotetemplate.replace("<|>", self.authornote)
used_all_tokens = False
#we're going to split our actions by sentence for better context. We'll add in which actions the sentence covers
if self.actions.action_count >= 0:
action_text = str(self.actions)
action_text_split = [[x+" ", []] for x in re.split("(?<=[.!?])\s+", action_text)]
action_text_split[-1][0] = action_text_split[-1][0][:-1]
Action_Position = [0, len(self.actions[0])] #First element is the action item, second is how much text is left
Sentence_Position = [0, len(action_text_split[0][0])]
while True:
advance_action = False
advance_sentence = False
if Action_Position[1] <= Sentence_Position[1]:
#We have enough text in the sentence to completely cover the action. Advance it to the next action
advance_action = True
if Sentence_Position[1] <= Action_Position[1]:
advance_sentence = True
if Action_Position[0] not in action_text_split[Sentence_Position[0]][1]:
#Since this action is in the sentence, add it to the list if it's not already there
action_text_split[Sentence_Position[0]][1].append(Action_Position[0])
#Fix the text length leftovers first since they interact with each other
if not advance_action:
Action_Position[1] -= Sentence_Position[1]
if not advance_sentence:
Sentence_Position[1] -= Action_Position[1]
if advance_action:
Action_Position[0] += 1
if Action_Position[0] >= len(self.actions):
break
Action_Position[1] = len(self.actions[Action_Position[0]])
if advance_sentence:
Sentence_Position[0] += 1
if Sentence_Position[0] >= len(action_text_split):
break
Sentence_Position[1] = len(action_text_split[Sentence_Position[0]][0])
#OK, action_text_split now contains a list of [sentence including trailing space if needed, [action IDs that sentence includes]]
for action in range(len(self.actions)):
self.actions.set_action_in_ai(action, used=False)
for i in range(len(action_text_split)-1, -1, -1):
if len(action_text_split) - i - 1 == self.andepth and self.authornote != "":
game_text = "{}{}".format(authors_note_final, game_text)
game_context.insert(0, {"type": "authors_note", "text": authors_note_final})
length = 0 if self.tokenizer is None else len(self.tokenizer.encode(action_text_split[i][0]))
if length+used_tokens <= token_budget and not used_all_tokens:
used_tokens += length
selected_text = action_text_split[i][0]
game_text = "{}{}".format(selected_text, game_text)
game_context.insert(0, {"type": "action", "text": selected_text})
for action in action_text_split[i][1]:
self.actions.set_action_in_ai(action)
#Now we need to check for used world info entries
for wi in self.worldinfo_v2:
if wi['uid'] not in used_world_info:
#Check to see if we have the keys/secondary keys in the text so far
match = False
for key in wi['key']:
if key in selected_text:
match = True
break
if wi['selective'] and match:
match = False
for key in wi['keysecondary']:
if key in selected_text:
match=True
break
if match:
if used_tokens+0 if 'token_length' not in wi or wi['token_length'] is None else wi['token_length'] <= token_budget:
used_tokens+=wi['token_length']
used_world_info.append(wi['uid'])
wi_text = wi["content"]
game_text = "{}{}".format(wi_text, game_text)
game_context.insert(0, {"type": "world_info", "text": wi_text})
self.worldinfo_v2.set_world_info_used(wi['uid'])
else:
used_all_tokens = True
else:
action_text_split = []
#if we don't have enough actions to get to author's note depth then we just add it right before the game text
if len(action_text_split) < self.andepth and self.authornote != "":
game_text = "{}{}".format(authors_note_final, game_text)
game_context.insert(0, {"type": "authors_note", "text": authors_note_final})
if not self.useprompt:
prompt_length = self.max_prompt_length if self.prompt_length > self.max_prompt_length else self.prompt_length
if prompt_length + used_tokens < token_budget:
used_tokens += self.max_prompt_length if self.prompt_length > self.max_prompt_length else self.prompt_length
#Find World Info entries in prompt
for wi in self.worldinfo_v2:
if wi['uid'] not in used_world_info:
#Check to see if we have the keys/secondary keys in the text so far
match = False
for key in wi['key']:
if key in self.prompt:
match = True
break
if wi['selective'] and match:
match = False
for key in wi['keysecondary']:
if key in self.prompt:
match=True
break
if match:
if used_tokens+0 if 'token_length' not in wi or wi['token_length'] is None else wi['token_length'] <= token_budget:
used_tokens+=wi['token_length']
used_world_info.append(wi['uid'])
wi_text = wi["content"]
text += wi_text
context.append({"type": "world_info", "text": wi_text})
self.worldinfo_v2.set_world_info_used(wi['uid'])
self.prompt_in_ai = True
prompt_text = self.prompt
if self.tokenizer and self.prompt_length > self.max_prompt_length:
prompt_text = self.tokenizer.decode(self.tokenizer.encode(self.prompt)[-self.max_prompt_length-1:])
else:
self.prompt_in_ai = False
prompt_text = ""
text += prompt_text
context.append({"type": "prompt", "text": prompt_text})
text += game_text
context += game_context