community: creation score_threshold in flashrank_rerank.py (#24016)

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 <benjamin.bernard@openpathview.fr>
Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com>
Co-authored-by: Chester Curme <chester.curme@gmail.com>
This commit is contained in:
Diverrez morgan 2024-07-26 15:34:39 +02:00 committed by GitHub
parent 190988d93e
commit c4d2a53f18
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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."""
@ -69,11 +73,14 @@ class FlashrankRerank(BaseDocumentCompressor):
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