stark

Stark：在文本和關係知識基礎上基準測試LLM檢索

消息

• [2024年10月]排行榜：我們在unggingface上構建了官方排行榜，我們呼籲提交！
• [2024年10月]攝像機就緒紙已經熄滅了！我們添加了多個檢索模型，包括BM25，ColbertV2，Gritlm。
• [2024年9月] Stark被接受為2024 Neurips Dataset＆Benchmark Track！
• [2024年6月]我們將基準作為pip包stark-qa做出。您現在可以直接從軟件包中加載數據！
• [2024年6月]我們將數據遷移到擁抱面孔！您無需更改任何內容，數據將自動下載。
• [2024年5月]我們通過三個可訪問的高質量人類生成的查詢數據集增強了基準。在我們更新的Arxiv中查看更多詳細信息！
• [2024年5月9日]我們發布了Stark SKB Explorer，這是一個交互式界面，供您探索我們的知識基礎！
• [2024年5月7日]我們在2024年斯坦福年度會員會議和2024年斯坦福數據科學會議上介紹了斯塔克。
• [2024年5月5日]據報導，在MarketPost和智源社區baai上報導了Stark。感謝您寫有關我們的工作的文章！
• [4月21日2024年]我們發布了Stark基準。

什麼是史塔克？

Stark是關於文本和關係知識庫的大規模半結構化檢索基準，涵蓋了產品搜索，學術紙搜索和生物醫學查詢中的應用。

Stark採用需要特定於上下文的推理的多樣，自然和實用的查詢，為評估由LLMS驅動的現實世界檢索系統設定了新的標準，並為未來的研究帶來了重大挑戰。

查看我們的網站以獲取更多概述！

訪問基準數據

1）設置

從PIP（推薦）

使用Python> = 3.8和<3.12

pip install stark-qa

來自來源

使用Python> = 3.8和<3.12創建一個conda env，並在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'

# Load the retrieval dataset
qa_dataset = load_qa ( dataset_name )
idx_split = qa_dataset . get_idx_split ()

# Load the semi-structured knowledge base
skb = load_skb ( dataset_name , download_processed = True , root = None )

LOAD_SKB的根參數指定存儲SKB數據的位置。對於默認值None ，數據將存儲在HuggingFace緩存中。

檢索任務的數據

檢索任務的問題答案對將自動在data/{dataset}/stark_qa中下載。我們提供了data/{dataset}/split中的官方拆分。

知識庫的數據

有兩種加載知識庫數據的方法：

• （建議）即時下載：設置download_processed=True時，所有三個基準的知識庫數據將自動下載和加載。
• RAW的過程數據：我們還提供了所有預處理代碼以實現透明度。因此，您可以通過設置download_processed=False處理原始數據。在這種情況下，Stark-Primekg大約需要5分鐘才能下載並加載處理後的數據。 Stark-Amazon和Stark-Mag可能需要大約一個小時的時間來處理原始數據。

3）評估基準

如果您正在運行評估，則可以安裝以下軟件包：

pip install llm2vec gritlm bm25

• 我們的評估要求將節點文檔嵌入到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 ：要生成嵌入的LLM名稱，例如text-embedding-ada-002 ， text-embedding-3-large ， voyage-large-2-instruct GritLM/GritLM-7B McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp
  • 有關其他參數，請參見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 VSS --emb_dir emb/ --output_dir output/ --emb_model GritLM/GritLM-7B --split test-0.1 --save_pred 

  python eval.py --dataset amazon --model LLMReranker --emb_dir emb/ --output_dir output/ --emb_model text-embedding-ada-002 --split human_generated_eval --llm_model gpt-4-1106-preview --save_pred

  鍵Args：

  • dataset ： amazon ， mag或prime的數據集之一。
  • model ：要評估的模型BM25 ， Colbertv2 ， VSS ， MultiVSS ， LLMReranker之一。
    • 請用參數--emb_model指定嵌入模型的名稱。
    • 如果您使用的是LLMReranker ，請用參數--llm_model指定LLM名稱。
    • 在命令行中指定API鍵

       export ANTHROPIC_API_KEY=YOUR_API_KEY
      

      或者

       export OPENAI_API_KEY=YOUR_API_KEY
      export OPENAI_ORG=YOUR_ORGANIZATION
      

      或者

       export VOYAGE_API_KEY=YOUR_API_KEY
      

  • emb_dir ：存儲嵌入的目錄。
  • split ：拆分以評估train ， val ， test ， test-0.1 （隨機樣本為10％）和human_generated_eval （在人類生成的查詢數據集上進行評估）。
  • output_dir ：存儲評估輸出的目錄。
  • surfix ：指定何時存儲的嵌入在文件夾doc{surfix}或query{surfix}中，例如，_no_compact，

參考

如果您在工作中使用我們的基准或代碼，請考慮引用我們的論文：

 @inproceedings{wu24stark,
    title        = {STaRK: Benchmarking LLM Retrieval on Textual and Relational Knowledge Bases},
    author       = {
        Shirley Wu and Shiyu Zhao and 
        Michihiro Yasunaga and Kexin Huang and 
        Kaidi Cao and Qian Huang and 
        Vassilis N. Ioannidis and Karthik Subbian and 
        James Zou and Jure Leskovec
    },
    booktitle    = {NeurIPS Datasets and Benchmarks Track},
    year         = {2024}
}