AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-08-12 11:09:07 +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

40
graphrag/mind_map_extractor.py
View File

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