Feature request: Option to disable cross encoder models

bin/config.ts

@@ -18,6 +18,7 @@ export function getConfig(): SystemConfig {
     },
     llms: {
       // sagemaker: [SupportedSageMakerModels.FalconLite]
+      enableSagemakerModels: false,
       sagemaker: [],
     },
     rag: {
@@ -64,6 +65,7 @@ export function getConfig(): SystemConfig {
           default: true,
         },
       ],
+      crossEncodingEnabled: false,
     },
   };
 }

cli/magic-create.ts

@@ -146,6 +146,18 @@ async function processCreateOptions(options: any): Promise<void> {
       },
       initial: options.bedrockRoleArn || "",
     },
+    {
+      type: "multiselect",
+      name: "selectedEmbeddingModels",
+      hint: "SPACE to select, ENTER to confirm selection",
+      message: "Which Embedding Models do you want to enable",
+      choices: embeddingModels.map((m) => ({ name: m.name, value: m })),
+      validate(choices: any) {
+        return (this as any).skipped || choices.length > 0
+          ? true
+          : "You need to select at least one model";
+      },
+    },
     {
       type: "confirm",
       name: "enableSagemakerModels",
@@ -180,6 +192,12 @@ async function processCreateOptions(options: any): Promise<void> {
       message: "Do you want to enable RAG",
       initial: options.enableRag || false,
     },
+    {
+      type: "confirm",
+      name: "enableCrossEncoding",
+      message: "Do you want to enable Cross-Encoding",
+      initial: options.enableCrossEncoding || false,
+    },
     {
       type: "multiselect",
       name: "ragsToEnable",
@@ -319,6 +337,7 @@ async function processCreateOptions(options: any): Promise<void> {
         }
       : undefined,
     llms: {
+      enableSagemakerModels: answers.enableSagemakerModels,
       sagemaker: answers.sagemakerModels,
     },
     rag: {
@@ -339,6 +358,7 @@ async function processCreateOptions(options: any): Promise<void> {
       },
       embeddingsModels: [{}],
       crossEncoderModels: [{}],
+      crossEncodingEnabled: answers.enableCrossEncoding,
     },
   };
 
@@ -347,12 +367,20 @@ async function processCreateOptions(options: any): Promise<void> {
     models.defaultEmbedding = embeddingModels[0].name;
   }
 
-  config.rag.crossEncoderModels[0] = {
-    provider: "sagemaker",
-    name: "cross-encoder/ms-marco-MiniLM-L-12-v2",
-    default: true,
-  };
-  config.rag.embeddingsModels = embeddingModels;
+  if (answers.enableCrossEncoding && answers.sagemakerModels.length > 0) {
+    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,
+    };
+  }
+  config.rag.embeddingsModels = embeddingModels.filter((model) => answers.selectedEmbeddingModels.includes(model.name));
   config.rag.embeddingsModels.forEach((m: any) => {
     if (m.name === models.defaultEmbedding) {
       m.default = true;

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

@@ -151,8 +151,7 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
       identityPool: authentication.identityPool,
       api: chatBotApi,
       chatbotFilesBucket: chatBotApi.filesBucket,
-      crossEncodersEnabled:
-        typeof ragEngines?.sageMakerRagModels?.model !== "undefined",
+      crossEncodersEnabled: props.config.rag.crossEncodingEnabled,
       sagemakerEmbeddingsEnabled:
         typeof ragEngines?.sageMakerRagModels?.model !== "undefined",
     });

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,12 +38,13 @@ 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 (config["rag"]["crossEncodingEnabled"]):
+        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")
+        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]
@@ -185,78 +187,69 @@ 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
-        )
-
-        for i in range(len(unique_items)):
-            score = passage_scores[i]
-            unique_items[i]["score"] = score
-            score_dict[unique_items[i]["chunk_id"]] = score
-
-    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 (config["rag"]["crossEncodingEnabled"]):
+        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 full_response:
-        unique_items = unique_items[:limit]
-        ret_value = {
-            "engine": "aurora",
-            "query_language": language_name,
-            "supported_languages": languages,
-            "detected_languages": detected_languages,
-            "items": convert_types(unique_items),
-            "vector_search_metric": metric,
-            "vector_search_items": convert_types(vector_search_records),
-            "keyword_search_items": convert_types(keyword_search_records),
-        }
-    else:
-        ret_items = list(filter(lambda val: val["score"] > threshold, unique_items))[
-            :limit
-        ]
-        if len(ret_items) < limit:
-            # inner product metric is negative hence we sort ascending
-            if metric == "inner":
-                unique_items = sorted(
-                    unique_items,
-                    key=lambda x: x["vector_search_score"] or 1,
-                    reverse=False,
-                )
-                ret_items = ret_items + (
-                    list(
-                        filter(
-                            lambda val: (val["vector_search_score"] or 1) < -0.5,
-                            unique_items,
-                        )
-                    )[: (limit - len(ret_items))]
-                )
-            else:
+            for i in range(len(unique_items)):
+                score = passage_scores[i]
+                unique_items[i]["score"] = score
+                score_dict[unique_items[i]["chunk_id"]] = score
+
+        unique_items = sorted(unique_items, key=lambda x: x["score"], reverse=True)
+
+    if (config["rag"]["crossEncodingEnabled"]):
+        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:
+            ret_value = {
+                "engine": "aurora",
+                "query_language": language_name,
+                "supported_languages": languages,
+                "detected_languages": detected_languages,
+                "items": convert_types(unique_items),
+                "vector_search_metric": metric,
+                "vector_search_items": convert_types(vector_search_records),
+                "keyword_search_items": convert_types(keyword_search_records),
+            }
+        else:
+            ret_items = list(filter(lambda val: val["score"] > threshold, unique_items))
+            if len(ret_items) < limit:
                 unique_items = sorted(
-                    unique_items,
-                    key=lambda x: x["vector_search_score"] or -1,
-                    reverse=True,
+                    unique_items, key=lambda x: x["vector_search_score"], reverse=True
                 )
                 ret_items = ret_items + (
                     list(
-                        filter(
-                            lambda val: (val["vector_search_score"] or -1) > 0.5,
-                            unique_items,
-                        )
+                        filter(lambda val: val["vector_search_score"] > 0.5, unique_items)
                     )[: (limit - len(ret_items))]
                 )
 
+            ret_value = {
+                "engine": "aurora",
+                "query_language": language_name,
+                "supported_languages": languages,
+                "detected_languages": detected_languages,
+                "items": convert_types(
+                    list(filter(lambda val: val["score"] > 0, unique_items))
+                ),
+            }
+    else:
         ret_value = {
             "engine": "aurora",
             "query_language": language_name,
             "supported_languages": languages,
             "detected_languages": detected_languages,
-            "items": convert_types(ret_items),
+            "items": convert_types(unique_items),
+            "vector_search_metric": metric,
+            "vector_search_items": convert_types(vector_search_records),
+            "keyword_search_items": convert_types(keyword_search_records),
         }
 
     logger.info(ret_value)