Add a model named "openai" and "--model_list" option (#392)

* Add a model named "openai" and "--model_list" option

* Remove old code

* Fix whitespace issues

README.md

@@ -25,8 +25,9 @@ Find more info here for using liteLLM: https://github.com/BerriAI/litellm/blob/m
    Or, just set environment variable `BBM_OPENAI_API_KEY` instead.
 - A sample book, `test_books/animal_farm.epub`, is provided for testing purposes.
 - The default underlying model is [GPT-3.5-turbo](https://openai.com/blog/introducing-chatgpt-and-whisper-apis), which is used by ChatGPT currently. Use `--model gpt4` to change the underlying model to `GPT4`.
-   If using `GPT4`, you can add `--use_context` to add a context paragraph to each passage sent to the model for translation (see below)
-- support DeepL model [DeepL Translator](https://rapidapi.com/splintPRO/api/dpl-translator) need pay to get the token use `--model deepl --deepl_key ${deepl_key}`
+  - Important to note that `gpt-4` is significantly more expensive than `gpt-4-turbo`, but to avoid bumping into rate limits, we automatically balance queries across `gpt-4-1106-preview`, `gpt-4`, `gpt-4-32k`, `gpt-4-0613`,`gpt-4-32k-0613`.
+    - If you want to use a specific model alias with OpenAI (eg `gpt-4-1106-preview` or `gpt-3.5-turbo-0125`), you can use `--model openai --model_list gpt-4-1106-preview,gpt-3.5-turbo-0125`. `--model_list` takes a comma-separated list of model aliases.
+  - If using `GPT4`, you can add `--use_context` to add a context paragraph to each passage sent to the model for translation (see below).- support DeepL model [DeepL Translator](https://rapidapi.com/splintPRO/api/dpl-translator) need pay to get the token use `--model deepl --deepl_key ${deepl_key}`
 - support DeepL free model `--model deeplfree`
 - support Google [Gemini](https://makersuite.google.com/app/apikey) model `--model gemini --gemini_key ${gemini_key}`
 - Support [Claude](https://console.anthropic.com/docs) model, use `--model claude --claude_key ${claude_key}`
@@ -83,10 +84,15 @@ export OPENAI_API_KEY=${your_api_key}
 # Use the GPT-4 model with context to Japanese
 python3 make_book.py --book_name test_books/animal_farm.epub --model gpt4 --use_context --language ja
 
+# Use a specific OpenAI model alias
+python3 make_book.py --book_name test_books/animal_farm.epub --model openai --model_list gpt-4-1106-preview --openai_key ${openai_key}
+
+# Use a specific list of OpenAI model aliases
+python3 make_book.py --book_name test_books/animal_farm.epub --model openai --model_list gpt-4-1106-preview,gpt-4-0125-preview,gpt-3.5-turbo-0125 --openai_key ${openai_key}
+
 # Use the DeepL model with Japanese
 python3 make_book.py --book_name test_books/animal_farm.epub --model deepl --deepl_key ${deepl_key} --language ja
 
-
 # Use the Claude model with Japanese
 python3 make_book.py --book_name test_books/animal_farm.epub --model claude --claude_key ${claude_key} --language ja
 

book_maker/cli.py

@@ -275,6 +275,12 @@ So you are close to reaching the limit. You have to choose your own value, there
         default=-1,
         help="merge multiple paragraphs into one block, may increase accuracy and speed up the process, but disturb the original format, must be used with `--single_translate`",
     )
+    parser.add_argument(
+        "--model_list",
+        type=str,
+        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`)",
+    )
 
     options = parser.parse_args()
 
@@ -290,7 +296,7 @@ So you are close to reaching the limit. You have to choose your own value, there
     translate_model = MODEL_DICT.get(options.model)
     assert translate_model is not None, "unsupported model"
     API_KEY = ""
-    if options.model in ["chatgptapi", "gpt4"]:
+    if options.model in ["openai", "chatgptapi", "gpt4"]:
         if OPENAI_API_KEY := (
             options.openai_key
             or env.get(
@@ -402,6 +408,14 @@ So you are close to reaching the limit. You have to choose your own value, there
         if not options.api_base:
             raise ValueError("`api_base` must be provided when using `deployment_id`")
         e.translate_model.set_deployment_id(options.deployment_id)
+    if options.model == "openai":
+        # Currently only supports `openai` when you also have --model_list set
+        if options.model_list:
+            e.translate_model.set_model_list(options.model_list.split(","))
+        else:
+            raise ValueError(
+                "When using `openai` model, you must also provide `--model_list`. For default model sets use `--model chatgptapi` or `--model gpt4`",
+            )
     # TODO refactor, quick fix for gpt4 model
     if options.model == "chatgptapi":
         e.translate_model.set_gpt35_models()

book_maker/translator/__init__.py

@@ -9,12 +9,13 @@ from book_maker.translator.tencent_transmart_translator import TencentTranSmart
 from book_maker.translator.custom_api_translator import CustomAPI
 
 MODEL_DICT = {
+    "openai": ChatGPTAPI,
     "chatgptapi": ChatGPTAPI,
+    "gpt4": ChatGPTAPI,
     "google": Google,
     "caiyun": Caiyun,
     "deepl": DeepL,
     "deeplfree": DeepLFree,
-    "gpt4": ChatGPTAPI,
     "claude": Claude,
     "gemini": Gemini,
     "tencentransmart": TencentTranSmart,

book_maker/translator/chatgptapi_translator.py

@@ -330,3 +330,8 @@ class ChatGPTAPI(Base):
             model_list = list(set(my_model_list) & set(GPT4_MODEL_LIST))
             print(f"Using model list {model_list}")
             self.model_list = cycle(model_list)
+
+    def set_model_list(self, model_list):
+        model_list = list(set(model_list))
+        print(f"Using model list {model_list}")
+        self.model_list = cycle(model_list)