Explorar el Código

refac/fix: reranking function

Timothy Jaeryang Baek hace 2 meses
padre
commit
abe280f0a3
Se han modificado 2 ficheros con 18 adiciones y 10 borrados
  1. 14 6
      backend/open_webui/routers/retrieval.py
  2. 4 4
      backend/open_webui/utils/middleware.py

+ 14 - 6
backend/open_webui/routers/retrieval.py

@@ -2050,11 +2050,13 @@ def query_doc_handler(
                 ),
                 k=form_data.k if form_data.k else request.app.state.config.TOP_K,
                 reranking_function=(
-                    lambda sentences: (
-                        request.app.state.RERANKING_FUNCTION(sentences, user=user)
-                        if request.app.state.RERANKING_FUNCTION
-                        else None
+                    (
+                        lambda sentences: request.app.state.RERANKING_FUNCTION(
+                            sentences, user=user
+                        )
                     )
+                    if request.app.state.RERANKING_FUNCTION
+                    else None
                 ),
                 k_reranker=form_data.k_reranker
                 or request.app.state.config.TOP_K_RERANKER,
@@ -2112,8 +2114,14 @@ def query_collection_handler(
                     query, prefix=prefix, user=user
                 ),
                 k=form_data.k if form_data.k else request.app.state.config.TOP_K,
-                reranking_function=lambda sentences: request.app.state.RERANKING_FUNCTION(
-                    sentences, user=user
+                reranking_function=(
+                    (
+                        lambda sentences: request.app.state.RERANKING_FUNCTION(
+                            sentences, user=user
+                        )
+                    )
+                    if request.app.state.RERANKING_FUNCTION
+                    else None
                 ),
                 k_reranker=form_data.k_reranker
                 or request.app.state.config.TOP_K_RERANKER,

+ 4 - 4
backend/open_webui/utils/middleware.py

@@ -653,13 +653,13 @@ async def chat_completion_files_handler(
                         ),
                         k=request.app.state.config.TOP_K,
                         reranking_function=(
-                            lambda sentences: (
-                                request.app.state.RERANKING_FUNCTION(
+                            (
+                                lambda sentences: request.app.state.RERANKING_FUNCTION(
                                     sentences, user=user
                                 )
-                                if request.app.state.RERANKING_FUNCTION
-                                else None
                             )
+                            if request.app.state.RERANKING_FUNCTION
+                            else None
                         ),
                         k_reranker=request.app.state.config.TOP_K_RERANKER,
                         r=request.app.state.config.RELEVANCE_THRESHOLD,