Fixed Generation

predibase · Apr 2, 2024 · 8dc9325 · 8dc9325
1 parent 8d5561c
commit 8dc9325
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 10 deletions.
diff --git a/server/lorax_server/models/causal_lm.py b/server/lorax_server/models/causal_lm.py
@@ -8,6 +8,7 @@
 from lorax_server.models import Model
 from lorax_server.models.types import (
     Batch,
+    NextTokens,
     PrefillTokens,
     AlternativeTokens,
     Generation,
@@ -733,11 +734,13 @@ def generate_token(
                     request.id,
                     prefill_tokens,
                     prefill_tokens_length,
-                    None,
-                    next_token_id_squeezed,
-                    next_token_logprob,
-                    next_token_text,
-                    next_token_id_squeezed.item() in self.all_special_ids,
+                    NextTokens(
+                        [next_token_id_squeezed],
+                        [next_token_logprob],
+                        [next_token_text],
+                        [next_token_id_squeezed.item() in self.all_special_ids],
+                        None,
+                    ),
                     generated_text,
                 )
 

diff --git a/server/lorax_server/models/seq2seq_lm.py b/server/lorax_server/models/seq2seq_lm.py
@@ -11,6 +11,7 @@
     GeneratedText,
     Batch,
     Generation,
+    NextTokens,
     PrefillTokens,
 )
 from lorax_server.pb import generate_pb2
@@ -717,11 +718,13 @@ def generate_token(
                     request.id,
                     prefill_tokens,
                     prefill_tokens_length,
-                    None,
-                    next_token_id_squeezed,
-                    next_token_logprob,
-                    next_token_text,
-                    next_token_id_squeezed.item() in self.all_special_ids,
+                    NextTokens(
+                        [next_token_id_squeezed],
+                        [next_token_logprob],
+                        [next_token_text],
+                        [next_token_id_squeezed.item() in self.all_special_ids],
+                        None,
+                    ),
                     generated_text,
                 )