diff --git a/graphrag/entity_resolution.py b/graphrag/entity_resolution.py index a8e044a9f..12b2bc9a6 100644 --- a/graphrag/entity_resolution.py +++ b/graphrag/entity_resolution.py @@ -99,7 +99,7 @@ class EntityResolution(Extractor): for candidate_resolution_i in candidate_resolution.items(): if not candidate_resolution_i[1]: continue - nursery.start_soon(self._resolve_candidate(candidate_resolution_i, resolution_result)) + nursery.start_soon(lambda: self._resolve_candidate(candidate_resolution_i, resolution_result)) connect_graph = nx.Graph() removed_entities = [] diff --git a/graphrag/general/extractor.py b/graphrag/general/extractor.py index 2a956615d..5576101c7 100644 --- a/graphrag/general/extractor.py +++ b/graphrag/general/extractor.py @@ -101,7 +101,7 @@ class Extractor: async with trio.open_nursery() as nursery: for i, (cid, ck) in enumerate(chunks): ck = truncate(ck, int(self._llm.max_length*0.8)) - nursery.start_soon(self._process_single_content, (cid, ck), i, len(chunks), out_results) + nursery.start_soon(lambda: self._process_single_content((cid, ck), i, len(chunks), out_results)) maybe_nodes = defaultdict(list) maybe_edges = defaultdict(list) @@ -120,7 +120,7 @@ class Extractor: all_entities_data = [] async with trio.open_nursery() as nursery: for en_nm, ents in maybe_nodes.items(): - nursery.start_soon(self._merge_nodes, en_nm, ents, all_entities_data) + nursery.start_soon(lambda: self._merge_nodes(en_nm, ents, all_entities_data)) now = trio.current_time() if callback: callback(msg = f"Entities merging done, {now-start_ts:.2f}s.") @@ -130,7 +130,7 @@ class Extractor: all_relationships_data = [] async with trio.open_nursery() as nursery: for (src, tgt), rels in maybe_edges.items(): - nursery.start_soon(self._merge_edges, src, tgt, rels, all_relationships_data) + nursery.start_soon(lambda: self._merge_edges(src, tgt, rels, all_relationships_data)) now = trio.current_time() if callback: callback(msg = f"Relationships merging done, {now-start_ts:.2f}s.") diff --git a/graphrag/general/mind_map_extractor.py b/graphrag/general/mind_map_extractor.py index b4ee6343e..c9ac2a64b 100644 --- a/graphrag/general/mind_map_extractor.py +++ b/graphrag/general/mind_map_extractor.py @@ -93,13 +93,13 @@ class MindMapExtractor(Extractor): for i in range(len(sections)): section_cnt = num_tokens_from_string(sections[i]) if cnt + section_cnt >= token_count and texts: - nursery.start_soon(self._process_document, "".join(texts), prompt_variables, res) + nursery.start_soon(lambda: self._process_document("".join(texts), prompt_variables, res)) texts = [] cnt = 0 texts.append(sections[i]) cnt += section_cnt if texts: - nursery.start_soon(self._process_document, "".join(texts), prompt_variables, res) + nursery.start_soon(lambda: self._process_document("".join(texts), prompt_variables, res)) if not res: return MindMapResult(output={"id": "root", "children": []}) merge_json = reduce(self._merge, res) diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index b03732b48..d035b87c9 100644 --- a/rag/svr/task_executor.py +++ b/rag/svr/task_executor.py @@ -297,7 +297,7 @@ async def build_chunks(task, progress_callback): return async with trio.open_nursery() as nursery: for d in docs: - nursery.start_soon(doc_keyword_extraction, chat_mdl, d, task["parser_config"]["auto_keywords"]) + nursery.start_soon(lambda: doc_keyword_extraction(chat_mdl, d, task["parser_config"]["auto_keywords"])) progress_callback(msg="Keywords generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st)) if task["parser_config"].get("auto_questions", 0): @@ -316,7 +316,7 @@ async def build_chunks(task, progress_callback): d["question_tks"] = rag_tokenizer.tokenize("\n".join(d["question_kwd"])) async with trio.open_nursery() as nursery: for d in docs: - nursery.start_soon(doc_question_proposal, chat_mdl, d, task["parser_config"]["auto_questions"]) + nursery.start_soon(lambda: doc_question_proposal(chat_mdl, d, task["parser_config"]["auto_questions"])) progress_callback(msg="Question generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st)) if task["kb_parser_config"].get("tag_kb_ids", []): @@ -356,7 +356,7 @@ async def build_chunks(task, progress_callback): d[TAG_FLD] = json.loads(cached) async with trio.open_nursery() as nursery: for d in docs_to_tag: - nursery.start_soon(doc_content_tagging, chat_mdl, d, topn_tags) + nursery.start_soon(lambda: doc_content_tagging(chat_mdl, d, topn_tags)) progress_callback(msg="Tagging {} chunks completed in {:.2f}s".format(len(docs), timer() - st)) return docs