If you’re looking for something out of the box, I can’t help you unfortunately. I do ML at a company where we have our own in-house tooling. I agree with your

If you’re looking for something out of the box, I can’t help you unfortunately. I do ML at a company where we have our own in-house tooling. I agree with your algorithm, it is equivalent to cosine similarity. But you might also consider whether you need all the shards at all. Maybe you can get away with a half or even a tenth of them. If you have some metric like ndcg you can measure the drop in performance and consider the trade off.