[ColBERT] idea: return top_n best scoring tokens

[fdurant]

I'm experimenting with RAGStack ColBERT and have a feature request.

In order to be able to produce a query-passage scoring interpretability visualization like this, it would be handy if the result of ColbertVectorStore.add_texts also included the top-n list of most contributing tokens, each with a normalized score that would be trivial to color-code in a UI. This could be achieved via an extra parameter include_token_scores: int = 0

[cbornet]

There is no scoring when adding documents. It only happens when retrieving.
Chunk scores are part of ColbertRetriever::text_search. Does it answer your need ?

[cbornet]

Also I don't think we keep track of the tokens. Only their embeddings. And the Chunk score is the max of the embeddings scores which are not exposed either,
@zzzming can you confirm ?

[epinzur]

@cbornet you are correct. We don't store the tokens for the text... only the embeddings of the tokens.

[rehan243]

Fwiw, I was trying to achieve something similar with ColBERT and ended up modifying the ColbertVectorStore class to return the token scores directly. I added a simple arg to the add_texts method, kinda like you're suggesting with include_token_scores. Here's what I did:

def add_texts(self, texts, include_token_scores=False):
    results = self._index_texts(texts)
    if include_token_scores:
        return [(result, self._get_top_tokens(result, n=5)) for result in results]
    return results

I've been using this to debug my RAG pipeline and it's been pretty helpful. The _get_top_tokens method isn't shown here, but it's just a simple sort on the token scores.