enable support for HNSW in dataprep and retriever for better performance (#1779)

* enable support vector-schema in dataprep redis

Signed-off-by: leslieluyu <yu1.lu@intel.com>

* enable support HNSW in retriever

Signed-off-by: leslieluyu <yu1.lu@intel.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: leslieluyu <yu1.lu@intel.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: xiguiw <111278656+xiguiw@users.noreply.github.com>

Author: 3 people

comps/dataprep/src/integrations/redis.py

@@ -52,6 +52,10 @@
 TIMEOUT_SECONDS = int(os.getenv("TIMEOUT_SECONDS", 600))
 SEARCH_BATCH_SIZE = int(os.getenv("SEARCH_BATCH_SIZE", 10))
 
+# Vector Schema Configuration
+DEFAULT_VECTOR_SCHEMA = {"algorithm": "HNSW", "m": 16, "ef_construction": 200}
+VECTOR_SCHEMA = os.getenv("VECTOR_SCHEMA", json.dumps(DEFAULT_VECTOR_SCHEMA))
+
 # Redis Connection Information
 REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
 REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
@@ -200,6 +204,12 @@ async def ingest_chunks_to_redis(file_name: str, chunks: List, embedder, index_n
 
     # if data will be saved to a different index name than the default one
     ingest_index_name = index_name if index_name else INDEX_NAME
+    # Parse vector schema
+    try:
+        vector_schema = json.loads(VECTOR_SCHEMA)
+    except json.JSONDecodeError as e:
+        logger.error(f"Invalid VECTOR_SCHEMA format: {e}")
+        vector_schema = DEFAULT_VECTOR_SCHEMA
 
     file_ids = []
     for i in range(0, num_chunks, batch_size):
@@ -214,6 +224,7 @@ async def ingest_chunks_to_redis(file_name: str, chunks: List, embedder, index_n
             embedding=embedder,
             index_name=ingest_index_name,
             redis_url=REDIS_URL,
+            vector_schema=vector_schema,
         )
         if logflag:
             logger.info(f"[ redis ingest chunks ] keys: {keys}")

comps/retrievers/src/integrations/config.py

@@ -46,7 +46,9 @@ def get_boolean_env_var(var_name, default_value=False):
 LOCAL_EMBEDDING_MODEL = os.getenv("LOCAL_EMBEDDING_MODEL", "maidalun1020/bce-embedding-base_v1")
 TEI_EMBEDDING_ENDPOINT = os.getenv("TEI_EMBEDDING_ENDPOINT", "")
 BRIDGE_TOWER_EMBEDDING = os.getenv("BRIDGE_TOWER_EMBEDDING", False)
+
 HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN", "")
+ENABLE_SCHEMA = get_boolean_env_var("ENABLE_SCHEMA", False)
 
 # OpenAI
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

comps/retrievers/src/integrations/redis.py

@@ -26,6 +26,7 @@
 from .config import (
     BRIDGE_TOWER_EMBEDDING,
     EMBED_MODEL,
+    ENABLE_SCHEMA,
     HF_TOKEN,
     INDEX_NAME,
     INDEX_SCHEMA,
@@ -100,7 +101,13 @@ async def _initialize_client(self, index_name=INDEX_NAME) -> Redis:
                 client = Redis(
                     embedding=self.embeddings, index_name=index_name, index_schema=INDEX_SCHEMA, redis_url=REDIS_URL
                 )
+            elif ENABLE_SCHEMA:
+                logger.info(f"generate redis instance with index_schema:{INDEX_SCHEMA}")
+                client = Redis(
+                    embedding=self.embeddings, index_name=index_name, index_schema=INDEX_SCHEMA, redis_url=REDIS_URL
+                )
             else:
+                logger.info(f"generate redis instance with index_name:{INDEX_NAME}")
                 client = Redis(embedding=self.embeddings, index_name=index_name, redis_url=REDIS_URL)
             return client
         except Exception as e:

comps/retrievers/src/integrations/redis_schema_hnsw.yml

@@ -0,0 +1,9 @@
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+vector:
+  - name: content_vector
+    algorithm: HNSW
+    datatype: FLOAT32
+    dims: 768
+    distance_metric: COSINE