From 405578f2b3616a98ef5fa3b66252e77e0ce1cf28 Mon Sep 17 00:00:00 2001 From: ebolam Date: Mon, 24 Oct 2022 20:28:56 -0400 Subject: [PATCH] Colab Update --- aiserver.py | 3 ++- tpu_mtj_backend.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/aiserver.py b/aiserver.py index 41dff108..47e7aea0 100644 --- a/aiserver.py +++ b/aiserver.py @@ -1441,6 +1441,7 @@ def general_startup(override_args=None): #setup socketio relay queue koboldai_settings.queue = multiprocessing.Queue() + socketio.start_background_task(socket_io_relay, koboldai_settings.queue, socketio) @@ -3080,7 +3081,7 @@ def load_model(use_gpu=True, gpu_layers=None, disk_layers=None, initial_load=Fal koboldai_vars.allowsp = True loadmodelsettings() loadsettings() - tpu_mtj_backend.load_model(koboldai_vars.custmodpth, hf_checkpoint=koboldai_vars.model not in ("TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX") and koboldai_vars.use_colab_tpu, **koboldai_vars.modelconfig) + tpu_mtj_backend.load_model(koboldai_vars.custmodpth, hf_checkpoint=koboldai_vars.model not in ("TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX") and koboldai_vars.use_colab_tpu, socketio_queue=koboldai_settings.queue, **koboldai_vars.modelconfig) #tpool.execute(tpu_mtj_backend.load_model, koboldai_vars.custmodpth, hf_checkpoint=koboldai_vars.model not in ("TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX") and koboldai_vars.use_colab_tpu, **koboldai_vars.modelconfig) koboldai_vars.modeldim = int(tpu_mtj_backend.params.get("d_embed", tpu_mtj_backend.params["d_model"])) tokenizer = tpu_mtj_backend.tokenizer diff --git a/tpu_mtj_backend.py b/tpu_mtj_backend.py index 1ddfb652..45073296 100644 --- a/tpu_mtj_backend.py +++ b/tpu_mtj_backend.py @@ -55,7 +55,8 @@ from mesh_transformer.util import to_bf16 import time -tqdm_print = None +socketio = None +queue = None params: Dict[str, Any] = {}