Skip to content

Commit

Permalink
community: creation score_threshold in flashrank_rerank.py (#24016)
Browse files Browse the repository at this point in the history
Description: 
add a optional score relevance threshold for select only coherent
document, it's in complement of top_n

Discussion:
add relevance score threshold in flashrank_rerank document compressors
#24013

Dependencies:
 no dependencies

---------

Co-authored-by: Benjamin BERNARD <[email protected]>
Co-authored-by: Bagatur <[email protected]>
Co-authored-by: Chester Curme <[email protected]>
  • Loading branch information
4 people authored Jul 26, 2024
1 parent 190988d commit c4d2a53
Showing 1 changed file with 14 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,12 @@ class FlashrankRerank(BaseDocumentCompressor):
"""Flashrank client to use for compressing documents"""
top_n: int = 3
"""Number of documents to return."""
score_threshold: float = 0.0
"""Minimum relevance threshold to return."""
model: Optional[str] = None
"""Model to use for reranking."""
prefix_metadata: str = ""
"""Prefix for flashrank_rerank metadata keys"""

class Config:
"""Configuration for this pydantic object."""
Expand Down Expand Up @@ -69,11 +73,14 @@ def compress_documents(
final_results = []

for r in rerank_response:
metadata = r["meta"]
metadata["relevance_score"] = r["score"]
doc = Document(
page_content=r["text"],
metadata=metadata,
)
final_results.append(doc)
if r["score"] >= self.score_threshold:
doc = Document(
page_content=r["text"],
metadata={
self.prefix_metadata + "id": r["id"],
self.prefix_metadata + "relevance_score": r["score"],
**r["meta"],
},
)
final_results.append(doc)
return final_results

0 comments on commit c4d2a53

Please sign in to comment.