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

Add individual configuration files for userscripts

This commit is contained in:
Gnome Ann 2021-12-22 02:59:31 -05:00
parent 8c63f7c034
commit d538782b1e
2 changed files with 139 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
aiserver.py
View File

@ -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,

141
bridge.lua
View File

@ -153,6 +153,8 @@ return function(_python, _bridged)
---@field generated_cols integer
---@field outputs table<integer, string>
---@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<string, file*>
local config_file_filename_map = {} ---@type table<file*, string>
---@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<integer, string>
koboldbridge.userscripts = {} ---@type table<integer, KoboldUserScriptModule>
koboldbridge.userscriptmodule_filename_map = {} ---@type table<KoboldUserScriptModule, string>
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