Skip to content

Commit

Permalink
Fix BGERerankFunction call return type. (#14)
Browse files Browse the repository at this point in the history
- Changing BGERerankFunction call return type from `List[float]` to
`List[RerankResult]`.
  • Loading branch information
muazhari authored Apr 9, 2024
1 parent 1edfa8c commit 8202e8f
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion milvus_model/reranker/bgereranker.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ def __init__(
def _batchify(self, texts: List[str], batch_size: int) -> List[List[str]]:
return [texts[i : i + batch_size] for i in range(0, len(texts), batch_size)]

def __call__(self, query: str, documents: List[str], top_k: int = 5) -> List[float]:
def __call__(self, query: str, documents: List[str], top_k: int = 5) -> List[RerankResult]:
query_document_pairs = [[query, doc] for doc in documents]
batched_texts = self._batchify(documents, self.batch_size)
scores = []
Expand Down

0 comments on commit 8202e8f

Please sign in to comment.