How to access individual dense and sparse scores in hybrid retrieval with Milvus?

[shalini0311]

Hi Team,

I'm currently using langchain with Milvus as the vector store, and I'm implementing a hybrid search setup that combines both dense and sparse retrieval.

This is a sample code I'm using:

from milvus_model.sparse import BM25EmbeddingFunction
class MyBM25SparseEmbedding(BM25SparseEmbedding):
    def __init__(self, corpus: List[str] = None):
        if corpus is not None:
            super().__init__(corpus=corpus)

    def load(self, save_path: str):
        self.bm25_ef = BM25EmbeddingFunction()
        self.bm25_ef.load(save_path)

    def save(self, save_path: str):
        self.bm25_ef.save(save_path)

For search

sparse_search_params = {"metric_type": "IP"}
dense_search_params = {"metric_type": "IP", "params": {}}

sparse_embedding_func = MyBM25SparseEmbedding()
sparse_embedding_func.load(bm25_store_file)
print(sparse_embedding_func.bm25_ef.idf)

collection = Collection(collection_name)
retriever = MilvusCollectionHybridSearchRetriever(
    collection=collection,
    rerank=WeightedRanker(0.5, 0.5),
    anns_fields=[dense_field, sparse_field],
    field_embeddings=[dense_embedding_func, sparse_embedding_func],
    field_search_params=[dense_search_params, sparse_search_params],
    top_k=3,
    text_field=text_field,
)

question = "How to start learning programming?"
print(question)
documents = retriever.get_relevant_documents(question)
print(f'Similar docs: {len(documents)}')
for document in documents:
    print(document)

I'd to like know about individual similarity scores using both embeddings to understand more about ranking strategy and assess confidence in retrieval.
Is there any way to retrieve both dense and sparse similarity scores in the above implementation?

Any pointers or workarounds would be appreciated.

Thanks!

[yhmo]

MilvusCollectionHybridSearchRetriever is langchain class that implements the langchain.BaseRetriever class. It is impossible to get the scores from the interfaces of langchain.BaseRetriever.

You can use pymilvus interfaces to achieve it.

Basically,

retriever = MilvusCollectionHybridSearchRetriever(
    collection=collection,
    rerank=WeightedRanker(0.5, 0.5),
    anns_fields=[dense_field, sparse_field],
    field_embeddings=[dense_embedding_func, sparse_embedding_func],
    field_search_params=[dense_search_params, sparse_search_params],
    top_k=3,
    text_field=text_field,
)
question = "How to start learning programming?"
print(question)
documents = retriever.get_relevant_documents(question)

is almost equal to:

sparse_vec = sparse_embedding_func.embed_documents([question])[0]
dense_vec = dense_embedding_func.embed_documents([question])[0]

collection = Collection(collection_name)

search_param = {
      "data": [dense_vec],
      "anns_field": dense_field,
      "param": dense_search_params,
      "limit": 3*2,
}
req_1 = AnnSearchRequest(**search_param)

search_param = {
      "data": [sparse_vec],
      "anns_field": sparse_field,
      "param": sparse_search_params,
      "limit": 3*2,
}
req_2 = AnnSearchRequest(**search_param)
hybrid_res = collection.hybrid_search([req_1, req_2], WeightedRanker(0.5, 0.5), 3, output_fields=[text_field])

In milvus, the two AnnSearchRequests are executed separately, the two results are combined by the WeightedRanker.
So, if you want to see the scores of each sub-request, you can:

search_param = {
    "data": [dense_vec],
    "anns_field": dense_field,
    "param": dense_search_params,
    "limit": 6,
}
results_1 = collection.search(**search_param)
print(result_1)

search_param = {
    "data": [sparse_vec],
    "anns_field": sparse_field,
    "param": sparse_search_params,
    "limit": 6,
}
results_2 = collection.search(**search_param)
print(result_2)

Note: I set the "limit=6" for sub-request, because the lanchain implementation set this value by:
self.field_limits = [self.top_k] * len(self.anns_fields)
The "self.top_k" is the limit value for hybrid_search() interface. The "self.field_limits" is the limit value for sub-request. The "self.anns_fields" is a list of all vector fields. In your case, len(self.anns_fields) = 2.
https://github.com/langchain-ai/langchain-milvus/blob/7d548dc259cf8b5b2fac50c0579ef68d138abb90/libs/milvus/langchain_milvus/retrievers/milvus_hybrid_search.py#L59

[shalini0311]

The "results_1" and "results_2" doesn't holds score as key, neither in the metadata. What to do?

[xiaofan-luan]

The "results_1" and "results_2" doesn't holds score as key, neither in the metadata. What to do?

see Milvus document about hybrid search.
You can use RRF as an alternative

[yhmo]

The "results_1" and "results_2" doesn't holds score as key, neither in the metadata. What to do?

I don't know what you mean "doesn't holds score as key". Each search interface will return a result with id-score(or id-distance) pairs. "distance" is "score".
You can try this script and observe the output:

from pymilvus import (
    connections,
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection,
    AnnSearchRequest, RRFRanker, WeightedRanker,
)

PK_FIELD = "id"
DENSE_FIELD = "dense_vector"
SPARSE_FIELD = "sparse_vector"
TEXT_FIELD = "text"

DENSE_DIM = 4

DENSE_METRIC = "COSINE"
SPARSE_METRIC = "IP"

connections.connect(host="localhost", port="19530")

collection_name = "hybrid_test"

