Can we retrieve the matching tokens for a document?

[mrisher]

We're trying to display some highlighting in the UI based on which tokens match. I can't find any method that returns the matching tokens from a query, but wanted to confirm? Ideally we'd have something in the results object from the call results, scores = retriever.retrieve(...) so we would know that document ID #123 had tokens ["fish", "cat"]?

Thank you

[xhluca]

This is not built in, but i think it should be easy to verify? E.g.:

import bm25s

query = "what are cats?"
doc = "Cats are animals"

query_tokens = set(bm25s.tokenize(query, return_ids=False)[0])
doc_tokens = bm25s.tokenize(doc, return_ids=False)[0]

matches = [1 if dt in query_tokens else 0 for dt in doc_tokens]

[xhluca]

Code will be more complex for other cases. however, it should be fast since oyu are only interested in the top-k instances for highlighting.

[mrisher]

Thank you. Does that also give me a way to get token positions? I would like to boost the score for exact phrase matches and don't really understand the internals of the token index.

[xhluca]

No, but I think it would be straightforward to build a scorer for reranking bm25s results. Since this would be modular and separate from bm25s, i think it would be more appropriate as a separate tool, perhaps in the utils module or in a separate library.

[greenersharp]

I would like to second the addition of this functionality.

In my use case I need to make sure the documents returned by retrieve match some of my query tokens.
I need to make 100% sure some query tokens are matched.

I think it would be very useful to have the retrieve method return the matching query tokens for each result.

Thank you.

[xhluca]

Does this solution work for you?

#89 (comment)

I think this shouldn't be in retrieve as retrieve is already complex. utils package would be better, and if it's a popular feature we could add it into retrieve in a future refactor of the retrieve function.

[xhluca]

Btw this is the retrieve function:

bm25s/bm25s/__init__.py

