Fixed tests

predibase · Jan 9, 2024 · f23a973 · f23a973
1 parent 09c1c6e
commit f23a973
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/server/lorax_server/models/causal_lm.py b/server/lorax_server/models/causal_lm.py
@@ -655,12 +655,15 @@ def generate_token(
                     prefill_tokens = PrefillTokens(
                         prefill_token_ids, prefill_logprobs, prefill_texts
                     )
+                    prefill_tokens_length = len(prefill_tokens.token_ids)
                 else:
                     prefill_tokens = None
+                    prefill_tokens_length = 0
 
                 generation = Generation(
                     request.id,
                     prefill_tokens,
+                    prefill_tokens_length,
                     next_token_id_squeezed,
                     next_token_logprob,
                     next_token_text,

diff --git a/server/lorax_server/models/seq2seq_lm.py b/server/lorax_server/models/seq2seq_lm.py
@@ -699,12 +699,15 @@ def generate_token(
                         [float("nan")],
                         [self.tokenizer.bos_token],
                     )
+                    prefill_tokens_length = len(prefill_tokens.token_ids)
                 else:
                     prefill_tokens = None
+                    prefill_tokens_length = 0
 
                 generation = Generation(
                     request.id,
                     prefill_tokens,
+                    prefill_tokens_length,
                     next_token_id_squeezed,
                     next_token_logprob,
                     next_token_text,