Warm up reranker on model switch (onyx-dot-app#2408)

* warm up reranker on model switch * properly type * fix issue * Update search_settings.py
UiPath · Oct 2, 2024 · a8af95e · a8af95e
1 parent 774f678
commit a8af95e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/backend/danswer/db/search_settings.py b/backend/danswer/db/search_settings.py
@@ -20,6 +20,7 @@
 from danswer.db.models import SearchSettings
 from danswer.indexing.models import IndexingSetting
 from danswer.natural_language_processing.search_nlp_models import clean_model_name
+from danswer.natural_language_processing.search_nlp_models import warm_up_cross_encoder
 from danswer.search.models import SavedSearchSettings
 from danswer.server.manage.embedding.models import (
     CloudEmbeddingProvider as ServerCloudEmbeddingProvider,
@@ -180,6 +181,14 @@ def update_current_search_settings(
         logger.warning("No current search settings found to update")
         return
 
+    # Whenever we update the current search settings, we should ensure that the local reranking model is warmed up.
+    if (
+        current_settings.provider_type is None
+        and search_settings.rerank_model_name is not None
+        and current_settings.rerank_model_name != search_settings.rerank_model_name
+    ):
+        warm_up_cross_encoder(search_settings.rerank_model_name)
+
     update_search_settings(current_settings, search_settings, preserved_fields)
     db_session.commit()
     logger.info("Current search settings updated successfully")

diff --git a/web/src/app/admin/embeddings/RerankingFormPage.tsx b/web/src/app/admin/embeddings/RerankingFormPage.tsx
@@ -66,6 +66,7 @@ const RerankingDetailsForm = forwardRef<
       >
         {({ values, setFieldValue, resetForm }) => {
           const resetRerankingValues = () => {
+            setRerankingDetails(originalRerankingDetails);
             resetForm();
           };
 
@@ -191,7 +192,8 @@ const RerankingDetailsForm = forwardRef<
                             {card.rerank_provider_type ===
                             RerankerProvider.LITELLM ? (
                               <LiteLLMIcon size={24} className="mr-2" />
-                            ) : RerankerProvider.COHERE ? (
+                            ) : card.rerank_provider_type ===
+                              RerankerProvider.COHERE ? (
                               <CohereIcon size={24} className="mr-2" />
                             ) : (
                               <MixedBreadIcon size={24} className="mr-2" />