diff --git a/settings.cfg.example b/settings.cfg.example index 02b7ffc..18bd0e1 100644 --- a/settings.cfg.example +++ b/settings.cfg.example @@ -1,5 +1,6 @@ [option] -#API key for OpenAI API +#API URL and key for OpenAI API +openai-url = https://api.openai.com/v1 openai-apikey = sk- #Target language for translation, e.g. "English", "Chinese", "Japanese" diff --git a/srt_translation.py b/srt_translation.py index 32fca9f..36fae43 100644 --- a/srt_translation.py +++ b/srt_translation.py @@ -33,10 +33,12 @@ config.read_string(config_text) # 获取openai_apikey和language +openai_url = config.get('option', 'openai-url') openai_apikey = config.get('option', 'openai-apikey') language_name = config.get('option', 'target-language') # 设置openai的API密钥 +openai.api_base = openai_url openai.api_key = openai_apikey import argparse @@ -90,16 +92,26 @@ def split_text(text): def is_translation_valid(original_text, translated_text): def get_index_lines(text): lines = text.split('\n') - index_lines = [line for line in lines if re.match(r'^\d+$', line.strip())] + index_lines = [line.strip() for line in lines if re.match(r'^\d+$', line.strip())] return index_lines + def get_subtitle_lines(text): + paras = re.split(r'(\n\s*\n)', text.strip()) + # 判断每段是否至少三行 + subtitle_lines = [para for para in paras if len(para.split('\n')) >= 3] + return subtitle_lines + original_index_lines = get_index_lines(original_text) translated_index_lines = get_index_lines(translated_text) + lines_match = original_index_lines == translated_index_lines + + original_index_para = get_subtitle_lines(original_text) + translated_index_para = get_subtitle_lines(translated_text) + # 判断两个数组个数相等 + paras_match = len(original_index_para) == len(translated_index_para) - print(original_text, original_index_lines) - print(translated_text, translated_index_lines) + return lines_match and paras_match - return original_index_lines == translated_index_lines def translate_text(text): max_retries = 3 retries = 0 @@ -107,11 +119,15 @@ def translate_text(text): while retries < max_retries: try: completion = openai.ChatCompletion.create( - model="gpt-3.5-turbo", + model="gpt-4o-mini", messages=[ + { + "role": "system", + "content": f"Translate the following subtitle text into {language_name}, but keep the subtitle number and timeline unchanged. Keep the format and line breaks. Please pay attention to the context to make whole dialog consistent. Use natural language and avoid word-for-word translation.", + }, { "role": "user", - "content": f"Translate the following subtitle text into {language_name}, but keep the subtitle number and timeline unchanged: \n{text}", + "content": text, } ], )