Browse Source

Fix entity_types. Close #6287 and #6608 (#6632)

### What problem does this PR solve?

Fix entity_types. Close #6287 and #6608

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
tags/v0.18.0
Zhichang Yu 7 months ago
parent
commit
d32a35d8fd
No account linked to committer's email address
1 changed files with 2 additions and 3 deletions
  1. 2
    3
      graphrag/general/graph_extractor.py

+ 2
- 3
graphrag/general/graph_extractor.py View File

@@ -11,7 +11,7 @@ from dataclasses import dataclass
import tiktoken
import trio

from graphrag.general.extractor import Extractor, ENTITY_EXTRACTION_MAX_GLEANINGS, DEFAULT_ENTITY_TYPES
from graphrag.general.extractor import Extractor, ENTITY_EXTRACTION_MAX_GLEANINGS
from graphrag.general.graph_prompt import GRAPH_EXTRACTION_PROMPT, CONTINUE_PROMPT, LOOP_PROMPT
from graphrag.utils import ErrorHandlerFn, perform_variable_replacements, chat_limiter, split_string_by_multi_markers
from rag.llm.chat_model import Base as CompletionLLM
@@ -91,11 +91,10 @@ class GraphExtractor(Extractor):

# Wire defaults into the prompt variables
self._prompt_variables = {
"entity_types": entity_types,
self._tuple_delimiter_key: DEFAULT_TUPLE_DELIMITER,
self._record_delimiter_key: DEFAULT_RECORD_DELIMITER,
self._completion_delimiter_key: DEFAULT_COMPLETION_DELIMITER,
self._entity_types_key: ",".join(DEFAULT_ENTITY_TYPES),
self._entity_types_key: entity_types,
}

async def _process_single_content(self, chunk_key_dp: tuple[str, str], chunk_seq: int, num_chunks: int, out_results):

Loading…
Cancel
Save