Skip to content

Commit

Permalink
fix: reformat
Browse files Browse the repository at this point in the history
  • Loading branch information
Anhui-tqhuang committed Feb 20, 2024
1 parent 9d1afad commit 831e438
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 4 deletions.
6 changes: 3 additions & 3 deletions private_gpt/components/reranker/reranker.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from llama_index.postprocessor.types import BaseNodePostprocessor
from private_gpt.settings.settings import Settings


@singleton
class RerankerComponent(BaseNodePostprocessor):
"""
Expand All @@ -17,6 +18,7 @@ class RerankerComponent(BaseNodePostprocessor):
If the number of nodes with score > cut_off is <= top_n, then return top_n nodes.
Otherwise, return all nodes with score > cut_off.
"""

reranker: FlagReranker = Field(description="Reranker class.")
top_n: int = Field(description="Top N nodes to return.")
cut_off: float = Field(description="Cut off score for nodes.")
Expand Down Expand Up @@ -66,6 +68,4 @@ def _postprocess_nodes(
if len(res) > self.top_n:
return res

return sorted(nodes, key=lambda x: x.score or 0.0, reverse=True)[
: self.top_n
]
return sorted(nodes, key=lambda x: x.score or 0.0, reverse=True)[: self.top_n]
2 changes: 1 addition & 1 deletion private_gpt/settings/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ class RerankerSettings(BaseModel):
)
hf_model_name: str = Field(
"BAAI/bge-reranker-large",
description="Name of the HuggingFace model to use for reranking"
description="Name of the HuggingFace model to use for reranking",
)
top_n: int = Field(
5,
Expand Down

0 comments on commit 831e438

Please sign in to comment.