dale_yxc/ragflow_python
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Let ThreadPool exit gracefully. (#3653)

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

#3646

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
This commit is contained in:
Kevin Hu
2024-11-26 16:31:07 +08:00
committed by GitHub
parent 5c59651bda
commit 0891a393d7
5 changed files with 37 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
graphrag/mind_map_extractor.py
View File

@@ -88,26 +88,26 @@ class MindMapExtractor:
prompt_variables = {}
try:
max_workers = int(os.environ.get('MINDMAP_EXTRACTOR_MAX_WORKERS', 12))
exe = ThreadPoolExecutor(max_workers=max_workers)
threads = []
token_count = max(self._llm.max_length * 0.8, self._llm.max_length - 512)
texts = []
res = []
cnt = 0
for i in range(len(sections)):
section_cnt = num_tokens_from_string(sections[i])
if cnt + section_cnt >= token_count and texts:
max_workers = int(os.environ.get('MINDMAP_EXTRACTOR_MAX_WORKERS', 12))
with ThreadPoolExecutor(max_workers=max_workers) as exe:
threads = []
token_count = max(self._llm.max_length * 0.8, self._llm.max_length - 512)
texts = []
cnt = 0
for i in range(len(sections)):
section_cnt = num_tokens_from_string(sections[i])
if cnt + section_cnt >= token_count and texts:
threads.append(exe.submit(self._process_document, "".join(texts), prompt_variables))
texts = []
cnt = 0
texts.append(sections[i])
cnt += section_cnt
if texts:
threads.append(exe.submit(self._process_document, "".join(texts), prompt_variables))
texts = []
cnt = 0
texts.append(sections[i])
cnt += section_cnt
if texts:
threads.append(exe.submit(self._process_document, "".join(texts), prompt_variables))
for i, _ in enumerate(threads):
res.append(_.result())
for i, _ in enumerate(threads):
res.append(_.result())
if not res:
return MindMapResult(output={"id": "root", "children": []})