Ran poe format

Author: Gabriel Nieves

graphrag/api/prompt_tune.py

@@ -54,7 +54,9 @@ async def generate_indexing_prompts(
     logger: ProgressLogger,
     root: str,
     chunk_size: PositiveInt = graphrag_config_defaults.chunks.size,
-    overlap: Annotated[int, annotated_types.Gt(-1)] = graphrag_config_defaults.chunks.overlap,
+    overlap: Annotated[
+        int, annotated_types.Gt(-1)
+    ] = graphrag_config_defaults.chunks.overlap,
     limit: PositiveInt = 15,
     selection_method: DocSelectionType = DocSelectionType.RANDOM,
     domain: str | None = None,
@@ -113,7 +115,9 @@ async def generate_indexing_prompts(
         default_llm_settings.max_retries = min(
             len(doc_list), language_model_defaults.max_retries
         )
-        logger.warning(f"max_retries not set, using default value: {default_llm_settings.max_retries}")  # noqa: G004
+        logger.warning(
+            f"max_retries not set, using default value: {default_llm_settings.max_retries}"
+        )  # noqa: G004
 
     logger.info("Creating language model...")
     llm = ModelManager().register_chat(

graphrag/cli/main.py

@@ -313,10 +313,16 @@ def _prompt_tune_cli(
         ),
     ] = 2,
     chunk_size: Annotated[
-        int, typer.Option(help="The size of each example text chunk. Overrides chunks.size in the configuration file.")
+        int,
+        typer.Option(
+            help="The size of each example text chunk. Overrides chunks.size in the configuration file."
+        ),
     ] = graphrag_config_defaults.chunks.size,
     overlap: Annotated[
-        int, typer.Option(help="The overlap size for chunking documents. Overrides chunks.overlap in the configuration file")
+        int,
+        typer.Option(
+            help="The overlap size for chunking documents. Overrides chunks.overlap in the configuration file"
+        ),
     ] = graphrag_config_defaults.chunks.overlap,
     language: Annotated[
         str | None,

graphrag/cli/prompt_tune.py

@@ -66,17 +66,15 @@ async def prompt_tune(
     # override chunking config in the configuration
     if chunk_size != graph_config.chunks.size:
         graph_config.chunks.size = chunk_size
-    
+
     if overlap != graph_config.chunks.overlap:
         graph_config.chunks.overlap = overlap
 
     progress_logger = LoggerFactory().create_logger(logger)
     info, error, success = _logger(progress_logger)
 
     enabled_logging, log_path = enable_logging_with_config(
-        graph_config,
-        verbose,
-        filename="prompt-tune.log"
+        graph_config, verbose, filename="prompt-tune.log"
     )
     if enabled_logging:
         info(f"Logging enabled at {log_path}", verbose)

graphrag/config/logging.py

@@ -34,9 +34,7 @@ def enable_logging(log_filepath: str | Path, verbose: bool = False) -> None:
 
 
 def enable_logging_with_config(
-    config: GraphRagConfig,
-    verbose: bool = False,
-    filename: str = "indexing-engine.log"
+    config: GraphRagConfig, verbose: bool = False, filename: str = "indexing-engine.log"
 ) -> tuple[bool, str]:
     """Enable logging to a file based on the config.
 

graphrag/prompt_tune/loader/input.py

@@ -28,9 +28,7 @@ async def _embed_chunks(
 ) -> tuple[pd.DataFrame, np.ndarray]:
     """Convert text chunks into dense text embeddings."""
     sampled_text_chunks = text_chunks.sample(n=min(n_subset_max, len(text_chunks)))
-    embeddings = await embedding_llm.aembed_batch(
-        sampled_text_chunks["text"].tolist()
-    )
+    embeddings = await embedding_llm.aembed_batch(sampled_text_chunks["text"].tolist())
     return text_chunks, np.array(embeddings)
 
 
@@ -66,20 +64,20 @@ async def load_docs_in_chunks(
     dataset = await create_input(config.input, logger, root)
     chunk_config = config.chunks
     chunks_df = create_base_text_units(
-        documents = dataset,
-        callbacks = NoopWorkflowCallbacks(),
-        group_by_columns = chunk_config.group_by_columns,
-        size = chunk_size,
-        overlap = overlap,
-        encoding_model = chunk_config.encoding_model,
-        strategy = chunk_config.strategy,
-        prepend_metadata = chunk_config.prepend_metadata,
-        chunk_size_includes_metadata = chunk_config.chunk_size_includes_metadata,
+        documents=dataset,
+        callbacks=NoopWorkflowCallbacks(),
+        group_by_columns=chunk_config.group_by_columns,
+        size=chunk_size,
+        overlap=overlap,
+        encoding_model=chunk_config.encoding_model,
+        strategy=chunk_config.strategy,
+        prepend_metadata=chunk_config.prepend_metadata,
+        chunk_size_includes_metadata=chunk_config.chunk_size_includes_metadata,
     )
 
     # Depending on the select method, build the dataset
     if limit <= 0 or limit > len(chunks_df):
-        logger.warning(f"Limit out of range, using default number of chunks: {LIMIT}") # noqa: G004
+        logger.warning(f"Limit out of range, using default number of chunks: {LIMIT}")  # noqa: G004
         limit = LIMIT
 
     if select_method == DocSelectionType.TOP: