Fix: filter_suffix in TritonClient (#150)

InternLM · Feb 4, 2024 · 6a54476 · 6a54476
1 parent a2c23ef
commit 6a54476
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lagent/llms/lmdepoly_wrapper.py b/lagent/llms/lmdepoly_wrapper.py
@@ -159,10 +159,11 @@ def stream_chat(self,
                 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 +