Browse Source

fix: 扩展文档的代码逻辑 (#227)

* fix: 扩展文档的代码逻辑

* Update local_doc_qa.py

---------

Co-authored-by: imClumsyPanda <littlepanda0716@gmail.com>
Winter 2 years ago
parent
commit
47922d2ee3
1 changed files with 2 additions and 1 deletions
  1. 2 1
      chains/local_doc_qa.py

+ 2 - 1
chains/local_doc_qa.py

@@ -68,6 +68,7 @@ def similarity_search_with_score_by_vector(
     scores, indices = self.index.search(np.array([embedding], dtype=np.float32), k)
     docs = []
     id_set = set()
+    store_len = len(self.index_to_docstore_id)
     for j, i in enumerate(indices[0]):
         if i == -1:
             # This happens when not enough docs are returned.
@@ -76,7 +77,7 @@ def similarity_search_with_score_by_vector(
         doc = self.docstore.search(_id)
         id_set.add(i)
         docs_len = len(doc.page_content)
-        for k in range(1, max(i, len(docs) - i)):
+        for k in range(1, max(i, store_len-i)):
             break_flag = False
             for l in [i + k, i - k]:
                 if 0 <= l < len(self.index_to_docstore_id):