Adding Neural Sparse Search preset

common/constants.ts

@@ -155,6 +155,15 @@ export const OPENAI_CONFIGS = {
   } as RemoteEmbeddingModelConfig,
 };
 
+// Neural Sparse 
+export const NEURAL_SPARSE_CONFIGS = {
+  [`opensearch-neural-sparse-encoding-v2-distill`]: {
+    dimension: 30522,
+    fieldName: 'passage_embedding',
+  } as RemoteEmbeddingModelConfig,
+  // TODO: Add More Neural Sparse Models
+};
+
 /**
  * Various constants pertaining to Workflow configs
  */
@@ -166,6 +175,7 @@ export enum WORKFLOW_TYPE {
   SEMANTIC_SEARCH = 'Semantic Search',
   MULTIMODAL_SEARCH = 'Multimodal Search',
   HYBRID_SEARCH = 'Hybrid Search',
+  NEURAL_SPARSE_SEARCH = 'Neural Sparse Search',
   VECTOR_SEARCH_WITH_RAG = 'RAG with Vector Retrieval',
   HYBRID_SEARCH_WITH_RAG = 'RAG with Hybrid Search',
   CUSTOM = 'Custom Search',
@@ -445,6 +455,22 @@ export const KNN_QUERY = {
   },
   size: 10,
 };
+export const NEURAL_SPARSE_SEARCH_TEMPLATE_QUERY = {
+  ext: {
+    ml_inference: {
+      text: QUERY_TEXT_PATTERN,
+    },
+  },
+  query: {
+    template: {
+      neural_sparse: {
+        [VECTOR_FIELD_PATTERN]: {
+          query_tokens: VECTOR_TEMPLATE_PLACEHOLDER,
+        },
+      },
+    },
+  },
+};
 export const SEMANTIC_SEARCH_QUERY_NEURAL = {
   _source: {
     excludes: [VECTOR_FIELD_PATTERN],
@@ -638,6 +664,10 @@ export const QUERY_PRESETS = [
     name: 'Basic k-NN',
     query: customStringify(KNN_QUERY),
   },
+  {
+    name: 'Neural Sparse Search Query',
+    query: customStringify(NEURAL_SPARSE_SEARCH_TEMPLATE_QUERY),
+  },
   {
     name: WORKFLOW_TYPE.MULTIMODAL_SEARCH,
     query: customStringify(MULTIMODAL_SEARCH_QUERY_BOOL),

common/interfaces.ts

@@ -46,7 +46,7 @@ export type ConfigFieldValue = string | {};
 export interface IConfigField {
   type: ConfigFieldType;
   id: string;
-  value?: ConfigFieldValue;
+  value?: ConfigFieldValue | null;
   selectOptions?: ConfigFieldValue[];
 }
 

common/utils.ts

@@ -51,6 +51,7 @@ export function isVectorSearchUseCase(workflowType?: WORKFLOW_TYPE): boolean {
       WORKFLOW_TYPE.SEMANTIC_SEARCH,
       WORKFLOW_TYPE.MULTIMODAL_SEARCH,
       WORKFLOW_TYPE.HYBRID_SEARCH,
+      WORKFLOW_TYPE.NEURAL_SPARSE_SEARCH,
       WORKFLOW_TYPE.VECTOR_SEARCH_WITH_RAG,
       WORKFLOW_TYPE.HYBRID_SEARCH_WITH_RAG,
     ].includes(workflowType)

documentation/models.md

@@ -473,6 +473,130 @@ POST /_plugins/_ml/models/_register
 }
 ```
 
+### Neural Sparse Encoding
+
+Deploy a sparse encoding model from the Hugging Face Model Hub to a SageMaker real-time inference endpoint using this [guide](https://github.com/zhichao-aws/opensearch-neural-sparse-sample/tree/main/examples/deploy_on_sagemaker).
+
+Connector:
+
+```
+POST /_plugins/_ml/connectors/_create
+{
+  "name": "Neural Sparse Encoding",
+  "description": "Test connector for Sagemaker model",
+  "version": 1,
+  "protocol": "aws_sigv4",
+  "credential": {
+        "access_key": "",
+        "secret_key": "",
+        "session_token": ""
+    },
+  "parameters": {
+    "region": "us-east-1",
+    "service_name": "sagemaker",
+    "model": "opensearch-neural-sparse-encoding-v2-distill"
+  },
+  "actions": [
+    {
+      "action_type": "predict",
+      "method": "POST",
+      "headers": {
+        "content-type": "application/json"
+      },
+      "url": "https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/ns-handler-3/invocations",
+     "request_body": "[\"${parameters.text_doc}\"]"
+
+    }
+  ]
+}
+```
+
+Model:
+
+```
+POST /_plugins/_ml/models/_register
+{ "name": "Neural Sparse Encoding Model",
+  "function_name": "remote",
+  "version": "1.0.0",
+  "connector_id": "<connector-id>",
+  "description": "Test connector for Sagemaker model",
+  "interface": {
+    "input": {
+        "type": "object",
+        "properties": {
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "text_doc": {
+                        "type": "string"
+                    }
+                },
+                "additionalProperties": true,
+                "required": [
+                    "text_doc"
+                ]
+            }
+        }
+    },
+    "output": {
+        "type": "object",
+        "properties": {
+            "inference_results": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "output": {
+                            "type": "array",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "name": {
+                                        "type": "string"
+                                    },
+                                    "dataAsMap": {
+                                        "type": "object",
+                                        "properties": {
+                                            "response": {
+                                                "type": "array",
+                                                "items": {
+                                                    "type": "object",
+                                                    "additionalProperties": {
+                                                        "type": "number"
+                                                    }
+                                                }
+                                            }
+                                        },
+                                        "required": [
+                                            "response"
+                                        ]
+                                    }
+                                },
+                                "required": [
+                                    "name",
+                                    "dataAsMap"
+                                ]
+                            }
+                        },
+                        "status_code": {
+                            "type": "integer"
+                        }
+                    },
+                    "required": [
+                        "output",
+                        "status_code"
+                    ]
+                }
+            }
+        },
+        "required": [
+            "inference_results"
+        ]
+    }
+  }
+}
+```
+
 ## Generative models
 
 ### Claude 3 Sonnet (hosted on Amazon Bedrock)

documentation/tutorial.md

@@ -491,3 +491,46 @@ Override the query to a knn query, including the embedding output. For example:
     }
 }
 ```
