yqkcn
commited on
Commit
·
5f0746a
1
Parent(s):
94ac15e
format mind_map_extractor code (#2686)
Browse files### What problem does this PR solve?
format mind_map_extractor code
### Type of change
- [x] Refactoring
- graphrag/mind_map_extractor.py +25 -15
graphrag/mind_map_extractor.py
CHANGED
@@ -65,17 +65,20 @@ class MindMapExtractor:
|
|
65 |
if isinstance(obj, str):
|
66 |
obj = [obj]
|
67 |
if isinstance(obj, list):
|
68 |
-
|
69 |
-
|
|
|
70 |
arr = []
|
71 |
for k, v in obj.items():
|
72 |
k = self._key(k)
|
73 |
-
if
|
74 |
-
|
75 |
-
|
76 |
-
|
77 |
-
|
78 |
-
|
|
|
|
|
79 |
return arr
|
80 |
|
81 |
def __call__(
|
@@ -110,15 +113,22 @@ class MindMapExtractor:
|
|
110 |
return MindMapResult(output={"id": "root", "children": []})
|
111 |
|
112 |
merge_json = reduce(self._merge, res)
|
113 |
-
if len(merge_json
|
114 |
-
|
115 |
-
|
116 |
-
merge_json = {
|
117 |
-
|
118 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
119 |
else:
|
120 |
k = self._key(list(merge_json.keys())[0])
|
121 |
-
merge_json = {"id": k, "children": self._be_children(list(merge_json.items())[0][1],
|
122 |
|
123 |
except Exception as e:
|
124 |
logging.exception("error mind graph")
|
|
|
65 |
if isinstance(obj, str):
|
66 |
obj = [obj]
|
67 |
if isinstance(obj, list):
|
68 |
+
keyset.update(obj)
|
69 |
+
obj = [re.sub(r"\*+", "", i) for i in obj]
|
70 |
+
return [{"id": i, "children": []} for i in obj if i]
|
71 |
arr = []
|
72 |
for k, v in obj.items():
|
73 |
k = self._key(k)
|
74 |
+
if k and k not in keyset:
|
75 |
+
keyset.add(k)
|
76 |
+
arr.append(
|
77 |
+
{
|
78 |
+
"id": k,
|
79 |
+
"children": self._be_children(v, keyset)
|
80 |
+
}
|
81 |
+
)
|
82 |
return arr
|
83 |
|
84 |
def __call__(
|
|
|
113 |
return MindMapResult(output={"id": "root", "children": []})
|
114 |
|
115 |
merge_json = reduce(self._merge, res)
|
116 |
+
if len(merge_json) > 1:
|
117 |
+
keys = [re.sub(r"\*+", "", k) for k, v in merge_json.items() if isinstance(v, dict)]
|
118 |
+
keyset = set(i for i in keys if i)
|
119 |
+
merge_json = {
|
120 |
+
"id": "root",
|
121 |
+
"children": [
|
122 |
+
{
|
123 |
+
"id": self._key(k),
|
124 |
+
"children": self._be_children(v, keyset)
|
125 |
+
}
|
126 |
+
for k, v in merge_json.items() if isinstance(v, dict) and self._key(k)
|
127 |
+
]
|
128 |
+
}
|
129 |
else:
|
130 |
k = self._key(list(merge_json.keys())[0])
|
131 |
+
merge_json = {"id": k, "children": self._be_children(list(merge_json.items())[0][1], {k})}
|
132 |
|
133 |
except Exception as e:
|
134 |
logging.exception("error mind graph")
|