def create_collection():
    if utility.has_collection(collection_name):
        utility.drop_collection(collection_name)

    fields = [
        FieldSchema(name=PK_FIELD, dtype=DataType.INT64, is_primary=True, auto_id=False),
        FieldSchema(name=DENSE_FIELD, dtype=DataType.FLOAT_VECTOR, dim=DENSE_DIM),
        FieldSchema(name=SPARSE_FIELD, dtype=DataType.SPARSE_FLOAT_VECTOR),
        FieldSchema(name=TEXT_FIELD, dtype=DataType.VARCHAR, max_length=100),
    ]

    schema = CollectionSchema(fields)
    collection = Collection(collection_name, schema)
    print("collection created")
    return collection

def create_index():
    collection = Collection(collection_name)
    collection.create_index(
        field_name=DENSE_FIELD,
        index_params={
            'metric_type': DENSE_METRIC,
            'index_type': "FLAT",
            'params': {},
        }
    )

    collection.create_index(
        field_name=SPARSE_FIELD,
        index_params={
            'metric_type': SPARSE_METRIC,
            'index_type': "SPARSE_INVERTED_INDEX",
            'params': {},
        }
    )

    collection.load()
    print("index created")

def insert():
    collection = Collection(collection_name)

    collection.insert({"id": 1, DENSE_FIELD: [1.0, 1.0, 1.0, 1.0], SPARSE_FIELD: {2: 0.5, 4: 0.8}, TEXT_FIELD: "text 1"})
    collection.insert({"id": 2, DENSE_FIELD: [2.0, 2.0, 2.0, 2.0], SPARSE_FIELD: {1: 0.2, 2: 0.6, 5: 0.7}, TEXT_FIELD: "text 2"})
    collection.insert({"id": 3, DENSE_FIELD: [3.0, 3.0, 3.0, 3.0], SPARSE_FIELD: {3: 0.3, 4: 0.2}, TEXT_FIELD: "text 3"})
    collection.insert({"id": 4, DENSE_FIELD: [4.0, 4.0, 4.0, 4.0], SPARSE_FIELD: {2: 0.2}, TEXT_FIELD: "text 4"})
    collection.insert({"id": 5, DENSE_FIELD: [5.0, 5.0, 5.0, 5.0], SPARSE_FIELD: {1: 0.5, 2: 0.6, 5: 0.9}, TEXT_FIELD: "text 5"})
    print("inserted")

    collection.flush()

def search():
    collection = Collection(collection_name)

    dense_vector = [2.0, 3.0, 4.0, 5.0]
    sparse_vector = {2: 0.3, 4: 0.5}

    print("\n================= normal search on float vector =================")
    res = collection.search(
        data=[dense_vector],
        anns_field=DENSE_FIELD,
        output_fields=[PK_FIELD, TEXT_FIELD],
        param={},
        limit=5,
    )
    for hits in res:
        for hit in hits:
            print(hit.to_dict())

    print("\n================= normal search on sparse vector =================")
    res = collection.search(
        data=[sparse_vector],
        anns_field=SPARSE_FIELD,
        output_fields=[PK_FIELD, TEXT_FIELD],
        param={},
        limit=5,
    )
    for hits in res:
        for hit in hits:
            print(hit.to_dict())

    print("\n================= hybrid search =================")
    req_list = []
    req_list.append(AnnSearchRequest(
        data=[dense_vector],
        anns_field=DENSE_FIELD,
        param={},
        limit=5,
    ))

    req_list.append(AnnSearchRequest(
        data=[sparse_vector],
        anns_field=SPARSE_FIELD,
        param={},
        limit=5,
    ))

    results = collection.hybrid_search(reqs=req_list,
                                       # rerank=RRFRanker(k=2),
                                       rerank=WeightedRanker(0.5, 0.5),
                                       output_fields=[PK_FIELD, TEXT_FIELD],
                                       limit=3,
                                       )
    for hits in results:
        for hit in hits:
            print(hit.to_dict())


if __name__ == '__main__':
    create_collection()
    insert()
    create_index()
    search()

The output:

collection created
inserted
index created

================= normal search on float vector =================
{'id': 5, 'distance': 0.9525793790817261, 'entity': {'text': 'text 5', 'id': 5}}
{'id': 4, 'distance': 0.9525793790817261, 'entity': {'text': 'text 4', 'id': 4}}
{'id': 2, 'distance': 0.9525793790817261, 'entity': {'text': 'text 2', 'id': 2}}
{'id': 1, 'distance': 0.9525793790817261, 'entity': {'text': 'text 1', 'id': 1}}
{'id': 3, 'distance': 0.9525793194770813, 'entity': {'text': 'text 3', 'id': 3}}

================= normal search on sparse vector =================
{'id': 1, 'distance': 0.550000011920929, 'entity': {'text': 'text 1', 'id': 1}}
{'id': 5, 'distance': 0.18000000715255737, 'entity': {'text': 'text 5', 'id': 5}}
{'id': 2, 'distance': 0.18000000715255737, 'entity': {'text': 'text 2', 'id': 2}}
{'id': 3, 'distance': 0.10000000149011612, 'entity': {'text': 'text 3', 'id': 3}}
{'id': 4, 'distance': 0.06000000238418579, 'entity': {'text': 'text 4', 'id': 4}}

================= hybrid search =================
{'id': 1, 'distance': 0.8181748390197754, 'entity': {'text': 'text 1', 'id': 1}}
{'id': 2, 'distance': 0.7664892077445984, 'entity': {'text': 'text 2', 'id': 2}}
{'id': 5, 'distance': 0.7664892077445984, 'entity': {'text': 'text 5', 'id': 5}}

The "distance" is the "score".