Add stopper hooks suppport to exllama

This commit is contained in:
Llama
2023-08-22 23:06:16 -07:00
parent 22fd49937a
commit b96d5d8646

View File

@@ -91,7 +91,7 @@ class model_backend(InferenceModel):
self.capabilties = ModelCapabilities(
embedding_manipulation=False,
post_token_hooks=True,
stopper_hooks=False,
stopper_hooks=True,
post_token_probs=False,
)
@@ -305,6 +305,15 @@ class model_backend(InferenceModel):
trim_count = 1
break
# Apply stoppers
do_stop = False
for stopper in self.stopper_hooks:
do_stop = stopper(self, self.generator.sequence)
if do_stop:
break
if do_stop:
break
utils.koboldai_vars.generated_tkns = max_new - trim_count
if trim_count > 0:
seq = self.generator.sequence[:, gen_in.size(1):-trim_count]