Wiki-all 数据集

wiki-all 数据集旨在通过大量向量和维度来测试大规模向量搜索算法的压力。整个数据集包含 8800 万个向量，每个向量有 768 个维度，用于测试在检索增强生成 (RAG) 工作负载中通常会遇到的向量类型。完整数据集大小约为 251GB，这有意大于 typical GPU 的内存大小。如此大规模的数据集是为了促进使用压缩和高效的核外（out-of-core）方法进行索引和搜索。

该数据集由来自 Kaggle 的英文 Wiki 文本和来自 Cohere Wikipedia 的多语言 Wiki 文本组成。

Cohere 的英文文本较旧（2022 年）且比 Kaggle 的英文 Wiki 文本（2023 年）小，因此已完全从 Cohere 数据集中移除。最终的 Wiki 文本包括来自 Kaggle 的英文 Wiki 和来自 Cohere 的其他语言。英文文本占总文本大小的 50%。

为了形成最终数据集，Wiki 文本被分块成 8500 万个 128-token 的片段。作为参考，Cohere 将 Wiki 文本分块成 104-token 的片段。最后，使用 paraphrase-multilingual-mpnet-base-v2 嵌入模型计算每个片段的嵌入。生成的数据集是一个大小为 8800 万 x 768 的嵌入矩阵。数据集还包含一个查询文件（含 1 万个查询向量）和一个用于评估最近邻算法的 groundtruth 文件。

获取数据集

完整数据集

数据集的一个版本以二进制格式提供，可直接由 cuvs-bench 工具使用。完整的 88M 数据集约为 251GB，下面的下载链接包含已分割成多个部分的 tar 包。

以下命令将下载所有 10 个部分并将它们解压到一个名为 wiki_all_88M 的目录中

上述方法有一个不幸的缺点，即如果命令因任何原因失败，所有部分都需要重新下载。文件也可以单独下载，然后解压到目录中。每个文件大约 27GB，共有 10 个文件。

curl -s https://data.rapids.ai/raft/datasets/wiki_all/wiki_all.tar.00
...
curl -s https://data.rapids.ai/raft/datasets/wiki_all/wiki_all.tar.09

cat wiki_all.tar.* | tar -xf - -C wiki_all_88M/

1M 和 10M 子集

还提供完整数据集的 1M 和 10M 子集，大小分别为 2.9GB 和 29GB。这些子集也包含 1 万个向量的查询集和相应的 groundtruth 文件。

curl -s https://data.rapids.ai/raft/datasets/wiki_all_1M/wiki_all_1M.tar
curl -s https://data.rapids.ai/raft/datasets/wiki_all_10M/wiki_all_10M.tar

使用数据集

下载数据集并将其解压到 wiki_all_88M 目录（或 wiki_all_1M/wiki_all_10M，取决于使用的是子集），然后就可以在基准测试工具中使用这些文件了。数据集名称为 wiki_all（或 wiki_all_1M/wiki_all_10M），可以通过在脚本中指定相应的名称 --dataset wiki_all_88M 来使用基准测试工具。

许可信息

Kaggle 上提供的英文 wiki 文本遵循 CC BY-NCSA 4.0 许可，Cohere wikipedia 数据集遵循 Apache 2.0 许可。