Feature request: Option to disable cross encoder models

bin/config.ts

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

cli/magic-create.ts

@@ -203,6 +203,24 @@ 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: "confirm",
+      name: "enableCrossEncoding",
+      message: "Do you want to enable Cross-Encoding",
+      initial: options.enableCrossEncoding || false,
+      skip(): boolean {
+        return !(this as any).state.answers.enableRag;
+      },
+    },
     {
       type: "multiselect",
       name: "ragsToEnable",
@@ -349,10 +367,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"),
@@ -369,6 +390,7 @@ async function processCreateOptions(options: any): Promise<void> {
       },
       embeddingsModels: [{}],
       crossEncoderModels: [{}],
+      crossEncodingEnabled: answers.enableCrossEncoding,
     },
   };
 
@@ -377,12 +399,24 @@ 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.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;
+  }
   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,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

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

@@ -17,6 +17,7 @@ def query_workspace_open_search(
 ):
     index_name = workspace_id.replace("-", "")
 
+    config = genai_core.parameters.get_config()
     embeddings_model_provider = workspace["embeddings_model_provider"]
     embeddings_model_name = workspace["embeddings_model_name"]
     cross_encoder_model_provider = workspace["cross_encoder_model_provider"]
@@ -36,13 +37,6 @@ def query_workspace_open_search(
     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]
@@ -95,23 +89,32 @@ def query_workspace_open_search(
                 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
-    unique_items = sorted(unique_items, key=lambda x: x["score"], reverse=True)
+        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
+            )
 
-    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"]]
+            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 full_response:
         unique_items = unique_items[:limit]
@@ -124,9 +127,11 @@ def query_workspace_open_search(
             "keyword_search_items": 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:
             unique_items = sorted(
                 unique_items, key=lambda x: x["vector_search_score"] or -1, reverse=True
@@ -208,4 +213,4 @@ def keyword_query(client, index_name: str, text: str, size: int = 25):
     ret_value = response["hits"]["hits"]
     ret_value = ret_value if ret_value is not None else []
 
-    return ret_value
+    return ret_value

lib/shared/types.ts

@@ -83,6 +83,7 @@ export interface SystemConfig {
     roleArn?: string;
   };
   llms: {
+    enableSagemakerModels: boolean;
     sagemaker: SupportedSageMakerModels[];
   };
   rag: {
@@ -117,6 +118,7 @@ export interface SystemConfig {
       name: string;
       default?: boolean;
     }[];
+    crossEncodingEnabled: boolean;
   };
 }