Create translate_epub.py

translate_epub.py

@@ -0,0 +1,248 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+from argparse import ArgumentParser
+import time
+import os, re
+import fnmatch
+import glob
+import shutil
+import zipfile
+from tqdm import tqdm
+
+def find_all_htmls(root_dir):
+    html_files = []
+    for foldername, subfolders, filenames in os.walk(root_dir):
+        for extension in ['*.html', '*.xhtml', '*.htm']:
+            for filename in fnmatch.filter(filenames, extension):
+                file_path = os.path.join(foldername, filename)
+                html_files.append(file_path)
+    return html_files
+
+def get_html_text_list(epub_path, text_length):
+    data_list = []
+
+    def clean_text(text):
+        text=re.sub(r'<rt[^>]*?>.*?</rt>', '', text)
+        text=re.sub(r'<[^>]*>|\n', '', text)
+        return text
+
+    with open(epub_path, 'r', encoding='utf-8') as f:
+        file_text = f.read()
+        matches = re.finditer(r'<(h[1-6]|p).*?>(.+?)</\1>', file_text, flags=re.DOTALL)
+        if not matches:
+            print("perhaps this file is a struct file")
+            return data_list, file_text
+        groups = []
+        text = ''
+        pre_end = 0
+        for match in matches:
+            if len(text + match.group(2)) <= text_length:
+                new_text = clean_text(match.group(2))
+                if new_text:
+                    groups.append(match)
+                    text += '\n' + new_text
+            else:
+                data_list.append((text, groups, pre_end))
+                pre_end = groups[-1].end()
+                new_text = clean_text(match.group(2))
+                if new_text:
+                    groups = [match]
+                    text = clean_text(match.group(2))
+                else:
+                    groups = []
+                    text = ''
+                    
+        if text:
+            data_list.append((text, groups, pre_end))
+    # TEST:
+    # for d in data_list:
+    #     print(f"{len(d[0])}", end=" ")
+    return data_list, file_text
+
+def get_prompt(input, model_version):
+    if model_version == '0.5' or model_version == '0.8':
+        prompt = "<reserved_106>将下面的日文文本翻译成中文：" + input + "<reserved_107>"
+        return prompt
+    if model_version == '0.7':
+        prompt = f"<|im_start|>user\n将下面的日文文本翻译成中文：{input}<|im_end|>\n<|im_start|>assistant\n"
+        return prompt
+    if model_version == '0.1':
+        prompt = "Human: \n将下面的日文文本翻译成中文：" + input + "\n\nAssistant: \n"
+        return prompt
+    if model_version == '0.4':
+        prompt = "User: 将下面的日文文本翻译成中文：" + input + "\nAssistant: "
+        return prompt
+
+    raise ValueError(f"Wrong model version{model_version}, please view https://huggingface.co/sakuraumi/Sakura-13B-Galgame")
+
+def split_response(response, model_version):
+    response = response.replace("</s>", "")
+    if model_version == '0.5' or model_version == '0.8':
+        output = response.split("<reserved_107>")[1]
+        return output
+    if model_version == '0.7':
+        output = response.split("<|im_start|>assistant\n")[1]
+        return output
+    if model_version == '0.1':
+        output = response.split("\n\nAssistant: \n")[1]
+        return output
+    if model_version == '0.4':
+        output = response.split("\nAssistant: ")[1]
+        return output
+    
+    raise ValueError(f"Wrong model version{model_version}, please view https://huggingface.co/sakuraumi/Sakura-13B-Galgame")
+
+def detect_degeneration(generation: list, model_version):
+    if model_version != "0.8":
+        return False
+    i = generation.index(196)
+    generation = generation[i+1:]
+    if len(generation) >= 1023:
+        print("model degeneration detected, retrying...")
+        return True
+    else:
+        return False
+
+def get_model_response(model: AutoModelForCausalLM, tokenizer: AutoTokenizer, prompt: str, model_version: str, generation_config: GenerationConfig, text_length: int):
+
+    backup_generation_config_stage2 = GenerationConfig(
+            temperature=1,
+            top_p=0.6,
+            top_k=40,
+            num_beams=1,
+            bos_token_id=1,
+            eos_token_id=2,
+            pad_token_id=0,
+            max_new_tokens=1024,
+            min_new_tokens=1,
+            do_sample=True
+        )
+
+    backup_generation_config_stage3 = GenerationConfig(
+            top_k=5,
+            num_beams=1,
+            bos_token_id=1,
+            eos_token_id=2,
+            pad_token_id=0,
+            max_new_tokens=1024,
+            min_new_tokens=1,
+            penalty_alpha=0.3
+        )
+
+    backup_generation_config = [backup_generation_config_stage2, backup_generation_config_stage3]
+
+    generation = model.generate(**tokenizer(prompt, return_tensors="pt").to(model.device), generation_config=generation_config)[0]
+    if len(generation) > 2 * text_length:
+        stage = 0
+        while detect_degeneration(list(generation), model_version):
+            stage += 1
+            if stage > 2:
+                print("model degeneration cannot be avoided.")
+                break
+            generation = model.generate(**tokenizer(prompt, return_tensors="pt").to(model.device), generation_config=backup_generation_config[stage-1])[0]
+    response = tokenizer.decode(generation)
+    output = split_response(response, model_version)
+    return output
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument("--model_name_or_path", type=str, default="SakuraLLM/Sakura-13B-LNovel-v0.8", help="model huggingface id or local path.")
+    parser.add_argument("--use_gptq_model", action="store_true", help="whether your model is gptq quantized.")
+    parser.add_argument("--model_version", type=str, default="0.8", help="model version written on huggingface readme, now we have ['0.1', '0.4', '0.5', '0.7', '0.8']")
+    parser.add_argument("--data_path", type=str, default="", help="file path of the epub you want to translate.")
+    parser.add_argument("--data_folder", type=str, default="", help="folder path of the epubs you want to translate.")
+    parser.add_argument("--output_folder", type=str, default="", help="save folder path of the epubs model translated.")
+    parser.add_argument("--text_length", type=int, default=512, help="input max length in each inference.")
+    args = parser.parse_args()
+
+    if args.use_gptq_model:
+        from auto_gptq import AutoGPTQForCausalLM
+
+    generation_config = GenerationConfig(
+        temperature=0.1,
+        top_p=0.3,
+        top_k=40,
+        num_beams=1,
+        bos_token_id=1,
+        eos_token_id=2,
+        pad_token_id=0,
+        max_new_tokens=1024,
+        min_new_tokens=1,
+        do_sample=True
+    )
+
+    print("Loading model...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=False, trust_remote_code=True)
+
+    if args.use_gptq_model:
+        model = AutoGPTQForCausalLM.from_quantized(args.model_name_or_path, device="cuda:0", trust_remote_code=True)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path, device_map="auto", trust_remote_code=True)
+
+    print("Start translating...")
+    start = time.time()
+    
+    epub_list = []
+    save_list = []
+    if args.data_path:
+        epub_list.append(args.data_path)
+        save_list.append(os.path.join(args.output_folder, os.path.basename(args.data_path)))
+    if args.data_folder:
+        os.makedirs(args.output_folder, exist_ok=True)
+        for f in os.listdir(args.data_folder):
+            if f.endswith(".epub"):
+                epub_list.append(os.path.join(args.data_folder, f))
+                save_list.append(os.path.join(args.output_folder, f))
+    
+    for epub_path, save_path in zip(epub_list, save_list):
+        print(f"translating {epub_path}...")
+        start_epub = time.time()
+
+        if os.path.exists('./temp'):
+            shutil.rmtree('./temp')
+        with zipfile.ZipFile(epub_path, 'r') as f:
+            f.extractall('./temp')
+        
+        for html_path in find_all_htmls('./temp'):
+            print(f"\ttranslating {html_path}...")
+            start_html = time.time()
+
+            translated = ''
+            data_list, file_text = get_html_text_list(html_path, args.text_length)
+            if len(data_list) == 0:
+                    continue
+            for text, groups, pre_end in tqdm(data_list):
+                prompt = get_prompt(text, args.model_version)
+                output = get_model_response(model, tokenizer, prompt, args.model_version, generation_config, args.text_length)
+                texts = output.strip().split('\n')
+                if len(texts) < len(groups):
+                    texts += [''] * (len(groups) - len(texts))
+                else:
+                    texts = texts[:len(groups)-1] + ['<br/>'.join(texts[len(groups)-1:])]
+                for t, match in zip(texts, groups):
+                    t = match.group(0).replace(match.group(2), t)
+                    translated += file_text[pre_end:match.start()] + t
+                    pre_end = match.end()
+
+            translated += file_text[data_list[-1][1][-1].end():]
+            with open(html_path, 'w', encoding='utf-8') as f:
+                f.write(translated)
+            
+            end_html = time.time()
+            print(f"\t{html_path} translated, used time: ", end_html-start_html)
+        
+        with zipfile.ZipFile(save_path, 'w', zipfile.ZIP_DEFLATED) as f:
+            for file_path in glob.glob(f'./temp/**', recursive=True):
+                if not os.path.isdir(file_path):
+                    relative_path = os.path.relpath(file_path, './temp')
+                    f.write(file_path, relative_path)
+        shutil.rmtree('./temp')
+                    
+        end_epub = time.time()
+        print(f"{epub_path} translated, used time: ", end_epub-start_epub)
+
+    end = time.time()
+    print("translation completed, used time: ", end-start)
+
+if __name__ == "__main__":
+    main()