Feature request: Option to disable cross encoder models

bin/config.ts

@@ -20,10 +20,12 @@ export function getConfig(): SystemConfig {
     },
     llms: {
       // sagemaker: [SupportedSageMakerModels.FalconLite]
+      enableSagemakerModels: false,
       sagemaker: [],
     },
     rag: {
       enabled: false,
+      enableEmbeddingModelsViaSagemaker: false,
       engines: {
         aurora: {
           enabled: false,
@@ -42,11 +44,13 @@ export function getConfig(): SystemConfig {
           provider: "sagemaker",
           name: "intfloat/multilingual-e5-large",
           dimensions: 1024,
+          default: false,
         },
         {
           provider: "sagemaker",
           name: "sentence-transformers/all-MiniLM-L6-v2",
           dimensions: 384,
+          default: false,
         },
         {
           provider: "bedrock",
@@ -58,8 +62,10 @@ export function getConfig(): SystemConfig {
           provider: "openai",
           name: "text-embedding-ada-002",
           dimensions: 1536,
+          default: false,
         },
       ],
+      crossEncodingEnabled: false,
       crossEncoderModels: [
         {
           provider: "sagemaker",

cli/magic-config.ts

@@ -22,21 +22,25 @@ const embeddingModels = [
     provider: "sagemaker",
     name: "intfloat/multilingual-e5-large",
     dimensions: 1024,
+    default: false,
   },
   {
     provider: "sagemaker",
     name: "sentence-transformers/all-MiniLM-L6-v2",
     dimensions: 384,
+    default: false,
   },
   {
     provider: "bedrock",
     name: "amazon.titan-embed-text-v1",
     dimensions: 1536,
+    default: false,
   },
   {
     provider: "openai",
     name: "text-embedding-ada-002",
     dimensions: 1536,
+    default: false,
   },
 ];
 
@@ -203,6 +207,15 @@ async function processCreateOptions(options: any): Promise<void> {
       message: "Do you want to enable RAG",
       initial: options.enableRag || false,
     },
+    {
+      type: "confirm",
+      name: "enableEmbeddingModelsViaSagemaker",
+      message: "Do you want to enable embedding models via SageMaker?",
+      initial: options.enableEmbeddingModelsViaSagemaker || false,
+      skip(): boolean {
+        return !(this as any).state.answers.enableRag;
+      },
+    },
     {
       type: "multiselect",
       name: "ragsToEnable",
@@ -328,7 +341,7 @@ async function processCreateOptions(options: any): Promise<void> {
       choices: embeddingModels.map((m) => ({ name: m.name, value: m })),
       initial: options.defaultEmbedding || undefined,
       skip(): boolean {
-        return !(this as any).state.answers.enableRag;
+        return !answers.enableRag;
       },
     },
   ];
@@ -349,10 +362,13 @@ async function processCreateOptions(options: any): Promise<void> {
         }
       : undefined,
     llms: {
+      enableSagemakerModels: answers.enableSagemakerModels,
       sagemaker: answers.sagemakerModels,
     },
     rag: {
       enabled: answers.enableRag,
+      enableEmbeddingModelsViaSagemaker:
+        answers.enableEmbeddingModelsViaSagemaker,
       engines: {
         aurora: {
           enabled: answers.ragsToEnable.includes("aurora"),
@@ -367,28 +383,44 @@ async function processCreateOptions(options: any): Promise<void> {
           enterprise: false,
         },
       },
+      crossEncodingEnabled: answers.enableEmbeddingModelsViaSagemaker,
       embeddingsModels: [{}],
       crossEncoderModels: [{}],
     },
   };
 
+  if (
+    answers.enableEmbeddingModelsViaSagemaker &&
+    answers.enableSagemakerModels
+  ) {
+    config.rag.crossEncoderModels[0] = {
+      provider: "sagemaker",
+      name: "cross-encoder/ms-marco-MiniLM-L-12-v2",
+      default: true,
+    };
+  } else {
+    config.rag.crossEncoderModels[0] = {
+      provider: "None",
+      name: "None",
+      default: true,
+    };
+  }
+  if (!config.rag.enableEmbeddingModelsViaSagemaker) {
+    config.rag.embeddingsModels = embeddingModels.filter(model => model.provider !== "sagemaker");
+  } else {
+    config.rag.embeddingsModels = embeddingModels;
+  }
   // If we have not enabled rag the default embedding is set to the first model
   if (!answers.enableRag) {
-    models.defaultEmbedding = embeddingModels[0].name;
+    (config.rag.embeddingsModels[0] as any).default = true;
+  } else {
+    config.rag.embeddingsModels.forEach((m: any) => {
+      if (m.name === models.defaultEmbedding) {
+        m.default = true;
+      }
+    });
   }
 
-  config.rag.crossEncoderModels[0] = {
-    provider: "sagemaker",
-    name: "cross-encoder/ms-marco-MiniLM-L-12-v2",
-    default: true,
-  };
-  config.rag.embeddingsModels = embeddingModels;
-  config.rag.embeddingsModels.forEach((m: any) => {
-    if (m.name === models.defaultEmbedding) {
-      m.default = true;
-    }
-  });
-
   config.rag.engines.kendra.createIndex =
     answers.ragsToEnable.includes("kendra");
   config.rag.engines.kendra.enabled =

lib/aws-genai-llm-chatbot-stack.ts

@@ -151,10 +151,8 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
       identityPool: authentication.identityPool,
       api: chatBotApi,
       chatbotFilesBucket: chatBotApi.filesBucket,
-      crossEncodersEnabled:
-        typeof ragEngines?.sageMakerRagModels?.model !== "undefined",
-      sagemakerEmbeddingsEnabled:
-        typeof ragEngines?.sageMakerRagModels?.model !== "undefined",
+      crossEncodersEnabled: props.config.rag.crossEncodingEnabled,
+      sagemakerEmbeddingsEnabled: props.config.rag.enableEmbeddingModelsViaSagemaker,
     });
 
     /**
@@ -283,10 +281,7 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
         ]
       );
 
-      if (
-        props.config.rag.engines.aurora.enabled ||
-        props.config.rag.engines.opensearch.enabled
-      ) {
+      if (props.config.llms.enableSagemakerModels) {
         NagSuppressions.addResourceSuppressionsByPath(
           this,
           [

lib/chatbot-api/functions/api-handler/routes/cross_encoders.py

@@ -30,14 +30,19 @@ def models():
 @tracer.capture_method
 def cross_encoders(input: dict):
     request = CrossEncodersRequest(**input)
-    selected_model = genai_core.cross_encoder.get_cross_encoder_model(
-        request.provider, request.model
-    )
-
-    if selected_model is None:
-        raise genai_core.types.CommonError("Model not found")
-
-    ret_value = genai_core.cross_encoder.rank_passages(
-        selected_model, request.reference, request.passages
-    )
-    return [{"score": v, "passage": p} for v, p in zip(ret_value, request.passages)]
+    config = genai_core.parameters.get_config()
+    crossEncodingEnabled = config["rag"]["crossEncodingEnabled"]
+    if (crossEncodingEnabled):
+        selected_model = genai_core.cross_encoder.get_cross_encoder_model(
+            request.provider, request.model
+        )
+
+        if selected_model is None:
+            raise genai_core.types.CommonError("Model not found")
+
+        ret_value = genai_core.cross_encoder.rank_passages(
+            selected_model, request.reference, request.passages
+        )
+        return [{"score": v, "passage": p} for v, p in zip(ret_value, request.passages)]
+
+    return [{"score": 0, "passage": p} for p in request.passages]

lib/rag-engines/data-import/index.ts

@@ -131,7 +131,7 @@ export class DataImport extends Construct {
         processingBucket,
         auroraDatabase: props.auroraDatabase,
         ragDynamoDBTables: props.ragDynamoDBTables,
-        sageMakerRagModelsEndpoint: props.sageMakerRagModels?.model.endpoint,
+        sageMakerRagModelsEndpoint: props.sageMakerRagModels?.model?.endpoint,
         openSearchVector: props.openSearchVector,
       }
     );

lib/rag-engines/index.ts

@@ -41,10 +41,7 @@ export class RagEngines extends Construct {
     const tables = new RagDynamoDBTables(this, "RagDynamoDBTables");
 
     let sageMakerRagModels: SageMakerRagModels | null = null;
-    if (
-      props.config.rag.engines.aurora.enabled ||
-      props.config.rag.engines.opensearch.enabled
-    ) {
+    if (props.config.llms.enableSagemakerModels) {
       sageMakerRagModels = new SageMakerRagModels(this, "SageMaker", {
         shared: props.shared,
         config: props.config,

lib/rag-engines/sagemaker-rag-models/index.ts

@@ -25,20 +25,22 @@ export class SageMakerRagModels extends Construct {
       .filter((c) => c.provider === "sagemaker")
       .map((c) => c.name);
 
-    const model = new SageMakerModel(this, "Model", {
-      vpc: props.shared.vpc,
-      region: cdk.Aws.REGION,
-      model: {
-        type: DeploymentType.CustomInferenceScript,
-        modelId: [
-          ...sageMakerEmbeddingsModelIds,
-          ...sageMakerCrossEncoderModelIds,
-        ],
-        codeFolder: path.join(__dirname, "./model"),
-        instanceType: "ml.g4dn.xlarge",
-      },
-    });
+    if (sageMakerEmbeddingsModelIds?.length > 0 || sageMakerCrossEncoderModelIds?.length > 0) {
+      const model = new SageMakerModel(this, "Model", {
+        vpc: props.shared.vpc,
+        region: cdk.Aws.REGION,
+        model: {
+          type: DeploymentType.CustomInferenceScript,
+          modelId: [
+            ...sageMakerEmbeddingsModelIds,
+            ...sageMakerCrossEncoderModelIds,
+          ],
+          codeFolder: path.join(__dirname, "./model"),
+          instanceType: "ml.g4dn.xlarge",
+        },
+      });
 
-    this.model = model;
+      this.model = model;
+    }
   }
 }

lib/shared/layers/python-sdk/python/genai_core/aurora/query.py

@@ -19,6 +19,7 @@ def query_workspace_aurora(
     full_response: bool,
     threshold: int = 0,
 ):
+    config = genai_core.parameters.get_config()
     table_name = sql.Identifier(workspace_id.replace("-", ""))
     embeddings_model_provider = workspace["embeddings_model_provider"]
     embeddings_model_name = workspace["embeddings_model_name"]
@@ -37,13 +38,6 @@ def query_workspace_aurora(
     if selected_model is None:
         raise genai_core.types.CommonError("Embeddings model not found")
 
-    cross_encoder_model = genai_core.cross_encoder.get_cross_encoder_model(
-        cross_encoder_model_provider, cross_encoder_model_name
-    )
-
-    if cross_encoder_model is None:
-        raise genai_core.types.CommonError("Cross encoder model not found")
-
     query_embeddings = genai_core.embeddings.generate_embeddings(
         selected_model, [query]
     )[0]
@@ -185,24 +179,33 @@ def query_workspace_aurora(
                 item["keyword_search_score"] = current["keyword_search_score"]
 
     unique_items = list(unique_items.values())
-    score_dict = dict({})
-    if len(unique_items) > 0:
-        passages = [record["content"] for record in unique_items]
-        passage_scores = genai_core.cross_encoder.rank_passages(
-            cross_encoder_model, query, passages
+
+    if (config["rag"]["crossEncodingEnabled"]):
+        cross_encoder_model = genai_core.cross_encoder.get_cross_encoder_model(
+            cross_encoder_model_provider, cross_encoder_model_name
         )
 
-        for i in range(len(unique_items)):
-            score = passage_scores[i]
-            unique_items[i]["score"] = score
-            score_dict[unique_items[i]["chunk_id"]] = score
+        if cross_encoder_model is None:
+            raise genai_core.types.CommonError("Cross encoder model not found")
+
+        score_dict = dict({})
+        if len(unique_items) > 0:
+            passages = [record["content"] for record in unique_items]
+            passage_scores = genai_core.cross_encoder.rank_passages(
+                cross_encoder_model, query, passages
+            )
 
-    unique_items = sorted(unique_items, key=lambda x: x["score"], reverse=True)
+            for i in range(len(unique_items)):
+                score = passage_scores[i]
+                unique_items[i]["score"] = score
+                score_dict[unique_items[i]["chunk_id"]] = score
 
-    for record in vector_search_records:
-        record["score"] = score_dict[record["chunk_id"]]
-    for record in keyword_search_records:
-        record["score"] = score_dict[record["chunk_id"]]
+        unique_items = sorted(unique_items, key=lambda x: x["score"], reverse=True)
+
+        for record in vector_search_records:
+            record["score"] = score_dict[record["chunk_id"]]
+        for record in keyword_search_records:
+            record["score"] = score_dict[record["chunk_id"]]
 
     if full_response:
         unique_items = unique_items[:limit]
@@ -217,9 +220,11 @@ def query_workspace_aurora(
             "keyword_search_items": convert_types(keyword_search_records),
         }
     else:
-        ret_items = list(filter(lambda val: val["score"] > threshold, unique_items))[
-            :limit
-        ]
+        if config["rag"]["crossEncodingEnabled"]:
+            ret_items = list(filter(lambda val: val["score"] > threshold, unique_items))[:limit]
+        else:
+            ret_items = unique_items[:limit]
+
         if len(ret_items) < limit:
             # inner product metric is negative hence we sort ascending
             if metric == "inner":
@@ -295,4 +300,4 @@ def _convert_records(source: str, records: List[dict]):
 
         converted_records.append(converted)
 
-    return converted_records
+    return converted_records