From 97e7d03d09bb49b55b061f7c50884a878f26086e Mon Sep 17 00:00:00 2001 From: daijun1 Date: Tue, 26 Dec 2023 16:32:33 +0800 Subject: [PATCH] fixed the issue that the HF model spontaneously conducted multiple rounds of Q&A and stream_chat method generates garbled characters Signed-off-by: daijun1 --- tools/transformers/modeling_internlm.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tools/transformers/modeling_internlm.py b/tools/transformers/modeling_internlm.py index 37f50d1..b817f69 100644 --- a/tools/transformers/modeling_internlm.py +++ b/tools/transformers/modeling_internlm.py @@ -844,6 +844,7 @@ class InternLMForCausalLM(InternLMPreTrainedModel): self.query = query self.history = history self.response = "" + self.chche = [] self.received_inputs = False self.queue.put((self.response, history + [(self.query, self.response)])) @@ -858,11 +859,18 @@ class InternLMForCausalLM(InternLMPreTrainedModel): self.received_inputs = True return - token = self.tokenizer.decode([value[-1]], skip_special_tokens=True) + self.chche.extend(value.tolist()) + token = self.tokenizer.decode(self.chche, skip_special_tokens=True) + if " " in token and len(token) <= 5: + return + if token.strip() != "": self.response = self.response + token history = self.history + [(self.query, self.response)] self.queue.put((self.response, history)) + self.chche = [] + else: + self.end() def end(self): self.queue.put(None)