Lines 694 to 874 in 3cab213

	num_docs = self.scores["num_docs"]
	if k > num_docs:
	raise ValueError(
	f"k of {k} is larger than the number of available scores"
	f", which is {num_docs} (corpus size should be larger than top-k)."
	f" Please set with a smaller k or increase the size of corpus."
	)
	allowed_return_as = ["tuple", "documents"]
	if return_as not in allowed_return_as:
	raise ValueError("`return_as` must be either 'tuple' or 'documents'")
	else:
	pass
	if n_threads == -1:
	n_threads = os.cpu_count()
	# if it's a list of list of tokens ids (int), we remove any integer not in the vocab_dict
	if is_list_of_list_of_type(query_tokens, type_=int):
	if not hasattr(self, "unique_token_ids_set") or self.unique_token_ids_set is None:
	raise ValueError(
	"The unique_token_ids_set attribute is not found. Please run the `index` method first, or make sure"
	"run retriever.load(load_vocab=True) before calling retrieve on list of list of token IDs (int)."
	)
	query_tokens_filtered = []
	for query in query_tokens:
	query_filtered = [
	token_id
	for token_id in query
	if token_id in self.unique_token_ids_set
	]
	if len(query_filtered) == 0:
	if "" not in self.vocab_dict:
	raise ValueError(
	"The query does not contain any tokens that are in the vocabulary. "
	"Please provide a query that contains at least one token that is in the vocabulary. "
	"Alternatively, you can set `create_empty_token=True` when calling `index` to add an empty token to the vocabulary. "
	"You can also manually add an empty token to the vocabulary by setting `retriever.vocab_dict[''] = max(retriever.vocab_dict.values()) + 1`. "
	"Then, run `retriever.unique_token_ids_set = set(retriever.vocab_dict.values())` to update the unique token IDs."
	)
	query_filtered = [self.vocab_dict[""]]
	query_tokens_filtered.append(query_filtered)
	query_tokens = query_tokens_filtered
	if isinstance(query_tokens, tuple) and not _is_tuple_of_list_of_tokens(
	query_tokens
	):
	if len(query_tokens) != 2:
	msg = (
	"Expected a list of string or a tuple of two elements: the first element is the "
	"list of unique token IDs, "
	"and the second element is the list of token IDs for each document."
	f"Found {len(query_tokens)} elements instead."
	)
	raise ValueError(msg)
	else:
	ids, vocab = query_tokens
	if not isinstance(ids, Iterable):
	raise ValueError(
	"The first element of the tuple passed to retrieve must be an iterable."
	)
	if not isinstance(vocab, dict):
	raise ValueError(
	"The second element of the tuple passed to retrieve must be a dictionary."
	)
	query_tokens = tokenization.Tokenized(ids=ids, vocab=vocab)
	if isinstance(query_tokens, tokenization.Tokenized):
	query_tokens = tokenization.convert_tokenized_to_string_list(query_tokens)
	corpus = corpus if corpus is not None else self.corpus
	if weight_mask is not None:
	if not isinstance(weight_mask, np.ndarray):
	raise ValueError("weight_mask must be a numpy array.")
	# check if weight_mask is a 1D array, if not raise an error
	if weight_mask.ndim != 1:
	raise ValueError("weight_mask must be a 1D array.")
	# check if the length of the weight_mask is the same as the length of the corpus
	if len(weight_mask) != self.scores["num_docs"]:
	raise ValueError(
	"The length of the weight_mask must be the same as the length of the corpus."
	)
	if self.backend == "numba":
	if _retrieve_numba_functional is None:
	raise ImportError(
	"Numba is not installed. Please install numba wiith `pip install numba` to use the numba backend."
	)
	backend_selection = (
	"numba" if backend_selection == "auto" else backend_selection
	)
	# if is list of list of int
	if is_list_of_list_of_type(query_tokens, type_=int):
	query_tokens_ids = query_tokens
	elif is_list_of_list_of_type(query_tokens, type_=str):
	query_tokens_ids = [self.get_tokens_ids(q) for q in query_tokens]
	else:
	raise ValueError(
	"The query_tokens must be a list of list of tokens (str for stemmed words, int for token ids matching corpus) or a tuple of two lists: the first list is the list of unique token IDs, and the second list is the list of token IDs for each document."
	)
	res = _retrieve_numba_functional(
	query_tokens_ids=query_tokens_ids,
	scores=self.scores,
	corpus=corpus,
	k=k,
	sorted=sorted,
	return_as=return_as,
	show_progress=show_progress,
	leave_progress=leave_progress,
	n_threads=n_threads,
	chunksize=None, # chunksize is ignored in the numba backend
	backend_selection=backend_selection, # backend_selection is ignored in the numba backend
	dtype=self.dtype,
	int_dtype=self.int_dtype,
	nonoccurrence_array=self.nonoccurrence_array,
	weight_mask=weight_mask,
	)
	if return_as == "tuple":
	return Results(documents=res[0], scores=res[1])
	else:
	return res
	tqdm_kwargs = {
	"total": len(query_tokens),
	"desc": "BM25S Retrieve",
	"leave": leave_progress,
	"disable": not show_progress,
	}
	topk_fn = partial(
	self._get_top_k_results,
	k=k,
	sorted=sorted,
	backend=backend_selection,
	weight_mask=weight_mask,
	)
	if n_threads == 0:
	# Use a simple map function to retrieve the results
	out = tqdm(map(topk_fn, query_tokens), **tqdm_kwargs)
	else:
	# Use concurrent.futures.ProcessPoolExecutor to parallelize the computation
	with ThreadPoolExecutor(max_workers=n_threads) as executor:
	process_map = executor.map(
	topk_fn,
	query_tokens,
	chunksize=chunksize,
	)
	out = list(tqdm(process_map, **tqdm_kwargs))
	scores, indices = zip(*out)
	scores, indices = np.array(scores), np.array(indices)
	corpus = corpus if corpus is not None else self.corpus
	if corpus is None:
	retrieved_docs = indices
	else:
	# if it is a JsonlCorpus object, we do not need to convert it to a list
	if isinstance(corpus, utils.corpus.JsonlCorpus):
	retrieved_docs = corpus[indices]
	elif isinstance(corpus, np.ndarray) and corpus.ndim == 1:
	retrieved_docs = corpus[indices]
	else:
	index_flat = indices.flatten().tolist()
	results = [corpus[i] for i in index_flat]
	retrieved_docs = np.array(results).reshape(indices.shape)
	if return_as == "tuple":
	return Results(documents=retrieved_docs, scores=scores)
	elif return_as == "documents":
	return retrieved_docs
	else:
	raise ValueError("`return_as` must be either 'tuple' or 'documents'")

almost 200 lines without counting the internal function like _retrieve_numba_functional. that's why i think it's cleaner to have a utils.

feel free to contribute this feature if you believe it's important to add it.

[greenersharp]

Thanks, I looked thru the code (very nice code btw) and it seems like _compute_relevance_from_scores is currently returning scores only, but it can be modified to return a tuple of (scores,query_index_ids). Then it's just a matter of modifying a few other functions to bubble the new tuple up to the retrieve function.

I'll work on adding this feature and submit the PR request sometime soon.

Take care.

[xhluca]

thanks.

i'm not sure if it makes sense to do change the signature of the outputs of internal functions, as they may be used elsewhere as well (it's been a while so don't remember on top of my head). also, you'd need to be careful with changing the signature of the output of retrieve as it may be a break changing.

Finally, it will be important to think about how this change will affect the performance of the end result, as the extra operations may increase the time it takes to compute the scores, as the _compute_relevance_from_scores_jit_ready function is carefully written to maximize speed and seemingly simple operations may add a lot of time. To address this, i would only feel comfortable if the PR also includes the results on the bm25-benchmarks compared to the latest release (perhaps on both linux and mac). If the performance end up being lower, we would need another function called _compute_relevance_from_scores_jit_ready_with_query_index_ids that is only used when we want the indexes, but that'd increase the complexity of the code (=harder to maintain/debug).

Overall, i'm not sure this is the best way to approach the problem, whereas an utils function that is applied posthoc, which only looks at the tokens that are matched, may make more sense.

I will look at the PR either way but i will consider the performance and maintanability seriously before merging, as _compute_relevance_from_scores and _compute_relevance_from_scores_jit_ready are low-level function that are not designed to be often modified for the reason above.