Kevin Hu commited on
Commit
140db5b
·
1 Parent(s): d102a8b

fix index of range (#3279)

Browse files

### What problem does this PR solve?

#3273

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

Files changed (1) hide show
  1. graphrag/mind_map_extractor.py +3 -4
graphrag/mind_map_extractor.py CHANGED
@@ -39,7 +39,6 @@ class MindMapResult:
39
 
40
 
41
  class MindMapExtractor:
42
-
43
  _llm: CompletionLLM
44
  _input_text_key: str
45
  _mind_map_prompt: str
@@ -93,7 +92,7 @@ class MindMapExtractor:
93
  max_workers = int(os.environ.get('MINDMAP_EXTRACTOR_MAX_WORKERS', 12))
94
  exe = ThreadPoolExecutor(max_workers=max_workers)
95
  threads = []
96
- token_count = max(self._llm.max_length * 0.8, self._llm.max_length-512)
97
  texts = []
98
  res = []
99
  cnt = 0
@@ -163,14 +162,14 @@ class MindMapExtractor:
163
  elif isinstance(value, list):
164
  new_value = {}
165
  for i in range(len(value)):
166
- if isinstance(value[i], list):
167
  new_value[value[i - 1]] = value[i][0]
168
  data[key] = new_value
169
  else:
170
  continue
171
  return data
172
 
173
- def _todict(self, layer:collections.OrderedDict):
174
  to_ret = layer
175
  if isinstance(layer, collections.OrderedDict):
176
  to_ret = dict(layer)
 
39
 
40
 
41
  class MindMapExtractor:
 
42
  _llm: CompletionLLM
43
  _input_text_key: str
44
  _mind_map_prompt: str
 
92
  max_workers = int(os.environ.get('MINDMAP_EXTRACTOR_MAX_WORKERS', 12))
93
  exe = ThreadPoolExecutor(max_workers=max_workers)
94
  threads = []
95
+ token_count = max(self._llm.max_length * 0.8, self._llm.max_length - 512)
96
  texts = []
97
  res = []
98
  cnt = 0
 
162
  elif isinstance(value, list):
163
  new_value = {}
164
  for i in range(len(value)):
165
+ if isinstance(value[i], list) and i > 0:
166
  new_value[value[i - 1]] = value[i][0]
167
  data[key] = new_value
168
  else:
169
  continue
170
  return data
171
 
172
+ def _todict(self, layer: collections.OrderedDict):
173
  to_ret = layer
174
  if isinstance(layer, collections.OrderedDict):
175
  to_ret = dict(layer)