diff --git a/pdf2zh/translator.py b/pdf2zh/translator.py index 5046f6e3..c514bec6 100644 --- a/pdf2zh/translator.py +++ b/pdf2zh/translator.py @@ -216,12 +216,34 @@ def translate(self, text): print(len(self.prompt(text, self.prompttext))) print(self.prompt(text, self.prompttext)[0]) print(self.prompt(text, self.prompttext)[1]) - response = self.client.chat( - model=self.model, - options=self.options, - messages=self.prompt(text, self.prompttext), - ) - return response["message"]["content"].strip() + maxlen = max(2000, len(text) * 3) + for model in self.model.split(";"): + for i in range(2): + if i: + print("[Retry]") + response = "" + try: + stream = self.client.chat( + model=model, + options=self.options, + messages=self.prompt(text, self.prompttext), + stream=True, + ) + for chunk in stream: + chunk = chunk["message"]["content"] + print(chunk, end="", flush=True) + response += chunk + if len(response) > maxlen: + raise Exception("Response too long") + if not response.endswith("\n"): + print() + return response.strip() + except Exception as e: + print() + print(e) + print("[Aborted.]") + # if translation fails after multiple retries, return the original text to prevent hang-up + return text class OpenAITranslator(BaseTranslator):