add tutorial for semantic search with byte quantized vector and Cohere embedding model #2127
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,278 @@ | ||
| # Topic | ||
|
|
||
| > This tutorial doesn't explain byte-quantized vectors in detail. For more information, see [Byte-quantized vectors in OpenSearch](https://opensearch.org/blog/byte-quantized-vectors-in-opensearch/). | ||
|
|
||
| This tutorial shows how to build semantic search using the [Cohere Embed model](https://docs.cohere.com/reference/embed) and byte-quantized vectors in OpenSearch. | ||
|
|
||
| The Cohere Embed v3 model supports several `embedding_types`. This tutorial uses the `int8` type for byte-quantized vectors. | ||
|
|
||
| Note: Replace the placeholders that start with `your_` with your own values. | ||
|
|
||
| # Steps | ||
|
|
||
| The Cohere Embed v3 model supports several input types. This tutorial uses the following input types (from the Cohere [documentation](https://docs.cohere.com/reference/embed)): | ||
| > - `search_document`: Used for embeddings stored in a vector database for search use cases. | ||
ylwu-amzn marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| > - `"search_query"`: Used for embeddings of search queries run against a vector DB to find relevant documents. | ||
|
There was a problem hiding this comment. IMHO, these descriptions from Cohere are a little hard to understand. Can it be something like:
There was a problem hiding this comment. This is copied from Cohere document directly. We should keep consistent with Cohere doc. @tianjing-li , can you help review ? There was a problem hiding this comment. This is a snippet from our official docs that could help clarify: There was a problem hiding this comment. Thanks @tianjing-li , can you share the doc link ? There was a problem hiding this comment. |
||
|
|
||
| You will create two models in this tutorial: | ||
| - A model used for ingestion with the `search_document` input type | ||
| - A model used for search with the `search_query` input type | ||
|
|
||
| ## 1. Create embedding model for ingestion | ||
|
|
||
| Create a connector with the `search_document` input type: | ||
|
|
||
| ``` | ||
| POST /_plugins/_ml/connectors/_create | ||
| { | ||
| "name": "Cohere embedding connector with int8 embedding type for ingestion", | ||
| "description": "Test connector for Cohere embedding model", | ||
| "version": 1, | ||
| "protocol": "http", | ||
| "credential": { | ||
| "cohere_key": "your_cohere_api_key" | ||
| }, | ||
| "parameters": { | ||
| "model": "embed-english-v3.0", | ||
| "embedding_types": ["int8"], | ||
| "input_type": "search_document" | ||
| }, | ||
| "actions": [ | ||
| { | ||
| "action_type": "predict", | ||
| "method": "POST", | ||
| "headers": { | ||
| "Authorization": "Bearer ${credential.cohere_key}", | ||
| "Request-Source": "unspecified:opensearch" | ||
| }, | ||
| "url": "https://api.cohere.ai/v1/embed", | ||
| "request_body": "{ \"model\": \"${parameters.model}\", \"texts\": ${parameters.texts}, \"input_type\":\"${parameters.input_type}\", \"embedding_types\": ${parameters.embedding_types} }", | ||
| "pre_process_function": "connector.pre_process.cohere.embedding", | ||
| "post_process_function": "\n def name = \"sentence_embedding\";\n def data_type = \"FLOAT32\";\n def result;\n if (params.embeddings.int8 != null) {\n data_type = \"INT8\";\n result = params.embeddings.int8;\n } else if (params.embeddings.uint8 != null) {\n data_type = \"UINT8\";\n result = params.embeddings.uint8;\n } else if (params.embeddings.float != null) {\n data_type = \"FLOAT32\";\n result = params.embeddings.float;\n }\n \n if (result == null) {\n return \"Invalid embedding result\";\n }\n \n def embedding_list = new StringBuilder(\"[\");\n \n for (int m=0; m<result.length; m++) {\n def embedding_size = result[m].length;\n def embedding = new StringBuilder(\"[\");\n def shape = [embedding_size];\n for (int i=0; i<embedding_size; i++) {\n def val;\n if (\"FLOAT32\".equals(data_type)) {\n val = result[m][i].floatValue();\n } else if (\"INT8\".equals(data_type) || \"UINT8\".equals(data_type)) {\n val = result[m][i].intValue();\n }\n embedding.append(val);\n if (i < embedding_size - 1) {\n embedding.append(\",\"); \n }\n }\n embedding.append(\"]\"); \n \n // workaround for compatible with neural-search\n def dummy_data_type = 'FLOAT32';\n \n def json = '{' +\n '\"name\":\"' + name + '\",' +\n '\"data_type\":\"' + dummy_data_type + '\",' +\n '\"shape\":' + shape + ',' +\n '\"data\":' + embedding +\n '}';\n embedding_list.append(json);\n if (m < result.length - 1) {\n embedding_list.append(\",\"); \n }\n }\n embedding_list.append(\"]\"); \n return embedding_list.toString();\n " | ||
| } | ||
| ] | ||
| } | ||
| ``` | ||
| Use the connector ID from the response to create a model: | ||
| ``` | ||
| POST /_plugins/_ml/models/_register?deploy=true | ||
| { | ||
| "name": "Cohere embedding model for INT8 with search_document input type", | ||
| "function_name": "remote", | ||
| "description": "test model", | ||
| "connector_id": "your_connector_id" | ||
| } | ||
| ``` | ||
| Note the model ID; you'll use it in step 2.1. | ||
ylwu-amzn marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| Test the model: | ||
| ``` | ||
| POST /_plugins/_ml/models/your_embedding_model_id/_predict | ||
| { | ||
| "parameters": { | ||
| "texts": ["hello", "goodbye"] | ||
| } | ||
| } | ||
| ``` | ||
| Sample response: | ||
|
|
||
| ``` | ||
| { | ||
| "inference_results": [ | ||
| { | ||
| "output": [ | ||
| { | ||
| "name": "sentence_embedding", | ||
| "data_type": "FLOAT32", | ||
| "shape": [ | ||
| 1024 | ||
| ], | ||
| "data": [ | ||
| 20, | ||
| -11, | ||
| -60, | ||
| -91, | ||
| ... | ||
| ] | ||
| }, | ||
| { | ||
| "name": "sentence_embedding", | ||
|
There was a problem hiding this comment. I think it's worth mentioning in the explanation that even though the |
||
| "data_type": "FLOAT32", | ||
| "shape": [ | ||
| 1024 | ||
| ], | ||
| "data": [ | ||
| 58, | ||
| -30, | ||
| 9, | ||
| -51, | ||
| ... | ||
| ] | ||
| } | ||
| ], | ||
| "status_code": 200 | ||
| } | ||
| ] | ||
| } | ||
| ``` | ||
|
|
||
| ## 2. Ingest data | ||
|
|
||
| ### 2.1 Create ingest pipeline | ||
|
|
||
| ``` | ||
| PUT /_ingest/pipeline/pipeline-cohere | ||
| { | ||
| "description": "Cohere embedding ingest pipeline", | ||
| "processors": [ | ||
| { | ||
| "text_embedding": { | ||
| "model_id": "your_embedding_model_id_created_in_step1", | ||
| "field_map": { | ||
| "passage_text": "passage_embedding" | ||
| } | ||
| } | ||
| } | ||
| ] | ||
| } | ||
| ``` | ||
|
|
||
ylwu-amzn marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| ### 2.2 Create KNN index with byte-quantized vector | ||
| For more information, refer to [this blog](https://opensearch.org/blog/byte-quantized-vectors-in-opensearch/). | ||
|
|
||
| ``` | ||
| PUT my_test_data | ||
| { | ||
| "settings": { | ||
| "index": { | ||
| "knn": true, | ||
| "knn.algo_param.ef_search": 100, | ||
| "default_pipeline": "pipeline-cohere" | ||
| } | ||
| }, | ||
| "mappings": { | ||
| "properties": { | ||
| "passage_text": { | ||
| "type": "text" | ||
| }, | ||
| "passage_embedding": { | ||
| "type": "knn_vector", | ||
| "dimension": 1024, | ||
| "data_type": "byte", | ||
| "method": { | ||
| "name": "hnsw", | ||
| "space_type": "l2", | ||
| "engine": "lucene", | ||
| "parameters": { | ||
| "ef_construction": 128, | ||
| "m": 24 | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| ``` | ||
|
|
||
| Ingest test data: | ||
|
|
||
| ``` | ||
| POST _bulk | ||
| { "index" : { "_index" : "my_test_data" } } | ||
| { "passage_text" : "OpenSearch is the flexible, scalable, open-source way to build solutions for data-intensive applications. Explore, enrich, and visualize your data with built-in performance, developer-friendly tools, and powerful integrations for machine learning, data processing, and more." } | ||
| { "index" : { "_index" : "my_test_data"} } | ||
| { "passage_text" : "BM25 is a keyword-based algorithm that performs well on queries containing keywords but fails to capture the semantic meaning of the query terms. Semantic search, unlike keyword-based search, takes into account the meaning of the query in the search context. Thus, semantic search performs well when a query requires natural language understanding." } | ||
|
|
||
| ``` | ||
|
|
||
| ## 3. Semantic search | ||
|
|
||
| Create another embedding model with the `search_query` input type: | ||
| ``` | ||
| POST /_plugins/_ml/connectors/_create | ||
| { | ||
| "name": "Cohere embedding connector with int8 embedding type for search", | ||
| "description": "Test connector for Cohere embedding model. Use this connector for search.", | ||
| "version": 1, | ||
| "protocol": "http", | ||
| "credential": { | ||
| "cohere_key": "your_cohere_api_key" | ||
| }, | ||
| "parameters": { | ||
| "model": "embed-english-v3.0", | ||
| "embedding_types": ["int8"], | ||
| "input_type": "search_query" | ||
| }, | ||
| "actions": [ | ||
| { | ||
| "action_type": "predict", | ||
| "method": "POST", | ||
| "headers": { | ||
| "Authorization": "Bearer ${credential.cohere_key}", | ||
| "Request-Source": "unspecified:opensearch" | ||
| }, | ||
| "url": "https://api.cohere.ai/v1/embed", | ||
| "request_body": "{ \"model\": \"${parameters.model}\", \"texts\": ${parameters.texts}, \"input_type\":\"${parameters.input_type}\", \"embedding_types\": ${parameters.embedding_types} }", | ||
| "pre_process_function": "connector.pre_process.cohere.embedding", | ||
| "post_process_function": "\n def name = \"sentence_embedding\";\n def data_type = \"FLOAT32\";\n def result;\n if (params.embeddings.int8 != null) {\n data_type = \"INT8\";\n result = params.embeddings.int8;\n } else if (params.embeddings.uint8 != null) {\n data_type = \"UINT8\";\n result = params.embeddings.uint8;\n } else if (params.embeddings.float != null) {\n data_type = \"FLOAT32\";\n result = params.embeddings.float;\n }\n \n if (result == null) {\n return \"Invalid embedding result\";\n }\n \n def embedding_list = new StringBuilder(\"[\");\n \n for (int m=0; m<result.length; m++) {\n def embedding_size = result[m].length;\n def embedding = new StringBuilder(\"[\");\n def shape = [embedding_size];\n for (int i=0; i<embedding_size; i++) {\n def val;\n if (\"FLOAT32\".equals(data_type)) {\n val = result[m][i].floatValue();\n } else if (\"INT8\".equals(data_type) || \"UINT8\".equals(data_type)) {\n val = result[m][i].intValue();\n }\n embedding.append(val);\n if (i < embedding_size - 1) {\n embedding.append(\",\"); \n }\n }\n embedding.append(\"]\"); \n \n // workaround for compatible with neural-search\n def dummy_data_type = 'FLOAT32';\n \n def json = '{' +\n '\"name\":\"' + name + '\",' +\n '\"data_type\":\"' + dummy_data_type + '\",' +\n '\"shape\":' + shape + ',' +\n '\"data\":' + embedding +\n '}';\n embedding_list.append(json);\n if (m < result.length - 1) {\n embedding_list.append(\",\"); \n }\n }\n embedding_list.append(\"]\"); \n return embedding_list.toString();\n " | ||
| } | ||
| ] | ||
| } | ||
| ``` | ||
| Use the connector ID from the response to create a model: | ||
| ``` | ||
| POST /_plugins/_ml/models/_register?deploy=true | ||
| { | ||
| "name": "Cohere embedding model for INT8 with search_document input type", | ||
| "function_name": "remote", | ||
| "description": "test model", | ||
| "connector_id": "your_connector_id" | ||
| } | ||
| ``` | ||
| Then you can use the model id from response to run neural search query: | ||
| ``` | ||
| POST /my_test_data/_search | ||
| { | ||
| "query": { | ||
| "neural": { | ||
| "passage_embedding": { | ||
| "query_text": "semantic search", | ||
| "model_id": "your_embedding_model_id", | ||
| "k": 100 | ||
| } | ||
| } | ||
| }, | ||
| "size": "1", | ||
| "_source": ["passage_text"] | ||
| } | ||
| ``` | ||
| Sample response | ||
| ``` | ||
| { | ||
| "took": 143, | ||
| "timed_out": false, | ||
| "_shards": { | ||
| "total": 1, | ||
| "successful": 1, | ||
| "skipped": 0, | ||
| "failed": 0 | ||
| }, | ||
| "hits": { | ||
| "total": { | ||
| "value": 2, | ||
| "relation": "eq" | ||
| }, | ||
| "max_score": 9.345969e-7, | ||
| "hits": [ | ||
| { | ||
| "_index": "my_test_data", | ||
| "_id": "_IXCuY0BJr_OiKWden7i", | ||
| "_score": 9.345969e-7, | ||
| "_source": { | ||
| "passage_text": "BM25 is a keyword-based algorithm that performs well on queries containing keywords but fails to capture the semantic meaning of the query terms. Semantic search, unlike keyword-based search, takes into account the meaning of the query in the search context. Thus, semantic search performs well when a query requires natural language understanding." | ||
| } | ||
| } | ||
| ] | ||
| } | ||
| } | ||
| ``` | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I would suggest that the your_xxx placeholders be emphasized somehow in the code samples, e.g., in italic or boldface or something.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It will pollute the REST API sample request by doing this. I think we should be good to leave as is.