Skip to content

Commit

Permalink
support md file type
Browse files Browse the repository at this point in the history
  • Loading branch information
colin4k committed Dec 22, 2024
1 parent 09a3aae commit 5f7d526
Show file tree
Hide file tree
Showing 5 changed files with 186 additions and 4 deletions.
2 changes: 2 additions & 0 deletions book_maker/loader/__init__.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
from book_maker.loader.epub_loader import EPUBBookLoader
from book_maker.loader.txt_loader import TXTBookLoader
from book_maker.loader.srt_loader import SRTBookLoader
from book_maker.loader.md_loader import MarkdownBookLoader

BOOK_LOADER_DICT = {
"epub": EPUBBookLoader,
"txt": TXTBookLoader,
"srt": SRTBookLoader,
"md": MarkdownBookLoader,
# TODO add more here
}
176 changes: 176 additions & 0 deletions book_maker/loader/md_loader.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,176 @@
import sys
from pathlib import Path

from book_maker.utils import prompt_config_to_kwargs

from .base_loader import BaseBookLoader


class MarkdownBookLoader(BaseBookLoader):
def __init__(
self,
md_name,
model,
key,
resume,
language,
model_api_base=None,
is_test=False,
test_num=5,
prompt_config=None,
single_translate=False,
context_flag=False,
context_paragraph_limit=0,
temperature=1.0,
) -> None:
self.md_name = md_name
self.translate_model = model(
key,
language,
api_base=model_api_base,
temperature=temperature,
**prompt_config_to_kwargs(prompt_config),
)
self.is_test = is_test
self.p_to_save = []
self.bilingual_result = []
self.bilingual_temp_result = []
self.test_num = test_num
self.batch_size = 10
self.single_translate = single_translate
self.md_paragraphs = []

try:
with open(f"{md_name}", encoding="utf-8") as f:
self.origin_book = f.read().splitlines()

except Exception as e:
raise Exception("can not load file") from e

self.resume = resume
self.bin_path = f"{Path(md_name).parent}/.{Path(md_name).stem}.temp.bin"
if self.resume:
self.load_state()

self.process_markdown_content()

def process_markdown_content(self):
"""将原始内容处理成 markdown 段落"""
current_paragraph = []
for line in self.origin_book:
# 如果是空行且当前段落不为空,保存当前段落
if not line.strip() and current_paragraph:
self.md_paragraphs.append('\n'.join(current_paragraph))
current_paragraph = []
# 如果是标题行,单独作为一个段落
elif line.strip().startswith('#'):
if current_paragraph:
self.md_paragraphs.append('\n'.join(current_paragraph))
current_paragraph = []
self.md_paragraphs.append(line)
# 其他情况,添加到当前段落
else:
current_paragraph.append(line)

# 处理最后一个段落
if current_paragraph:
self.md_paragraphs.append('\n'.join(current_paragraph))

@staticmethod
def _is_special_text(text):
return text.isdigit() or text.isspace() or len(text) == 0

def _make_new_book(self, book):
pass

def make_bilingual_book(self):
index = 0
p_to_save_len = len(self.p_to_save)

try:
sliced_list = [
self.md_paragraphs[i : i + self.batch_size]
for i in range(0, len(self.md_paragraphs), self.batch_size)
]
for paragraphs in sliced_list:
batch_text = '\n\n'.join(paragraphs)
if self._is_special_text(batch_text):
continue
if not self.resume or index >= p_to_save_len:
try:
max_retries = 3
retry_count = 0
while retry_count < max_retries:
try:
temp = self.translate_model.translate(batch_text)
break
except AttributeError as ae:
print(f"翻译出错: {ae}")
retry_count += 1
if retry_count == max_retries:
raise Exception("翻译模型初始化失败") from ae
except Exception as e:
print(f"翻译过程中出错: {e}")
raise Exception("翻译过程中出现错误") from e

self.p_to_save.append(temp)
if not self.single_translate:
self.bilingual_result.append(batch_text)
self.bilingual_result.append(temp)
index += self.batch_size
if self.is_test and index > self.test_num:
break

self.save_file(
f"{Path(self.md_name).parent}/{Path(self.md_name).stem}_bilingual.md",
self.bilingual_result,
)

except (KeyboardInterrupt, Exception) as e:
print(f"发生错误: {e}")
print("程序将保存进度,您可以稍后继续")
self._save_progress()
self._save_temp_book()
sys.exit(1) # 使用非零退出码表示错误

def _save_temp_book(self):
index = 0
sliced_list = [
self.origin_book[i : i + self.batch_size]
for i in range(0, len(self.origin_book), self.batch_size)
]

for i in range(len(sliced_list)):
batch_text = "".join(sliced_list[i])
self.bilingual_temp_result.append(batch_text)
if self._is_special_text(self.origin_book[i]):
continue
if index < len(self.p_to_save):
self.bilingual_temp_result.append(self.p_to_save[index])
index += 1

self.save_file(
f"{Path(self.md_name).parent}/{Path(self.md_name).stem}_bilingual_temp.txt",
self.bilingual_temp_result,
)

def _save_progress(self):
try:
with open(self.bin_path, "w", encoding="utf-8") as f:
f.write("\n".join(self.p_to_save))
except:
raise Exception("can not save resume file")

def load_state(self):
try:
with open(self.bin_path, encoding="utf-8") as f:
self.p_to_save = f.read().splitlines()
except Exception as e:
raise Exception("can not load resume file") from e

def save_file(self, book_path, content):
try:
with open(book_path, "w", encoding="utf-8") as f:
f.write("\n".join(content))
except:
raise Exception("can not save file")
2 changes: 1 addition & 1 deletion book_maker/translator/gemini_translator.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ def __init__(
or environ.get(PROMPT_ENV_MAP["system"])
or None # Allow None, but not empty string
)

