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")
|