chore: update preset models (#159)

* chore: update preset models * fix: add chat model * fix: ebturbo pro
baidubce · Dec 28, 2023 · 7d50cfe · 7d50cfe
1 parent 21b2059
commit 7d50cfe
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 0 deletions.
diff --git a/src/qianfan/resources/images/text2image.py b/src/qianfan/resources/images/text2image.py
@@ -57,6 +57,9 @@ def _supported_models(cls) -> Dict[str, QfLLMInfo]:
                     "steps",
                     "sampler_index",
                     "user_id",
+                    "seed",
+                    "cfg_scale",
+                    "style",
                 },
             ),
             UNSPECIFIED_MODEL: QfLLMInfo(

diff --git a/src/qianfan/resources/llm/chat_completion.py b/src/qianfan/resources/llm/chat_completion.py
@@ -113,6 +113,20 @@ def _supported_models(cls) -> Dict[str, QfLLMInfo]:
                     "user_id",
                 },
             ),
+            "ERNIE-Bot-turbo-pro": QfLLMInfo(
+                endpoint="/chat/eb_turbo_pro",
+                required_keys={"messages"},
+                optional_keys={
+                    "stream",
+                    "temperature",
+                    "top_p",
+                    "penalty_score",
+                    "user_id",
+                    "tools",
+                    "tool_choice",
+                    "system",
+                },
+            ),
             "ERNIE-Bot-turbo-AI": QfLLMInfo(
                 endpoint="/chat/ai_apaas",
                 required_keys={"messages"},

diff --git a/src/qianfan/resources/llm/completion.py b/src/qianfan/resources/llm/completion.py
@@ -113,6 +113,20 @@ def _supported_models(cls) -> Dict[str, QfLLMInfo]:
                     "user_id",
                 },
             ),
+            "ERNIE-Bot-turbo-pro": QfLLMInfo(
+                endpoint="/chat/eb_turbo_pro",
+                required_keys={"messages"},
+                optional_keys={
+                    "stream",
+                    "temperature",
+                    "top_p",
+                    "penalty_score",
+                    "user_id",
+                    "tools",
+                    "tool_choice",
+                    "system",
+                },
+            ),
             "EB-turbo-AppBuilder": QfLLMInfo(
                 endpoint="/chat/ai_apaas",
                 required_keys={"messages"},
@@ -291,6 +305,19 @@ def _supported_models(cls) -> Dict[str, QfLLMInfo]:
                     "stop",
                 },
             ),
+            "Yi-34B-Chat": QfLLMInfo(
+                endpoint="/chat/yi_34b_chat",
+                required_keys={"messages"},
+                optional_keys={
+                    "stream",
+                    "user_id",
+                    "temperature",
+                    "top_k",
+                    "top_p",
+                    "penalty_score",
+                    "stop",
+                },
+            ),
             UNSPECIFIED_MODEL: QfLLMInfo(
                 endpoint="",
                 required_keys={"prompt"},

diff --git a/src/qianfan/resources/llm/embedding.py b/src/qianfan/resources/llm/embedding.py
@@ -62,6 +62,11 @@ def _supported_models(cls) -> Dict[str, QfLLMInfo]:
                 required_keys={"input"},
                 optional_keys={"user_id"},
             ),
+            "tao-8k": QfLLMInfo(
+                endpoint="/embeddings/tao_8k",
+                required_keys={"input"},
+                optional_keys={"user_id"},
+            ),
             UNSPECIFIED_MODEL: QfLLMInfo(
                 endpoint="", required_keys={"input"}, optional_keys=set()
             ),