fix: sieve of documents was not getting updated if top_k changed

Author: thiswillbeyourgithub

WDoc/WDoc.py

@@ -1426,7 +1426,7 @@ def retrieve_documents(inputs):
                 }
                 rag_chain = (
                     retrieve_documents
-                    | sieve_documents(top_k=self.top_k, max_top_k=self.max_top_k)
+                    | sieve_documents(instance=self)
                     | refilter_documents
                 )
                 tried_top_k = []
@@ -1576,7 +1576,7 @@ def retrieve_documents(inputs):
 
             rag_chain = (
                 retrieve_documents
-                | sieve_documents(top_k=self.top_k, max_top_k=self.max_top_k)
+                | sieve_documents(instance=self)
                 | pbar_chain(
                         llm=self.eval_llm,
                         len_func="len(inputs['unfiltered_docs'])",

WDoc/utils/tasks/query.py

@@ -48,22 +48,25 @@ def check_intermediate_answer(ans: str) -> bool:
 
 
 @optional_typecheck
-def sieve_documents(top_k: int, max_top_k: int) -> RunnableLambda:
+def sieve_documents(instance) -> RunnableLambda:
     """cap the number of retrieved documents as if multiple retrievers are used
     we can end up with a lot more document!
     """
-    assert max_top_k >= top_k
     @chain
     @optional_typecheck
     def _sieve(inputs: dict) -> dict:
         assert "question_to_answer" in inputs, inputs.keys()
         assert "unfiltered_docs" in inputs, inputs.keys()
-        if len(inputs) > top_k:
+        # we have to pass an instance otherwise we can't know if the top_k got updated
+        assert hasattr(instance, "top_k")
+        assert hasattr(instance, "max_top_k")
+        assert instance.max_top_k >= instance.top_k
+        if len(inputs) > instance.top_k:
             red(
                 "Number of documents found via embeddings was "
-                f"'{inputs['unfiltered_docs']}' which is > top_k ({top_k}) "
+                f"'{inputs['unfiltered_docs']}' which is > top_k ({instance.top_k}) "
                 "so we crop")
-        inputs["unfiltered_docs"] = inputs["unfiltered_docs"][:top_k]
+        inputs["unfiltered_docs"] = inputs["unfiltered_docs"][:instance.top_k]
         return inputs
     return _sieve