SPLADE Neural Retrieval Model

SPLADE uses Facebook Hydra for modular configuration management. This allows you to define complex configurations using hierarchical YAML files and override any parameter or entire configuration package via the command line or environment variables.

Key Concepts

• Overriding a parameter: Use the parameter's package path. For example, if a parameter is nested under init_dict, use init_dict.parameter_name=value.
• Overriding a package: Use the package path to swap out entire configuration modules (e.g., train/data=msmarco).
• Config Selection: You can select a specific configuration file by setting the SPLADE_CONFIG_NAME environment variable or by providing a full path via SPLADE_CONFIG_FULLPATH.

Before starting the experiment, download and extract the required dataset using the following commands:

wget https://www.dropbox.com/s/kjk9scpku3mrqnn/data.tar.gz?dl=0 -O pruning_data.tar.gz
tar xzvf pruning_data.tar.gz

Training is performed using the splade.hf_train module via distributed PyTorch launch. You can choose between two training configurations located in main_config/two_msmarco:

1. Without titles: Use --config-name=splade_default.
2. With titles: Use --config-name=splade_titles.

Replace NUMGPU with the number of GPUs available on your node.

# Training without titles
python -m torch.distributed.launch --use_env --nproc_per_node NUMGPU -m splade.hf_train --config-name=splade_default --config-dir=main_config/two_msmarco

# Training with titles
python -m torch.distributed.launch --use_env --nproc_per_node NUMGPU -m splade.hf_train --config-name=splade_titles --config-dir=main_config/two_msmarco

Download the required pisa_index.tar.gz file and extract it to your current working directory to prepare the data for retrieval tasks.

wget https://www.dropbox.com/s/odkkbgg8lopcduk/pisa_index.tar.gz?dl=0 -O pisa_index.tar.gz
tar xzvf pisa_index.tar.gz

To use SPLADE, you need torch and the Hugging Face transformers library, along with the splade package. Ensure you have a working Python environment with these dependencies installed.

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer
from splade.models.transformer_rep import Splade

To perform retrieval across all indexes, use the run_all.sh script with the argument 3.

Important: If the indexes were not pruned in the previous step, this command will only query the base model.

Supported Model Names:

• eff_v_large
• eff_v_medium
• eff_v_small
• msmarco-deepimpact
• msmarco-unicoil-tilde

bash run_all.sh MODELNAME 3

The repository provides modules to manage the entire lifecycle of SPLADE models. You can run the full pipeline (training, indexing, and retrieval) using splade.all.

Experiments are managed via Hydra. You can override configuration parameters directly in the command line.

Quick start with toy data:

1. Activate the environment.
2. Set PYTHONPATH to the current directory.
3. Set SPLADE_CONFIG_NAME to your desired config (e.g., config_default.yaml).
4. Run the splade.all module.

conda activate splade_env
export PYTHONPATH=$PYTHONPATH:$(pwd)
export SPLADE_CONFIG_NAME="config_default.yaml"
python3 -m splade.all \
  config.checkpoint_dir=experiments/debug/checkpoint \
  config.index_dir=experiments/debug/index \
  config.out_dir=experiments/debug/out

To execute the entire workflow—training the model, indexing data, performing retrieval, and evaluating results—run the src.all module while specifying the desired configuration name via the SPLADE_CONFIG_NAME environment variable.

SPLADE_CONFIG_NAME=config_splade python -m src.all

Before training efficient SPLADE models using the provided configuration files, you must download and extract the middle-trained Pre-trained Language Models (PLMs) to the root of the repository.

Run the following commands to fetch and unpack the models:

wget https://www.dropbox.com/s/hir60b9yj194dv7/mlm_flops.tar.gz?dl=0
tar -xzvf mlm_flops.tar.gz?dl=0

You can use the Jupyter notebook inference_splade.ipynb to load a trained model and perform inference. This is useful for inspecting the predicted "bag-of-expanded-words" generated by the model.

Available models on Hugging Face include:

• naver/splade_v2_max (v2)
• naver/splade_v2_distil (v2)
• naver/splade-cocondenser-selfdistil (SPLADE++)
• naver/splade-cocondenser-ensembledistil (SPLADE++)
• naver/efficient-splade-V-large-doc + naver/efficient-splade-V-large-query (efficient SPLADE)
• naver/efficient-splade-VI-BT-large-doc + efficient-splade-VI-BT-large-query (efficient SPLADE)

To use the HuggingFace-based training version of SPLADE, install the following dependencies via pip:

pip install torch transformers==4.29.2 hydra-core faiss-cpu pytest numba h5py pytrec_eval tensorboard accelerate matplotlib

To perform all pruning operations for a specific model, use the run_all.sh script with the argument 1. This script creates the necessary directories and executes prune_all.sh.

Supported Model Names:

• eff_v_large
• eff_v_medium
• eff_v_small
• msmarco-deepimpact
• msmarco-unicoil-tilde

bash run_all.sh MODELNAME 1