AVB/bedrock-access-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support of Mistral Large

Browse Source
This commit is contained in:
Aiden Dai
2024-04-03 11:46:10 +08:00
parent 268c5ef0f1
commit 02095c030f
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/api/models/bedrock.py
View File

@@ -49,6 +49,7 @@ SUPPORTED_BEDROCK_MODELS = {
"meta.llama2-70b-chat-v1": "Llama 2 Chat 70B", "meta.llama2-70b-chat-v1": "Llama 2 Chat 70B",
"mistral.mistral-7b-instruct-v0:2": "Mistral 7B Instruct", "mistral.mistral-7b-instruct-v0:2": "Mistral 7B Instruct",
"mistral.mixtral-8x7b-instruct-v0:1": "Mixtral 8x7B Instruct", "mistral.mixtral-8x7b-instruct-v0:1": "Mixtral 8x7B Instruct",
"mistral.mistral-large-2402-v1:0": "Mistral Large",
} }
SUPPORTED_BEDROCK_EMBEDDING_MODELS = { SUPPORTED_BEDROCK_EMBEDDING_MODELS = {
@@ -657,7 +658,7 @@ def get_model(model_id: str) -> BedrockModel:
return ClaudeModel() return ClaudeModel()
elif model_name in ["Llama 2 Chat 13B", "Llama 2 Chat 70B"]: elif model_name in ["Llama 2 Chat 13B", "Llama 2 Chat 70B"]:
return Llama2Model() return Llama2Model()
elif model_name in ["Mistral 7B Instruct", "Mixtral 8x7B Instruct"]: elif model_name in ["Mistral 7B Instruct", "Mixtral 8x7B Instruct", "Mistral Large"]:
return MistralModel() return MistralModel()
else: else:
logger.error("Unsupported model id " + model_id) logger.error("Unsupported model id " + model_id)