+
+---
+
+## 9. Neural sparse search
+
+### ML resources
+Create and deploy a [Neural Sparse Encoding model](https://github.com/opensearch-project/dashboards-flow-framework/blob/main/documentation/models.md#neural-sparse-encoding).
+
+### Index
+
+Ensure the index mappings have a `rank_features` field - something like the following:
+
+```
+"<embedding_field_name>": {
+    "type": "rank_features"
+}
+```
+
+### Ingest pipeline
+
+Single ML inference processor. Map your input text to the `text_doc` model input field. Optionally map the output `response` to a new document field. Transform the response if needed using JSONPath expression. 
+
+
+### Search pipeline
+
+Single ML inference **search request** processor. Map the query field containing the input text to the `text_doc` model input field. Optionally map the output `response` to a new field. Transform the response if needed using JSONPath expression. Override the query to a neural sparse query. For example:
+
+```
+{
+    "_source": {
+        "excludes": [
+            "<embedding_field>"
+        ]
+    },
+    "query": {
+        "neural_sparse": {
+            "<embedding_field>": {
+                "query_tokens": ${response},
+            }
+        }
+    }
+}
+```

public/configs/search_request_processors/ml_search_request_processor.ts

@@ -3,22 +3,25 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { ConfigFieldType } from 'common';
 import { generateId } from '../../utils';
 import { MLProcessor } from '../ml_processor';
 
 /**
  * The ML processor in the context of search request
  */
 export class MLSearchRequestProcessor extends MLProcessor {
-  constructor() {
+  constructor(includeQueryTemplate: boolean = true) {
     super();
     this.id = generateId('ml_processor_search_request');
     this.optionalFields = [
-      {
-        id: 'query_template',
-        type: 'jsonString',
-      },
-      ...(this.optionalFields || []),
+      ...(includeQueryTemplate
+        ? [{
+            id: 'query_template',
+            type: 'jsonString' as ConfigFieldType,
+          }]
+        : []),
+        ...(this.optionalFields || []),
     ];
   }
 }

public/pages/workflow_detail/workflow_inputs/ingest_inputs/advanced_settings.tsx

@@ -14,7 +14,7 @@ import {
 } from '@elastic/eui';
 import { JsonField } from '../input_fields';
 import { getIn, useFormikContext } from 'formik';
-import { WorkflowFormValues } from '../../../../../common';
+import { WORKFLOW_TYPE, WorkflowFormValues } from '../../../../../common';
 import { AppState } from '../../../../store';
 import {
   getEmbeddingField,
@@ -28,6 +28,7 @@ import {
 
 interface AdvancedSettingsProps {
   setHasInvalidDimensions: (hasInvalidDimensions: boolean) => void;
+  workflowType: WORKFLOW_TYPE | undefined;
 }
 
 /**
@@ -64,7 +65,7 @@ export function AdvancedSettings(props: AdvancedSettingsProps) {
 
           // If a dimension is found, it is a known embedding model.
           // Ensure the index is configured to be knn-enabled.
-          if (dimension !== undefined) {
+          if (dimension !== undefined && props.workflowType !== WORKFLOW_TYPE.NEURAL_SPARSE_SEARCH) {
             if (!isKnnIndex(curSettings)) {
               setFieldValue(
                 indexSettingsPath,

public/pages/workflow_detail/workflow_inputs/ingest_inputs/ingest_data.tsx

@@ -13,9 +13,9 @@ import {
 } from '@elastic/eui';
 import { TextField } from '../input_fields';
 import { AdvancedSettings } from './advanced_settings';
-import { KNN_VECTOR_DOCS_LINK } from '../../../../../common';
+import { KNN_VECTOR_DOCS_LINK, WORKFLOW_TYPE } from '../../../../../common';
 
-interface IngestDataProps {}
+interface IngestDataProps {workflowType: WORKFLOW_TYPE | undefined;}
 
 /**
  * Input component for configuring the data ingest (the OpenSearch index)
@@ -56,7 +56,7 @@ export function IngestData(props: IngestDataProps) {
         />
       </EuiFlexItem>
       <EuiFlexItem>
-        <AdvancedSettings setHasInvalidDimensions={setHasInvalidDimensions} />
+        <AdvancedSettings setHasInvalidDimensions={setHasInvalidDimensions} workflowType={props.workflowType} />
       </EuiFlexItem>
     </EuiFlexGroup>
   );

public/pages/workflow_detail/workflow_inputs/ingest_inputs/ingest_inputs.tsx

@@ -51,7 +51,7 @@ export function IngestInputs(props: IngestInputsProps) {
         <EuiHorizontalRule margin="none" />
       </EuiFlexItem>
       <EuiFlexItem grow={false}>
-        <IngestData />
+        <IngestData workflowType={props.workflow?.ui_metadata?.type}/>
       </EuiFlexItem>
     </EuiFlexGroup>
   );

public/pages/workflows/new_workflow/quick_configure_modal.tsx

@@ -434,6 +434,7 @@ function injectQuickConfigureFields(
     switch (workflow.ui_metadata?.type) {
       case WORKFLOW_TYPE.SEMANTIC_SEARCH:
       case WORKFLOW_TYPE.HYBRID_SEARCH:
+      case WORKFLOW_TYPE.NEURAL_SPARSE_SEARCH:
       case WORKFLOW_TYPE.MULTIMODAL_SEARCH: {
         if (!isEmpty(quickConfigureFields) && workflow.ui_metadata?.config) {
           workflow.ui_metadata.config = updateIngestProcessors(
@@ -444,6 +445,7 @@ function injectQuickConfigureFields(
           );
           workflow.ui_metadata.config = updateIndexConfig(
             workflow.ui_metadata.config,
+            workflow.ui_metadata?.type,
             quickConfigureFields
           );
           workflow.ui_metadata.config.search.request.value = injectPlaceholderValues(
@@ -470,6 +472,7 @@ function injectQuickConfigureFields(
           );
           workflow.ui_metadata.config = updateIndexConfig(
             workflow.ui_metadata.config,
+            workflow.ui_metadata?.type,
             quickConfigureFields
           );
           workflow.ui_metadata.config.search.request.value = injectPlaceholderValues(
@@ -795,6 +798,7 @@ function updateRAGSearchResponseProcessors(
 // prefill index mappings/settings, if applicable
 function updateIndexConfig(
   config: WorkflowConfig,
+  workflow_type: WORKFLOW_TYPE,
   fields: QuickConfigureFields
 ): WorkflowConfig {
   if (
@@ -823,10 +827,10 @@ function updateIndexConfig(
       };
     }
     if (fields.vectorField) {
-      properties[fields.vectorField] = {
-        type: 'knn_vector',
-        dimension: fields.embeddingLength || '',
-      };
+      properties[fields.vectorField] = 
+        workflow_type !== WORKFLOW_TYPE.NEURAL_SPARSE_SEARCH
+          ? { type: 'knn_vector', dimension: fields.embeddingLength || '' }
+          : { type: 'rank_features' };
     }
     if (fields.labelField) {
       properties[fields.labelField] = {