Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

added top_k argument in the run function of ElasticSearcBM25Retriever #130

Merged
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):
def run(self, query: str, top_k: int=None):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

top_k should be typed as Optional

docs = self._document_store._bm25_retrieval(
query=query,
filters=self._filters,
fuzziness=self._fuzziness,
top_k=self._top_k,
top_k=self._top_k if top_k == None else top_k,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you can simplify this statement with 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,17 +64,18 @@ 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]):
def run(self, query_embedding: List[float], top_k:int = None):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above, top_k is optional

"""
Retrieve documents using a vector similarity metric.

:param query_embedding: Embedding of the query.
:param top_k: Maximum number of Documents to return
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for fixing the docs! It's out of the scope of this PR, but would you mind adding a similar docstring to the run method of the other retriever component?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@masci

The changes you suggested will be done soon.

Question: When you mentioned other retriever components you mean in the main haystack right?

Copy link
Contributor

@masci masci Dec 23, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

:return: List of Document similar to `query_embedding`.
"""
docs = self._document_store._embedding_retrieval(
query_embedding=query_embedding,
filters=self._filters,
top_k=self._top_k,
top_k=self._top_k if top_k == None else top_k,
num_candidates=self._num_candidates,
)
return {"documents": docs}
Loading