Implemented abort function in model classes.

Horde stopper not working properly
This commit is contained in:
ebolam
2023-09-27 12:43:42 -04:00
parent 2f3b24eba7
commit ac0f627616
4 changed files with 54 additions and 20 deletions

View File

@@ -171,6 +171,7 @@ class InferenceModel:
"""Root class for all models."""
def __init__(self) -> None:
self.abort = False
self.gen_state = {}
self.post_token_hooks = []
self.stopper_hooks = []
@@ -669,6 +670,9 @@ class InferenceModel:
for hook in self.post_token_hooks:
hook(self, input_ids)
def abort_generation(self, abort=True):
self.abort=abort
def get_supported_gen_modes(self) -> List[GenerationMode]:
"""Returns a list of compatible `GenerationMode`s for the current model.