This commit is contained in:
ebolam
2022-11-19 19:04:20 -05:00
8 changed files with 749 additions and 127 deletions

View File

@@ -187,6 +187,9 @@ class koboldai_vars(object):
def reset_model(self):
self._model_settings.reset_for_model_load()
def is_chat_v2(self):
return self.chat_style > 0 and self.chatmode
def get_token_representation(self, text: Union[str, list, None]) -> list:
if not self.tokenizer or not text:
return []
@@ -206,6 +209,7 @@ class koboldai_vars(object):
return ""
return [], 0, 0+self.genamt, []
if self.alt_gen:
method = 2
else:
@@ -273,6 +277,15 @@ class koboldai_vars(object):
######################################### Get Action Text by Sentence ########################################################
action_text_split = self.actions.to_sentences(submitted_text=submitted_text)
if action_text_split == []:
if return_text:
return ""
return [], 0, 0+self.genamt, []
# Always add newlines on chat v2
if self.is_chat_v2():
action_text_split[-1][0] = action_text_split[-1][0].strip() + "\n"
######################################### Prompt ########################################################
@@ -808,6 +821,7 @@ class story_settings(settings):
# bias experiment
self.memory_attn_bias = 1
self.an_attn_bias = 1
self.chat_style = 0
################### must be at bottom #########################
@@ -1281,7 +1295,7 @@ class KoboldStoryRegister(object):
old_text = None
old_length = None
old = None
self.actions[i] = {"Selected Text": text, "Probabilities": [], "Options": []}
self.actions[i] = {"Selected Text": text, "Probabilities": [], "Options": [], "Time": int(time.time())}
self.story_settings.assign_world_info_to_actions(action_id=i, no_transmit=True)
process_variable_changes(self.socketio, "story", 'actions', {"id": i, 'action': self.actions[i]}, old)
@@ -1317,6 +1331,9 @@ class KoboldStoryRegister(object):
temp = {}
data_to_send = []
for item in json_data['actions']:
if "Time" not in json_data["actions"][item]:
json_data["actions"][item]["Time"] = int(time.time())
temp[int(item)] = json_data['actions'][item]
if int(item) >= self.action_count-100: #sending last 100 items to UI
data_to_send.append({"id": item, 'action': temp[int(item)]})
@@ -1348,6 +1365,7 @@ class KoboldStoryRegister(object):
if self.actions[action_id]["Selected Text"] != text:
self.actions[action_id]["Selected Text"] = text
self.actions[action_id]["Probabilities"] = []
self.actions[action_id]["Time"] = self.actions[action_id].get("Time", int(time.time()))
selected_text_length = 0
self.actions[action_id]["Selected Text Length"] = selected_text_length
for item in self.actions[action_id]["Options"]:
@@ -1358,8 +1376,13 @@ class KoboldStoryRegister(object):
else:
selected_text_length = 0
self.actions[action_id] = {"Selected Text": text, "Selected Text Length": selected_text_length,
"Options": [], "Probabilities": []}
self.actions[action_id] = {
"Selected Text": text,
"Selected Text Length": selected_text_length,
"Options": [],
"Probabilities": [],
"Time": int(time.time()),
}
if self.story_settings is not None:
self.story_settings.assign_world_info_to_actions(action_id=action_id, no_transmit=True)
@@ -1395,14 +1418,23 @@ class KoboldStoryRegister(object):
self.actions[self.action_count+1]['Options'].append({"text": option, "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": []})
else:
old_options = None
self.actions[self.action_count+1] = {"Selected Text": "", "Selected Text Length": 0, "Options": [{"text": x, "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": []} for x in option_list]}
self.actions[self.action_count+1] = {
"Selected Text": "",
"Selected Text Length": 0,
"Options": [{"text": x, "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": []} for x in option_list],
"Time": int(time.time()),
}
process_variable_changes(self.socketio, "story", 'actions', {"id": self.action_count+1, 'action': self.actions[self.action_count+1]}, None)
self.set_game_saved()
def set_options(self, option_list, action_id):
if action_id not in self.actions:
old_options = None
self.actions[action_id] = {"Selected Text": "", "Options": option_list}
self.actions[action_id] = {
"Selected Text": "",
"Options": option_list,
"Time": int(time.time()),
}
else:
old_options = self.actions[action_id]["Options"]
self.actions[action_id]["Options"] = []
@@ -1556,7 +1588,7 @@ class KoboldStoryRegister(object):
if not found:
self.actions[self.action_count+1]['Options'].append({"text": text_list[i], "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": [], "stream_id": i})
else:
self.actions[self.action_count+1] = {"Selected Text": "", "Selected Text Length": 0, "Options": []}
self.actions[self.action_count+1] = {"Selected Text": "", "Selected Text Length": 0, "Options": [], "Time": int(time.time())}
for i in range(len(text_list)):
self.actions[self.action_count+1]['Options'].append({"text": text_list[i], "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": [], "stream_id": i})
@@ -1581,7 +1613,7 @@ class KoboldStoryRegister(object):
selected_text_length = len(self.koboldai_vars.tokenizer.encode(text_list[0]))
else:
selected_text_length = 0
self.actions[self.action_count+1] = {"Selected Text": text_list[0], "Selected Text Length": selected_text_length, "Options": []}
self.actions[self.action_count+1] = {"Selected Text": text_list[0], "Selected Text Length": selected_text_length, "Options": [], "Time": int(time.time())}
@@ -1730,6 +1762,7 @@ class KoboldWorldInfo(object):
self.world_info = {}
self.world_info_folder = OrderedDict()
self.world_info_folder['root'] = []
self.image_store = {}
self.story_settings = story_settings
def reset(self):
@@ -1805,8 +1838,9 @@ class KoboldWorldInfo(object):
if self.socketio is not None:
self.socketio.emit("world_info_folder", {x: self.world_info_folder[x] for x in self.world_info_folder}, broadcast=True, room="UI_2")
def add_item(self, title, key, keysecondary, folder, constant, manual_text,
comment, use_wpp=False, wpp={'name': "", 'type': "", 'format': "W++", 'attributes': {}},
def add_item(self, title, key, keysecondary, folder, constant, manual_text,
comment, wi_type="wi", use_wpp=False,
wpp={'name': "", 'type': "", 'format': "W++", 'attributes': {}},
v1_uid=None, recalc=True, sync=True, send_to_ui=True):
if len(self.world_info) == 0:
uid = 0
@@ -1851,6 +1885,7 @@ class KoboldWorldInfo(object):
'manual_text': manual_text,
"content": content,
"comment": comment,
"type": wi_type,
"token_length": token_length,
"selective": len(keysecondary) > 0,
"used_in_game": constant,
@@ -1880,7 +1915,7 @@ class KoboldWorldInfo(object):
ignore = self.koboldai_vars.calc_ai_text()
return uid
def edit_item(self, uid, title, key, keysecondary, folder, constant, manual_text, comment, use_wpp=False, before=None, wpp={'name': "", 'type': "", 'format': "W++", 'attributes': {}}):
def edit_item(self, uid, title, key, keysecondary, folder, constant, manual_text, comment, wi_type, use_wpp=False, before=None, wpp={'name': "", 'type': "", 'format': "W++", 'attributes': {}}):
logger.debug("Editing World Info {}: {}".format(uid, title))
old_folder = self.world_info[uid]['folder']
#move the world info entry if the folder changed or if there is a new order requested
@@ -1915,6 +1950,7 @@ class KoboldWorldInfo(object):
'manual_text': manual_text,
"content": content,
"comment": comment,
"type": wi_type,
"token_length": token_length,
"selective": len(keysecondary) > 0,
"used_in_game": constant,
@@ -1934,10 +1970,15 @@ class KoboldWorldInfo(object):
def delete(self, uid):
del self.world_info[uid]
if uid in self.image_store:
del self.image_store[uid]
for folder in self.world_info_folder:
if uid in self.world_info_folder[folder]:
self.world_info_folder[folder].remove(uid)
self.story_settings.gamesaved = False
self.sync_world_info_to_old_format()
if self.socketio is not None:
@@ -1989,19 +2030,38 @@ class KoboldWorldInfo(object):
if folder is None:
return {
"folders": {x: self.world_info_folder[x] for x in self.world_info_folder},
"entries": self.world_info
"entries": self.world_info,
"images": self.image_store
}
else:
return {
"folders": {x: self.world_info_folder[x] for x in self.world_info_folder if x == folder},
"entries": {x: self.world_info[x] for x in self.world_info if self.world_info[x]['folder'] == folder}
"entries": {x: self.world_info[x] for x in self.world_info if self.world_info[x]['folder'] == folder},
"images": self.image_store
}
def upgrade_entry(self, wi_entry: dict) -> dict:
# If we do not have a type, or it is incorrect, set to WI.
if wi_entry.get("type") not in ["constant", "chatcharacter", "wi"]:
wi_entry["type"] = "wi"
return wi_entry
def load_json(self, data, folder=None):
if "images" in data:
self.image_store = data["images"]
data["entries"] = {k: self.upgrade_entry(v) for k,v in data["entries"].items()}
if folder is None:
self.world_info = {int(x): data['entries'][x] for x in data['entries']}
self.world_info_folder = data['folders']
#Add the item
start_time = time.time()
for uid, item in data['entries'].items():
start_time = time.time()
self.add_item(item['title'] if 'title' in item else item['key'][0],
item['key'] if 'key' in item else [],
item['keysecondary'] if 'keysecondary' in item else [],
@@ -2009,13 +2069,14 @@ class KoboldWorldInfo(object):
item['constant'] if 'constant' in item else False,
item['manual_text'] if 'manual_text' in item else item['content'],
item['comment'] if 'comment' in item else '',
wi_type=item["type"],
use_wpp=item['use_wpp'] if 'use_wpp' in item else False,
wpp=item['wpp'] if 'wpp' in item else {'name': "", 'type': "", 'format': "W++", 'attributes': {}},
recalc=False, sync=False)
if folder is None:
#self.world_info = {int(x): data['entries'][x] for x in data['entries']}
self.world_info_folder = data['folders']
logger.debug("Load World Info {} took {}s".format(uid, time.time()-start_time))
logger.debug("Load World Info took {}s".format(time.time()-start_time))
try:
start_time = time.time()
self.sync_world_info_to_old_format()