From 0d188dd37b77426f3643793b3601302904a66fae Mon Sep 17 00:00:00 2001 From: Billie Rinaldi Date: Wed, 28 May 2025 16:40:26 -0400 Subject: [PATCH] Fix prompt tuning sampling error when LIMIT is greater than number of chunks --- graphrag/prompt_tune/loader/input.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/graphrag/prompt_tune/loader/input.py b/graphrag/prompt_tune/loader/input.py index fa49ebeeb9..7cdc3006dc 100644 --- a/graphrag/prompt_tune/loader/input.py +++ b/graphrag/prompt_tune/loader/input.py @@ -67,8 +67,8 @@ async def load_docs_in_chunks( # 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 - limit = LIMIT + limit = min(LIMIT, len(chunks_df)) + logger.warning(f"Limit out of range, using default number of chunks: {limit}") # noqa: G004 if select_method == DocSelectionType.TOP: chunks_df = chunks_df[:limit]