Fix usage of get_tokenizer.

Author: dworthen

graphrag/index/text_splitting/text_splitting.py

@@ -85,12 +85,12 @@ class TokenTextSplitter(TextSplitter):
 
     def __init__(
         self,
-        tokenizer: Tokenizer = get_tokenizer(),
+        tokenizer: Tokenizer | None = None,
         **kwargs: Any,
     ):
         """Init method definition."""
         super().__init__(**kwargs)
-        self._tokenizer = tokenizer
+        self._tokenizer = tokenizer or get_tokenizer()
 
     def num_tokens(self, text: str) -> int:
         """Return the number of tokens in a string."""

graphrag/prompt_tune/generator/extract_graph_prompt.py

@@ -24,7 +24,7 @@ def create_extract_graph_prompt(
     examples: list[str],
     language: str,
     max_token_count: int,
-    tokenizer: Tokenizer = get_tokenizer(),
+    tokenizer: Tokenizer | None = None,
     json_mode: bool = False,
     output_path: Path | None = None,
     min_examples_required: int = 2,
@@ -56,6 +56,8 @@ def create_extract_graph_prompt(
     if isinstance(entity_types, list):
         entity_types = ", ".join(map(str, entity_types))
 
+    tokenizer = tokenizer or get_tokenizer()
+
     tokens_left = (
         max_token_count
         - tokenizer.num_tokens(prompt)

graphrag/query/context_builder/community_context.py

@@ -24,7 +24,7 @@
 def build_community_context(
     community_reports: list[CommunityReport],
     entities: list[Entity] | None = None,
-    tokenizer: Tokenizer = get_tokenizer(),
+    tokenizer: Tokenizer | None = None,
     use_community_summary: bool = True,
     column_delimiter: str = "|",
     shuffle_data: bool = True,
@@ -46,6 +46,7 @@ def build_community_context(
 
     The calculated weight is added as an attribute to the community reports and added to the context data table.
     """
+    tokenizer = tokenizer or get_tokenizer()
 
     def _is_included(report: CommunityReport) -> bool:
         return report.rank is not None and report.rank >= min_community_rank

graphrag/query/context_builder/conversation_history.py

@@ -148,7 +148,7 @@ def get_user_turns(self, max_user_turns: int | None = 1) -> list[str]:
 
     def build_context(
         self,
-        tokenizer: Tokenizer = get_tokenizer(),
+        tokenizer: Tokenizer | None = None,
         include_user_turns_only: bool = True,
         max_qa_turns: int | None = 5,
         max_context_tokens: int = 8000,
@@ -168,6 +168,7 @@ def build_context(
             context_name: Name of the context, default is "Conversation History".
 
         """
+        tokenizer = tokenizer or get_tokenizer()
         qa_turns = self.to_qa_turns()
         if include_user_turns_only:
             qa_turns = [

graphrag/query/context_builder/local_context.py

@@ -29,14 +29,16 @@
 
 def build_entity_context(
     selected_entities: list[Entity],
-    tokenizer: Tokenizer = get_tokenizer(),
+    tokenizer: Tokenizer | None = None,
     max_context_tokens: int = 8000,
     include_entity_rank: bool = True,
     rank_description: str = "number of relationships",
     column_delimiter: str = "|",
     context_name="Entities",
 ) -> tuple[str, pd.DataFrame]:
     """Prepare entity data table as context data for system prompt."""
+    tokenizer = tokenizer or get_tokenizer()
+
     if len(selected_entities) == 0:
         return "", pd.DataFrame()
 
@@ -91,12 +93,13 @@ def build_entity_context(
 def build_covariates_context(
     selected_entities: list[Entity],
     covariates: list[Covariate],
-    tokenizer: Tokenizer = get_tokenizer(),
+    tokenizer: Tokenizer | None = None,
     max_context_tokens: int = 8000,
     column_delimiter: str = "|",
     context_name: str = "Covariates",
 ) -> tuple[str, pd.DataFrame]:
     """Prepare covariate data tables as context data for system prompt."""
+    tokenizer = tokenizer or get_tokenizer()
     # create an empty list of covariates
     if len(selected_entities) == 0 or len(covariates) == 0:
         return "", pd.DataFrame()
@@ -155,7 +158,7 @@ def build_covariates_context(
 def build_relationship_context(
     selected_entities: list[Entity],
     relationships: list[Relationship],
-    tokenizer: Tokenizer = get_tokenizer(),
+    tokenizer: Tokenizer | None = None,
     include_relationship_weight: bool = False,
     max_context_tokens: int = 8000,
     top_k_relationships: int = 10,
@@ -164,6 +167,7 @@ def build_relationship_context(
     context_name: str = "Relationships",
 ) -> tuple[str, pd.DataFrame]:
     """Prepare relationship data tables as context data for system prompt."""
+    tokenizer = tokenizer or get_tokenizer()
     selected_relationships = _filter_relationships(
         selected_entities=selected_entities,
         relationships=relationships,

graphrag/query/context_builder/source_context.py

@@ -20,14 +20,15 @@
 
 def build_text_unit_context(
     text_units: list[TextUnit],
-    tokenizer: Tokenizer = get_tokenizer(),
+    tokenizer: Tokenizer | None = None,
     column_delimiter: str = "|",
     shuffle_data: bool = True,
     max_context_tokens: int = 8000,
     context_name: str = "Sources",
     random_state: int = 86,
 ) -> tuple[str, dict[str, pd.DataFrame]]:
     """Prepare text-unit data table as context data for system prompt."""
+    tokenizer = tokenizer or get_tokenizer()
     if text_units is None or len(text_units) == 0:
         return ("", {})
 

graphrag/query/structured_search/base.py

@@ -59,13 +59,13 @@ def __init__(
         self,
         model: ChatModel,
         context_builder: T,
-        tokenizer: Tokenizer = get_tokenizer(),
+        tokenizer: Tokenizer | None = None,
         model_params: dict[str, Any] | None = None,
         context_builder_params: dict[str, Any] | None = None,
     ):
         self.model = model
         self.context_builder = context_builder
-        self.tokenizer = tokenizer
+        self.tokenizer = tokenizer or get_tokenizer()
         self.model_params = model_params or {}
         self.context_builder_params = context_builder_params or {}
 

graphrag/query/structured_search/basic_search/basic_context.py

@@ -30,11 +30,11 @@ def __init__(
         text_embedder: EmbeddingModel,
         text_unit_embeddings: BaseVectorStore,
         text_units: list[TextUnit] | None = None,
-        tokenizer: Tokenizer = get_tokenizer(),
+        tokenizer: Tokenizer | None = None,
         embedding_vectorstore_key: str = "id",
     ):
         self.text_embedder = text_embedder
-        self.tokenizer = tokenizer
+        self.tokenizer = tokenizer or get_tokenizer()
         self.text_units = text_units
         self.text_unit_embeddings = text_unit_embeddings
         self.embedding_vectorstore_key = embedding_vectorstore_key

graphrag/query/structured_search/basic_search/search.py

@@ -16,7 +16,6 @@
 from graphrag.query.context_builder.builders import BasicContextBuilder
 from graphrag.query.context_builder.conversation_history import ConversationHistory
 from graphrag.query.structured_search.base import BaseSearch, SearchResult
-from graphrag.tokenizer.get_tokenizer import get_tokenizer
 from graphrag.tokenizer.tokenizer import Tokenizer
 
 logger = logging.getLogger(__name__)
@@ -32,7 +31,7 @@ def __init__(
         self,
         model: ChatModel,
         context_builder: BasicContextBuilder,
-        tokenizer: Tokenizer = get_tokenizer(),
+        tokenizer: Tokenizer | None = None,
         system_prompt: str | None = None,
         response_type: str = "multiple paragraphs",
         callbacks: list[QueryCallbacks] | None = None,

graphrag/query/structured_search/drift_search/drift_context.py

@@ -47,7 +47,7 @@ def __init__(
         reports: list[CommunityReport] | None = None,
         relationships: list[Relationship] | None = None,
         covariates: dict[str, list[Covariate]] | None = None,
-        tokenizer: Tokenizer = get_tokenizer(),
+        tokenizer: Tokenizer | None = None,
         embedding_vectorstore_key: str = EntityVectorStoreKey.ID,
         config: DRIFTSearchConfig | None = None,
         local_system_prompt: str | None = None,
@@ -59,7 +59,7 @@ def __init__(
         self.config = config or DRIFTSearchConfig()
         self.model = model
         self.text_embedder = text_embedder
-        self.tokenizer = tokenizer
+        self.tokenizer = tokenizer or get_tokenizer()
         self.local_system_prompt = local_system_prompt or DRIFT_LOCAL_SYSTEM_PROMPT
         self.reduce_system_prompt = reduce_system_prompt or DRIFT_REDUCE_PROMPT