supprt openai batch api (#423)

* feat: support batch api

book_maker/cli.py

@@ -304,6 +304,18 @@ def main():
         dest="model_list",
         help="Rather than using our preset lists of models, specify exactly the models you want as a comma separated list `gpt-4-32k,gpt-3.5-turbo-0125` (Currently only supports: `openai`)",
     )
+    parser.add_argument(
+        "--batch",
+        dest="batch_flag",
+        action="store_true",
+        help="Enable batch translation using ChatGPT's batch API for improved efficiency",
+    )
+    parser.add_argument(
+        "--batch-use",
+        dest="batch_use_flag",
+        action="store_true",
+        help="Use pre-generated batch translations to create files. Run with --batch first before using this option",
+    )
 
     options = parser.parse_args()
 
@@ -461,6 +473,10 @@ def main():
         e.translate_model.set_gpt4omini_models()
     if options.block_size > 0:
         e.block_size = options.block_size
+    if options.batch_flag:
+        e.batch_flag = options.batch_flag
+    if options.batch_use_flag:
+        e.batch_use_flag = options.batch_use_flag
 
     e.make_bilingual_book()
 

book_maker/config.py

@@ -0,0 +1,8 @@
+config = {
+    "translator": {
+        "chatgptapi": {
+            "context_paragraph_limit": 3,
+            "batch_context_update_interval": 50,
+        }
+    },
+}

book_maker/loader/epub_loader.py

@@ -2,6 +2,7 @@
 import pickle
 import string
 import sys
+import time
 from copy import copy
 from pathlib import Path
 
@@ -65,6 +66,8 @@ def __init__(
         self.only_filelist = ""
         self.single_translate = single_translate
         self.block_size = -1
+        self.batch_use_flag = False
+        self.batch_flag = False
 
         # monkey patch for # 173
         def _write_items_patch(obj):
@@ -142,11 +145,18 @@ def _process_paragraph(self, p, new_p, index, p_to_save_len):
         if self.resume and index < p_to_save_len:
             p.string = self.p_to_save[index]
         else:
+            t_text = ""
+            if self.batch_flag:
+                self.translate_model.add_to_batch_translate_queue(index, new_p.text)
+            elif self.batch_use_flag:
+                t_text = self.translate_model.batch_translate(index)
+            else:
+                t_text = self.translate_model.translate(new_p.text)
             if type(p) == NavigableString:
-                new_p = self.translate_model.translate(new_p.text)
+                new_p = t_text
                 self.p_to_save.append(new_p)
             else:
-                new_p.string = self.translate_model.translate(new_p.text)
+                new_p.string = t_text
                 self.p_to_save.append(new_p.text)
 
         self.helper.insert_trans(
@@ -456,13 +466,26 @@ def process_item(
 
         return index
 
+    def batch_init_then_wait(self):
+        name, _ = os.path.splitext(self.epub_name)
+        if self.batch_flag or self.batch_use_flag:
+            self.translate_model.batch_init(name)
+            if self.batch_use_flag:
+                start_time = time.time()
+                while not self.translate_model.is_completed_batch():
+                    print("Batch translation is not completed yet")
+                    time.sleep(2)
+                    if time.time() - start_time > 300:  # 5 minutes
+                        raise Exception("Batch translation timed out after 5 minutes")
+
     def make_bilingual_book(self):
         self.helper = EPUBBookLoaderHelper(
             self.translate_model,
             self.accumulated_num,
             self.translation_style,
             self.context_flag,
         )
+        self.batch_init_then_wait()
         new_book = self._make_new_book(self.origin_book)
         all_items = list(self.origin_book.get_items())
         trans_taglist = self.translate_tags.split(",")
@@ -520,7 +543,10 @@ def make_bilingual_book(self):
                     name, _ = os.path.splitext(self.epub_name)
                     epub.write_epub(f"{name}_bilingual.epub", new_book, {})
             name, _ = os.path.splitext(self.epub_name)
-            epub.write_epub(f"{name}_bilingual.epub", new_book, {})
+            if self.batch_flag:
+                self.translate_model.batch()
+            else:
+                epub.write_epub(f"{name}_bilingual.epub", new_book, {})
             if self.accumulated_num == 1:
                 pbar.close()
         except (KeyboardInterrupt, Exception) as e: