From a1df22668b2cf5e81fa393680851e84b6058f51c Mon Sep 17 00:00:00 2001 From: DocShotgun <126566557+DocShotgun@users.noreply.github.com> Date: Fri, 10 May 2024 12:30:17 -0700 Subject: [PATCH 1/4] API: Add min_tokens Bans the EOS token until the generation reaches a minimum length. This will not prevent the model from otherwise ending the generation early by outputting other stop conditions. --- backends/exllamav2/model.py | 25 ++++++++++++++++++++++++- common/sampling.py | 6 ++++++ sampler_overrides/sample_preset.yml | 3 +++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py index 2df3bd5..89834f9 100644 --- a/backends/exllamav2/model.py +++ b/backends/exllamav2/model.py @@ -905,6 +905,9 @@ class ExllamaV2Container: kwargs.get("max_tokens"), self.config.max_seq_len - prompt_tokens ) + # Set min_tokens to generate while keeping EOS banned + min_tokens = unwrap(kwargs.get("min_tokens"), 0) + # This is an inverse of skip_special_tokens decode_special_tokens = unwrap(not kwargs.get("skip_special_tokens"), False) @@ -941,10 +944,27 @@ class ExllamaV2Container: "installed ExLlamaV2 version." ) + # Check if temporary token bans are supported + # TODO: Remove when a new version of ExllamaV2 is released + if min_tokens: + stream_signature = signature(self.generator.stream_ex) + + try: + _bound_vars = stream_signature.bind_partial( + ban_tokens=[] + ) + except TypeError: + logger.warning( + "min_tokens is not supported by the currently " + "installed ExLlamaV2 version." + ) + min_tokens = 0 + # Log generation options to console # Some options are too large, so log the args instead log_generation_params( max_tokens=max_tokens, + min_tokens=min_tokens, stream=kwargs.get("stream"), **gen_settings_log_dict, token_healing=token_healing, @@ -997,7 +1017,10 @@ class ExllamaV2Container: # Run dict generation # Guarantees return of chunk, eos, and chunk_token_ids - raw_generation = self.generator.stream_ex() + if generated_tokens < min_tokens: + raw_generation = self.generator.stream_ex(ban_tokens=eos_tokens) + else: + raw_generation = self.generator.stream_ex() if token_healing: # Extract healed token diff --git a/common/sampling.py b/common/sampling.py index 6c22bfd..9e90a0f 100644 --- a/common/sampling.py +++ b/common/sampling.py @@ -18,6 +18,11 @@ class BaseSamplerRequest(BaseModel): examples=[150], ) + min_tokens: Optional[int] = Field( + default_factory=lambda: get_default_sampler_value("min_tokens", 0), + examples=[0], + ) + generate_window: Optional[int] = Field( default_factory=lambda: get_default_sampler_value("generate_window"), examples=[512], @@ -260,6 +265,7 @@ class BaseSamplerRequest(BaseModel): gen_params = { "max_tokens": self.max_tokens, + "min_tokens": self.min_tokens, "generate_window": self.generate_window, "stop": self.stop, "add_bos_token": self.add_bos_token, diff --git a/sampler_overrides/sample_preset.yml b/sampler_overrides/sample_preset.yml index d91b1a1..ee8c32e 100644 --- a/sampler_overrides/sample_preset.yml +++ b/sampler_overrides/sample_preset.yml @@ -11,6 +11,9 @@ max_tokens: override: 150 force: false +min_tokens: + override: 0 + force: false stop: override: [] force: false From c0b631ba929f6241480ac5a8029be5768e154f2e Mon Sep 17 00:00:00 2001 From: DocShotgun <126566557+DocShotgun@users.noreply.github.com> Date: Fri, 10 May 2024 13:53:55 -0700 Subject: [PATCH 2/4] API: Add banned_strings From exllamav2: List of strings that the generator will refuse to output. As soon as a partial match happens, a checkpoint is saved that the generator can rewind to if need be. Subsequent tokens are then held until the full string is resolved (match or no match) and either emitted or discarded, accordingly. --- backends/exllamav2/model.py | 18 ++++++++++++++++++ common/sampling.py | 9 +++++++++ sampler_overrides/sample_preset.yml | 4 ++++ 3 files changed, 31 insertions(+) diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py index 89834f9..aa3c1af 100644 --- a/backends/exllamav2/model.py +++ b/backends/exllamav2/model.py @@ -791,6 +791,7 @@ class ExllamaV2Container: ) stop_conditions: List[Union[str, int]] = unwrap(kwargs.get("stop"), []) + banned_strings: List[str] = unwrap(kwargs.get("banned_strings"), []) add_bos_token = unwrap(kwargs.get("add_bos_token"), True) ban_eos_token = unwrap(kwargs.get("ban_eos_token"), False) logit_bias = kwargs.get("logit_bias") @@ -960,6 +961,22 @@ class ExllamaV2Container: ) min_tokens = 0 + # Check if banned_strings is supported + # TODO: Remove when a new version of ExllamaV2 is released + if banned_strings: + begin_stream_signature = signature(self.generator.begin_stream_ex) + + try: + _bound_vars = begin_stream_signature.bind_partial( + banned_strings=[] + ) + begin_stream_args["banned_strings"] = banned_strings + except TypeError: + logger.warning( + "banned_strings is not supported by the currently " + "installed ExLlamaV2 version." + ) + # Log generation options to console # Some options are too large, so log the args instead log_generation_params( @@ -979,6 +996,7 @@ class ExllamaV2Container: logprobs=request_logprobs, stop_conditions=stop_conditions, banned_tokens=banned_tokens, + banned_strings=banned_strings, logit_bias=logit_bias, ) diff --git a/common/sampling.py b/common/sampling.py index 9e90a0f..10ef93c 100644 --- a/common/sampling.py +++ b/common/sampling.py @@ -32,6 +32,10 @@ class BaseSamplerRequest(BaseModel): default_factory=lambda: get_default_sampler_value("stop", []) ) + banned_strings: Optional[Union[str, List[str]]] = Field( + default_factory=lambda: get_default_sampler_value("banned_strings", []) + ) + token_healing: Optional[bool] = Field( default_factory=lambda: get_default_sampler_value("token_healing", False) ) @@ -257,6 +261,10 @@ class BaseSamplerRequest(BaseModel): if self.stop and isinstance(self.stop, str): self.stop = [self.stop] + # Convert banned_strings to an array of strings + if self.banned_strings and isinstance(self.banned_strings, str): + self.banned_strings = [self.banned_strings] + # Convert string banned tokens to an integer list if self.banned_tokens and isinstance(self.banned_tokens, str): self.banned_tokens = [ @@ -268,6 +276,7 @@ class BaseSamplerRequest(BaseModel): "min_tokens": self.min_tokens, "generate_window": self.generate_window, "stop": self.stop, + "banned_strings": self.banned_strings, "add_bos_token": self.add_bos_token, "ban_eos_token": self.ban_eos_token, "skip_special_tokens": self.skip_special_tokens, diff --git a/sampler_overrides/sample_preset.yml b/sampler_overrides/sample_preset.yml index ee8c32e..f3dac71 100644 --- a/sampler_overrides/sample_preset.yml +++ b/sampler_overrides/sample_preset.yml @@ -18,6 +18,10 @@ stop: override: [] force: false additive: false +banned_strings: + override: [] + force: false + additive: false token_healing: override: false force: false From 7eee936a3fc5f3f09c9c4e89642b21e6c899e7b4 Mon Sep 17 00:00:00 2001 From: kingbri Date: Fri, 10 May 2024 21:20:00 -0400 Subject: [PATCH 3/4] Model: Remove old code and fix API handling skip_special_tokens is in stable exl2. Also default the parameters if they are not present in the function signature. Signed-off-by: kingbri --- backends/exllamav2/model.py | 19 ++----------------- 1 file changed, 2 insertions(+), 17 deletions(-) diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py index aa3c1af..57313ee 100644 --- a/backends/exllamav2/model.py +++ b/backends/exllamav2/model.py @@ -929,24 +929,9 @@ class ExllamaV2Container: } ) - # Check if decode_special_tokens is supported - # TODO: Remove when a new version of ExllamaV2 is released - if decode_special_tokens: - begin_stream_signature = signature(self.generator.begin_stream_ex) - - try: - _bound_vars = begin_stream_signature.bind_partial( - decode_special_tokens=True - ) - begin_stream_args["decode_special_tokens"] = decode_special_tokens - except TypeError: - logger.warning( - "skip_special_tokens is not supported by the currently " - "installed ExLlamaV2 version." - ) + # MARK: Function signature checks. Not needed in newer ExllamaV2 versions # Check if temporary token bans are supported - # TODO: Remove when a new version of ExllamaV2 is released if min_tokens: stream_signature = signature(self.generator.stream_ex) @@ -962,7 +947,6 @@ class ExllamaV2Container: min_tokens = 0 # Check if banned_strings is supported - # TODO: Remove when a new version of ExllamaV2 is released if banned_strings: begin_stream_signature = signature(self.generator.begin_stream_ex) @@ -976,6 +960,7 @@ class ExllamaV2Container: "banned_strings is not supported by the currently " "installed ExLlamaV2 version." ) + banned_strings = [] # Log generation options to console # Some options are too large, so log the args instead From 366d57cf4561311485621784ab76120e9daff436 Mon Sep 17 00:00:00 2001 From: kingbri Date: Fri, 10 May 2024 21:20:41 -0400 Subject: [PATCH 4/4] Tree: Format Signed-off-by: kingbri --- backends/exllamav2/model.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py index 57313ee..f911e5e 100644 --- a/backends/exllamav2/model.py +++ b/backends/exllamav2/model.py @@ -936,9 +936,7 @@ class ExllamaV2Container: stream_signature = signature(self.generator.stream_ex) try: - _bound_vars = stream_signature.bind_partial( - ban_tokens=[] - ) + _bound_vars = stream_signature.bind_partial(ban_tokens=[]) except TypeError: logger.warning( "min_tokens is not supported by the currently " @@ -951,9 +949,7 @@ class ExllamaV2Container: begin_stream_signature = signature(self.generator.begin_stream_ex) try: - _bound_vars = begin_stream_signature.bind_partial( - banned_strings=[] - ) + _bound_vars = begin_stream_signature.bind_partial(banned_strings=[]) begin_stream_args["banned_strings"] = banned_strings except TypeError: logger.warning(