Skip to content

Commit

Permalink
added top_k arguement in the the run function of ElasticSearchBM25Ret…
Browse files Browse the repository at this point in the history
…riever
  • Loading branch information
sahusiddharth committed Dec 22, 2023
1 parent 82d2af3 commit 35eb043
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -48,12 +48,12 @@ def from_dict(cls, data: Dict[str, Any]) -> "ElasticsearchBM25Retriever":
return default_from_dict(cls, data)

@component.output_types(documents=List[Document])
def run(self, query: str, top_k: int=None):
def run(self, query: str, top_k: Optional[int] = None):
docs = self._document_store._bm25_retrieval(
query=query,
filters=self._filters,
fuzziness=self._fuzziness,
top_k=self._top_k if top_k == None else top_k,
top_k=top_k or self._top_k,
scale_score=self._scale_score,
)
return {"documents": docs}
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ElasticsearchEmbeddingRetriever":
return default_from_dict(cls, data)

@component.output_types(documents=List[Document])
def run(self, query_embedding: List[float], top_k:int = None):
def run(self, query_embedding: List[float], top_k: Optional[int] = None):
"""
Retrieve documents using a vector similarity metric.
Expand All @@ -75,7 +75,7 @@ def run(self, query_embedding: List[float], top_k:int = None):
docs = self._document_store._embedding_retrieval(
query_embedding=query_embedding,
filters=self._filters,
top_k=self._top_k if top_k == None else top_k,
top_k=top_k or self._top_k,
num_candidates=self._num_candidates,
)
return {"documents": docs}

0 comments on commit 35eb043

Please sign in to comment.