From d538782b1e78bcb7a43720c9f1982544e73c838f Mon Sep 17 00:00:00 2001 From: Gnome Ann <> Date: Wed, 22 Dec 2021 02:59:31 -0500 Subject: [PATCH] Add individual configuration files for userscripts --- aiserver.py | 4 +- bridge.lua | 141 ++++++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 139 insertions(+), 6 deletions(-) diff --git a/aiserver.py b/aiserver.py index e31eedd8..f971cab3 100644 --- a/aiserver.py +++ b/aiserver.py @@ -7,7 +7,6 @@ # External packages import eventlet -from transformers.generation_logits_process import RepetitionPenaltyLogitsProcessor eventlet.monkey_patch() import os os.environ['EVENTLET_THREADPOOL_SIZE'] = '1' @@ -617,7 +616,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly", "TPUMeshTransforme # Patch transformers to use our custom logit warpers - from transformers import LogitsProcessorList, LogitsWarper, LogitsProcessor, TopKLogitsWarper, TopPLogitsWarper, TemperatureLogitsWarper + from transformers import LogitsProcessorList, LogitsWarper, LogitsProcessor, TopKLogitsWarper, TopPLogitsWarper, TemperatureLogitsWarper, RepetitionPenaltyLogitsProcessor def dynamic_processor_wrap(cls, field_name, var_name): old_call = cls.__call__ @@ -1461,6 +1460,7 @@ vars.lua_state = lupa.LuaRuntime(unpack_returned_tuples=True) bridged = { "corescript_path": os.path.join(os.path.dirname(os.path.realpath(__file__)), "cores"), "userscript_path": os.path.join(os.path.dirname(os.path.realpath(__file__)), "userscripts"), + "config_path": os.path.join(os.path.dirname(os.path.realpath(__file__)), "userscripts"), "lib_path": os.path.join(os.path.dirname(os.path.realpath(__file__)), "extern", "lualibs"), "load_callback": load_callback, "print": lua_print, diff --git a/bridge.lua b/bridge.lua index 8e5955f3..2542aedc 100644 --- a/bridge.lua +++ b/bridge.lua @@ -153,6 +153,8 @@ return function(_python, _bridged) ---@field generated_cols integer ---@field outputs table ---@field num_outputs integer + ---@field feedback string + ---@field is_config_file_open boolean local kobold = setmetatable({}, metawrapper) local KoboldLib_mt = setmetatable({}, metawrapper) local KoboldLib_getters = setmetatable({}, metawrapper) @@ -216,6 +218,11 @@ return function(_python, _bridged) koboldbridge.outputs = {} koboldbridge.feedback = nil ---@type string? + function koboldbridge:clear_userscript_metadata() + self.logging_name = nil + self.filename = nil + end + ---@return nil local function maybe_require_regeneration() if koboldbridge.userstate == "genmod" or koboldbridge.userstate == "outmod" then @@ -224,6 +231,84 @@ return function(_python, _bridged) end + --========================================================================== + -- Userscript API: Configuration + --========================================================================== + + local config_files = {} ---@type table + local config_file_filename_map = {} ---@type table + + ---@return file*? + local function open_and_handle_errors(...) + local file, err_msg = io.open(...) + if err_msg ~= nil then + koboldbridge.obliterate_multiverse() + error(err_msg) + return + end + return file + end + + ---@param file file* + local function new_close_pre(file) + if file == nil then + file = io.output() + end + local filename = config_file_filename_map[file] + if filename ~= nil then + config_file_filename_map[file] = nil + config_files[filename] = nil + end + end + + ---@param f fun(file?: file*) + local function _new_close(f) + ---@param file file*? + return function(file) + new_close_pre(file) + return f(file) + end + end + debug.getmetatable(io.stdout).__index.close = _new_close(io.stdout.close) + + ---@param filename string + ---@return boolean + local function is_config_file_open(filename) + return config_files[filename] ~= nil + end + + ---@param filename string + ---@param clear? boolean + ---@return file* + local function get_config_file(filename, clear) + if not is_config_file_open(filename) then + local config_filepath = join_folder_and_filename(bridged.config_path, filename .. ".conf") + open_and_handle_errors(config_filepath, "a"):close() + config_files[filename] = open_and_handle_errors(config_filepath, clear and "w+b" or "r+b") + config_file_filename_map[config_files[filename]] = filename + end + return config_files[filename] + end + + ---@param clear? boolean + ---@return file* + function kobold.get_config_file(clear) + return get_config_file(koboldbridge.filename, clear) + end + + ---@param t KoboldLib + ---@return boolean + function KoboldLib_getters.is_config_file_open(t) + return is_config_file_open(koboldbridge.filename) + end + + ---@param t KoboldLib + ---@param v boolean + function KoboldLib_setters.is_config_file_open(t, v) + error("`KoboldLib.is_config_file_open` is a read-only attribute") + end + + --========================================================================== -- Userscript API: World Info --========================================================================== @@ -1257,7 +1342,8 @@ return function(_python, _bridged) -- Core script API --========================================================================== - koboldbridge.userscripts = {} ---@type table + koboldbridge.userscripts = {} ---@type table + koboldbridge.userscriptmodule_filename_map = {} ---@type table koboldbridge.num_userscripts = 0 koboldbridge.inmod = nil ---@type function? koboldbridge.genmod = nil ---@type function? @@ -1267,6 +1353,7 @@ return function(_python, _bridged) ---@field inmod function? ---@field genmod function? ---@field outmod function? + ---@field config file* ---@class KoboldCoreScript ---@field inmod function? @@ -1280,6 +1367,7 @@ return function(_python, _bridged) ---@field filename string ---@field modulename string ---@field description string + ---@field is_config_file_open boolean ---@field inmod function? ---@field genmod function? ---@field outmod function? @@ -1297,6 +1385,12 @@ return function(_python, _bridged) outmod = false, } + ---@param clear? boolean + ---@return file* + function KoboldUserScriptModule:get_config_file(clear) + return get_config_file(koboldbridge.userscriptmodule_filename_map[self], clear) + end + ---@generic K ---@param t KoboldUserScriptModule ---@param k K @@ -1316,6 +1410,8 @@ return function(_python, _bridged) function KoboldUserScriptModule_mt.__index(t, k) if type(k) == "string" and KoboldUserScriptModule_fields[k] ~= nil then return rawget(t, "_" .. k) + elseif k == "is_config_file_open" then + return is_config_file_open(koboldbridge.userscriptmodule_filename_map[t]) end return rawget(t, k) end @@ -1405,6 +1501,7 @@ return function(_python, _bridged) local envs = {} koboldbridge.logging_name = nil + koboldbridge.filename = nil local old_load = load local function _safe_load(_g) @@ -1610,6 +1707,8 @@ return function(_python, _bridged) output = io.output, read = io.read, write = io.write, + close = _new_close(io.close), + lines = io.lines, flush = io.flush, type = io.type, }, @@ -1654,7 +1753,11 @@ return function(_python, _bridged) end function koboldbridge.obliterate_multiverse() + for k, v in pairs(config_files) do + pcall(v.close, v) + end envs = {} + koboldbridge.userscripts = {} koboldbridge.num_userscripts = 0 koboldbridge.inmod = nil koboldbridge.genmod = nil @@ -1669,22 +1772,49 @@ return function(_python, _bridged) ---@return nil function koboldbridge.load_userscripts(filenames, modulenames, descriptions) set_require_path(bridged.userscript_path) + config_files = {} + config_file_filename_map = {} koboldbridge.userscripts = {} + koboldbridge.userscriptmodule_filename_map = {} koboldbridge.num_userscripts = 0 for i, filename in _python.enumerate(filenames) do bridged.load_callback(filename, modulenames[i]) koboldbridge.logging_name = modulenames[i] + koboldbridge.filename = filename ---@type KoboldUserScript local _userscript = old_loadfile(join_folder_and_filename(bridged.userscript_path, filename), "t", koboldbridge.get_universe(filename))() koboldbridge.logging_name = nil + koboldbridge.filename = nil local userscript = deepcopy(KoboldUserScriptModule) - rawset(userscript, "_inmod", function() koboldbridge.logging_name = modulenames[i]; if _userscript.inmod ~= nil then _userscript.inmod() end end) - rawset(userscript, "_genmod", function() koboldbridge.logging_name = modulenames[i]; if _userscript.genmod ~= nil then _userscript.genmod() end end) - rawset(userscript, "_outmod", function() koboldbridge.logging_name = modulenames[i]; if _userscript.outmod ~= nil then _userscript.outmod() end end) + rawset(userscript, "_inmod", function() + koboldbridge.logging_name = modulenames[i] + koboldbridge.filename = filename + if _userscript.inmod ~= nil then + _userscript.inmod() + end + koboldbridge:clear_userscript_metadata() + end) + rawset(userscript, "_genmod", function() + koboldbridge.logging_name = modulenames[i] + koboldbridge.filename = filename + if _userscript.genmod ~= nil then + _userscript.genmod() + end + koboldbridge:clear_userscript_metadata() + end) + rawset(userscript, "_outmod", function() + koboldbridge.logging_name = modulenames[i] + koboldbridge.filename = filename + if _userscript.outmod ~= nil then + _userscript.outmod() + end + koboldbridge:clear_userscript_metadata() + end) rawset(userscript, "_filename", filename) rawset(userscript, "_modulename", modulenames[i]) rawset(userscript, "_description", descriptions[i]) koboldbridge.userscripts[i+1] = userscript + koboldbridge.userscriptmodule_filename_map[userscript] = filename koboldbridge.num_userscripts = i + 1 end end @@ -1700,6 +1830,7 @@ return function(_python, _bridged) function koboldbridge.execute_inmod() local r + koboldbridge:clear_userscript_metadata() koboldbridge.restart_sequence = nil koboldbridge.userstate = "inmod" koboldbridge.regeneration_required = false @@ -1722,6 +1853,7 @@ return function(_python, _bridged) ---@return any, boolean function koboldbridge.execute_genmod() local r + koboldbridge:clear_userscript_metadata() koboldbridge.generating = true koboldbridge.userstate = "genmod" if koboldbridge.genmod ~= nil then @@ -1758,6 +1890,7 @@ return function(_python, _bridged) function koboldbridge.execute_outmod() local r + koboldbridge:clear_userscript_metadata() koboldbridge.generating = false koboldbridge.userstate = "outmod" koboldbridge.num_outputs = kobold.settings.numseqs