STaRK: 在文本和关系知识库上的大模型检索基准
新闻
- [2024年6月] 我们将基准作为一个pip包发布 stark-qa。你现在可以直接从包中加载数据了!
- [2024年6月] 我们将数据迁移到了 Hugging Face! 你不需要更改任何内容,数据会自动下载。
- [2024年5月] 我们增强了基准数据集,新增了三个高质量的人类生成的查询数据集,现已开放访问。详情请见我们的更新的arxiv!
- [2024年5月9日] 我们发布了 STaRK SKB Explorer,一个交互界面供你探索我们的知识库!
- [2024年5月7日] 我们在2024年斯坦福年度赞助商会议和2024年斯坦福数据科学大会上展示了STaRK。
- [2024年5月5日] STaRK被Marketpost和智源社区 BAAI报道。感谢对我们工作的报道!
- [2024年4月21日] 我们发布了STaRK基准。
什么是STaRK?
STaRK是一个基于文本和关系知识库的大规模半结构化信息检索基准,覆盖了产品搜索、学术论文搜索和生物医学查询等应用。
STaRK特有多样性、自然语言风格以及实用性的查询,这些查询需要特定情境下的推理能力,因而为评估受大模型驱动的真实世界检索系统设立了新标准,并且为未来的研究提出了显著的挑战。
🔥 查看更多概述,请访问我们的网站: website!
访问基准数据
1) 环境设置
从pip安装 (推荐)
使用python >=3.8且<3.12
pip install stark-qa
从源码安装
创建一个包含python >=3.8且<3.12的conda环境,并安装requirements.txt
中的必要包。
conda create -n stark python=3.11
conda activate stark
pip install -r requirements.txt
2) 数据加载
from stark_qa import load_qa, load_skb
dataset_name = 'amazon'
# 加载检索数据集
qa_dataset = load_qa(dataset_name)
idx_split = qa_dataset.get_idx_split()
# 加载半结构化知识库
skb = load_skb(dataset_name, download_processed=True, root=None)
load_skb
的root
参数指定了存储SKB数据的位置。默认值为None
时,数据将存储在huggingface缓存中。
检索任务数据
默认情况下,检索任务的问题和答案对将会自动下载到data/{dataset}/stark_qa
中。我们提供了官方分割数据位于data/{dataset}/split
中。
知识库数据
有两种方式加载知识库数据:
- (推荐) 即时下载:当设置
download_processed=True
时,所有三种基准的知识库数据将被自动下载和加载。 - 从原始数据处理:出于透明性考虑,我们提供了所有的预处理代码。因此,你可以通过设置
download_processed=False
从头处理原始数据。在这种情况下,处理STaRK-PrimeKG大约需要花费5分钟时间,而处理STaRK-Amazon和STaRK-MAG从原始数据到加载大约需要一个小时。
3) 基准评估
-
我们的评估需要将节点文档嵌入到
candidate_emb_dict.pt
中,这是一个字典node_id -> torch.Tensor
。如果没有可用的查询嵌入,将会自动生成查询嵌入。你可以运行以下Python脚本下载由text-embedding-ada-002
生成的查询嵌入和文档嵌入。 (我们提供这些嵌入,以便你可以立即在我们的基准上运行。)python emb_download.py --dataset amazon --emb_dir emb/
或者你可以运行以下代码生成查询或文档嵌入。例如,
python emb_generate.py --dataset amazon --mode query --emb_dir emb/ --emb_model text-embedding-ada-002
dataset
:amazon
、mag
或prime
之一。mode
: 嵌入的内容,query
或doc
(节点文档)之一。emb_dir
: 存储嵌入的目录。emb_model
: 生成嵌入的大模型名称,如text-embedding-ada-002
、text-embedding-3-large
。- 其他参数见
emb_generate.py
。
-
运行Python脚本进行评估。例如,
python eval.py --dataset amazon --model VSS --emb_dir emb/ --output_dir output/ --emb_model text-embedding-ada-002 --split test --save_pred
python eval.py --dataset amazon --model LLMReranker --emb_dir emb/ --output_dir output/ --emb_model text-embedding-ada-002 --split test --llm_model gpt-4-1106-preview --save_pred
dataset
: 要评估的数据集,amazon
、mag
或prime
之一。model
: 要评估的模型,VSS
、MultiVSS
、LLMReranker
之一。- 请使用参数
--emb_model
指定嵌入模型名称。 - 如果你使用的是
LLMReranker
,请使用参数--llm_model
指定大模型名称。 - 在命令行中指定API密钥
或export ANTHROPIC_API_KEY=YOUR_API_KEY
export OPENAI_API_KEY=YOUR_API_KEY export OPENAI_ORG=YOUR_ORGANIZATION
config/openai_api_key.txt
或config/claude_api_key.txt
本地文件中指定- 请使用参数
emb_dir
: 存储嵌入的目录。split
: 要评估的数据集切分,train
、val
、test
和human_generated_eval
之一(在人工生成的查询数据集上评估)。output_dir
: 存储评估输出的目录。
参考文献
如果你在你的工作中使用了我们的基准或代码,请考虑引用我们的文章:
@article{wu24stark,
title = {STaRK: 在文本和关系知识库上的大模型检索基准},
author = {
Shirley Wu, Shiyu Zhao,
Michihiro Yasunaga, Kexin Huang,
Kaidi Cao, Qian Huang,
Vassilis N. Ioannidis, Karthik Subbian,
James Zou, Jure Leskovec
},
eprinttype = {arXiv},
eprint = {2404.13207},
year = {2024}
}