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

Fix a bug with soft prompts when using transformers XGLM

This commit is contained in:
Gnome Ann 2022-02-25 12:35:23 -05:00
parent 1fc173890e
commit a22d59e191
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
aiserver.py
View File

@ -1005,7 +1005,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly", "TPUMeshTransforme
if(hasattr(self, "transformer")): if(hasattr(self, "transformer")):
inputs_embeds = self.transformer.wte(input_ids) inputs_embeds = self.transformer.wte(input_ids)
else: else:
inputs_embeds = self.model.embed_tokens(input_ids) * self.model.embed_scale inputs_embeds = self.model.embed_tokens(input_ids)
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(
@ -1013,6 +1013,8 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly", "TPUMeshTransforme
vars.sp[shifted_input_ids.clamp(min=0)], vars.sp[shifted_input_ids.clamp(min=0)],
inputs_embeds, inputs_embeds,
) )
if(not hasattr(self, "transformer")):
inputs_embeds *= self.model.embed_scale
kwargs['inputs_embeds'] = inputs_embeds kwargs['inputs_embeds'] = inputs_embeds
return old_forward(self, *args, **kwargs) return old_forward(self, *args, **kwargs)
cls.forward = new_causallm_forward cls.forward = new_causallm_forward