Warm up models before first document indexed (#333)

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()

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))