From a353d883a08dd7ed7d6eec6c8306ec0872598294 Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Fri, 5 Jul 2024 16:07:42 +0530 Subject: [PATCH] Make it optional to set the encoder, cross-encoder configs via admin UI Upgrade sentence-transformer, add einops dependency for some sentence transformer models like nomic --- pyproject.toml | 3 ++- src/khoj/database/models/__init__.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b3b4af35..de82586f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,8 @@ dependencies = [ "pyyaml ~= 6.0", "rich >= 13.3.1", "schedule == 1.1.0", - "sentence-transformers == 2.5.1", + "sentence-transformers == 3.0.1", + "einops == 0.8.0", "transformers >= 4.28.0", "torch == 2.2.2", "uvicorn == 0.17.6", diff --git a/src/khoj/database/models/__init__.py b/src/khoj/database/models/__init__.py index 62afdd2b..096d14bc 100644 --- a/src/khoj/database/models/__init__.py +++ b/src/khoj/database/models/__init__.py @@ -215,11 +215,11 @@ class SearchModelConfig(BaseModel): # Bi-encoder model of sentence-transformer type to load from HuggingFace bi_encoder = models.CharField(max_length=200, default="thenlper/gte-small") # Config passed to the sentence-transformer model constructor. E.g. device="cuda:0", trust_remote_server=True etc. - bi_encoder_model_config = models.JSONField(default=dict) + bi_encoder_model_config = models.JSONField(default=dict, blank=True) # Query encode configs like prompt, precision, normalize_embeddings, etc. for sentence-transformer models - bi_encoder_query_encode_config = models.JSONField(default=dict) + bi_encoder_query_encode_config = models.JSONField(default=dict, blank=True) # Docs encode configs like prompt, precision, normalize_embeddings, etc. for sentence-transformer models - bi_encoder_docs_encode_config = models.JSONField(default=dict) + bi_encoder_docs_encode_config = models.JSONField(default=dict, blank=True) # Cross-encoder model of sentence-transformer type to load from HuggingFace cross_encoder = models.CharField(max_length=200, default="mixedbread-ai/mxbai-rerank-xsmall-v1") # Inference server API endpoint to use for embeddings inference. Bi-encoder model should be hosted on this server