Minor improvements for loader and translator (#103)

* Minor improvements for loader and translator

* use rotate key

* fix: typo

---------

Co-authored-by: yihong0618 <zouzou0208@gmail.com>

book_maker/__main__.py

@@ -1 +1,4 @@
 from cli import main
+
+if __name__ == "__main__":
+    main()

book_maker/loader/base_loader.py

@@ -1,20 +1,7 @@
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 
 
-class BaseBookLoader:
-    def __init__(
-        self,
-        epub_name,
-        model,
-        key,
-        resume,
-        language,
-        model_api_base=None,
-        is_test=False,
-        test_num=5,
-    ):
-        pass
-
+class BaseBookLoader(ABC):
     @staticmethod
     def _is_special_text(text):
         return text.isdigit() or text.isspace()

book_maker/loader/epub_loader.py

@@ -1,4 +1,3 @@
-import argparse
 import os
 import pickle
 import sys
@@ -33,7 +32,7 @@ class EPUBBookLoader(BaseBookLoader):
 
         try:
             self.origin_book = epub.read_epub(self.epub_name)
-        except:
+        except Exception:
             # tricky for #71 if you don't know why please check the issue and ignore this
             # when upstream change will TODO fix this
             def _load_spine(self):
@@ -119,7 +118,7 @@ class EPUBBookLoader(BaseBookLoader):
         try:
             with open(self.bin_path, "rb") as f:
                 self.p_to_save = pickle.load(f)
-        except:
+        except Exception:
             raise Exception("can not load resume file")
 
     def _save_temp_book(self):
@@ -164,5 +163,5 @@ class EPUBBookLoader(BaseBookLoader):
         try:
             with open(self.bin_path, "wb") as f:
                 pickle.dump(self.p_to_save, f)
-        except:
+        except Exception:
             raise Exception("can not save resume file")

book_maker/translator/base_translator.py

@@ -1,17 +1,15 @@
-from abc import abstractmethod
+import itertools
+from abc import ABC, abstractmethod
 
 
-class Base:
-    def __init__(self, key, language, api_base=None):
-        self.key = key
+class Base(ABC):
+    def __init__(self, key, language):
+        self.keys = itertools.cycle(key.split(","))
         self.language = language
-        self.current_key_index = 0
 
-    def get_key(self, key_str):
-        keys = key_str.split(",")
-        key = keys[self.current_key_index]
-        self.current_key_index = (self.current_key_index + 1) % len(keys)
-        return key
+    @abstractmethod
+    def rotate_key(self):
+        pass
 
     @abstractmethod
     def translate(self, text):

book_maker/translator/chatgptapi_translator.py

@@ -7,15 +7,17 @@ from .base_translator import Base
 
 class ChatGPTAPI(Base):
     def __init__(self, key, language, api_base=None):
-        super().__init__(key, language, api_base=api_base)
-        self.key = key
-        self.language = language
+        super().__init__(key, language)
+        self.key_len = len(key.split(","))
         if api_base:
             openai.api_base = api_base
 
+    def rotate_key(self):
+        openai.api_key = next(self.keys)
+
     def translate(self, text):
         print(text)
-        openai.api_key = self.get_key(self.key)
+        self.rotate_key()
         try:
             completion = openai.ChatCompletion.create(
                 model="gpt-3.5-turbo",
@@ -36,11 +38,10 @@ class ChatGPTAPI(Base):
             )
         except Exception as e:
             # TIME LIMIT for open api please pay
-            key_len = self.key.count(",") + 1
-            sleep_time = int(60 / key_len)
+            sleep_time = int(60 / self.key_len)
             time.sleep(sleep_time)
             print(e, f"will sleep  {sleep_time} seconds")
-            openai.api_key = self.get_key(self.key)
+            self.rotate_key()
             completion = openai.ChatCompletion.create(
                 model="gpt-3.5-turbo",
                 messages=[

book_maker/translator/deepl_translator.py

@@ -2,8 +2,4 @@ from .base_translator import Base
 
 
 class DeepL(Base):
-    def __init__(self, session, key, api_base=None):
-        super().__init__(session, key, api_base=api_base)
-
-    def translate(self, text):
-        return super().translate(text)
+    pass

book_maker/translator/gpt3_translator.py

@@ -7,7 +7,6 @@ from .base_translator import Base
 class GPT3(Base):
     def __init__(self, key, language, api_base=None):
         super().__init__(key, language)
-        self.api_key = key
         self.api_url = (
             f"{api_base}v1/completions"
             if api_base
@@ -27,9 +26,12 @@ class GPT3(Base):
         self.session = requests.session()
         self.language = language
 
+    def rotate_key(self):
+        self.headers["Authorization"] = f"Bearer {next(self.keys)}"
+
     def translate(self, text):
         print(text)
-        self.headers["Authorization"] = f"Bearer {self.get_key(self.api_key)}"
+        self.rotate_key()
         self.data["prompt"] = f"Please help me to translate，`{text}` to {self.language}"
         r = self.session.post(self.api_url, headers=self.headers, json=self.data)
         if not r.ok: