From 13fef1e5d3a263644420e882fabb03a21ca4c894 Mon Sep 17 00:00:00 2001 From: Hunsmore Date: Thu, 31 Aug 2023 15:38:55 +0800 Subject: [PATCH] add bloomz_7b, llama-2-7b, llama-2-13b, llama-2-70b to ErnieBotChat (#10024) - Description: Add bloomz_7b, llama-2-7b, llama-2-13b, llama-2-70b to ErnieBotChat, which only supported ERNIE-Bot-turbo and ERNIE-Bot. - Issue: #10022, - Dependencies: no extra dependencies --------- Co-authored-by: hetianfeng --- libs/langchain/langchain/chat_models/ernie.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/libs/langchain/langchain/chat_models/ernie.py b/libs/langchain/langchain/chat_models/ernie.py index d3fdce5c31f..367341c11f3 100644 --- a/libs/langchain/langchain/chat_models/ernie.py +++ b/libs/langchain/langchain/chat_models/ernie.py @@ -98,12 +98,19 @@ class ErnieBotChat(BaseChatModel): def _chat(self, payload: object) -> dict: base_url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat" - if self.model_name == "ERNIE-Bot-turbo": - url = f"{base_url}/eb-instant" - elif self.model_name == "ERNIE-Bot": - url = f"{base_url}/completions" + model_paths = { + "ERNIE-Bot-turbo": "eb-instant", + "ERNIE-Bot": "completions", + "BLOOMZ-7B": "bloomz_7b1", + "Llama-2-7b-chat": "llama_2_7b", + "Llama-2-13b-chat": "llama_2_13b", + "Llama-2-70b-chat": "llama_2_70b", + } + if self.model_name in model_paths: + url = f"{base_url}/{model_paths[self.model_name]}" else: raise ValueError(f"Got unknown model_name {self.model_name}") + resp = requests.post( url, timeout=self.request_timeout,