From 75809803841c940a2e41ab4514fc5f116294c23c Mon Sep 17 00:00:00 2001 From: "Lumberbot (aka Jack)" <39504233+meeseeksmachine@users.noreply.github.com> Date: Wed, 8 May 2024 23:36:33 +0200 Subject: [PATCH] Backport PR #778: Add Titan embedding model v2 (#779) Co-authored-by: Sanjiv Das --- .../jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py index 3b187e1c2..3dea152e7 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py @@ -98,10 +98,11 @@ class HfHubEmbeddingsProvider(BaseEmbeddingsProvider, HuggingFaceHubEmbeddings): registry = True +# See model ID list here: https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html class BedrockEmbeddingsProvider(BaseEmbeddingsProvider, BedrockEmbeddings): id = "bedrock" name = "Bedrock" - models = ["amazon.titan-embed-text-v1"] + models = ["amazon.titan-embed-text-v1", "amazon.titan-embed-text-v2:0"] model_id_key = "model_id" pypi_package_deps = ["boto3"] auth_strategy = AwsAuthStrategy()