self.interval = 3
genai.configure(api_key=next(self.keys))
generation_config["temperature"] = temperature

Expand Down
4 changes: 4 additions & 0 deletions prompt_md.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"system": "You are a highly skilled translator responsible for translating the content of books in Markdown format from English into Chinese.",
"user": "## Strategies\nYou will follow a three-step translation process:\n### 1. Translate the input content from English into Chinese, respect the intention of the original text, keep the original Markdown format unchanged, and do not delete or omit any content, nor add additional explanations or remarks.\n### 2. Read the original text and the translation carefully, and then put forward constructive criticism and helpful suggestions to improve the translation. The final style and tone of the translation should conform to the Chinese language style.\nYou must strictly follow the rules below.\n- Never change the Markdown markup structure. Don't add or remove links. Do not change any URL.\n- Never touch or change the contents of code blocks even if they appear to have a bug.\n- Always preserve the original line breaks. Do not add or remove blank lines.\n- Never touch any permalink at the end of each heading.\n- Never touch HTML-like tags such as `<Notes>`.\nWhen writing suggestions, pay attention to whether there are ways to improve the translation in terms of:\n- Accuracy (by correcting errors such as additions, mistranslations, omissions or untranslated text).\n- Fluency (by applying the rules of Chinese grammar, spelling and punctuation, and ensuring there is no unnecessary repetition).\n- Conciseness and abbreviation (please appropriately simplify and abbreviate the translation result while keeping the original meaning unchanged to avoid the translation being too lengthy).\n### 3. Based on the results of steps 1 and 2, refine and polish the translation, and do not add additional explanations or remarks.\n## Output\nFor each step of the translation process, output the results within the appropriate XML tags:\n<step1_initial_translation>\n[Insert your initial translation here.]\n</step1_initial_translation>\n<step2_reflection>\n[Insert your reflection on the translation and put forward specific here, useful and constructive suggestions to improve the translation. Each suggestion should target a specific part of the translation.]\n</step2_reflection>\n<step3_refined_translation>\n[Insert your refined and polished translation here.]\n</step3_refined_translation>\n## Input\nThe following is the content of the book that needs to be translated within the <INPUT> tag:\n<INPUT>{text}</INPUT>"
}
6 changes: 3 additions & 3 deletions prompt_template_sample.json
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"system": "您是一位技术精湛的翻译人员,负责将书籍的内容从其他语言翻译成中文。",
"user": "## 策略\n您将遵循三步翻译流程:\n1. 将输入的内容翻译成中文,尊重原文意图,保持原始SRT格式不变,不删除或省略任何内容,也不添加额外的解释或说明。\n2. 仔细阅读原文和译文,然后提出建设性的批评和有益的建议以改进翻译。翻译的最终风格和语气应符合中文语言风格。写建议时,注意是否有办法改进翻译的\n- 准确性 (通过纠正添加、误译、遗漏或未翻译文本的错误)\n- 流畅性 (通过应用中文语法、拼写和标点符号规则,并确保没有不必要的重复)\n- 精简缩写 (请在保持原意不变的情况下尽量精简缩写翻译结果,避免翻译结果冗长)\n3. 根据步骤1和2的结果,完善和润色翻译,不要添加额外的解释、说明。\n\n## 输出\n\n对于翻译过程的每个步骤,在适当的 XML 标签中输出结果:\n\n<step1_initial_translation>\n[在此插入您的初始翻译]\n</step1_initial_translation>\n\n<step2_reflection>\n[插入您对翻译的反思,提出具体、有用且建设性的建议,以改进翻译。每个建议都应针对翻译的一个特定部分。]\n</step2_reflection>\n\n<step3_refined_translation>\n[在此处插入您精炼和润色后的翻译]\n</step3_refined_translation>\n\n以下<INPUT>标签内是需要翻译的书籍内容:\n<INPUT>{text}</INPUT>"
}
"system": "You are a highly skilled academic translator. Please complete the translation task according to the following instructions and provide only the final polished translation.",
"user": "## Strategies\nYou will follow a three-step translation process:\n### Step.1 Initial Direct Translation: Translate the content from English to Chinese sentence by sentence, respecting the original intent without deleting, omitting, or adding any extra explanations or notes.\n ### Step.2 Reflection and Revision: Carefully review both the input content and the initial direct translation from Step 1. Check if the translation conveys the original meaning, if the grammatical structure is correct, if word choices are appropriate, and if there are any ambiguities or polysemous words. The final style and tone should conform to Chinese language conventions. \nYou must strictly follow the rules below.\n- Don't add or remove links. Do not change any URL.\n- Do not translate the reference list.\n- Never touch,change or translate the mathematical formulas.\n- Never touch,change or translate the contents of code blocks even if they appear to have a bug.\n- Always preserve the original line breaks. Do not add or remove blank lines.\nProvide constructive criticism and helpful suggestions to improve: \n- translation accuracy (correct additions, mistranslations, omissions, or untranslated text errors),\n- fluency (apply Chinese grammar, spelling, and punctuation rules, and ensure no unnecessary repetition), \n- conciseness (streamline the translation results while maintaining the original meaning, avoiding wordiness).\n ### Step.3 Polish and Optimize: Based on the results from Steps 1 and 2, refine and polish the translation, ensuring the final translation adheres to Chinese style without additional explanations or notes. The content to be translated is wrapped in the following <INPUT> tags:\n\n<INPUT>{text}</INPUT>. \n\nPlease write and output only the final polished translation here: "
}

0 comments on commit 5f7d526

Please sign in to comment.