diff --git a/graphrag/mind_map_extractor.py b/graphrag/mind_map_extractor.py index d25b24c24..6705fa6e6 100644 --- a/graphrag/mind_map_extractor.py +++ b/graphrag/mind_map_extractor.py @@ -39,7 +39,6 @@ class MindMapResult: class MindMapExtractor: - _llm: CompletionLLM _input_text_key: str _mind_map_prompt: str @@ -93,7 +92,7 @@ class MindMapExtractor: max_workers = int(os.environ.get('MINDMAP_EXTRACTOR_MAX_WORKERS', 12)) exe = ThreadPoolExecutor(max_workers=max_workers) threads = [] - token_count = max(self._llm.max_length * 0.8, self._llm.max_length-512) + token_count = max(self._llm.max_length * 0.8, self._llm.max_length - 512) texts = [] res = [] cnt = 0 @@ -163,14 +162,14 @@ class MindMapExtractor: elif isinstance(value, list): new_value = {} for i in range(len(value)): - if isinstance(value[i], list): + if isinstance(value[i], list) and i > 0: new_value[value[i - 1]] = value[i][0] data[key] = new_value else: continue return data - def _todict(self, layer:collections.OrderedDict): + def _todict(self, layer: collections.OrderedDict): to_ret = layer if isinstance(layer, collections.OrderedDict): to_ret = dict(layer)