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

Fix XGLM soft prompts

Browse Source
This commit is contained in:
Gnome Ann 2022-02-02 22:45:16 -05:00
parent d847d04605
commit 78f52063c7
1 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
aiserver.py
View File

@ -752,7 +752,10 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly", "TPUMeshTransforme
if(vars.sp is not None): if(vars.sp is not None):
shifted_input_ids = input_ids - self.config.vocab_size shifted_input_ids = input_ids - self.config.vocab_size
input_ids.clamp_(max=self.config.vocab_size-1) input_ids.clamp_(max=self.config.vocab_size-1)
inputs_embeds = self.transformer.wte(input_ids) if(hasattr(self, "transformer")):
inputs_embeds = self.transformer.wte(input_ids)
else:
inputs_embeds = self.model.embed_tokens(input_ids) * self.embed_scale
if(vars.sp is not None): if(vars.sp is not None):
vars.sp = vars.sp.to(inputs_embeds.dtype).to(inputs_embeds.device) vars.sp = vars.sp.to(inputs_embeds.dtype).to(inputs_embeds.device)
inputs_embeds = torch.where( inputs_embeds = torch.where(
@ -765,11 +768,11 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly", "TPUMeshTransforme
cls.forward = new_causallm_forward cls.forward = new_causallm_forward
for cls in (GPT2LMHeadModel, GPTNeoForCausalLM): for cls in (GPT2LMHeadModel, GPTNeoForCausalLM):
patch_causallm(cls) patch_causallm(cls)
try: for c in ("GPTJForCausalLM", "XGLMForCausalLM"):
from transformers import GPTJForCausalLM try:
patch_causallm(GPTJForCausalLM) patch_causallm(getattr(__import__("transformers"), c))
except: except:
pass pass
# Patch transformers to use our custom logit warpers # Patch transformers to use our custom logit warpers