fixed token buffer in case of no end sequece

LAION-AI · Mar 18, 2023 · fa72a04 · fa72a04
1 parent 5e0a04b
commit fa72a04
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/inference/worker/utils.py b/inference/worker/utils.py
@@ -39,10 +39,15 @@ def add(self, token: interface.Token):
     def finish(self, reason: Literal["length", "eos_token", "stop_sequence"]) -> Iterable[interface.Token]:
         if reason == "stop_sequence":
             end_sequence = ""
+            end_tokens = []
             while self.tokens:
-                end_sequence = self.tokens.pop().text + end_sequence
+                token = self.tokens.pop()
+                end_tokens.append(token)
+                end_sequence = token.text + end_sequence
                 if end_sequence in self.stop_sequences:
                     break
+            else:
+                self.tokens.extend(reversed(end_tokens))
             yield from self.tokens
         else:
             yield from self.tokens

diff --git a/inference/worker/work.py b/inference/worker/work.py
@@ -106,6 +106,7 @@ def handle_work_request(
             send_token.to_token_response(request_id=work_request.id),
         )
 
+    logger.debug(f"Generated text: {stream_response.generated_text}")
     utils.send_response(
         ws,
         inference.GeneratedTextResponse(