first commit
This commit is contained in:
9
reranker/__init__.py
Normal file
9
reranker/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
Reranker implementations for mem0 search functionality.
|
||||
"""
|
||||
|
||||
from .base import BaseReranker
|
||||
from .cohere_reranker import CohereReranker
|
||||
from .sentence_transformer_reranker import SentenceTransformerReranker
|
||||
|
||||
__all__ = ["BaseReranker", "CohereReranker", "SentenceTransformerReranker"]
|
||||
20
reranker/base.py
Normal file
20
reranker/base.py
Normal file
@@ -0,0 +1,20 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Dict, Any
|
||||
|
||||
class BaseReranker(ABC):
|
||||
"""Abstract base class for all rerankers."""
|
||||
|
||||
@abstractmethod
|
||||
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Rerank documents based on relevance to the query.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
documents: List of documents to rerank, each with 'memory' field
|
||||
top_k: Number of top documents to return (None = return all)
|
||||
|
||||
Returns:
|
||||
List of reranked documents with added 'rerank_score' field
|
||||
"""
|
||||
pass
|
||||
85
reranker/cohere_reranker.py
Normal file
85
reranker/cohere_reranker.py
Normal file
@@ -0,0 +1,85 @@
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from mem0.reranker.base import BaseReranker
|
||||
|
||||
try:
|
||||
import cohere
|
||||
COHERE_AVAILABLE = True
|
||||
except ImportError:
|
||||
COHERE_AVAILABLE = False
|
||||
|
||||
|
||||
class CohereReranker(BaseReranker):
|
||||
"""Cohere-based reranker implementation."""
|
||||
|
||||
def __init__(self, config):
|
||||
"""
|
||||
Initialize Cohere reranker.
|
||||
|
||||
Args:
|
||||
config: CohereRerankerConfig object with configuration parameters
|
||||
"""
|
||||
if not COHERE_AVAILABLE:
|
||||
raise ImportError("cohere package is required for CohereReranker. Install with: pip install cohere")
|
||||
|
||||
self.config = config
|
||||
self.api_key = config.api_key or os.getenv("COHERE_API_KEY")
|
||||
if not self.api_key:
|
||||
raise ValueError("Cohere API key is required. Set COHERE_API_KEY environment variable or pass api_key in config.")
|
||||
|
||||
self.model = config.model
|
||||
self.client = cohere.Client(self.api_key)
|
||||
|
||||
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Rerank documents using Cohere's rerank API.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
documents: List of documents to rerank
|
||||
top_k: Number of top documents to return
|
||||
|
||||
Returns:
|
||||
List of reranked documents with rerank_score
|
||||
"""
|
||||
if not documents:
|
||||
return documents
|
||||
|
||||
# Extract text content for reranking
|
||||
doc_texts = []
|
||||
for doc in documents:
|
||||
if 'memory' in doc:
|
||||
doc_texts.append(doc['memory'])
|
||||
elif 'text' in doc:
|
||||
doc_texts.append(doc['text'])
|
||||
elif 'content' in doc:
|
||||
doc_texts.append(doc['content'])
|
||||
else:
|
||||
doc_texts.append(str(doc))
|
||||
|
||||
try:
|
||||
# Call Cohere rerank API
|
||||
response = self.client.rerank(
|
||||
model=self.model,
|
||||
query=query,
|
||||
documents=doc_texts,
|
||||
top_n=top_k or self.config.top_k or len(documents),
|
||||
return_documents=self.config.return_documents,
|
||||
max_chunks_per_doc=self.config.max_chunks_per_doc,
|
||||
)
|
||||
|
||||
# Create reranked results
|
||||
reranked_docs = []
|
||||
for result in response.results:
|
||||
original_doc = documents[result.index].copy()
|
||||
original_doc['rerank_score'] = result.relevance_score
|
||||
reranked_docs.append(original_doc)
|
||||
|
||||
return reranked_docs
|
||||
|
||||
except Exception:
|
||||
# Fallback to original order if reranking fails
|
||||
for doc in documents:
|
||||
doc['rerank_score'] = 0.0
|
||||
return documents[:top_k] if top_k else documents
|
||||
147
reranker/huggingface_reranker.py
Normal file
147
reranker/huggingface_reranker.py
Normal file
@@ -0,0 +1,147 @@
|
||||
from typing import List, Dict, Any, Union
|
||||
import numpy as np
|
||||
|
||||
from mem0.reranker.base import BaseReranker
|
||||
from mem0.configs.rerankers.base import BaseRerankerConfig
|
||||
from mem0.configs.rerankers.huggingface import HuggingFaceRerankerConfig
|
||||
|
||||
try:
|
||||
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
||||
import torch
|
||||
TRANSFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
TRANSFORMERS_AVAILABLE = False
|
||||
|
||||
|
||||
class HuggingFaceReranker(BaseReranker):
|
||||
"""HuggingFace Transformers based reranker implementation."""
|
||||
|
||||
def __init__(self, config: Union[BaseRerankerConfig, HuggingFaceRerankerConfig, Dict]):
|
||||
"""
|
||||
Initialize HuggingFace reranker.
|
||||
|
||||
Args:
|
||||
config: Configuration object with reranker parameters
|
||||
"""
|
||||
if not TRANSFORMERS_AVAILABLE:
|
||||
raise ImportError("transformers package is required for HuggingFaceReranker. Install with: pip install transformers torch")
|
||||
|
||||
# Convert to HuggingFaceRerankerConfig if needed
|
||||
if isinstance(config, dict):
|
||||
config = HuggingFaceRerankerConfig(**config)
|
||||
elif isinstance(config, BaseRerankerConfig) and not isinstance(config, HuggingFaceRerankerConfig):
|
||||
# Convert BaseRerankerConfig to HuggingFaceRerankerConfig with defaults
|
||||
config = HuggingFaceRerankerConfig(
|
||||
provider=getattr(config, 'provider', 'huggingface'),
|
||||
model=getattr(config, 'model', 'BAAI/bge-reranker-base'),
|
||||
api_key=getattr(config, 'api_key', None),
|
||||
top_k=getattr(config, 'top_k', None),
|
||||
device=None, # Will auto-detect
|
||||
batch_size=32, # Default
|
||||
max_length=512, # Default
|
||||
normalize=True, # Default
|
||||
)
|
||||
|
||||
self.config = config
|
||||
|
||||
# Set device
|
||||
if self.config.device is None:
|
||||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
else:
|
||||
self.device = self.config.device
|
||||
|
||||
# Load model and tokenizer
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.config.model)
|
||||
self.model = AutoModelForSequenceClassification.from_pretrained(self.config.model)
|
||||
self.model.to(self.device)
|
||||
self.model.eval()
|
||||
|
||||
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Rerank documents using HuggingFace cross-encoder model.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
documents: List of documents to rerank
|
||||
top_k: Number of top documents to return
|
||||
|
||||
Returns:
|
||||
List of reranked documents with rerank_score
|
||||
"""
|
||||
if not documents:
|
||||
return documents
|
||||
|
||||
# Extract text content for reranking
|
||||
doc_texts = []
|
||||
for doc in documents:
|
||||
if 'memory' in doc:
|
||||
doc_texts.append(doc['memory'])
|
||||
elif 'text' in doc:
|
||||
doc_texts.append(doc['text'])
|
||||
elif 'content' in doc:
|
||||
doc_texts.append(doc['content'])
|
||||
else:
|
||||
doc_texts.append(str(doc))
|
||||
|
||||
try:
|
||||
scores = []
|
||||
|
||||
# Process documents in batches
|
||||
for i in range(0, len(doc_texts), self.config.batch_size):
|
||||
batch_docs = doc_texts[i:i + self.config.batch_size]
|
||||
batch_pairs = [[query, doc] for doc in batch_docs]
|
||||
|
||||
# Tokenize batch
|
||||
inputs = self.tokenizer(
|
||||
batch_pairs,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.config.max_length,
|
||||
return_tensors="pt"
|
||||
).to(self.device)
|
||||
|
||||
# Get scores
|
||||
with torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
batch_scores = outputs.logits.squeeze(-1).cpu().numpy()
|
||||
|
||||
# Handle single item case
|
||||
if batch_scores.ndim == 0:
|
||||
batch_scores = [float(batch_scores)]
|
||||
else:
|
||||
batch_scores = batch_scores.tolist()
|
||||
|
||||
scores.extend(batch_scores)
|
||||
|
||||
# Normalize scores if requested
|
||||
if self.config.normalize:
|
||||
scores = np.array(scores)
|
||||
scores = (scores - scores.min()) / (scores.max() - scores.min() + 1e-8)
|
||||
scores = scores.tolist()
|
||||
|
||||
# Combine documents with scores
|
||||
doc_score_pairs = list(zip(documents, scores))
|
||||
|
||||
# Sort by score (descending)
|
||||
doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Apply top_k limit
|
||||
final_top_k = top_k or self.config.top_k
|
||||
if final_top_k:
|
||||
doc_score_pairs = doc_score_pairs[:final_top_k]
|
||||
|
||||
# Create reranked results
|
||||
reranked_docs = []
|
||||
for doc, score in doc_score_pairs:
|
||||
reranked_doc = doc.copy()
|
||||
reranked_doc['rerank_score'] = float(score)
|
||||
reranked_docs.append(reranked_doc)
|
||||
|
||||
return reranked_docs
|
||||
|
||||
except Exception:
|
||||
# Fallback to original order if reranking fails
|
||||
for doc in documents:
|
||||
doc['rerank_score'] = 0.0
|
||||
final_top_k = top_k or self.config.top_k
|
||||
return documents[:final_top_k] if final_top_k else documents
|
||||
142
reranker/llm_reranker.py
Normal file
142
reranker/llm_reranker.py
Normal file
@@ -0,0 +1,142 @@
|
||||
import re
|
||||
from typing import List, Dict, Any, Union
|
||||
|
||||
from mem0.reranker.base import BaseReranker
|
||||
from mem0.utils.factory import LlmFactory
|
||||
from mem0.configs.rerankers.base import BaseRerankerConfig
|
||||
from mem0.configs.rerankers.llm import LLMRerankerConfig
|
||||
|
||||
|
||||
class LLMReranker(BaseReranker):
|
||||
"""LLM-based reranker implementation."""
|
||||
|
||||
def __init__(self, config: Union[BaseRerankerConfig, LLMRerankerConfig, Dict]):
|
||||
"""
|
||||
Initialize LLM reranker.
|
||||
|
||||
Args:
|
||||
config: Configuration object with reranker parameters
|
||||
"""
|
||||
# Convert to LLMRerankerConfig if needed
|
||||
if isinstance(config, dict):
|
||||
config = LLMRerankerConfig(**config)
|
||||
elif isinstance(config, BaseRerankerConfig) and not isinstance(config, LLMRerankerConfig):
|
||||
# Convert BaseRerankerConfig to LLMRerankerConfig with defaults
|
||||
config = LLMRerankerConfig(
|
||||
provider=getattr(config, 'provider', 'openai'),
|
||||
model=getattr(config, 'model', 'gpt-4o-mini'),
|
||||
api_key=getattr(config, 'api_key', None),
|
||||
top_k=getattr(config, 'top_k', None),
|
||||
temperature=0.0, # Default for reranking
|
||||
max_tokens=100, # Default for reranking
|
||||
)
|
||||
|
||||
self.config = config
|
||||
|
||||
# Create LLM configuration for the factory
|
||||
llm_config = {
|
||||
"model": self.config.model,
|
||||
"temperature": self.config.temperature,
|
||||
"max_tokens": self.config.max_tokens,
|
||||
}
|
||||
|
||||
# Add API key if provided
|
||||
if self.config.api_key:
|
||||
llm_config["api_key"] = self.config.api_key
|
||||
|
||||
# Initialize LLM using the factory
|
||||
self.llm = LlmFactory.create(self.config.provider, llm_config)
|
||||
|
||||
# Default scoring prompt
|
||||
self.scoring_prompt = getattr(self.config, 'scoring_prompt', None) or self._get_default_prompt()
|
||||
|
||||
def _get_default_prompt(self) -> str:
|
||||
"""Get the default scoring prompt template."""
|
||||
return """You are a relevance scoring assistant. Given a query and a document, you need to score how relevant the document is to the query.
|
||||
|
||||
Score the relevance on a scale from 0.0 to 1.0, where:
|
||||
- 1.0 = Perfectly relevant and directly answers the query
|
||||
- 0.8-0.9 = Highly relevant with good information
|
||||
- 0.6-0.7 = Moderately relevant with some useful information
|
||||
- 0.4-0.5 = Slightly relevant with limited useful information
|
||||
- 0.0-0.3 = Not relevant or no useful information
|
||||
|
||||
Query: "{query}"
|
||||
Document: "{document}"
|
||||
|
||||
Provide only a single numerical score between 0.0 and 1.0. Do not include any explanation or additional text."""
|
||||
|
||||
def _extract_score(self, response_text: str) -> float:
|
||||
"""Extract numerical score from LLM response."""
|
||||
# Look for decimal numbers between 0.0 and 1.0
|
||||
pattern = r'\b([01](?:\.\d+)?)\b'
|
||||
matches = re.findall(pattern, response_text)
|
||||
|
||||
if matches:
|
||||
score = float(matches[0])
|
||||
return min(max(score, 0.0), 1.0) # Clamp between 0.0 and 1.0
|
||||
|
||||
# Fallback: return 0.5 if no valid score found
|
||||
return 0.5
|
||||
|
||||
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Rerank documents using LLM scoring.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
documents: List of documents to rerank
|
||||
top_k: Number of top documents to return
|
||||
|
||||
Returns:
|
||||
List of reranked documents with rerank_score
|
||||
"""
|
||||
if not documents:
|
||||
return documents
|
||||
|
||||
scored_docs = []
|
||||
|
||||
for doc in documents:
|
||||
# Extract text content
|
||||
if 'memory' in doc:
|
||||
doc_text = doc['memory']
|
||||
elif 'text' in doc:
|
||||
doc_text = doc['text']
|
||||
elif 'content' in doc:
|
||||
doc_text = doc['content']
|
||||
else:
|
||||
doc_text = str(doc)
|
||||
|
||||
try:
|
||||
# Generate scoring prompt
|
||||
prompt = self.scoring_prompt.format(query=query, document=doc_text)
|
||||
|
||||
# Get LLM response
|
||||
response = self.llm.generate_response(
|
||||
messages=[{"role": "user", "content": prompt}]
|
||||
)
|
||||
|
||||
# Extract score from response
|
||||
score = self._extract_score(response)
|
||||
|
||||
# Create scored document
|
||||
scored_doc = doc.copy()
|
||||
scored_doc['rerank_score'] = score
|
||||
scored_docs.append(scored_doc)
|
||||
|
||||
except Exception:
|
||||
# Fallback: assign neutral score if scoring fails
|
||||
scored_doc = doc.copy()
|
||||
scored_doc['rerank_score'] = 0.5
|
||||
scored_docs.append(scored_doc)
|
||||
|
||||
# Sort by relevance score in descending order
|
||||
scored_docs.sort(key=lambda x: x['rerank_score'], reverse=True)
|
||||
|
||||
# Apply top_k limit
|
||||
if top_k:
|
||||
scored_docs = scored_docs[:top_k]
|
||||
elif self.config.top_k:
|
||||
scored_docs = scored_docs[:self.config.top_k]
|
||||
|
||||
return scored_docs
|
||||
107
reranker/sentence_transformer_reranker.py
Normal file
107
reranker/sentence_transformer_reranker.py
Normal file
@@ -0,0 +1,107 @@
|
||||
from typing import List, Dict, Any, Union
|
||||
import numpy as np
|
||||
|
||||
from mem0.reranker.base import BaseReranker
|
||||
from mem0.configs.rerankers.base import BaseRerankerConfig
|
||||
from mem0.configs.rerankers.sentence_transformer import SentenceTransformerRerankerConfig
|
||||
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
SENTENCE_TRANSFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
SENTENCE_TRANSFORMERS_AVAILABLE = False
|
||||
|
||||
|
||||
class SentenceTransformerReranker(BaseReranker):
|
||||
"""Sentence Transformer based reranker implementation."""
|
||||
|
||||
def __init__(self, config: Union[BaseRerankerConfig, SentenceTransformerRerankerConfig, Dict]):
|
||||
"""
|
||||
Initialize Sentence Transformer reranker.
|
||||
|
||||
Args:
|
||||
config: Configuration object with reranker parameters
|
||||
"""
|
||||
if not SENTENCE_TRANSFORMERS_AVAILABLE:
|
||||
raise ImportError("sentence-transformers package is required for SentenceTransformerReranker. Install with: pip install sentence-transformers")
|
||||
|
||||
# Convert to SentenceTransformerRerankerConfig if needed
|
||||
if isinstance(config, dict):
|
||||
config = SentenceTransformerRerankerConfig(**config)
|
||||
elif isinstance(config, BaseRerankerConfig) and not isinstance(config, SentenceTransformerRerankerConfig):
|
||||
# Convert BaseRerankerConfig to SentenceTransformerRerankerConfig with defaults
|
||||
config = SentenceTransformerRerankerConfig(
|
||||
provider=getattr(config, 'provider', 'sentence_transformer'),
|
||||
model=getattr(config, 'model', 'cross-encoder/ms-marco-MiniLM-L-6-v2'),
|
||||
api_key=getattr(config, 'api_key', None),
|
||||
top_k=getattr(config, 'top_k', None),
|
||||
device=None, # Will auto-detect
|
||||
batch_size=32, # Default
|
||||
show_progress_bar=False, # Default
|
||||
)
|
||||
|
||||
self.config = config
|
||||
self.model = SentenceTransformer(self.config.model, device=self.config.device)
|
||||
|
||||
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Rerank documents using sentence transformer cross-encoder.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
documents: List of documents to rerank
|
||||
top_k: Number of top documents to return
|
||||
|
||||
Returns:
|
||||
List of reranked documents with rerank_score
|
||||
"""
|
||||
if not documents:
|
||||
return documents
|
||||
|
||||
# Extract text content for reranking
|
||||
doc_texts = []
|
||||
for doc in documents:
|
||||
if 'memory' in doc:
|
||||
doc_texts.append(doc['memory'])
|
||||
elif 'text' in doc:
|
||||
doc_texts.append(doc['text'])
|
||||
elif 'content' in doc:
|
||||
doc_texts.append(doc['content'])
|
||||
else:
|
||||
doc_texts.append(str(doc))
|
||||
|
||||
try:
|
||||
# Create query-document pairs
|
||||
pairs = [[query, doc_text] for doc_text in doc_texts]
|
||||
|
||||
# Get similarity scores
|
||||
scores = self.model.predict(pairs)
|
||||
if isinstance(scores, np.ndarray):
|
||||
scores = scores.tolist()
|
||||
|
||||
# Combine documents with scores
|
||||
doc_score_pairs = list(zip(documents, scores))
|
||||
|
||||
# Sort by score (descending)
|
||||
doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Apply top_k limit
|
||||
final_top_k = top_k or self.config.top_k
|
||||
if final_top_k:
|
||||
doc_score_pairs = doc_score_pairs[:final_top_k]
|
||||
|
||||
# Create reranked results
|
||||
reranked_docs = []
|
||||
for doc, score in doc_score_pairs:
|
||||
reranked_doc = doc.copy()
|
||||
reranked_doc['rerank_score'] = float(score)
|
||||
reranked_docs.append(reranked_doc)
|
||||
|
||||
return reranked_docs
|
||||
|
||||
except Exception:
|
||||
# Fallback to original order if reranking fails
|
||||
for doc in documents:
|
||||
doc['rerank_score'] = 0.0
|
||||
final_top_k = top_k or self.config.top_k
|
||||
return documents[:final_top_k] if final_top_k else documents
|
||||
96
reranker/zero_entropy_reranker.py
Normal file
96
reranker/zero_entropy_reranker.py
Normal file
@@ -0,0 +1,96 @@
|
||||
import os
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from mem0.reranker.base import BaseReranker
|
||||
|
||||
try:
|
||||
from zeroentropy import ZeroEntropy
|
||||
ZERO_ENTROPY_AVAILABLE = True
|
||||
except ImportError:
|
||||
ZERO_ENTROPY_AVAILABLE = False
|
||||
|
||||
|
||||
class ZeroEntropyReranker(BaseReranker):
|
||||
"""Zero Entropy-based reranker implementation."""
|
||||
|
||||
def __init__(self, config):
|
||||
"""
|
||||
Initialize Zero Entropy reranker.
|
||||
|
||||
Args:
|
||||
config: ZeroEntropyRerankerConfig object with configuration parameters
|
||||
"""
|
||||
if not ZERO_ENTROPY_AVAILABLE:
|
||||
raise ImportError("zeroentropy package is required for ZeroEntropyReranker. Install with: pip install zeroentropy")
|
||||
|
||||
self.config = config
|
||||
self.api_key = config.api_key or os.getenv("ZERO_ENTROPY_API_KEY")
|
||||
if not self.api_key:
|
||||
raise ValueError("Zero Entropy API key is required. Set ZERO_ENTROPY_API_KEY environment variable or pass api_key in config.")
|
||||
|
||||
self.model = config.model or "zerank-1"
|
||||
|
||||
# Initialize Zero Entropy client
|
||||
if self.api_key:
|
||||
self.client = ZeroEntropy(api_key=self.api_key)
|
||||
else:
|
||||
self.client = ZeroEntropy() # Will use ZERO_ENTROPY_API_KEY from environment
|
||||
|
||||
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Rerank documents using Zero Entropy's rerank API.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
documents: List of documents to rerank
|
||||
top_k: Number of top documents to return
|
||||
|
||||
Returns:
|
||||
List of reranked documents with rerank_score
|
||||
"""
|
||||
if not documents:
|
||||
return documents
|
||||
|
||||
# Extract text content for reranking
|
||||
doc_texts = []
|
||||
for doc in documents:
|
||||
if 'memory' in doc:
|
||||
doc_texts.append(doc['memory'])
|
||||
elif 'text' in doc:
|
||||
doc_texts.append(doc['text'])
|
||||
elif 'content' in doc:
|
||||
doc_texts.append(doc['content'])
|
||||
else:
|
||||
doc_texts.append(str(doc))
|
||||
|
||||
try:
|
||||
# Call Zero Entropy rerank API
|
||||
response = self.client.models.rerank(
|
||||
model=self.model,
|
||||
query=query,
|
||||
documents=doc_texts,
|
||||
)
|
||||
|
||||
# Create reranked results
|
||||
reranked_docs = []
|
||||
for result in response.results:
|
||||
original_doc = documents[result.index].copy()
|
||||
original_doc['rerank_score'] = result.relevance_score
|
||||
reranked_docs.append(original_doc)
|
||||
|
||||
# Sort by relevance score in descending order
|
||||
reranked_docs.sort(key=lambda x: x['rerank_score'], reverse=True)
|
||||
|
||||
# Apply top_k limit
|
||||
if top_k:
|
||||
reranked_docs = reranked_docs[:top_k]
|
||||
elif self.config.top_k:
|
||||
reranked_docs = reranked_docs[:self.config.top_k]
|
||||
|
||||
return reranked_docs
|
||||
|
||||
except Exception:
|
||||
# Fallback to original order if reranking fails
|
||||
for doc in documents:
|
||||
doc['rerank_score'] = 0.0
|
||||
return documents[:top_k] if top_k else documents
|
||||
Reference in New Issue
Block a user