diff --git a/book_maker/__main__.py b/book_maker/__main__.py index 6fa7a56..2aa7140 100644 --- a/book_maker/__main__.py +++ b/book_maker/__main__.py @@ -1 +1,4 @@ from cli import main + +if __name__ == "__main__": + main() diff --git a/book_maker/loader/base_loader.py b/book_maker/loader/base_loader.py index 46986ff..ed6040d 100644 --- a/book_maker/loader/base_loader.py +++ b/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() diff --git a/book_maker/loader/epub_loader.py b/book_maker/loader/epub_loader.py index ee1127d..4f2df22 100644 --- a/book_maker/loader/epub_loader.py +++ b/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") diff --git a/book_maker/translator/base_translator.py b/book_maker/translator/base_translator.py index 2e5fe14..57bb2c1 100644 --- a/book_maker/translator/base_translator.py +++ b/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): diff --git a/book_maker/translator/chatgptapi_translator.py b/book_maker/translator/chatgptapi_translator.py index 6ab719d..49d75d9 100644 --- a/book_maker/translator/chatgptapi_translator.py +++ b/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=[ diff --git a/book_maker/translator/deepl_translator.py b/book_maker/translator/deepl_translator.py index eea19c0..b692769 100644 --- a/book_maker/translator/deepl_translator.py +++ b/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 diff --git a/book_maker/translator/gpt3_translator.py b/book_maker/translator/gpt3_translator.py index f277aa3..1794063 100644 --- a/book_maker/translator/gpt3_translator.py +++ b/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: