Reranking Model

integreat_chat/core/settings.py

@@ -97,10 +97,13 @@
 RAG_FALLBACK_LANGUAGE = "en"
 
 # SEARCH_MAX_DOCUMENTS - number of documents retrieved from the VDB
-SEARCH_MAX_DOCUMENTS = 15
+SEARCH_MAX_DOCUMENTS = 30
 SEARCH_SCORE_THRESHOLD = 0.1
 SEARCH_MAX_PAGES = 10
 SEARCH_EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
+OPENSEARCH_CROSSENCODER_MODEL_NAME = (
+    "huggingface/cross-encoders/ms-marco-MiniLM-L-6-v2"
+)
 OPENSEARCH_EMBEDDING_MODEL_NAME = (
     "huggingface/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
 )
@@ -115,7 +118,9 @@
 SEARCH_DENSE_WEIGHT = 0.7
 SEARCH_FALLBACK_LANGUAGE = "en"
 SEARCH_OPENSEARCH_MODEL_ID = config["OPENSEARCH"]["MODEL_ID"]
+SEARCH_OPENSEARCH_MODEL_ID_RERANKER = config["OPENSEARCH"]["MODEL_ID_RERANKER"]
 SEARCH_OPENSEARCH_MODEL_GROUP_ID = config["OPENSEARCH"]["MODEL_GROUP_ID"]
