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

Don't change the shape of transformer.wte

This commit is contained in:
Gnome Ann 2021-10-28 15:39:59 -04:00
parent a1ae11630a
commit 511817132a
1 changed files with 2 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
aiserver.py
View File

@ -515,16 +515,11 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
def patch_causallm(cls): def patch_causallm(cls):
old_forward = cls.forward old_forward = cls.forward
def new_causallm_forward(self, *args, **kwargs): def new_causallm_forward(self, *args, **kwargs):
num_embeddings = self.config.vocab_size
if(vars.sp is not None):
num_embeddings += vars.sp.shape[0]
if(self.transformer.wte.num_embeddings != num_embeddings):
self.resize_token_embeddings(num_embeddings)
input_ids = kwargs.get('input_ids').to(self.device) input_ids = kwargs.get('input_ids').to(self.device)
assert input_ids is not None assert input_ids is not None
kwargs['input_ids'] = None kwargs['input_ids'] = None
inputs_embeds = self.transformer.wte(input_ids) inputs_embeds = self.transformer.wte(input_ids.clamp(max=self.config.vocab_size-1))
input_ids -= self.config.vocab_size input_ids = input_ids - self.config.vocab_size # Don't use the -= operator here, you'll get a cryptic error message
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(