Error in SentenceTransformersDiversityRanker with MMR strategy when len(documents) < top_k

[AdrianoKF]

I'm trying to use SentenceTransformersDiversityRanker with the maximum margin relevance strategy, but I'm running into an error when the number of input documents is less than top_k.

SentenceTransformersDiversityRanker._maximum_margin_relevance raises a ValueError("No best document found, check if the documents list contains any documents."), since the while loop will select all available documents (which makes selected == documents, but len(selected) < top_k) and then run out of documents to select.

Am I missing something (entirely possible, since I'm new to Haystack), or should that method have a special case for len(documents) < top_k, such as:

-       while len(selected) < top_k:
+       while len(selected) < top_k and len(selected) < len(documents):
            # ...

            if best_idx is None:
                raise ValueError("No best document found, check if the documents list contains any documents.")

[anakin87]

Hello, @AdrianoKF!

I cannot reproduce the bug.

from haystack import Document
from haystack.components.rankers import SentenceTransformersDiversityRanker

ranker = SentenceTransformersDiversityRanker(model="sentence-transformers/all-MiniLM-L6-v2", similarity="cosine",
strategy="maximum_margin_relevance")
ranker.warm_up()

docs = [Document(content="Regular Exercise"), Document(content="Balanced Nutrition"), Document(content="Positive Mindset"), 
        Document(content="Eating Well"), Document(content="Doing physical activities"), Document(content="Thinking positively")]
        
query = "How can I maintain physical fitness?"
output = ranker.run(query=query, documents=docs, top_k=10)
docs = output["documents"]

print(docs)

This raises a reasonable error

  File ".../haystack/try.py", line 12, in <module>
    output = ranker.run(query=query, documents=docs, top_k=10)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../haystack/haystack/components/rankers/sentence_transformers_diversity.py", line 425, in run
    raise ValueError(f"top_k must be between 1 and {len(documents)}, but got {top_k}")
ValueError: top_k must be between 1 and 6, but got 10

Feel free to attach a reproducible example.

[AdrianoKF]

Thanks for your quick reply, @anakin87!

Indeed, with your code, I get the same error message (although one might debate if running with len(documents) < top_k is a valid use case after all).

After some more digging I found the difference: your code passes top_k to run(), while mine passed it as an init arg to the ranker.

When I modify your code as below, I get the original ValueError:

from haystack import Document
from haystack.components.rankers import SentenceTransformersDiversityRanker

# -- top_k passed here
ranker = SentenceTransformersDiversityRanker(
    model="sentence-transformers/all-MiniLM-L6-v2",
    similarity="cosine",
    strategy="maximum_margin_relevance",
    top_k=10,
)
ranker.warm_up()

# ...

query = "How can I maintain physical fitness?"
# -- no top_k arg here
output = ranker.run(query=query, documents=docs)
# --
docs = output["documents"]

print(docs)

So it appears the code in general handles the edge case already, however, not on all possible code paths. I can prepare a small PR, if there's interest.

[anakin87]

Thanks!

I opened an issue here: #9695

If you want to contribute, feel free to open a PR.