[Feature] Add an attribute api_key into TurboMindAPIModel default None (#1475)

Co-authored-by: Maxime <maximeshe@163.com>
Add an attribute api_key into TurboMindAPIModel default None then we can set the api_key while using lmdeploy to deploy the llm model
This commit is contained in:
Maxime SHE 2024-09-05 17:51:16 +08:00 committed by GitHub
parent 6c9cd9a260
commit 45efdc994d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 4 additions and 1 deletions

View File

@ -28,6 +28,7 @@ internlm_chat_20b = dict(
type=TurboMindAPIModel,
abbr='internlm-chat-20b-turbomind',
api_addr='http://0.0.0.0:23333',
api_key='internlm-chat-20b', # api_key
max_out_len=100,
max_seq_len=2048,
batch_size=8,
@ -40,6 +41,7 @@ internlm_chat_7b = dict(
type=TurboMindAPIModel,
abbr='internlm-chat-7b-turbomind',
api_addr='http://0.0.0.0:23333',
api_key='interlm-chat-7b', # api_key
max_out_len=100,
max_seq_len=2048,
batch_size=16,

View File

@ -40,6 +40,7 @@ class TurboMindAPIModel(BaseModel):
def __init__(self,
api_addr: str = 'http://0.0.0.0:23333',
api_key: str | None = None,
max_seq_len: int = 2048,
meta_template: Optional[Dict] = None,
end_str: Optional[str] = None,
@ -48,7 +49,7 @@ class TurboMindAPIModel(BaseModel):
max_seq_len=max_seq_len,
meta_template=meta_template)
from lmdeploy.serve.openai.api_client import APIClient
self.chatbot = APIClient(api_addr)
self.chatbot = APIClient(api_addr, api_key)
self.model_name = self.chatbot.available_models[0]
self.logger = get_logger()
self.template_parser = LMTemplateParser(meta_template)