Skip to content

Instantly share code, notes, and snippets.

@dtaivpp
Last active May 16, 2024 16:03
Show Gist options
  • Save dtaivpp/d7e8d8a3ee5debaf896ed2f45b915ad3 to your computer and use it in GitHub Desktop.
Save dtaivpp/d7e8d8a3ee5debaf896ed2f45b915ad3 to your computer and use it in GitHub Desktop.
Shows how to use semantic search locally on OpenSearch using Distillibert.

Local Semantic Search in OpenSearch!

If you'd like to watch this demo through it's available on YouTube or if you prefer reading there is a walkthrough on my blog.

Cluster settings for Amazon OpenSearch and Locally running

PUT /_cluster/settings
{
    "persistent": {
        "plugins.ml_commons.only_run_on_ml_node": false,
        "plugins.ml_commons.model_access_control_enabled": true,
        "plugins.ml_commons.native_memory_threshold": "99",
        "plugins.ml_commons.model_auto_redeploy.enable": true,
        "plugins.ml_commons.model_auto_redeploy.lifetime_retry_times": 3
    }
}

Register your model group

POST /_plugins/_ml/model_groups/_register
{
    "name": "Model_Group",
    "description": "Public ML Model Group",
    "access_mode": "public"
}
#MODEL_GROUP: 

Register your model

POST /_plugins/_ml/models/_register
{
    "name": "huggingface/sentence-transformers/msmarco-distilbert-base-tas-b",
    "version": "1.0.2",
    "model_group_id": "<MODEL_GROUP>",
    "model_format": "TORCH_SCRIPT"
}
# TASK_ID:

Check if model is downloaded

GET /_plugins/_ml/tasks/<TASK_ID>
# MODEL_ID:

Deploy your model

POST /_plugins/_ml/models/<MODEL_ID>/_deploy

Check that model was deployed

GET /_plugins/_ml/tasks/<TASK_ID>

Testing the model

POST /_plugins/_ml/_predict/text_embedding/<MODEL_ID>
{
  "text_docs":[ "This should get embedded"],
  "return_number": true,
  "target_response": ["sentence_embedding"]
}

Create and ingestion pipeline for embeddings

PUT _ingest/pipeline/embedding-ingest-pipeline
{
  "description": "Neural Search Pipeline",
  "processors" : [
    {
      "text_embedding": {
        "model_id": "<MODEL_ID>",
        "field_map": {
          "content": "content_embedding"
        }
      }
    }
  ]
}

Creating a hybrid search pipeline

## Put the search pipeline in place
PUT _search/pipeline/hybrid-search-pipeline
{
  "phase_results_processors": [
    {
      "normalization-processor": {
        "normalization": {
          "technique": "min_max"
        },
        "combination": {
          "technique": "arithmetic_mean",
          "parameters": {
            "weights": [
              0.3,
              0.7
            ]
          }
        }
      }
    }
  ]
}

Create the index

PUT /documents
{
    "settings": {
        "index.knn": true,
        "default_pipeline": "embedding-ingest-pipeline",
        "index.search.default_pipeline": "hybrid-search-pipeline"
    },
    "mappings": {
        "properties": {
            "content_embedding": {
                "type": "knn_vector",
                "dimension": 768,
                "method": {
                    "name": "hnsw",
                    "space_type": "innerproduct",
                    "engine": "nmslib"
                }
            },
            "content": {
                "type": "text"
            }
        }
    }
}

Upload documents

POST /documents/_bulk
{ "index": {"_id": "1234" } }
{ "content": "There once was a racecar driver that was super fast"}
{ "index": {"_id": "1235" } }
{ "content": "The golf driver used by tiger woods is the TaylorMade Qi10 LS prototype"}
{ "index": {"_id": "1236" } }
{ "content": "Some may say that supercar drivers dont really mind risk"}

Search the documents with lexical and vector search

GET /documents/_search
{
  "_source": {
    "exclude": [
      "content_embedding"
    ]
  },
  "query": {
    "hybrid": {
      "queries": [
        {
          "match": {
            "content": {
              "query": "sports automobile"
            }
          }
        },
        {
          "neural": {
            "content_embedding": {
              "query_text": "sports cars",
              "model_id": "<MODEL_ID>",
              "k": 5
            }
          }
        }
      ]
    }
  }
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment