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

Files changed (1) hide show
  1. 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
- for i in obj: keyset.add(i)
69
- return [{"id": re.sub(r"\*+", "", i), "children": []} for i in obj if re.sub(r"\*+", "", i)]
 
70
  arr = []
71
  for k, v in obj.items():
72
  k = self._key(k)
73
- if not k or k in keyset: continue
74
- keyset.add(k)
75
- arr.append({
76
- "id": k,
77
- "children": self._be_children(v, keyset)
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.keys()) > 1:
114
- keyset = set(
115
- [re.sub(r"\*+", "", k) for k, v in merge_json.items() if isinstance(v, dict) and re.sub(r"\*+", "", k)])
116
- merge_json = {"id": "root",
117
- "children": [{"id": self._key(k), "children": self._be_children(v, keyset)} for k, v in
118
- merge_json.items() if isinstance(v, dict) and self._key(k)]}
 
 
 
 
 
 
 
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], set([k]))}
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")