rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-02-17 12:10:49 +01:00
Code Issues Projects Releases Wiki Activity

TPU backend no longer needs to recompile after changing softprompt

Browse Source
This commit is contained in:
Gnome Ann 2021-12-05 02:49:15 -05:00
parent d2d338d314
commit 150ce033c9
1 changed files with 19 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
aiserver.py
View File

@ -1802,12 +1802,25 @@ def tpumtjgenerate(txt, minimum, maximum, found_entries=None):
raise ValueError("Dynamic world info scanning is not supported by the TPU backend yet") raise ValueError("Dynamic world info scanning is not supported by the TPU backend yet")
soft_tokens = None soft_tokens = None
if(vars.sp is not None): if(vars.sp is None):
soft_tokens = np.arange( global np
tpu_mtj_backend.params["n_vocab"] + tpu_mtj_backend.params["n_vocab_padding"], if 'np' not in globals():
tpu_mtj_backend.params["n_vocab"] + tpu_mtj_backend.params["n_vocab_padding"] + vars.sp_length, import numpy as np
dtype=np.uint32 tensor = np.zeros((1, tpu_mtj_backend.params["d_model"]), dtype=np.float32)
rows = tensor.shape[0]
padding_amount = tpu_mtj_backend.params["seq"] - (tpu_mtj_backend.params["seq"] % -tpu_mtj_backend.params["cores_per_replica"]) - rows
tensor = np.pad(tensor, ((0, padding_amount), (0, 0)))
tensor = tensor.reshape(
tpu_mtj_backend.params["cores_per_replica"],
-1,
tpu_mtj_backend.params["d_model"],
) )
vars.sp = tensor
soft_tokens = np.arange(
tpu_mtj_backend.params["n_vocab"] + tpu_mtj_backend.params["n_vocab_padding"],
tpu_mtj_backend.params["n_vocab"] + tpu_mtj_backend.params["n_vocab_padding"] + vars.sp_length,
dtype=np.uint32
)
genout = tpu_mtj_backend.infer( genout = tpu_mtj_backend.infer(
txt, txt,
@ -2676,7 +2689,7 @@ def spRequest(filename):
if(vars.model in ("TPUMeshTransformerGPTJ",)): if(vars.model in ("TPUMeshTransformerGPTJ",)):
rows = tensor.shape[0] rows = tensor.shape[0]
padding_amount = -(rows % -tpu_mtj_backend.params["cores_per_replica"]) padding_amount = tpu_mtj_backend.params["seq"] - (tpu_mtj_backend.params["seq"] % -tpu_mtj_backend.params["cores_per_replica"]) - rows
tensor = np.pad(tensor, ((0, padding_amount), (0, 0))) tensor = np.pad(tensor, ((0, padding_amount), (0, 0)))
tensor = tensor.reshape( tensor = tensor.reshape(
tpu_mtj_backend.params["cores_per_replica"], tpu_mtj_backend.params["cores_per_replica"],