Fix: filter_suffix in TritonClient (#150)

2024-02-04 17:36:52 +08:00
parent a2c23ef9dd
commit 6a5447663a
1 changed files with 2 additions and 1 deletions
--- a/lagent/llms/lmdepoly_wrapper.py
+++ b/lagent/llms/lmdepoly_wrapper.py
@@ -159,10 +159,11 @@ class TritonClient(BaseModel):
                self.chatbot._session, prompt, max_tokens, sequence_start,
                sequence_end):
            status = self.state_map.get(status)
+            # The stop symbol also appears in the output of the last STREAM_ING state.
+            res = filter_suffix(res, self.gen_params.get('stop_words'))
            if status < ModelStatusCode.END:
                return status, res, _
            elif status == ModelStatusCode.END:  # remove stop_words
-                res = filter_suffix(res, self.gen_params.get('stop_words'))
                self.chatbot._session.histories = (
                    self.chatbot._session.histories +
                    self.chatbot._session.prompt +