klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-07 21:29:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add search model config from khoj.yml to Postgres DB via migration script

Browse Source
This commit is contained in:
Debanjum Singh Solanky
2023-11-14 17:18:20 -08:00
parent 4af194d74b
commit 3d8d6145f2
1 changed files with 16 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
src/khoj/migrations/migrate_server_pg.py
View File

@@ -30,7 +30,7 @@ search-type:
encoder: sentence-transformers/all-MiniLM-L6-v2 encoder: sentence-transformers/all-MiniLM-L6-v2
encoder-type: null encoder-type: null
model-directory: ~/.khoj/search/symmetric model-directory: ~/.khoj/search/symmetric
version: 0.12.4 version: 0.14.0
The new version will looks like this: The new version will looks like this:
@@ -53,11 +53,7 @@ search-type:
asymmetric: asymmetric:
cross-encoder: cross-encoder/ms-marco-MiniLM-L-6-v2 cross-encoder: cross-encoder/ms-marco-MiniLM-L-6-v2
encoder: sentence-transformers/multi-qa-MiniLM-L6-cos-v1 encoder: sentence-transformers/multi-qa-MiniLM-L6-cos-v1
image: version: 0.15.0
encoder: sentence-transformers/clip-ViT-B-32
encoder-type: null
model-directory: /Users/si/.khoj/search/image
version: 0.12.4
""" """
import logging import logging
@@ -68,6 +64,7 @@ from database.models import (
OpenAIProcessorConversationConfig, OpenAIProcessorConversationConfig,
OfflineChatProcessorConversationConfig, OfflineChatProcessorConversationConfig,
ChatModelOptions, ChatModelOptions,
SearchModel,
) )
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -87,6 +84,19 @@ def migrate_server_pg(args):
if raw_config is None: if raw_config is None:
return args return args
if "search-type" in raw_config and raw_config["search-type"]:
if "asymmetric" in raw_config["search-type"]:
# Delete all existing search models
SearchModel.objects.filter(model_type=SearchModel.ModelType.TEXT).delete()
# Create new search model from existing Khoj YAML config
asymmetric_search = raw_config["search-type"]["asymmetric"]
SearchModel.objects.create(
name="default",
model_type=SearchModel.ModelType.TEXT,
bi_encoder=asymmetric_search.get("encoder"),
cross_encoder=asymmetric_search.get("cross-encoder"),
)
if "processor" in raw_config and raw_config["processor"] and "conversation" in raw_config["processor"]: if "processor" in raw_config and raw_config["processor"] and "conversation" in raw_config["processor"]:
processor_conversation = raw_config["processor"]["conversation"] processor_conversation = raw_config["processor"]["conversation"]