+
 OPENSEARCH_USER = (
     config["OPENSEARCH"]["USER"]
     if "USER" in config["OPENSEARCH"]

integreat_chat/search/management/commands/opensearch_setup.py

@@ -14,12 +14,13 @@ class Command(BaseCommand):
 
     def handle(self, *args, **options):
         oss = OpenSearchSetup(password=settings.OPENSEARCH_PASSWORD)
-        group_id, model_id = oss.setup()
+        group_id, model_id, model_id_ce = oss.setup()
         self.stdout.write(
             self.style.SUCCESS(  # pylint: disable=no-member
                 f'Successfully set up OpenSearch. Change the following settings '
                 f'in the OPENSEARCH section of your config:\n'
                 f'MODEL_GROUP_ID = {group_id}\n'
-                f'MODEL_ID = {model_id}\n'
+                f'MODEL_ID_EMBEDDING_MODEL = {model_id}\n'
+                f'MODEL_ID_CROSSENCODER_MODEL = {model_id_ce}\n'
             )
         )

integreat_chat/search/services/opensearch.py

@@ -2,15 +2,19 @@
 Setup and use of OpenSearch
 """
 
+import logging
 import hashlib
 import time
 from datetime import timedelta, datetime
+import math
 
 import requests
 from django.conf import settings
 from langchain_text_splitters import HTMLHeaderTextSplitter
 from integreat_chat.core.utils.integreat_cms import get_all_pages, get_parent_page_titles
 
+LOGGER = logging.getLogger("django")
+
 class OpenSearch:
     """
     Class for searching and updating documents in OpenSearch
@@ -36,6 +40,7 @@ def __init__(
         self.user = user
         self.password = password
         self.model_id = settings.SEARCH_OPENSEARCH_MODEL_ID
+        self.model_id_reranker = settings.SEARCH_OPENSEARCH_MODEL_ID_RERANKER
         self.model_group_id = settings.SEARCH_OPENSEARCH_MODEL_GROUP_ID
 
     def request(self, path: str, payload: dict, method: str = "GET") -> dict:
@@ -75,6 +80,7 @@ def reduce_search_result(
         """
         result = []
         found_urls = []
+        sigmoid = lambda x: 1 / (1 + math.exp(-x))
         if "hits" not in response:
             raise ValueError("Missing hits in result")
         for document in response["hits"]["hits"]:
@@ -87,7 +93,7 @@ def reduce_search_result(
                 "url": document["_source"]["url"],
                 "title": document["_source"]["title"],
                 "parent_titles": document["_source"]["parent_titles"],
-                "score": document["_score"],
+                "score": sigmoid(document["_score"]),
                 "chunk_text": document["_source"]["chunk_text"],
             })
             found_urls.append(document["_source"]["url"])
@@ -129,7 +135,7 @@ def search(self, region_slug: str, language_slug: str, message: str) -> dict:
                             "title_embedding": {
                                 "query_text": message,
                                 "model_id": self.model_id,
-                                "k": 5
+                                "k": 30
                             }
                         }
                         },
@@ -138,12 +144,19 @@ def search(self, region_slug: str, language_slug: str, message: str) -> dict:
                             "chunk_embedding": {
                                 "query_text": message,
                                 "model_id": self.model_id,
-                                "k": 5
+                                "k": 30
                             }
                         }
                         }
                     ]
                 }
+            },
+            "ext": {
+                "rerank": {
+                "query_context": {
+                    "query_text": message
+                }
+                }
             }
         }
         return self.request(
@@ -325,20 +338,26 @@ def setup(self) -> str:
         group_id = self.create_model_group()
         if not group_id:
             raise ValueError("Unexpected OpenSearch response while creating model group")
-        model_id = self.register_embedding_model(group_id)
-        if not model_id:
-            raise ValueError("Unexpected OpenSearch response while registering model")
-        self.deploy_model(model_id)
-        self.create_ingestion_pipeline(model_id)
+        model_id_embedding = self.register_embedding_model(group_id)
+        model_id_crossencoder = self.register_crossencoder_model(group_id)
+        if not model_id_embedding:
+            raise ValueError("Unexpected OpenSearch response while registering embedding model")
+        elif not model_id_crossencoder:
+            raise ValueError("Unexpected OpenSearch response while registering crossencoder model")
+        self.deploy_model(model_id_embedding)
+        self.deploy_model(model_id_crossencoder)
+        self.create_ingestion_pipeline(model_id_embedding)
         self.create_search_pipeline()
-        return group_id, model_id
+        return group_id, model_id_embedding, model_id_crossencoder
 
     def delete_model_group(self):
         """
         Delete previously created model group and model
         """
         self.request(f"/_plugins/_ml/models/{self.model_id}/_undeploy", {}, "POST")
         self.request(f"/_plugins/_ml/models/{self.model_id}", {}, "DELETE")
+        self.request(f"/_plugins/_ml/models/{self.model_id_reranker}/_undeploy", {}, "POST")
+        self.request(f"/_plugins/_ml/models/{self.model_id_reranker}", {}, "DELETE")
         self.request(f"/_plugins/_ml/model_groups/{self.model_group_id}", {}, "DELETE")
 
     def prepare_index(self, region_slug: str = "", language_slug: str = ""):
@@ -369,13 +388,37 @@ def create_model_group(self):
         Create model group
         """
         payload = {
-            "name": "integreat-chat-2025-01-31",
-            "description": "Integreat Chat embedding models"
+            "name": "integreat-chat-2025-06-05",
+            "description": "Integreat Chat model group"
         }
         response = self.request("/_plugins/_ml/model_groups/_register", payload, "POST")
+        LOGGER.debug(f"Model group response: {response}")
         if "model_group_id" in response:
             return response["model_group_id"]
         return False
+
+    def register_crossencoder_model(self, model_group_id: str) -> str:
+        """
+        Register crossencoder model
+        """
+        payload = {
+            "name": settings.OPENSEARCH_CROSSENCODER_MODEL_NAME,
+            "version": "1.0.2",
+            "model_group_id": model_group_id,
+            "model_format": "TORCH_SCRIPT"
+        }
+        register_response = self.request(
+            "/_plugins/_ml/models/_register", payload, "POST"
+        )
+        LOGGER.debug(f"Embedding model response: {register_response}")
+        if "task_id" in register_response:
+            for n in range(0, 10):
+                time.sleep(5)
+                if "model_id" in (task_response := self.request(
+                    f"/_plugins/_ml/tasks/{register_response['task_id']}", {}, "GET"
+                )):
+                    return task_response["model_id"]
+        return False
 
     def register_embedding_model(self, model_group_id: str) -> str:
         """
@@ -390,6 +433,7 @@ def register_embedding_model(self, model_group_id: str) -> str:
         register_response = self.request(
             "/_plugins/_ml/models/_register", payload, "POST"
         )
+        LOGGER.debug(f"Cross Encoder model response: {register_response}")
         if "task_id" in register_response:
             for n in range(0, 10):  # pylint: disable=W0612
                 time.sleep(5)
@@ -461,7 +505,20 @@ def create_search_pipeline(self):
                         }
                     }
                 }
-            ]
+            ],
+            "response_processors": [
+            {
+                "rerank": {
+                    "ml_opensearch": {
+                        "model_id": self.model_id_reranker
+                    },
+                    "context": {
+                    "document_fields": [
+                        "chunk_text"
+                    ]
+                    }
+                }
+            }]
         }
         self.request(f"/_search/pipeline/{self.search_pipeline_name}", payload, "PUT")
 
@@ -496,6 +553,7 @@ def create_index(self, index_slug: str):
             "settings": {
                 "index.knn": True,
                 "default_pipeline": self.ingest_pipeline_name,
+                "index.search.default_pipeline": self.search_pipeline_name,
             },
             "mappings": {
                 "properties": {