feat: Add rerank score filtering with configurable threshold
- Add DEFAULT_MIN_RERANK_SCORE constant (default: 0.0) - Add MIN_RERANK_SCORE environment variable support - Filter chunks with rerank scores below threshold in process_chunks_unified - Add info-level logging for filtering operations - Handle empty results gracefully after filtering - Maintain backward compatibility with non-reranked chunks
This commit is contained in:
parent
358fbd689f
commit
ebaff228aa
5 changed files with 42 additions and 2 deletions
|
|
@ -67,6 +67,9 @@ ENABLE_LLM_CACHE=true
|
||||||
|
|
||||||
### Reranker configuration (Set ENABLE_RERANK to true in reranking model is configed)
|
### Reranker configuration (Set ENABLE_RERANK to true in reranking model is configed)
|
||||||
# ENABLE_RERANK=True
|
# ENABLE_RERANK=True
|
||||||
|
### Minimum rerank score for document chunk exclusion (set to 0.0 to keep all chunks, 0.6 or above if LLM is not strong enought)
|
||||||
|
# MIN_RERANK_SCORE=0.0
|
||||||
|
### Rerank model configuration (required when ENABLE_RERANK=True)
|
||||||
# RERANK_MODEL=jina-reranker-v2-base-multilingual
|
# RERANK_MODEL=jina-reranker-v2-base-multilingual
|
||||||
# RERANK_BINDING_HOST=https://api.jina.ai/v1/rerank
|
# RERANK_BINDING_HOST=https://api.jina.ai/v1/rerank
|
||||||
# RERANK_BINDING_API_KEY=your_rerank_api_key_here
|
# RERANK_BINDING_API_KEY=your_rerank_api_key_here
|
||||||
|
|
|
||||||
|
|
@ -19,6 +19,7 @@ from lightrag.constants import (
|
||||||
DEFAULT_MAX_TOTAL_TOKENS,
|
DEFAULT_MAX_TOTAL_TOKENS,
|
||||||
DEFAULT_COSINE_THRESHOLD,
|
DEFAULT_COSINE_THRESHOLD,
|
||||||
DEFAULT_RELATED_CHUNK_NUMBER,
|
DEFAULT_RELATED_CHUNK_NUMBER,
|
||||||
|
DEFAULT_MIN_RERANK_SCORE,
|
||||||
)
|
)
|
||||||
|
|
||||||
# use the .env that is inside the current folder
|
# use the .env that is inside the current folder
|
||||||
|
|
@ -290,6 +291,11 @@ def parse_args() -> argparse.Namespace:
|
||||||
args.rerank_binding_host = get_env_value("RERANK_BINDING_HOST", None)
|
args.rerank_binding_host = get_env_value("RERANK_BINDING_HOST", None)
|
||||||
args.rerank_binding_api_key = get_env_value("RERANK_BINDING_API_KEY", None)
|
args.rerank_binding_api_key = get_env_value("RERANK_BINDING_API_KEY", None)
|
||||||
|
|
||||||
|
# Min rerank score configuration
|
||||||
|
args.min_rerank_score = get_env_value(
|
||||||
|
"MIN_RERANK_SCORE", DEFAULT_MIN_RERANK_SCORE, float
|
||||||
|
)
|
||||||
|
|
||||||
# Query configuration
|
# Query configuration
|
||||||
args.history_turns = get_env_value("HISTORY_TURNS", DEFAULT_HISTORY_TURNS, int)
|
args.history_turns = get_env_value("HISTORY_TURNS", DEFAULT_HISTORY_TURNS, int)
|
||||||
args.top_k = get_env_value("TOP_K", DEFAULT_TOP_K, int)
|
args.top_k = get_env_value("TOP_K", DEFAULT_TOP_K, int)
|
||||||
|
|
|
||||||
|
|
@ -23,6 +23,9 @@ DEFAULT_ENABLE_RERANK = True
|
||||||
DEFAULT_COSINE_THRESHOLD = 0.2
|
DEFAULT_COSINE_THRESHOLD = 0.2
|
||||||
DEFAULT_RELATED_CHUNK_NUMBER = 5
|
DEFAULT_RELATED_CHUNK_NUMBER = 5
|
||||||
|
|
||||||
|
# Rerank configuration defaults
|
||||||
|
DEFAULT_MIN_RERANK_SCORE = 0.0
|
||||||
|
|
||||||
# Separator for graph fields
|
# Separator for graph fields
|
||||||
GRAPH_FIELD_SEP = "<SEP>"
|
GRAPH_FIELD_SEP = "<SEP>"
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -31,6 +31,7 @@ from lightrag.constants import (
|
||||||
DEFAULT_MAX_TOTAL_TOKENS,
|
DEFAULT_MAX_TOTAL_TOKENS,
|
||||||
DEFAULT_COSINE_THRESHOLD,
|
DEFAULT_COSINE_THRESHOLD,
|
||||||
DEFAULT_RELATED_CHUNK_NUMBER,
|
DEFAULT_RELATED_CHUNK_NUMBER,
|
||||||
|
DEFAULT_MIN_RERANK_SCORE,
|
||||||
)
|
)
|
||||||
from lightrag.utils import get_env_value
|
from lightrag.utils import get_env_value
|
||||||
|
|
||||||
|
|
@ -284,6 +285,11 @@ class LightRAG:
|
||||||
rerank_model_func: Callable[..., object] | None = field(default=None)
|
rerank_model_func: Callable[..., object] | None = field(default=None)
|
||||||
"""Function for reranking retrieved documents. All rerank configurations (model name, API keys, top_k, etc.) should be included in this function. Optional."""
|
"""Function for reranking retrieved documents. All rerank configurations (model name, API keys, top_k, etc.) should be included in this function. Optional."""
|
||||||
|
|
||||||
|
min_rerank_score: float = field(
|
||||||
|
default=get_env_value("MIN_RERANK_SCORE", DEFAULT_MIN_RERANK_SCORE, float)
|
||||||
|
)
|
||||||
|
"""Minimum rerank score threshold for filtering chunks after reranking."""
|
||||||
|
|
||||||
# Storage
|
# Storage
|
||||||
# ---
|
# ---
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1871,7 +1871,29 @@ async def process_chunks_unified(
|
||||||
top_n=rerank_top_k,
|
top_n=rerank_top_k,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 2. Apply chunk_top_k limiting if specified
|
# 2. Filter by minimum rerank score if reranking is enabled
|
||||||
|
if query_param.enable_rerank and unique_chunks:
|
||||||
|
min_rerank_score = global_config.get("min_rerank_score", 0.5)
|
||||||
|
original_count = len(unique_chunks)
|
||||||
|
|
||||||
|
# Filter chunks with score below threshold
|
||||||
|
filtered_chunks = []
|
||||||
|
for chunk in unique_chunks:
|
||||||
|
rerank_score = chunk.get("rerank_score", 1.0) # Default to 1.0 if no score
|
||||||
|
if rerank_score >= min_rerank_score:
|
||||||
|
filtered_chunks.append(chunk)
|
||||||
|
|
||||||
|
unique_chunks = filtered_chunks
|
||||||
|
filtered_count = original_count - len(unique_chunks)
|
||||||
|
|
||||||
|
if filtered_count > 0:
|
||||||
|
logger.info(
|
||||||
|
f"Rerank filtering remained: {len(unique_chunks)} chunks (min rerank score: {min_rerank_score})"
|
||||||
|
)
|
||||||
|
if not unique_chunks:
|
||||||
|
return []
|
||||||
|
|
||||||
|
# 3. Apply chunk_top_k limiting if specified
|
||||||
if query_param.chunk_top_k is not None and query_param.chunk_top_k > 0:
|
if query_param.chunk_top_k is not None and query_param.chunk_top_k > 0:
|
||||||
if len(unique_chunks) > query_param.chunk_top_k:
|
if len(unique_chunks) > query_param.chunk_top_k:
|
||||||
unique_chunks = unique_chunks[: query_param.chunk_top_k]
|
unique_chunks = unique_chunks[: query_param.chunk_top_k]
|
||||||
|
|
@ -1879,7 +1901,7 @@ async def process_chunks_unified(
|
||||||
f"Kept chunk_top-k: {len(unique_chunks)} chunks (deduplicated original: {origin_count})"
|
f"Kept chunk_top-k: {len(unique_chunks)} chunks (deduplicated original: {origin_count})"
|
||||||
)
|
)
|
||||||
|
|
||||||
# 3. Token-based final truncation
|
# 4. Token-based final truncation
|
||||||
tokenizer = global_config.get("tokenizer")
|
tokenizer = global_config.get("tokenizer")
|
||||||
if tokenizer and unique_chunks:
|
if tokenizer and unique_chunks:
|
||||||
# Set default chunk_token_limit if not provided
|
# Set default chunk_token_limit if not provided
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue