基于矩阵乘法检索文档

Question

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

我的意思是，如何从一维火炬张量中获取最高 1000 个值的索引？p

Answer 1

在点积之后得到相似度分数。您可以通过以下方式获取前 1000 个指数

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

Answer 2

我想你正在寻找 torch.topk 它会 return 顶部 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]])

基于矩阵乘法检索文档

Retrieving documents based on matrix multiplication

information-retrieval

cosine-similarity

pytorch