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

format mind_map_extractor code (#2686)

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

format mind_map_extractor code

### Type of change

- [x] Refactoring
This commit is contained in:
yqkcn 2024-09-30 16:29:15 +08:00 committed by GitHub
parent 79b8210498
commit 9945988e44
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 25 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
graphrag/mind_map_extractor.py
View File

@ -65,17 +65,20 @@ class MindMapExtractor:
if isinstance(obj, str): if isinstance(obj, str):
obj = [obj] obj = [obj]
if isinstance(obj, list): if isinstance(obj, list):
for i in obj: keyset.add(i) keyset.update(obj)
return [{"id": re.sub(r"\*+", "", i), "children": []} for i in obj if re.sub(r"\*+", "", i)] obj = [re.sub(r"\*+", "", i) for i in obj]
return [{"id": i, "children": []} for i in obj if i]
arr = [] arr = []
for k, v in obj.items(): for k, v in obj.items():
k = self._key(k) k = self._key(k)
if not k or k in keyset: continue if k and k not in keyset:
keyset.add(k) keyset.add(k)
arr.append({ arr.append(
{
"id": k, "id": k,
"children": self._be_children(v, keyset) "children": self._be_children(v, keyset)
}) }
)
return arr return arr
def __call__( def __call__(
@ -110,15 +113,22 @@ class MindMapExtractor:
return MindMapResult(output={"id": "root", "children": []}) return MindMapResult(output={"id": "root", "children": []})
merge_json = reduce(self._merge, res) merge_json = reduce(self._merge, res)
if len(merge_json.keys()) > 1: if len(merge_json) > 1:
keyset = set( keys = [re.sub(r"\*+", "", k) for k, v in merge_json.items() if isinstance(v, dict)]
[re.sub(r"\*+", "", k) for k, v in merge_json.items() if isinstance(v, dict) and re.sub(r"\*+", "", k)]) keyset = set(i for i in keys if i)
merge_json = {"id": "root", merge_json = {
"children": [{"id": self._key(k), "children": self._be_children(v, keyset)} for k, v in "id": "root",
merge_json.items() if isinstance(v, dict) and self._key(k)]} "children": [
{
"id": self._key(k),
"children": self._be_children(v, keyset)
}
for k, v in merge_json.items() if isinstance(v, dict) and self._key(k)
]
}
else: else:
k = self._key(list(merge_json.keys())[0]) k = self._key(list(merge_json.keys())[0])
merge_json = {"id": k, "children": self._be_children(list(merge_json.items())[0][1], set([k]))} merge_json = {"id": k, "children": self._be_children(list(merge_json.items())[0][1], {k})}
except Exception as e: except Exception as e:
logging.exception("error mind graph") logging.exception("error mind graph")