From 384bf1befebb334f0d7e3acd1bf4f49402eab2b4 Mon Sep 17 00:00:00 2001 From: Yuhong Sun Date: Thu, 24 Aug 2023 20:01:50 -0700 Subject: [PATCH] Warm up models before first document indexed (#333) --- backend/danswer/background/update.py | 4 ++++ backend/danswer/search/search_utils.py | 6 +++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/backend/danswer/background/update.py b/backend/danswer/background/update.py index b0491a288b9..91bea0dbada 100755 --- a/backend/danswer/background/update.py +++ b/backend/danswer/background/update.py @@ -35,6 +35,7 @@ from danswer.db.models import Connector from danswer.db.models import IndexAttempt from danswer.db.models import IndexingStatus +from danswer.search.search_utils import warm_up_models from danswer.utils.logger import IndexAttemptSingleton from danswer.utils.logger import setup_logger @@ -417,4 +418,7 @@ def update_loop(delay: int = 10, num_workers: int = NUM_INDEXING_WORKERS) -> Non if __name__ == "__main__": + logger.info("Warming up Embedding Model(s)") + warm_up_models(indexer_only=True) + logger.info("Starting Indexing Loop") update_loop() diff --git a/backend/danswer/search/search_utils.py b/backend/danswer/search/search_utils.py index 4b896fcd6dd..4839f740850 100644 --- a/backend/danswer/search/search_utils.py +++ b/backend/danswer/search/search_utils.py @@ -61,10 +61,14 @@ def get_default_intent_model() -> TFDistilBertForSequenceClassification: return _INTENT_MODEL -def warm_up_models() -> None: +def warm_up_models(indexer_only: bool = False) -> None: warm_up_str = "Danswer is amazing" get_default_tokenizer()(warm_up_str) get_default_embedding_model().encode(warm_up_str) + + if indexer_only: + return + cross_encoders = get_default_reranking_model_ensemble() [ cross_encoder.predict((warm_up_str, warm_up_str))