基于矩阵乘法检索文件。

Question

我有一个在多维向量空间中表示文档集合的模型。因此，例如，对于100k个文档，我的模型以300维向量的形式来表示它们。所以，最后，我得到一个大小为 [100K, 300]. 为了根据与给定查询的相关性来检索这些文档，我进行矩阵乘法。例如，我将一个给定的查询表示为一个 [300, 1]. 然后我用矩阵乘法得到余弦相似度分数，如下图所示。[100K, 300]*[300, 1] = [100K, 1]. 现在我怎样才能从这个集合中检索出余弦相似度最高的前1000个文档。最简单的方法是根据余弦相似度进行排序，然后抓取前1000个文档。有没有什么方法可以用pytorch中的一些函数来检索这些文档？

我的意思是，如何从1D torch张量中获取最高1000个值的指数？p

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]])