stark
1.0.0
Stark est une référence de récupération semi-structurée à grande échelle sur les bases de connaissances textuelles et relationnelles, couvrant les applications de la recherche de produits, de la recherche de documents académiques et des demandes de biomédecine.
Avec des requêtes diverses, à consonance naturelle et pratiques qui nécessitent un raisonnement spécifique au contexte, Stark établit une nouvelle norme pour évaluer les systèmes de récupération du monde réel pilotés par les LLM et présente des défis importants pour les recherches futures.
Consultez notre site Web pour plus d'ensemble!
Avec Python> = 3,8 et <3.12
pip install stark-qa Créez un Conda Env avec Python> = 3,8 et <3.12 et installez les packages requis dans requirements.txt .
conda create -n stark python=3.11
conda activate stark
pip install -r requirements.txt 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 ) L'argument racine pour LOAD_SKB spécifie l'emplacement pour stocker les données SKB. Avec la valeur par défaut None , les données seront stockées dans HuggingFace Cache.
Les paires de réponses de question pour la tâche de récupération seront automatiquement téléchargées dans data/{dataset}/stark_qa par défaut. Nous avons fourni une scission officielle dans data/{dataset}/split .
Il existe deux façons de charger les données de base de connaissances:
download_processed=True .download_processed=False . Dans ce cas, Stark-Primekg met environ 5 minutes pour télécharger et charger les données traitées. Stark-Amazon et Stark-Mag peuvent prendre environ une heure pour traiter les données brutes. Si vous exécutez EVAL, vous pouvez installer les packages suivants:
pip install llm2vec gritlm bm25 Notre évaluation nécessite intégrer les documents de nœud dans candidate_emb_dict.pt , qui est un dictionnaire node_id -> torch.Tensor . Les incorporations de requête seront générées automatiquement si elles ne sont pas disponibles. Vous pouvez soit exécuter le script Python suivant pour télécharger des incorporations de requête et des intégres de documents générés par text-embedding-ada-002 . (Nous les fournissons afin que vous puissiez courir sur notre référence tout de suite.)
python emb_download.py --dataset amazon --emb_dir emb/Ou vous pouvez exécuter le code suivant pour générer vous-même la requête ou le document. Par exemple,
python emb_generate.py --dataset amazon --mode query --emb_dir emb/ --emb_model text-embedding-ada-002dataset : l'un d' amazon , mag ou prime .mode : Le contenu à intégrer, l'un des query ou doc (documents de nœud).emb_dir : Le répertoire pour stocker des intégres.emb_model : le nom LLM pour générer des intégres, tels GritLM/GritLM-7B text-embedding-ada-002 , text-embedding-3-large Embedding- voyage-large-2-instruct McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntpemb_generate.py pour d'autres arguments.Exécutez le script Python pour l'évaluation. Par exemple,
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_predArgs clés:
dataset : l'ensemble de données à évaluer, l'un d' amazon , mag ou prime .model : Le modèle à évaluer, l'un des BM25 , Colbertv2 , VSS , MultiVSS , LLMReranker .--emb_model .LLMReranker , veuillez spécifier le nom LLM avec Argument --llm_model . 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 : Le répertoire pour stocker des intégres.split : la division pour évaluer sur, une des train , val , test , test-0.1 (échantillon aléatoire à 10%) et human_generated_eval (à évaluer sur l'ensemble de données de requête généré par l'homme).output_dir : le répertoire pour stocker les sorties d'évaluation.surfix : Spécifiez lorsque les intégres stockés sont dans le dossier doc{surfix} ou query{surfix} , par exemple, _no_compact, Veuillez envisager de citer notre papier si vous utilisez notre référence ou notre code dans votre travail:
@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}
}
