From c074e9da4fd3bc73fb8c489212f9fee8f1c11492 Mon Sep 17 00:00:00 2001 From: Samuel Herman Date: Thu, 6 Mar 2025 15:20:15 -0800 Subject: [PATCH 1/2] add opensearch retriever for the BM25 and neural benchmarks Signed-off-by: Samuel Herman --- .gitignore | 3 + beir/retrieval/search/lexical/bm25_search.py | 2 +- beir/retrieval/search/opensearch/__init__.py | 5 + .../search/opensearch/dense/__init__.py | 5 + .../search/opensearch/dense/neural_search.py | 128 ++++ .../search/opensearch/lexical/__init__.py | 5 + .../search/opensearch/lexical/bm25_search.py | 102 +++ .../search/opensearch/opensearch_search.py | 580 ++++++++++++++++++ .../benchmarking/opensearch_benchmark_bm25.py | 78 +++ .../opensearch_benchmark_neural.py | 62 ++ requirements.txt | 6 + 11 files changed, 975 insertions(+), 1 deletion(-) create mode 100644 beir/retrieval/search/opensearch/__init__.py create mode 100644 beir/retrieval/search/opensearch/dense/__init__.py create mode 100644 beir/retrieval/search/opensearch/dense/neural_search.py create mode 100644 beir/retrieval/search/opensearch/lexical/__init__.py create mode 100644 beir/retrieval/search/opensearch/lexical/bm25_search.py create mode 100644 beir/retrieval/search/opensearch/opensearch_search.py create mode 100644 examples/benchmarking/opensearch_benchmark_bm25.py create mode 100644 examples/benchmarking/opensearch_benchmark_neural.py create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore index df2b2b25..0d0f5881 100644 --- a/.gitignore +++ b/.gitignore @@ -134,3 +134,6 @@ dmypy.json # Pyre type checker .pyre/ + +# IDE +.idea diff --git a/beir/retrieval/search/lexical/bm25_search.py b/beir/retrieval/search/lexical/bm25_search.py index 2eff4370..085e67d7 100644 --- a/beir/retrieval/search/lexical/bm25_search.py +++ b/beir/retrieval/search/lexical/bm25_search.py @@ -26,7 +26,7 @@ def __init__( maxsize: int = 24, number_of_shards: int = "default", initialize: bool = True, - sleep_for: int = 2, + sleep_for: int = 2 ): self.results = {} self.batch_size = batch_size diff --git a/beir/retrieval/search/opensearch/__init__.py b/beir/retrieval/search/opensearch/__init__.py new file mode 100644 index 00000000..b5f1baac --- /dev/null +++ b/beir/retrieval/search/opensearch/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from .opensearch_search import OpenSearchEngine + +__all__ = ["OpenSearchEngine"] diff --git a/beir/retrieval/search/opensearch/dense/__init__.py b/beir/retrieval/search/opensearch/dense/__init__.py new file mode 100644 index 00000000..d7bc7659 --- /dev/null +++ b/beir/retrieval/search/opensearch/dense/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from .neural_search import NeuralSearch + +__all__ = ["NeuralSearch"] \ No newline at end of file diff --git a/beir/retrieval/search/opensearch/dense/neural_search.py b/beir/retrieval/search/opensearch/dense/neural_search.py new file mode 100644 index 00000000..9d5f2927 --- /dev/null +++ b/beir/retrieval/search/opensearch/dense/neural_search.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import time + +import tqdm +import logging + +from ...base import BaseSearch +from ..opensearch_search import OpenSearchEngine +logger = logging.getLogger("NeuralSearch") + +def sleep(seconds): + if seconds: + time.sleep(seconds) + + +class NeuralSearch(BaseSearch): + def __init__( + self, + index_name: str, + hostname: str = "localhost", + keys: dict[str, str] = {"title": "title", "body": "txt", "embedding": "embedding"}, + language: str = "english", + batch_size: int = 128, + timeout: int = 100, + retry_on_timeout: bool = True, + maxsize: int = 24, + number_of_shards: int = "default", + initialize: bool = True, + sleep_for: int = 2 + ): + self.model_id = None + self.results = {} + self.batch_size = batch_size + self.initialize = initialize + self.sleep_for = sleep_for + self.config = { + "hostname": hostname, + "index_name": index_name, + "keys": keys, + "timeout": timeout, + "retry_on_timeout": retry_on_timeout, + "maxsize": maxsize, + "number_of_shards": number_of_shards, + "language": language, + } + # Initialize OpenSearch engine + self.os_engine = OpenSearchEngine(self.config) + if self.initialize: + self.initialise() + + def initialise(self): + """ + Initialise OpenSearch for neural search. + """ + # Setup ML infrastructure + self.os_engine.configure_ml_settings() + # Register model group and get ID + model_group_response = self.os_engine.register_model_group() + model_group_id = model_group_response["model_group_id"] + # Register model using group ID + model_register_response = self.os_engine.register_model(model_group_id=model_group_id) + logger.info(f"Model registration response: {model_register_response}") + self.model_id = self.os_engine.wait_for_model_deployment(task_id=model_register_response["task_id"]) # Use this ID in create_ingest_pipeline + logger.info(f"Model ID: {self.model_id}") + deploy_task_response = self.os_engine.deploy_model(self.model_id) + logger.info(f"Model deployment response: {deploy_task_response}") + self.os_engine.wait_for_model_deployment(task_id=deploy_task_response["task_id"]) + # Create pipeline and index + self.os_engine.create_ingest_pipeline(model_id=self.model_id) + self.os_engine.create_neural_search_index() + + def search( + self, + corpus: dict[str, dict[str, str]], + queries: dict[str, str], + top_k: int, + *args, + **kwargs, + ) -> dict[str, dict[str, float]]: + # Index the corpus within elastic-search + # False, if the corpus has been already indexed + if self.initialize: + self.index(corpus) + # Sleep for few seconds so that elastic-search indexes the docs properly + sleep(self.sleep_for) + + # retrieve neural search results from OpenSearch + query_ids = list(queries.keys()) + queries = [queries[qid] for qid in query_ids] + + for start_idx in tqdm.trange(0, len(queries), self.batch_size, desc="que"): + query_ids_batch = query_ids[start_idx : start_idx + self.batch_size] + results = self.os_engine.neural_multisearch( + texts=queries[start_idx : start_idx + self.batch_size], + model_id=self.model_id, + top_hits=top_k + 1, + ) # Add 1 extra if query is present with documents + + for query_id, hit in zip(query_ids_batch, results): + scores = {} + for corpus_id, score in hit["hits"]: + if corpus_id != query_id: # query doesnt return in results + scores[corpus_id] = score + self.results[query_id] = scores + + return self.results + + def index(self, corpus: dict[str, dict[str, str]]): + progress = tqdm.tqdm(unit="docs", total=len(corpus)) + # dictionary structure = {_id: {title_key: title, text_key: text}} + dictionary = { + idx: { + self.config["keys"]["title"]: corpus[idx].get("title", None), + self.config["keys"]["body"]: corpus[idx].get("text", None), + } + for idx in list(corpus.keys()) + } + self.os_engine.bulk_add_to_index( + generate_actions=self.os_engine.generate_actions(dictionary=dictionary, update=False), + progress=progress, + ) + + def cleanup(self): + self.os_engine.delete_index() + self.os_engine.delete_ingest_pipeline() + self.os_engine.undeploy_model(self.model_id) + self.os_engine.delete_model(self.model_id) diff --git a/beir/retrieval/search/opensearch/lexical/__init__.py b/beir/retrieval/search/opensearch/lexical/__init__.py new file mode 100644 index 00000000..31f7e932 --- /dev/null +++ b/beir/retrieval/search/opensearch/lexical/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from .bm25_search import BM25Search + +__all__ = ["BM25Search"] diff --git a/beir/retrieval/search/opensearch/lexical/bm25_search.py b/beir/retrieval/search/opensearch/lexical/bm25_search.py new file mode 100644 index 00000000..a076ee28 --- /dev/null +++ b/beir/retrieval/search/opensearch/lexical/bm25_search.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import time + +import tqdm + +from ...base import BaseSearch +from ..opensearch_search import OpenSearchEngine + + +def sleep(seconds): + if seconds: + time.sleep(seconds) + + +class BM25Search(BaseSearch): + def __init__( + self, + index_name: str, + hostname: str = "localhost", + keys: dict[str, str] = {"title": "title", "body": "txt"}, + language: str = "english", + batch_size: int = 128, + timeout: int = 100, + retry_on_timeout: bool = True, + maxsize: int = 24, + number_of_shards: int = "default", + initialize: bool = True, + sleep_for: int = 2 + ): + self.results = {} + self.batch_size = batch_size + self.initialize = initialize + self.sleep_for = sleep_for + self.config = { + "hostname": hostname, + "index_name": index_name, + "keys": keys, + "timeout": timeout, + "retry_on_timeout": retry_on_timeout, + "maxsize": maxsize, + "number_of_shards": number_of_shards, + "language": language, + } + self.es = OpenSearchEngine(self.config) + if self.initialize: + self.initialise() + + def initialise(self): + self.es.delete_index() + sleep(self.sleep_for) + self.es.create_index() + + def search( + self, + corpus: dict[str, dict[str, str]], + queries: dict[str, str], + top_k: int, + *args, + **kwargs, + ) -> dict[str, dict[str, float]]: + # Index the corpus within elastic-search + # False, if the corpus has been already indexed + if self.initialize: + self.index(corpus) + # Sleep for few seconds so that elastic-search indexes the docs properly + sleep(self.sleep_for) + + # retrieve results from BM25 + query_ids = list(queries.keys()) + queries = [queries[qid] for qid in query_ids] + + for start_idx in tqdm.trange(0, len(queries), self.batch_size, desc="que"): + query_ids_batch = query_ids[start_idx : start_idx + self.batch_size] + results = self.es.lexical_multisearch( + texts=queries[start_idx : start_idx + self.batch_size], + top_hits=top_k + 1, + ) # Add 1 extra if query is present with documents + + for query_id, hit in zip(query_ids_batch, results): + scores = {} + for corpus_id, score in hit["hits"]: + if corpus_id != query_id: # query doesnt return in results + scores[corpus_id] = score + self.results[query_id] = scores + + return self.results + + def index(self, corpus: dict[str, dict[str, str]]): + progress = tqdm.tqdm(unit="docs", total=len(corpus)) + # dictionary structure = {_id: {title_key: title, text_key: text}} + dictionary = { + idx: { + self.config["keys"]["title"]: corpus[idx].get("title", None), + self.config["keys"]["body"]: corpus[idx].get("text", None), + } + for idx in list(corpus.keys()) + } + self.es.bulk_add_to_index( + generate_actions=self.es.generate_actions(dictionary=dictionary, update=False), + progress=progress, + ) diff --git a/beir/retrieval/search/opensearch/opensearch_search.py b/beir/retrieval/search/opensearch/opensearch_search.py new file mode 100644 index 00000000..a26df292 --- /dev/null +++ b/beir/retrieval/search/opensearch/opensearch_search.py @@ -0,0 +1,580 @@ +from __future__ import annotations + +import logging +import time + +from opensearchpy import OpenSearch +from opensearchpy.helpers import streaming_bulk + +logger = logging.getLogger("OpenSearchEngine") +#tracer.setLevel(logging.CRITICAL) # suppressing INFO messages for opensearch + + +# We named it OpenSearchEngine to avoid conflict with OpenSearch class from the client library +class OpenSearchEngine: + def __init__(self, os_credentials: dict[str, object]): + logging.info("Activating OpenSearch....") + logging.info("OpenSearch Credentials: %s", os_credentials) + self.index_name = os_credentials["index_name"] + self.check_index_name() + + # Same language analyzers as Elasticsearch + self.languages = [ + "arabic", "armenian", "basque", "bengali", "brazilian", "bulgarian", + "catalan", "cjk", "czech", "danish", "dutch", "english", "estonian", + "finnish", "french", "galician", "german", "greek", "hindi", + "hungarian", "indonesian", "irish", "italian", "latvian", "lithuanian", + "norwegian", "persian", "portuguese", "romanian", "russian", "sorani", + "spanish", "swedish", "turkish", "thai" + ] + + self.language = os_credentials["language"] + self.check_language_supported() + + self.text_key = os_credentials["keys"]["body"] + self.embedding_key = os_credentials["keys"]["embedding"] + self.title_key = os_credentials["keys"]["title"] + self.number_of_shards = os_credentials["number_of_shards"] + + # OpenSearch client initialization + self.os = OpenSearch( + hosts=[os_credentials["hostname"]], + timeout=os_credentials["timeout"], + retry_on_timeout=os_credentials["retry_on_timeout"], + max_retries=3, + use_ssl=os_credentials.get("use_ssl", False), + verify_certs=os_credentials.get("verify_certs", False), + ssl_show_warn=os_credentials.get("ssl_show_warn", False), + http_auth=os_credentials.get("http_auth", None) + ) + + def check_language_supported(self): + if self.language.lower() not in self.languages: + raise ValueError( + f"Invalid Language: {self.language}, not supported by OpenSearch. Languages Supported: {self.languages}" + ) + + def check_index_name(self): + for char in r'#:\/*?"<>|,': + if char in self.index_name: + raise ValueError(r'Invalid OpenSearch Index, must not contain the characters ===> #:\/*?"<>|,') + + if self.index_name.startswith(("_", "-", "+")): + raise ValueError("Invalid OpenSearch Index, must not start with characters ===> _ or - or +") + + if self.index_name in [".", ".."]: + raise ValueError("Invalid OpenSearch Index, must not be . or ..") + + if not self.index_name.islower(): + raise ValueError("Invalid OpenSearch Index, must be lowercase") + + def create_index(self): + logging.info(f"Creating fresh OpenSearch-Index named - {self.index_name}") + + try: + mapping = { + "settings": { + "number_of_shards": self.number_of_shards if self.number_of_shards != "default" else 1, + "analysis": { + "analyzer": { + "default": {"type": self.language} + } + } + }, + "mappings": { + "properties": { + self.title_key: {"type": "text", "analyzer": self.language}, + self.text_key: {"type": "text", "analyzer": self.language} + } + } + } + + self.os.indices.create( + index=self.index_name, + body=mapping, + ignore=[400] # 400: IndexAlreadyExistsException + ) + except Exception as e: + logging.error(f"Unable to create Index in OpenSearch. Reason: {e}") + + def delete_index(self): + logging.info(f"Deleting previous OpenSearch-Index named - {self.index_name}") + try: + self.os.indices.delete( + index=self.index_name, + ignore=[400, 404] # 404: IndexDoesntExistException + ) + except Exception as e: + logging.error(f"Unable to delete Index in OpenSearch. Reason: {e}") + + def bulk_add_to_index(self, generate_actions, progress): + for ok, action in streaming_bulk( + client=self.os, + index=self.index_name, + actions=generate_actions, + ): + progress.update(1) + progress.reset() + progress.close() + + def lexical_search(self, text: str, top_hits: int, ids: list[str] = None, skip: int = 0) -> dict[str, object]: + req_body = { + "query": { + "multi_match": { + "query": text, + "type": "best_fields", + "fields": [self.text_key, self.title_key], + "tie_breaker": 0.5 + } + } + } + + if ids: + req_body = { + "query": { + "bool": { + "must": req_body["query"], + "filter": {"ids": {"values": ids}} + } + } + } + + res = self.os.search( + index=self.index_name, + body=req_body, + size=skip + top_hits, + search_type="dfs_query_then_fetch" + ) + + hits = [(hit["_id"], hit["_score"]) for hit in res["hits"]["hits"][skip:]] + return self.hit_template(os_res=res, hits=hits) + + def lexical_multisearch(self, texts: list[str], top_hits: int, skip: int = 0) -> list[dict[str, object]]: + assert skip + top_hits <= 10000, "OpenSearch Window too large, Max-Size = 10000" + + request = [] + for text in texts: + req_head = {"index": self.index_name, "search_type": "dfs_query_then_fetch"} + req_body = { + "_source": False, + "query": { + "multi_match": { + "query": text, + "type": "best_fields", + "fields": [self.title_key, self.text_key], + "tie_breaker": 0.5 + } + }, + "size": skip + top_hits + } + request.extend([req_head, req_body]) + + res = self.os.msearch(body=request) + + result = [] + for resp in res["responses"]: + hits = [(hit["_id"], hit["_score"]) for hit in resp["hits"]["hits"][skip:]] + result.append(self.hit_template(os_res=resp, hits=hits)) + return result + + def generate_actions(self, dictionary: dict[str, dict[str, str]], update: bool = False): + for _id, value in dictionary.items(): + if not update: + doc = { + "_id": str(_id), + "_op_type": "index", + "refresh": "wait_for", + self.text_key: value[self.text_key], + self.title_key: value[self.title_key] + } + else: + doc = { + "_id": str(_id), + "_op_type": "update", + "refresh": "wait_for", + "doc": { + self.text_key: value[self.text_key], + self.title_key: value[self.title_key] + } + } + yield doc + + def hit_template(self, os_res: dict[str, object], hits: list[tuple[str, float]]) -> dict[str, object]: + return { + "meta": { + "total": os_res["hits"]["total"]["value"], + "took": os_res["took"], + "num_hits": len(hits) + }, + "hits": hits + } + + def configure_ml_settings(self): + """Configure OpenSearch ML Commons settings for text embedding. + + Sets required cluster settings for running ML models: + - Allows running ML on any node + - Enables model access control + - Sets native memory threshold + """ + logging.info("Configuring ML cluster settings...") + try: + 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" + } + } + self.os.cluster.put_settings(body=settings) + except Exception as e: + logging.error(f"Unable to configure ML cluster settings. Reason: {e}") + + def register_model_group(self, name: str = "local_model_group", description: str = "A model group for local models"): + """Register a new model group for ML models. + + Args: + name: Name of the model group + description: Description of the model group's purpose + """ + logging.info(f"Registering model group: {name}") + try: + body = { + "name": name, + "description": description + } + + response = self.os.transport.perform_request( + method="POST", + url="/_plugins/_ml/model_groups/_register", + body=body + ) + return response + except Exception as e: + logging.error(f"Unable to register model group. Reason: {e}") + + def register_model( + self, + name: str = "huggingface/sentence-transformers/msmarco-distilbert-base-tas-b", + version: str = "1.0.2", + model_group_id: str = None, + model_format: str = "TORCH_SCRIPT" + ) -> dict: + """Register a new ML model. + + Args: + name: Name/path of the model + version: Version of the model + model_group_id: ID of the model group to associate with (from register_model_group response) + model_format: Format of the model (e.g., TORCH_SCRIPT) + + Returns: + dict: Response from the API containing task_id and status + """ + if not model_group_id: + raise ValueError("model_group_id is required. Get it from register_model_group() response.") + + logging.info(f"Registering model: {name}") + try: + body = { + "name": name, + "version": version, + "model_group_id": model_group_id, + "model_format": model_format + } + + response = self.os.transport.perform_request( + method="POST", + url="/_plugins/_ml/models/_register", + body=body + ) + return response + except Exception as e: + logging.error(f"Unable to register model. Reason: {e}") + raise + + def wait_for_model_registration(self, task_id: str, timeout: int = 300) -> str: + """Wait for the model registration to complete. + + Args: + task_id: ID of the task (from registration response) + timeout: Timeout in seconds + + Returns: + bool: model_id if the model is registered, throw exception otherwise + """ + logging.info(f"Waiting for model registration to complete: {task_id}") + try: + response = self.os.transport.perform_request( + method="GET", + url=f"/_plugins/_ml/tasks/{task_id}" + ) + start_time = time.time() + while response["state"] != "COMPLETED": + time.sleep(1) + response = self.os.transport.perform_request( + method="GET", + url=f"/_plugins/_ml/tasks/{task_id}" + ) + if time.time() - start_time > timeout: + raise TimeoutError("Model registration timed out") + return response["model_id"] + except Exception as e: + logging.error(f"Unable to wait for model registration. Reason: {e}") + raise + + def deploy_model(self, model_id: str) -> dict: + """Deploy a model to the cluster. + + Args: + model_id: ID of the model (from register_model response) + + Returns: + dict: Response from the API containing task_id and status + """ + logging.info(f"Deploying model: {model_id}") + try: + response = self.os.transport.perform_request( + method="POST", + url=f"/_plugins/_ml/models/{model_id}/_deploy", + body={} + ) + return response + except Exception as e: + logging.error(f"Unable to deploy model. Reason: {e}") + raise + + def wait_for_model_deployment(self, task_id: str, timeout: int = 300) -> str: + """Wait for the model deployment to complete. + + Args: + task_id: ID of the task (from deploy_model response) + timeout: Timeout in seconds + + Returns: + bool: model_id if the model is deployed, throw exception otherwise + """ + logging.info(f"Waiting for model deployment to complete: {task_id}") + try: + response = self.os.transport.perform_request( + method="GET", + url=f"/_plugins/_ml/tasks/{task_id}" + ) + start_time = time.time() + while response["state"] != "COMPLETED": + time.sleep(1) + response = self.os.transport.perform_request( + method="GET", + url=f"/_plugins/_ml/tasks/{task_id}" + ) + if time.time() - start_time > timeout: + raise TimeoutError("Model deployment timed out") + return response["model_id"] + except Exception as e: + logging.error(f"Unable to wait for model deployment. Reason: {e}") + raise + + def create_ingest_pipeline(self, model_id: str): + """Create an ingest pipeline for text embedding. + + Args: + model_id: ID of the model (from register_model response) + """ + logging.info("Creating Ingest Pipeline for Neural Search") + try: + pipeline_config = { + "description": "An NLP ingest pipeline", + "processors": [ + { + "text_embedding": { + "model_id": model_id, + "field_map": { + self.text_key: self.embedding_key + } + } + } + ] + } + self.os.ingest.put_pipeline( + id="nlp-ingest-pipeline", + body=pipeline_config + ) + except Exception as e: + logging.error(f"Unable to create Ingest Pipeline in OpenSearch. Reason: {e}") + + def create_neural_search_index(self, dimension: int = 768): + """ + Create neural search index with knn_vector field + """ + logging.info(f"Creating fresh OpenSearch-Index named - {self.index_name}") + try: + mapping = { + "settings": { + "number_of_shards": self.number_of_shards if self.number_of_shards != "default" else 1, + "analysis": { + "analyzer": { + "default": {"type": self.language} + } + }, + "index.knn": True, + "default_pipeline": "nlp-ingest-pipeline" + }, + "mappings": { + "properties": { + self.title_key: {"type": "text", "analyzer": self.language}, + self.text_key: {"type": "text", "analyzer": self.language}, + self.embedding_key: {"type": "knn_vector", "dimension": dimension, + "method": { + "name": "disk_ann", + "space_type": "l2", + "engine": "jvector" + } + } + } + } + } + + self.os.indices.create( + index=self.index_name, + body=mapping + ) + except Exception as e: + logging.error(f"Unable to create Index in OpenSearch. Reason: {e}") + + def neural_search( + self, + text: str, + model_id: str, + top_hits: int = 5, + exclude_embeddings: bool = True + ) -> dict[str, object]: + """Perform neural search using text embeddings. + + Args: + text: Query text to search for + model_id: ID of the model to use for embedding + top_hits: Number of top results to return (default: 5) + exclude_embeddings: Whether to exclude embeddings from response (default: True) + + Returns: + dict: Search results containing hits and metadata + """ + logging.info(f"Performing neural search for: {text}") + try: + body = { + "_source": { + "excludes": [self.embedding_key] if exclude_embeddings else [] + }, + "query": { + "neural": { + self.embedding_key: { + "query_text": text, + "model_id": model_id, + "k": top_hits + } + } + } + } + + response = self.os.search( + index=self.index_name, + body=body + ) + + hits = [(hit["_id"], hit["_score"]) for hit in response["hits"]["hits"]] + return self.hit_template(os_res=response, hits=hits) + + except Exception as e: + logging.error(f"Neural search failed. Reason: {e}") + raise + + def neural_multisearch( + self, + texts: list[str], + model_id: str, + top_hits: int = 5, + exclude_embeddings: bool = True + ) -> list[dict[str, object]]: + """Perform neural search for multiple queries in batch. + + Args: + texts: List of query texts to search for + model_id: ID of the model to use for embedding + top_hits: Number of top results to return per query + exclude_embeddings: Whether to exclude embeddings from response + + Returns: + list: List of search results, one per query + """ + logging.info(f"Performing batch neural search for {len(texts)} queries") + try: + request = [] + for text in texts: + req_head = {"index": self.index_name} + req_body = { + "_source": { + "excludes": [self.embedding_key] if exclude_embeddings else [] + }, + "query": { + "neural": { + self.embedding_key: { + "query_text": text, + "model_id": model_id, + "k": top_hits + } + } + } + } + request.extend([req_head, req_body]) + + response = self.os.msearch(body=request) + logger.info(f"Batch neural search response: {response}") + results = [] + for resp in response["responses"]: + hits = [(hit["_id"], hit["_score"]) for hit in resp["hits"]["hits"]] + results.append(self.hit_template(os_res=resp, hits=hits)) + return results + + except Exception as e: + logging.error(f"Batch neural search failed. Reason: {e}") + raise + + def delete_ingest_pipeline(self): + """Delete the ingest pipeline for text embedding.""" + logging.info("Deleting Ingest Pipeline for Neural Search") + try: + self.os.ingest.delete_pipeline( + id="nlp-ingest-pipeline" + ) + except Exception as e: + logging.error(f"Unable to delete Ingest Pipeline in OpenSearch. Reason: {e}") + + def undeploy_model(self, model_id: str): + """Undeploy a model from the cluster. + + Args: + model_id: ID of the model to undeploy + """ + logging.info(f"Undeploying model: {model_id}") + try: + self.os.transport.perform_request( + method="POST", + url=f"/_plugins/_ml/models/{model_id}/_undeploy" + ) + except Exception as e: + logging.error(f"Unable to undeploy model. Reason: {e}") + + def delete_model(self, model_id: str): + """Delete a model from the cluster. + + Args: + model_id: ID of the model to delete + """ + logging.info(f"Deleting model: {model_id}") + try: + self.os.transport.perform_request( + method="DELETE", + url=f"/_plugins/_ml/models/{model_id}" + ) + except Exception as e: + logging.error(f"Unable to delete model. Reason: {e}") diff --git a/examples/benchmarking/opensearch_benchmark_bm25.py b/examples/benchmarking/opensearch_benchmark_bm25.py new file mode 100644 index 00000000..d3eb2c0d --- /dev/null +++ b/examples/benchmarking/opensearch_benchmark_bm25.py @@ -0,0 +1,78 @@ +import datetime +import logging +import os +import pathlib +import random + +from beir import LoggingHandler, util +from beir.datasets.data_loader import GenericDataLoader +from beir.retrieval.evaluation import EvaluateRetrieval +from beir.retrieval.search.opensearch.lexical import BM25Search as BM25 + +random.seed(42) + +#### Just some code to print debug information to stdout +logging.basicConfig( + format="%(asctime)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + level=logging.INFO, + handlers=[LoggingHandler()], +) + +#### /print debug information to stdout + +#### Download dbpedia-entity.zip dataset and unzip the dataset +dataset = "dbpedia-entity" +url = f"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{dataset}.zip" +out_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "datasets") +data_path = util.download_and_unzip(url, out_dir) + +#### Loading test queries and corpus in DBPedia +corpus, queries, qrels = GenericDataLoader(data_path).load(split="test") +corpus_ids, query_ids = list(corpus), list(queries) + +#### Randomly sample 1M pairs from Original Corpus (4.63M pairs) +#### First include all relevant documents (i.e. present in qrels) +corpus_set = set() +for query_id in qrels: + corpus_set.update(list(qrels[query_id].keys())) +corpus_new = {corpus_id: corpus[corpus_id] for corpus_id in corpus_set} + +#### Remove already seen k relevant documents and sample (1M - k) docs randomly +remaining_corpus = list(set(corpus_ids) - corpus_set) +sample = 1000000 - len(corpus_set) + +for corpus_id in random.sample(remaining_corpus, sample): + corpus_new[corpus_id] = corpus[corpus_id] + +#### Provide parameters for OpenSearch +hostname = "localhost:9200" +index_name = dataset +model = BM25(index_name=index_name, hostname=hostname) +bm25 = EvaluateRetrieval(model) + +#### Index 1M passages into the index (seperately) +bm25.retriever.index(corpus_new) + +#### Saving benchmark times +''' +time_taken_all = {} + +for query_id in query_ids: + query = queries[query_id] + + #### Measure time to retrieve top-10 BM25 documents using single query latency + start = datetime.datetime.now() + results = bm25.retriever.es.lexical_search(text=query, top_hits=10) + end = datetime.datetime.now() + + #### Measuring time taken in ms (milliseconds) + time_taken = end - start + time_taken = time_taken.total_seconds() * 1000 + time_taken_all[query_id] = time_taken + logging.info(f"{query_id}: {query} {time_taken:.2f}ms") + +time_taken = list(time_taken_all.values()) +logging.info(f"Average time taken: {sum(time_taken) / len(time_taken_all):.2f}ms") +''' +bm25.evaluate(qrels=qrels, results=bm25.retrieve(corpus=corpus, queries=queries), k_values=[1, 3, 5, 10, 100]) diff --git a/examples/benchmarking/opensearch_benchmark_neural.py b/examples/benchmarking/opensearch_benchmark_neural.py new file mode 100644 index 00000000..0355866f --- /dev/null +++ b/examples/benchmarking/opensearch_benchmark_neural.py @@ -0,0 +1,62 @@ +import datetime +import logging +import os +import pathlib +import random + +from beir import LoggingHandler, util +from beir.datasets.data_loader import GenericDataLoader +from beir.retrieval.evaluation import EvaluateRetrieval +from beir.retrieval.search.opensearch.dense import NeuralSearch as Neural + +random.seed(42) + +#### Just some code to print debug information to stdout +logging.basicConfig( + format="%(asctime)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + level=logging.INFO, + handlers=[LoggingHandler()], +) + +#### /print debug information to stdout + +#### Download dbpedia-entity.zip dataset and unzip the dataset +dataset = "dbpedia-entity" +url = f"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{dataset}.zip" +out_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "datasets") +data_path = util.download_and_unzip(url, out_dir) + +#### Loading test queries and corpus in DBPedia +corpus, queries, qrels = GenericDataLoader(data_path).load(split="test") +corpus_ids, query_ids = list(corpus), list(queries) + +#### Randomly sample 1M pairs from Original Corpus (4.63M pairs) +#### First include all relevant documents (i.e. present in qrels) +corpus_set = set() +for query_id in qrels: + corpus_set.update(list(qrels[query_id].keys())) +# Take only first 50 relevant documents if there are more +corpus_set = set(list(corpus_set)[:50]) +corpus_new = {corpus_id: corpus[corpus_id] for corpus_id in corpus_set} + +#### Remove already seen k relevant documents and sample (100 - k) docs randomly +remaining_corpus = list(set(corpus_ids) - corpus_set) +target_size = 100 # Total corpus size of 100 documents +sample = target_size - len(corpus_set) + +for corpus_id in random.sample(remaining_corpus, sample): + corpus_new[corpus_id] = corpus[corpus_id] + +#### Provide parameters for OpenSearch +hostname = "localhost:9200" +index_name = dataset +model = Neural(index_name=index_name, hostname=hostname) +neural = EvaluateRetrieval(model) + +#### Index 1M passages into the index (separately) +#neural.retriever.index(corpus_new) + +#### Saving benchmark times +neural.evaluate(qrels=qrels, results=neural.retrieve(corpus=corpus_new, queries=queries), k_values=[1, 3, 5, 10, 100]) +neural.retriever.cleanup() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 00000000..ce31c677 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +sentence-transformers +pytrec-eval-terrier +datasets +elasticsearch +opensearch-py +faiss-cpu \ No newline at end of file From 492ece22e6a3179f39d78a68f1c98422dc50c60e Mon Sep 17 00:00:00 2001 From: Samuel Herman Date: Tue, 11 Mar 2025 08:04:25 -0700 Subject: [PATCH 2/2] add hybrid search retriever and benchmark Signed-off-by: Samuel Herman --- .../search/opensearch/hybrid/__init__.py | 5 + .../search/opensearch/hybrid/hybrid_search.py | 132 ++++++++++++ .../search/opensearch/opensearch_search.py | 191 ++++++++++++++++++ .../opensearch_benchmark_hybrid.py | 62 ++++++ .../opensearch_benchmark_neural.py | 2 +- 5 files changed, 391 insertions(+), 1 deletion(-) create mode 100644 beir/retrieval/search/opensearch/hybrid/__init__.py create mode 100644 beir/retrieval/search/opensearch/hybrid/hybrid_search.py create mode 100644 examples/benchmarking/opensearch_benchmark_hybrid.py diff --git a/beir/retrieval/search/opensearch/hybrid/__init__.py b/beir/retrieval/search/opensearch/hybrid/__init__.py new file mode 100644 index 00000000..8fb10666 --- /dev/null +++ b/beir/retrieval/search/opensearch/hybrid/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from .hybrid_search import HybridSearch + +__all__ = ["HybridSearch"] \ No newline at end of file diff --git a/beir/retrieval/search/opensearch/hybrid/hybrid_search.py b/beir/retrieval/search/opensearch/hybrid/hybrid_search.py new file mode 100644 index 00000000..3899b52e --- /dev/null +++ b/beir/retrieval/search/opensearch/hybrid/hybrid_search.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +import time + +import tqdm +import logging + +from ...base import BaseSearch +from ..opensearch_search import OpenSearchEngine +logger = logging.getLogger("NeuralSearch") + +def sleep(seconds): + if seconds: + time.sleep(seconds) + + +class HybridSearch(BaseSearch): + def __init__( + self, + index_name: str, + hostname: str = "localhost", + keys: dict[str, str] = {"title": "title", "body": "txt", "embedding": "embedding"}, + language: str = "english", + batch_size: int = 128, + timeout: int = 100, + retry_on_timeout: bool = True, + maxsize: int = 24, + number_of_shards: int = "default", + initialize: bool = True, + sleep_for: int = 2 + ): + self.model_id = None + self.results = {} + self.batch_size = batch_size + self.initialize = initialize + self.sleep_for = sleep_for + self.config = { + "hostname": hostname, + "index_name": index_name, + "keys": keys, + "timeout": timeout, + "retry_on_timeout": retry_on_timeout, + "maxsize": maxsize, + "number_of_shards": number_of_shards, + "language": language, + } + # Initialize OpenSearch engine + self.os_engine = OpenSearchEngine(self.config) + if self.initialize: + self.initialise() + + def initialise(self): + """ + Initialise OpenSearch for neural search. + """ + # Setup ML infrastructure + self.os_engine.configure_ml_settings() + # Register model group and get ID + model_group_response = self.os_engine.register_model_group() + model_group_id = model_group_response["model_group_id"] + # Register model using group ID + model_register_response = self.os_engine.register_model(model_group_id=model_group_id) + logger.info(f"Model registration response: {model_register_response}") + self.model_id = self.os_engine.wait_for_model_deployment(task_id=model_register_response["task_id"]) # Use this ID in create_ingest_pipeline + logger.info(f"Model ID: {self.model_id}") + deploy_task_response = self.os_engine.deploy_model(self.model_id) + logger.info(f"Model deployment response: {deploy_task_response}") + self.os_engine.wait_for_model_deployment(task_id=deploy_task_response["task_id"]) + # Create ingest pipeline + self.os_engine.create_ingest_pipeline(model_id=self.model_id) + # Create index + self.os_engine.create_neural_search_index() + # Create search pipeline + self.os_engine.create_search_pipeline() + + def search( + self, + corpus: dict[str, dict[str, str]], + queries: dict[str, str], + top_k: int, + *args, + **kwargs, + ) -> dict[str, dict[str, float]]: + # Index the corpus within elastic-search + # False, if the corpus has been already indexed + if self.initialize: + self.index(corpus) + # Sleep for few seconds so that elastic-search indexes the docs properly + sleep(self.sleep_for) + + # retrieve neural search results from OpenSearch + query_ids = list(queries.keys()) + queries = [queries[qid] for qid in query_ids] + + for start_idx in tqdm.trange(0, len(queries), self.batch_size, desc="que"): + query_ids_batch = query_ids[start_idx : start_idx + self.batch_size] + results = self.os_engine.hybrid_multisearch( + texts=queries[start_idx : start_idx + self.batch_size], + model_id=self.model_id, + top_hits=top_k + 1, + ) # Add 1 extra if query is present with documents + + for query_id, hit in zip(query_ids_batch, results): + scores = {} + for corpus_id, score in hit["hits"]: + if corpus_id != query_id: # query doesnt return in results + scores[corpus_id] = score + self.results[query_id] = scores + + return self.results + + def index(self, corpus: dict[str, dict[str, str]]): + progress = tqdm.tqdm(unit="docs", total=len(corpus)) + # dictionary structure = {_id: {title_key: title, text_key: text}} + dictionary = { + idx: { + self.config["keys"]["title"]: corpus[idx].get("title", None), + self.config["keys"]["body"]: corpus[idx].get("text", None), + } + for idx in list(corpus.keys()) + } + self.os_engine.bulk_add_to_index( + generate_actions=self.os_engine.generate_actions(dictionary=dictionary, update=False), + progress=progress, + ) + + def cleanup(self): + self.os_engine.delete_search_pipeline() + self.os_engine.delete_index() + self.os_engine.delete_ingest_pipeline() + self.os_engine.undeploy_model(self.model_id) + self.os_engine.delete_model(self.model_id) diff --git a/beir/retrieval/search/opensearch/opensearch_search.py b/beir/retrieval/search/opensearch/opensearch_search.py index a26df292..daa33d09 100644 --- a/beir/retrieval/search/opensearch/opensearch_search.py +++ b/beir/retrieval/search/opensearch/opensearch_search.py @@ -441,6 +441,51 @@ def create_neural_search_index(self, dimension: int = 768): except Exception as e: logging.error(f"Unable to create Index in OpenSearch. Reason: {e}") + def create_search_pipeline( + self, + pipeline_id: str = "nlp-search-pipeline", + weights: list[float] = [0.3, 0.7], + normalization: str = "min_max", + combination: str = "arithmetic_mean" + ) -> None: + """Create a search pipeline for hybrid search post-processing. + + Args: + pipeline_id: ID of the pipeline (default: "nlp-search-pipeline") + weights: List of weights for combining scores (default: [0.3, 0.7]) + normalization: Score normalization technique (default: "min_max") + combination: Score combination technique (default: "arithmetic_mean") + """ + logging.info(f"Creating Search Pipeline for Hybrid Search: {pipeline_id}") + try: + pipeline_config = { + "description": "Post processor for hybrid search", + "phase_results_processors": [ + { + "normalization-processor": { + "normalization": { + "technique": normalization + }, + "combination": { + "technique": combination, + "parameters": { + "weights": weights + } + } + } + } + ] + } + + self.os.transport.perform_request( + method="PUT", + url=f"/_search/pipeline/{pipeline_id}", + body=pipeline_config + ) + except Exception as e: + logging.error(f"Unable to create Search Pipeline in OpenSearch. Reason: {e}") + raise + def neural_search( self, text: str, @@ -539,6 +584,152 @@ def neural_multisearch( logging.error(f"Batch neural search failed. Reason: {e}") raise + def hybrid_search( + self, + text: str, + model_id: str, + top_hits: int = 5, + pipeline_id: str = "nlp-search-pipeline", + exclude_embeddings: bool = True + ) -> dict[str, object]: + """Perform hybrid search combining lexical and neural search. + + Args: + text: Query text to search for + model_id: ID of the model to use for neural embedding + top_hits: Number of top results to return (default: 5) + pipeline_id: ID of the search pipeline to use (default: "nlp-search-pipeline") + exclude_embeddings: Whether to exclude embeddings from response (default: True) + + Returns: + dict: Search results containing hits and metadata + """ + logging.info(f"Performing hybrid search for: {text}") + try: + body = { + "_source": { + "excludes": [self.embedding_key] if exclude_embeddings else [] + }, + "query": { + "hybrid": { + "queries": [ + { + "match": { + self.text_key: { + "query": text + } + } + }, + { + "neural": { + self.embedding_key: { + "query_text": text, + "model_id": model_id, + "k": top_hits + } + } + } + ] + } + } + } + + response = self.os.search( + index=self.index_name, + body=body, + params={"search_pipeline": pipeline_id} + ) + + hits = [(hit["_id"], hit["_score"]) for hit in response["hits"]["hits"]] + return self.hit_template(os_res=response, hits=hits) + + except Exception as e: + logging.error(f"Hybrid search failed. Reason: {e}") + raise + + def hybrid_multisearch( + self, + texts: list[str], + model_id: str, + top_hits: int = 5, + pipeline_id: str = "nlp-search-pipeline", + exclude_embeddings: bool = True + ) -> list[dict[str, object]]: + """Perform hybrid search for multiple queries in batch. + + Args: + texts: List of query texts to search for + model_id: ID of the model to use for embedding + top_hits: Number of top results to return per query + pipeline_id: ID of the search pipeline to use + exclude_embeddings: Whether to exclude embeddings from response + + Returns: + list: List of search results, one per query + """ + logging.info(f"Performing batch hybrid search for {len(texts)} queries") + try: + results = [] + for text in texts: + body = { + "_source": { + "excludes": [self.embedding_key] if exclude_embeddings else [] + }, + "query": { + "hybrid": { + "queries": [ + { + "match": { + self.text_key: { + "query": text + } + } + }, + { + "neural": { + self.embedding_key: { + "query_text": text, + "model_id": model_id, + "k": top_hits + } + } + } + ] + } + } + } + + response = self.os.search( + index=self.index_name, + body=body, + params={"search_pipeline": pipeline_id} + ) + + hits = [(hit["_id"], hit["_score"]) for hit in response["hits"]["hits"]] + results.append(self.hit_template(os_res=response, hits=hits)) + + return results + + except Exception as e: + logging.error(f"Batch hybrid search failed. Reason: {e}") + raise + + def delete_search_pipeline(self, pipeline_id: str = "nlp-search-pipeline") -> None: + """Delete a search pipeline. + + Args: + pipeline_id: ID of the pipeline to delete (default: "nlp-search-pipeline") + """ + logging.info(f"Deleting Search Pipeline: {pipeline_id}") + try: + self.os.transport.perform_request( + method="DELETE", + url=f"/_search/pipeline/{pipeline_id}" + ) + except Exception as e: + logging.error(f"Unable to delete Search Pipeline in OpenSearch. Reason: {e}") + raise + def delete_ingest_pipeline(self): """Delete the ingest pipeline for text embedding.""" logging.info("Deleting Ingest Pipeline for Neural Search") diff --git a/examples/benchmarking/opensearch_benchmark_hybrid.py b/examples/benchmarking/opensearch_benchmark_hybrid.py new file mode 100644 index 00000000..ea2cf111 --- /dev/null +++ b/examples/benchmarking/opensearch_benchmark_hybrid.py @@ -0,0 +1,62 @@ +import datetime +import logging +import os +import pathlib +import random + +from beir import LoggingHandler, util +from beir.datasets.data_loader import GenericDataLoader +from beir.retrieval.evaluation import EvaluateRetrieval +from beir.retrieval.search.opensearch.hybrid import HybridSearch as Hybrid + +random.seed(42) + +#### Just some code to print debug information to stdout +logging.basicConfig( + format="%(asctime)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + level=logging.INFO, + handlers=[LoggingHandler()], +) + +#### /print debug information to stdout + +#### Download dbpedia-entity.zip dataset and unzip the dataset +dataset = "dbpedia-entity" +url = f"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{dataset}.zip" +out_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "datasets") +data_path = util.download_and_unzip(url, out_dir) + +#### Loading test queries and corpus in DBPedia +corpus, queries, qrels = GenericDataLoader(data_path).load(split="test") +corpus_ids, query_ids = list(corpus), list(queries) + +#### Randomly sample 1M pairs from Original Corpus (4.63M pairs) +#### First include all relevant documents (i.e. present in qrels) +corpus_set = set() +for query_id in qrels: + corpus_set.update(list(qrels[query_id].keys())) +# Take only first 50 relevant documents if there are more +corpus_set = set(list(corpus_set)[:50]) +corpus_new = {corpus_id: corpus[corpus_id] for corpus_id in corpus_set} + +#### Remove already seen k relevant documents and sample (100 - k) docs randomly +remaining_corpus = list(set(corpus_ids) - corpus_set) +target_size = 100 # Total corpus size of 100 documents +sample = target_size - len(corpus_set) + +for corpus_id in random.sample(remaining_corpus, sample): + corpus_new[corpus_id] = corpus[corpus_id] + +#### Provide parameters for OpenSearch +hostname = "localhost:9200" +index_name = dataset +model = Hybrid(index_name=index_name, hostname=hostname) +hybrid = EvaluateRetrieval(model) + +#### Index 1M passages into the index (separately) +#neural.retriever.index(corpus_new) + +#### Saving benchmark times +hybrid.evaluate(qrels=qrels, results=hybrid.retrieve(corpus=corpus_new, queries=queries), k_values=[1, 3, 5, 10, 100]) +hybrid.retriever.cleanup() diff --git a/examples/benchmarking/opensearch_benchmark_neural.py b/examples/benchmarking/opensearch_benchmark_neural.py index 0355866f..4065f084 100644 --- a/examples/benchmarking/opensearch_benchmark_neural.py +++ b/examples/benchmarking/opensearch_benchmark_neural.py @@ -42,7 +42,7 @@ #### Remove already seen k relevant documents and sample (100 - k) docs randomly remaining_corpus = list(set(corpus_ids) - corpus_set) -target_size = 100 # Total corpus size of 100 documents +target_size = 1000000 # Total corpus size of 100 documents sample = target_size - len(corpus_set) for corpus_id in random.sample(remaining_corpus, sample):