diff --git a/integreat_chat/core/settings.py b/integreat_chat/core/settings.py index 9ff1f941..4943d9f2 100644 --- a/integreat_chat/core/settings.py +++ b/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"] diff --git a/integreat_chat/search/management/commands/opensearch_setup.py b/integreat_chat/search/management/commands/opensearch_setup.py index 602156dc..4f4dbb7f 100644 --- a/integreat_chat/search/management/commands/opensearch_setup.py +++ b/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' ) ) diff --git a/integreat_chat/search/services/opensearch.py b/integreat_chat/search/services/opensearch.py index 8505c620..3ccd3b33 100644 --- a/integreat_chat/search/services/opensearch.py +++ b/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,13 +338,17 @@ 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): """ @@ -339,6 +356,8 @@ def delete_model_group(self): """ 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": {