Merge pull request #8 from hugulas/main

support chatgpt as translate engine

README.md

@@ -58,8 +58,8 @@
 3. 若想使用其他翻译的API, 在`main.py` 中更改 `from translate_func import baidu_translate as net_translate`  
    - 有道翻译 `from translate_func import youdao_translate as net_translate` 
    - 谷歌翻译 `from translate_func import google_translate as net_translate`
-
-
+   - gpt翻译 效果更好 `from translate_func import gpt_translate as net_translate`
+     - 在secretKey里放自己的api key
 
 ### Step.3 批量翻译
 

main.py

@@ -16,14 +16,15 @@
 import time
 import shutil
 import inspect
+import traceback
 
 import fitz
 import requests
 from tqdm import tqdm
 from docx import Document
 from docx.shared import Inches
 from docx.oxml.ns import qn
-from translate_func import baidu_translate as net_translate
+from translate_func import gpt_translate as net_translate
 
 save_img = True
 save_docx = True
@@ -197,7 +198,7 @@ def main():
                                 # fitz.Rect(end[0],begin[1],end[2],end[3])为新扩展的矩形框坐标
                                 if begin[2] > end[2]:  # 如果起始点的右下角x坐标小于结束点的右下角x坐标
                                     new_page.insert_textbox(fitz.Rect(end[0], begin[1], begin[2], end[3]), res, fontname="song",
-                                                        fontfile=os.path.join(root,'EasyTrans',
+                                                        fontfile=os.path.join(root,'EasyTrans-mac',
                                                                               'SimSun.ttf'),
                                                         fontsize=fonts, align=text_pos)
                                 else:
@@ -263,7 +264,7 @@ def main():
         except Exception as e:
 
             print('翻译过程出现异常......')
-            print(e)
+            traceback.print_exc()
             new_file_name = os.path.join(root,'EasyTrans-mac', 'output_file',  file_name[:-4] + '_translated' + '.pdf')  # 翻译后的pdf保存路径
             print(new_file_name)
             new_docx_name = os.path.join(root,'EasyTrans-mac', 'output_file', file_name[:-4] + '_translated' + '.docx')  # 翻译后的docx保存路径

requirements.txt

@@ -9,3 +9,5 @@ requests
 six
 tqdm
 urllib3
+openai
+tenacity

translate_func.py

@@ -20,7 +20,12 @@
 from hashlib import md5
 from time import sleep, time
 import re
-
+import openai
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_random_exponential,
+)  # for exponential backoff
 from account import *
 
 
@@ -130,6 +135,80 @@ def google_translate(content):
     return res
 
 
+@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+def translate(key, target_language, text, use_azure=False, api_base="", deployment_name="", options=None):
+    # Set up OpenAI API version
+    if use_azure:
+        openai.api_type = "azure"
+        openai.api_version = AZURE_API_VERSION
+        openai.api_base = api_base
+    if api_base:
+        openai.api_base = api_base
+    # Set up OpenAI API key
+    openai.api_key = key[0]
+    if not text:
+        return ""
+    # lang
+
+    # Set up the prompt
+    messages = [{
+        'role': 'system',
+        'content': 'You are a translator assistant.'
+    }, {
+        "role":
+        "user",
+        "content":
+        f"Translate the following text into {target_language}. Retain the original format. Return only the translation and nothing else:\n{text}",
+    }]
+    if use_azure:
+        completion = openai.ChatCompletion.create(
+            # No need to specify model since deployment contain that information.
+            messages=messages,
+            deployment_id=deployment_name
+        )
+    else:
+        completion = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo-16k",
+            messages=messages,
+        )
+
+    t_text = (completion["choices"][0].get("message").get(
+        "content").encode("utf8").decode())
+
+    return t_text
+
+def divide_string_by_sentences(input_string, max_length):
+    sentence_regex = re.compile(r'.*?[.?!\n]+', re.DOTALL)
+    sentences = sentence_regex.findall(input_string)
+
+    parts = []
+    current_part = ''
+    for sentence in sentences:
+        if len(current_part) + len(sentence) <= max_length:
+            current_part += sentence
+        else:
+            parts.append(current_part.strip())
+            current_part = sentence
+
+    if current_part:
+        parts.append(current_part.strip())
+
+    return parts
+
+# chatgpt翻译
+def gpt_translate(content):
+    '''实现gpt的翻译'''
+
+    content = content.replace('\n','')
+    print(content)
+    max_length = 11000
+    if len(content) > max_length:
+        parts = divide_string_by_sentences(content, max_length)
+        return "".join([gpt_translate(part) for part in parts])
+    else:
+        return translate(secretKey,"Chinese", content, api_base="https://aigptx.top/v1")
+
+
 # 百度翻译方法
 @CallingCounter
 def baidu_translate(content,boundary=0):