基于矩阵乘法的文档检索

Question

我有一个模型，它表示多维向量空间中的文档集合。例如，对于100 k文档，我的模型以300维向量的形式表示它们。最后，我得到了一个大小为[100K, 300]的矩阵。为了根据给定查询的相关性检索这些文档，我执行矩阵乘法。例如，我将给定的查询表示为[300, 1]。然后，利用矩阵乘法得到余弦相似度分数如下：[100K, 300]*[300, 1] = [100K, 1]。现在，如何从这个具有最高余弦相似性的集合中检索前1000个文档。最简单的方法是根据余弦相似性进行排序，并获取前1000个文档。有什么方法可以用这种方式检索文档吗？

我的意思是，怎样才能从一维火炬张量中得到最高1000值的指数呢？

Answer 1

一旦你有了相似的分数后的点积。您可以获得前1000个索引如下

top_indices = torch.argsort(sims)[:1000]
similar_docs = sims[top_indices]

Answer 2

我想你是在寻找torch.topk，它将返回顶级k最大元素values和indices。

例如

x = torch.arange(100).view(-1,1)
x.shape
torch.Size([100, 1])

value, indices = x.topk(k=10, dim=0)
value
tensor([[99],
        [98],
        [97],
        [96],
        [95],
        [94],
        [93],
        [92],
        [91],
        [90]])
indices
tensor([[99],
        [98],
        [97],
        [96],
        [95],
        [94],
        [93],
        [92],
        [91],
        [90]])