support: gpt4o (#425)

* support: gpt4o * fix: remove wrong code * fix: fix ci error
yihong0618 · Aug 24, 2024 · 6912206 · 6912206
1 parent 9e4e7b5
commit 6912206
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 1 deletion.
diff --git a/book_maker/cli.py b/book_maker/cli.py
@@ -331,7 +331,7 @@ def main():
     translate_model = MODEL_DICT.get(options.model)
     assert translate_model is not None, "unsupported model"
     API_KEY = ""
-    if options.model in ["openai", "chatgptapi", "gpt4", "gpt4omini"]:
+    if options.model in ["openai", "chatgptapi", "gpt4", "gpt4omini", "gpt4o"]:
         if OPENAI_API_KEY := (
             options.openai_key
             or env.get(
@@ -449,6 +449,7 @@ def main():
             "chatgptapi",
             "gpt4",
             "gpt4omini",
+            "gpt4o",
         ], "only support chatgptapi for deployment_id"
         if not options.api_base:
             raise ValueError("`api_base` must be provided when using `deployment_id`")
@@ -471,6 +472,8 @@ def main():
         e.translate_model.set_gpt4_models()
     if options.model == "gpt4omini":
         e.translate_model.set_gpt4omini_models()
+    if options.model == "gpt4o":
+        e.translate_model.set_gpt4o_models()
     if options.block_size > 0:
         e.block_size = options.block_size
     if options.batch_flag:

diff --git a/book_maker/translator/__init__.py b/book_maker/translator/__init__.py
@@ -14,6 +14,7 @@
     "chatgptapi": ChatGPTAPI,
     "gpt4": ChatGPTAPI,
     "gpt4omini": ChatGPTAPI,
+    "gpt4o": ChatGPTAPI,
     "google": Google,
     "caiyun": Caiyun,
     "deepl": DeepL,

diff --git a/book_maker/translator/chatgptapi_translator.py b/book_maker/translator/chatgptapi_translator.py
@@ -42,6 +42,12 @@
     "gpt-4o-mini",
     "gpt-4o-mini-2024-07-18",
 ]
+GPT4o_MODEL_LIST = [
+    "gpt-4o",
+    "gpt-4o-2024-05-13",
+    "gpt-4o-2024-08-06",
+    "chatgpt-4o-latest",
+]
 
 
 class ChatGPTAPI(Base):
@@ -404,6 +410,18 @@ def set_gpt4omini_models(self):
             print(f"Using model list {model_list}")
             self.model_list = cycle(model_list)
 
+    def set_gpt4o_models(self):
+        # for issue #375 azure can not use model list
+        if self.deployment_id:
+            self.model_list = cycle(["gpt-4o"])
+        else:
+            my_model_list = [
+                i["id"] for i in self.openai_client.models.list().model_dump()["data"]
+            ]
+            model_list = list(set(my_model_list) & set(GPT4o_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}")