Fix.

huggingface · Dec 5, 2023 · cb8a168 · cb8a168
1 parent be481a4
commit cb8a168
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py
@@ -833,7 +833,7 @@ def generate_token(
             batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs
         )
 
-        speculative_length = speculative_ids.shape[1]
+        speculative_length = 0 if speculative_ids is None else speculative_ids.shape[1]
         if prefill:
             if len(batch) > 1 and prefill_logprobs:
                 # We create the prefill_tokens_indices tensor that will be used to gather prefill logprobs

diff --git a/server/text_generation_server/models/flash_mistral.py b/server/text_generation_server/models/flash_mistral.py
@@ -21,6 +21,7 @@
     FlashMistralForCausalLM,
     MistralConfig,
 )
+from text_generation_server.utils.speculate import get_speculate
 from text_generation_server.utils import (
     initialize_torch_distributed,
     weight_files,
@@ -132,8 +133,7 @@ def from_pb(
 
             # Paged attention
             # Remove one as the first token des not have a past
-            from text_generation_server.models import SPECULATE
-            speculative_length = SPECULATE
+            speculative_length = get_speculate()
             total_tokens = input_length + max_new_tokens - 1 + speculative_length
 
             # Needed blocks can not go over SLIDING_WINDOW_BLOCKS