CachedEmbeddingWrapper - separate embeddings extraction

[nseidan]

Hi,

While exploring CacheWrapper, I observed that embeddings for queries and corpus are put together in the same array when I try to get embeddings from a task.
For example,

vector_map = TextVectorMap("cache/MyRetrievalTask")
vector_map.load(name="MyRetrievalTask")
vectors = np.asarray(vector_map.vectors)

zero_mask = (vectors == 0).all(axis=1)
vectors = vectors[~zero_mask]
#vectors => a list of embeddings for both queries and corpus

Is there a way to separate the embeddings for queries and the embeddings for the corpus?
Ultimately, like this:

[<doc_id1, [embeddings]>, <doc_id2, [embeddings]>...]
[<query_id1, [embeddings]>, <query_id2, [embeddings]>...]

[orionw]

At the moment, no. You'd have to do it the way the CacheWrapper does it, hashing the text to find the mapping (i.e. iterate over all queries and find them, then you know the rest are documents)

If you're interested in making the changes to store ids as well, happy to take a PR!