AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-06-04 11:24:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make node merging parallel. (#5324)

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

#5314

### Type of change

- [x] Performance Improvement
This commit is contained in:
Kevin Hu 2025-02-25 12:02:44 +08:00 committed by GitHub
parent e2448fb6dd
commit ecf5f6976f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 27 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
graphrag/general/extractor.py
View File

@ -97,7 +97,7 @@ class Extractor:
): ):
results = [] results = []
max_workers = int(os.environ.get('GRAPH_EXTRACTOR_MAX_WORKERS', 50)) max_workers = int(os.environ.get('GRAPH_EXTRACTOR_MAX_WORKERS', 10))
with ThreadPoolExecutor(max_workers=max_workers) as exe: with ThreadPoolExecutor(max_workers=max_workers) as exe:
threads = [] threads = []
for i, (cid, ck) in enumerate(chunks): for i, (cid, ck) in enumerate(chunks):
@ -123,12 +123,21 @@ class Extractor:
maybe_edges[tuple(sorted(k))].extend(v) maybe_edges[tuple(sorted(k))].extend(v)
logging.info("Inserting entities into storage...") logging.info("Inserting entities into storage...")
all_entities_data = [] all_entities_data = []
for en_nm, ents in maybe_nodes.items(): with ThreadPoolExecutor(max_workers=max_workers) as exe:
all_entities_data.append(self._merge_nodes(en_nm, ents)) threads = []
for en_nm, ents in maybe_nodes.items():
threads.append(
exe.submit(self._merge_nodes, en_nm, ents))
for t in threads:
n = t.result()
if not isinstance(n, Exception):
all_entities_data.append(n)
elif callback:
callback(msg="Knowledge graph nodes merging error: {}".format(str(n)))
logging.info("Inserting relationships into storage...") logging.info("Inserting relationships into storage...")
all_relationships_data = [] all_relationships_data = []
for (src,tgt), rels in maybe_edges.items(): for (src, tgt), rels in maybe_edges.items():
all_relationships_data.append(self._merge_edges(src, tgt, rels)) all_relationships_data.append(self._merge_edges(src, tgt, rels))
if not len(all_entities_data) and not len(all_relationships_data): if not len(all_entities_data) and not len(all_relationships_data):
@ -167,17 +176,20 @@ class Extractor:
sorted(set([dp["description"] for dp in entities] + already_description)) sorted(set([dp["description"] for dp in entities] + already_description))
) )
already_source_ids = flat_uniq_list(entities, "source_id") already_source_ids = flat_uniq_list(entities, "source_id")
description = self._handle_entity_relation_summary( try:
entity_name, description description = self._handle_entity_relation_summary(
) entity_name, description
node_data = dict( )
entity_type=entity_type, node_data = dict(
description=description, entity_type=entity_type,
source_id=already_source_ids, description=description,
) source_id=already_source_ids,
node_data["entity_name"] = entity_name )
self._set_entity_(entity_name, node_data) node_data["entity_name"] = entity_name
return node_data self._set_entity_(entity_name, node_data)
return node_data
except Exception as e:
return e
def _merge_edges( def _merge_edges(
self, self,