AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-08-13 02:58:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

Let parallism of RAPTOR controlable. (#5379)

Browse Source 
### What problem does this PR solve?

#4874
### Type of change

- [x] Refactoring
This commit is contained in:
Kevin Hu 2025-02-26 15:58:06 +08:00 committed by GitHub
parent 5cab6c4ccb
commit 96e9d50060
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
rag/raptor.py
View File

@ -14,6 +14,7 @@
# limitations under the License. # limitations under the License.
# #
import logging import logging
import os
import re import re
from concurrent.futures import ThreadPoolExecutor, ALL_COMPLETED, wait from concurrent.futures import ThreadPoolExecutor, ALL_COMPLETED, wait
from threading import Lock from threading import Lock
@ -122,7 +123,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
lbls = [np.where(prob > self._threshold)[0] for prob in probs] lbls = [np.where(prob > self._threshold)[0] for prob in probs]
lbls = [lbl[0] if isinstance(lbl, np.ndarray) else lbl for lbl in lbls] lbls = [lbl[0] if isinstance(lbl, np.ndarray) else lbl for lbl in lbls]
lock = Lock() lock = Lock()
with ThreadPoolExecutor(max_workers=12) as executor: with ThreadPoolExecutor(max_workers=int(os.environ.get('GRAPH_EXTRACTOR_MAX_WORKERS', 10))) as executor:
threads = [] threads = []
for c in range(n_clusters): for c in range(n_clusters):
ck_idx = [i + start for i in range(len(lbls)) if lbls[i] == c] ck_idx = [i + start for i in range(len(lbls)) if lbls[i] == c]