### Complete Splade Training Workflow Example

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/training.md

A comprehensive Python script demonstrating the full training workflow for Splade models. It covers model initialization, data loading, optimizer and regularizer setup, and the training loop execution.

```python
import torch
from omegaconf import DictConfig
from splade.models.models_utils import get_model
from splade.losses.regularization import init_regularizer
from splade.datasets.dataloaders import SiamesePairsDataLoader
from splade.datasets.datasets import PairsDatasetPreLoad
from splade.optim.bert_optim import init_simple_bert_optim
from splade.tasks.transformer_trainer import SiameseTransformerTrainer
from splade.utils.utils import set_seed

# Set random seed
set_seed(42)

# Load configuration (simplified)
config = {
    "checkpoint_dir": "checkpoints",
    "nb_iterations": 100000,
    "lr": 2e-5,
    "warmup_steps": 10000,
    "weight_decay": 0.01,
    "fp16": True
}

init_dict = {
    "model_type_or_dir": "bert-base-uncased",
    "agg": "max"
}

# Create model
model = torch.nn.DataParallel(
    get_model({"matching_type": "splade"}, init_dict)
).cuda()

# Create data loaders
train_dataset = PairsDatasetPreLoad("data/train_pairs")
train_loader = SiamesePairsDataLoader(
    dataset=train_dataset,
    tokenizer_type="bert-base-uncased",
    max_length=512,
    batch_size=32,
    shuffle=True
)

# Initialize optimizer
optimizer, scheduler = init_simple_bert_optim(
    model,
    lr=config["lr"],
    warmup_steps=config["warmup_steps"],
    weight_decay=config["weight_decay"],
    num_training_steps=config["nb_iterations"]
)

# Initialize regularizer
regularizer = {
    "FLOPS": init_regularizer("FLOPS")
}

# Create trainer
from splade.tasks.base.trainer import TrainerIter
from splade.losses.pairwise import DistilMarginMSE

trainer = TrainerIter(
    iterations=(1, config["nb_iterations"]),
    model=model,
    loss=DistilMarginMSE(),
    optimizer=optimizer,
    config=config,
    train_loader=train_loader,
    scheduler=scheduler,
    regularizer=regularizer
)

# Train
trainer.train()
```

--------------------------------

### Install Pisa from the weight-queries branch

Source: https://github.com/naver/splade/blob/main/efficient_splade_pisa/README.md

Clone the Pisa repository, build it with CMake, and compile the release version. This is a prerequisite step.

```bash
git clone https://github.com/pisa-engine/pisa.git
cd pisa
mkdir build
cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
make
cd ../../ 

```

--------------------------------

### Splade Configuration File Example

Source: https://github.com/naver/splade/blob/main/conf/README.md

Example of a Hydra configuration file (conf/config_splade.yaml) for Splade. It defines defaults for training, indexing, and retrieval, and allows direct parameter setting.

```yaml
# FILES
defaults:
  ############## TRAIN ###################################
  - train/config: splade
  - train/data: msmarco
  - train/model: splade
  ############## INDEX ###################################
  - index: msmarco
  ############## RETRIEVE ################################
  - retrieve_evaluate: all

# Direct PARAMETER setting
config: # to be provided for each run
  checkpoint_dir: ???
  index_dir: ???
  out_dir: ???
  fp16: false

```

--------------------------------

### Splade Configuration File Example (train/config/splade.yaml)

Source: https://github.com/naver/splade/blob/main/conf/README.md

Example content of a Hydra configuration file for the Splade model. It defines global parameters for initialization, including model type, tokenizer, and FP16 settings.

```yaml
# @package _global_

init_dict:
  model_type_or_dir: distilbert-base-uncased
  model_type_or_dir_q: null
  freeze_d_model: 0
  agg: max
  fp16: true

config:
  tokenizer_type: distilbert-base-uncased

```

--------------------------------

### Training Script Usage Example

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/training.md

Demonstrates how to run the main training script and override configuration parameters. This includes setting output directories and training hyperparameters.

```bash
export SPLADE_CONFIG_NAME="config_splade++_ensembledistil.yaml"
python -m splade.train \
  config.checkpoint_dir=experiments/checkpoint \
  config.index_dir=experiments/index \
  config.out_dir=experiments/out
```

```bash
# Override training hyperparameters
python -m splade.train \
  config.nb_iterations=100000 \
  config.lr=1e-5 \
  config.warmup_steps=5000
```

```bash
# Override regularization
python -m splade.train \
  config.regularizer.FLOPS.lambda_q=0.05 \
  config.regularizer.FLOPS.lambda_d=0.01
```

--------------------------------

### FLOPs Configuration Example

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Sets the path to queries used for estimating FLOPs (floating point operations) for efficiency analysis.

```yaml
flops_queries: data/queries_for_flops
```

--------------------------------

### Configuration Composition for Training

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Example of setting an environment variable for a specific configuration name and overriding parameters for training SPLADE++ with ensemble distillation.

```bash
export SPLADE_CONFIG_NAME="config_splade++_ensembledistil.yaml"
python -m splade.train \
  config.checkpoint_dir=experiments/splade_v2 \
  config.index_dir=experiments/splade_v2/index \
  config.out_dir=experiments/splade_v2/out
```

--------------------------------

### Install SPLADE Environment

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Set up a Conda environment for SPLADE and activate it. This includes creating the environment and installing dependencies from a YAML file.

```bash
conda create -n splade_env python=3.9
conda activate splade_env
conda env create -f conda_splade_env.yml
```

--------------------------------

### Indexing Configuration Example

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Defines batch size for document encoding and the index directory. The index directory is required at runtime.

```yaml
index_retrieve_batch_size: 256      # Batch size for document encoding
index_dir: ???                       # Required at runtime
```

--------------------------------

### Install Dependencies for SPLADE Training

Source: https://github.com/naver/splade/blob/main/splade/hf/README.md

Installs necessary Python packages for training SPLADE models. Ensure you have PyTorch and HuggingFace Transformers installed.

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

--------------------------------

### Command-line Usage for HuggingFace Training

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/training.md

Example of how to launch the HuggingFace-based training script. Ensure the SPLADE_CONFIG_NAME environment variable is set before execution, and provide specific configuration overrides as needed.

```bash
export SPLADE_CONFIG_NAME="config_hf_splade_sigir23_32neg_distil.yaml"
python -m splade.hf_train config.checkpoint_dir=experiments/hf_checkpoint
```

--------------------------------

### Define example document

Source: https://github.com/naver/splade/blob/main/inference_splade.ipynb

Provides a sample document text from the MS MARCO passage collection for demonstration purposes.

```python
# example document from MS MARCO passage collection (doc_id = 8003157)

doc = "Glass and Thermal Stress. Thermal Stress is created when one area of a glass pane gets hotter than an adjacent area. If the stress is too great then the glass will crack. The stress level at which the glass will break is governed by several factors."
```

--------------------------------

### Data Configuration Example

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Specifies paths for training data, including triplets, validation qrels, query collections, and document collections. Supports optional hard negatives.

```yaml
TRAIN_PAIRS_PATH: data/toy_finetune/triplets
VALIDATION_QRELS: data/toy_finetune/val_qrels.json
Q_COLLECTION_PATH: [data/toy_finetune/val_queries]
DOCUMENT_PATH: data/toy_finetune/documents
```

--------------------------------

### Retrieval and Evaluation Configuration Example

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Specifies paths for query collections and evaluation qrels, along with retrieval parameters like top_k and threshold.

```yaml
Q_COLLECTION_PATH: [data/toy_finetune/val_queries]
EVAL_QREL_PATH: data/toy_finetune/val_qrels.json
top_k: 1000
threshold: 0
```

--------------------------------

### BCEWithLogitsLoss Example

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/losses.md

Use BCEWithLogitsLoss for binary cross-entropy. It treats positive examples as class 1 and negatives as class 0. Suitable for pointwise classification tasks.

```python
from splade.losses.pointwise import BCEWithLogitsLoss

loss_fn = BCEWithLogitsLoss()
loss = loss_fn({"pos_score": pos_scores, "neg_score": neg_scores})
```

--------------------------------

### DistilKLLoss Example

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/losses.md

Employ DistilKLLoss for KL divergence distillation. This loss aligns the student's score distribution with the teacher's. Input includes student and teacher positive and negative scores.

```python
from splade.losses.pairwise import DistilKLLoss

loss_fn = DistilKLLoss()
loss = loss_fn({
    "pos_score": student_pos,
    "neg_score": student_neg,
    "teacher_pos_score": teacher_pos,
    "teacher_neg_score": teacher_neg
})
```

--------------------------------

### Train a Reranker Model

Source: https://github.com/naver/splade/blob/main/splade/hf/README.md

Initiates the training of a reranker model using SPLADE output. This example uses a toy configuration for demonstration.

```bash
python -m splade.hf_train_reranker --config-name=config_reranker_train_toy
```

--------------------------------

### PairwiseNLL Loss Example

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/losses.md

Use PairwiseNLL for pairwise negative log-likelihood loss. It treats positive and negative examples as a softmax classification problem. Requires positive and negative scores as input.

```python
from splade.losses.pairwise import PairwiseNLL

loss_fn = PairwiseNLL()
pos_scores = torch.randn(8, 1)
neg_scores = torch.randn(8, 3)
loss = loss_fn({"pos_score": pos_scores, "neg_score": neg_scores})
```

--------------------------------

### Initialize Ranking Loss Functions

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Set up standard ranking loss functions like PairwiseNLL and DistilMarginMSE. Includes initialization for regularization and a weight scheduler for the training loop.

```python
from splade.losses.pairwise import PairwiseNLL, DistilMarginMSE
from splade.losses.regularization import init_regularizer, RegWeightScheduler

# Standard ranking loss
loss_fn = PairwiseNLL()

# With distillation
loss_fn = DistilMarginMSE()

# Add regularization
regularizer = init_regularizer("FLOPS")
reg_scheduler = RegWeightScheduler(lambda_=0.06, T=400000)

# Training loop
for step, batch in enumerate(train_loader):
    output = model(**batch)
    
    # Compute losses
    main_loss = loss_fn(output)
    lambda_t = reg_scheduler.step()
    reg_loss = lambda_t * regularizer(output["d_rep"])
    
    total_loss = main_loss + reg_loss
    total_loss.backward()
```

--------------------------------

### Download MSMARCO Data

Source: https://github.com/naver/splade/blob/main/main_config/two_msmarco/README.md

Download the collection data for the MSMARCO experiments. Ensure you have wget installed.

```bash
wget https://www.dropbox.com/s/te0qblvvbba76q3/collection_with_titles.tar.gz?dl=0 -O collection_with_titles.tar.gz
tar xzvf collection_with_titles.tar.gz
```

--------------------------------

### Get Total Number of Documents Indexed

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/indexing.md

Returns the total count of documents that have been added to the index.

```python
def nb_docs(self) -> int
```

```python
num_docs = index.nb_docs()
```

--------------------------------

### Query Collection Format (TSV)

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Example of the TSV format for query collection, mapping query IDs to their text.

```text
query_id	query_text
0	example query
1	another search query
2	third test query
```

--------------------------------

### Get Number of Terms in Index

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/indexing.md

Returns the count of unique vocabulary terms that have at least one document associated with them.

```python
def __len__(self) -> int
```

```python
num_terms = len(index)
```

--------------------------------

### Run Training with Custom Configuration

Source: https://github.com/naver/splade/blob/main/README.md

This command initiates the training process for SPLADE. It allows overriding specific configuration parameters directly from the command line.

```bash
python3 -m splade.train
```

--------------------------------

### Add New Parameters

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Demonstrates how to add new parameters to the configuration at runtime that may not be defined in the original config file.

```bash
# Add new parameters not in config
python -m splade.train +quantization_factor_document=100
```

--------------------------------

### Collection Data Format (TSV)

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Example of the TSV format for collection data, mapping document IDs to their text content.

```text
doc_id	document_text
0	First document content here
1	Second document with more text
2	Another document example
```

--------------------------------

### Create Anserini Readable Files with Custom Quantization

Source: https://github.com/naver/splade/blob/main/README.md

This command prepares data for Anserini, allowing custom quantization factors for documents and queries.

```bash
SPLADE_CONFIG_FULLPATH=/path/to/checkpoint/dir/config.yaml python3 -m splade.create_anserini +quantization_factor_document=100 +quantization_factor_query=100
```

--------------------------------

### Run Training Configuration

Source: https://github.com/naver/splade/blob/main/benchmarking_sigir23/README.md

Export the Python path and run a training configuration using a specified middletrained model, configuration file, and output directories.

```bash
export PYTHONPATH=$PYTHONPATH:$(pwd)
python -m splade.hf_train init_dict.model_type_or_dir=<YOUR MIDDLETRAINED MODEL>  --config-name <config> config.checkpoint_dir=<PATH>  config.index_dir=<PATHINDEX>  config.out_dir=<PATHRES>  
```

--------------------------------

### Create Directory if Not Exists

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/utilities.md

Use `makedir` to create a directory. It ensures the directory exists before proceeding.

```python
from splade.utils.utils import makedir

makedir("output/checkpoints")
```

--------------------------------

### Initialize EvalDatasetMonoT5 for MonoT5 Reranking

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/datasets.md

Initializes the EvalDatasetMonoT5 for monoT5-style reranking. Requires paths to the run file, document, and query collections, with optional top_k and finish_qrel parameters.

```python
class EvalDatasetMonoT5:
    def __init__(self, run_file, document_dir, query_dir, top_k=-1, finish_qrel=None)
```

--------------------------------

### Qrels Data Format (JSON)

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Example of the JSON format for relevance judgments (qrels), mapping query IDs to document IDs and their relevance scores.

```json
{
  "query_id_1": {
    "doc_id_1": 1,
    "doc_id_2": 0
  },
  "query_id_2": {
    "doc_id_3": 2,
    "doc_id_1": 1
  }
}
```

--------------------------------

### Main Training Script Entry Point

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/training.md

The main training script that uses Hydra for configuration management. It orchestrates model creation, optimizer initialization, checkpoint loading, and trainer execution.

```python
@hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base="1.2")
def train(exp_dict: DictConfig)
```

--------------------------------

### PairwiseBPR Loss Example

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/losses.md

Implement PairwiseBPR for Bayesian Personalized Ranking loss. This loss is suitable for ranking tasks and requires positive and negative scores.

```python
from splade.losses.pairwise import PairwiseBPR

loss_fn = PairwiseBPR()
loss = loss_fn({"pos_score": pos_scores, "neg_score": neg_scores})
```

--------------------------------

### Correct Initialization of SpladeDoc Model

Source: https://github.com/naver/splade/blob/main/_autodocs/errors.md

Demonstrates the correct way to initialize SpladeDoc, avoiding invalid configurations like specifying a separate query encoder or freezing the document model.

```python
# Incorrect
model = SpladeDoc(
    "bert-base-uncased",
    model_type_or_dir_q="bert-base-uncased",  # Error!
    freeze_d_model=True  # Error!
)

# Correct
model = SpladeDoc("bert-base-uncased")
```

--------------------------------

### HuggingFace Trainer with Hydra Configuration

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/training.md

This snippet shows the entry point for training with HuggingFace's Trainer, managed by Hydra for configuration. It requires setting the SPLADE_CONFIG_NAME environment variable and passing configuration arguments via the command line.

```python
@hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base="1.2")
def train_hf(exp_dict: DictConfig)
```

--------------------------------

### TextCollectionDataLoader Initialization and Usage

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/dataloaders.md

Initialize TextCollectionDataLoader with tokenizer and max length. Iterate through the loader to get batches containing tokenized inputs and original texts.

```python
from splade.datasets.dataloaders import TextCollectionDataLoader

loader = TextCollectionDataLoader(
    dataset=dataset,
    tokenizer_type="bert-base-uncased",
    max_length=512,
    batch_size=64
)

for batch in loader:
    tokens = batch["input_ids"]
    texts = batch["text"]
```

--------------------------------

### Enable FP16 Mixed Precision Training

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Use the command-line argument `config.fp16=true` to enable FP16 mixed precision training, which can significantly speed up training on compatible hardware.

```bash
python -m splade.train config.fp16=true
```

--------------------------------

### Training Data Format (TSV)

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Example of the TSV format for training data, containing query text, positive document text, and negative document text.

```text
query_text	positive_document_text	negative_document_text
how to cook pasta	pasta is made from wheat and water...	cooking requires heat...
what is machine learning
machine learning is a subset of AI...	python is a programming language...
```

--------------------------------

### Initialize SpladeTopK Model

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Instantiate the SpladeTopK model, optimized for efficiency with top-k pruning. Configure the number of dimensions for document and query pruning.

```python
from splade.models.transformer_rep import SpladeTopK

model = SpladeTopK("bert-base-uncased", top_d=32, top_q=5, agg="max")
```

--------------------------------

### Creating Directories with makedir

Source: https://github.com/naver/splade/blob/main/_autodocs/errors.md

Use the `makedir` function to pre-create necessary directories for checkpoints, index, or results if they do not exist.

```python
from splade.utils.utils import makedir

makedir("checkpoints")
makedir("index")
makedir("results")
```

--------------------------------

### DistilMarginMSE Loss Example

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/losses.md

Use DistilMarginMSE for margin distillation loss. It matches the student's score margin to the teacher's margin. Requires both student and teacher positive and negative scores.

```python
from splade.losses.pairwise import DistilMarginMSE

loss_fn = DistilMarginMSE()
loss = loss_fn({
    "pos_score": student_pos,
    "neg_score": student_neg,
    "teacher_pos_score": teacher_pos,
    "teacher_neg_score": teacher_neg
})
```

--------------------------------

### Run All Splade Components

Source: https://github.com/naver/splade/blob/main/conf/README.md

Execute training, indexing, retrieval, and evaluation in a single command. Ensure the SPLADE_CONFIG_NAME environment variable is set to 'config_splade'.

```bash
SPLADE_CONFIG_NAME=config_splade python -m src.all
```

--------------------------------

### Training Entry Point (train.py)

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/training.md

The main training script that orchestrates the training process using Hydra for configuration management. It handles configuration parsing, model creation, optimizer initialization, checkpoint loading, and trainer execution.

```APIDOC
## Training Entry Point

### train.py

Main training script using Hydra configuration.

```python
@hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base="1.2")
def train(exp_dict: DictConfig)
```

**Key steps:**

1. Parse configuration with `get_initialize_config()`
2. Create model with `get_model()`
3. Initialize optimizer and scheduler with `init_simple_bert_optim()`
4. Load checkpoint if resuming (step number, model state, optimizer state)
5. Create trainer instance
6. Call `trainer.train()`

**Usage:**

```bash
export SPLADE_CONFIG_NAME="config_splade++_ensembledistil.yaml"
python -m splade.train \
  config.checkpoint_dir=experiments/checkpoint \
  config.index_dir=experiments/index \
  config.out_dir=experiments/out
```

**Configuration overrides:**

```bash
# Override training hyperparameters
python -m splade.train \
  config.nb_iterations=100000 \
  config.lr=1e-5 \
  config.warmup_steps=5000

# Override regularization
python -m splade.train \
  config.regularizer.FLOPS.lambda_q=0.05 \
  config.regularizer.FLOPS.lambda_d=0.01
```
```

--------------------------------

### InBatchPairwiseNLL Loss Example

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/losses.md

Utilize InBatchPairwiseNLL for in-batch negative sampling. It considers negatives from other queries in the batch and optional hard negatives. Input includes positive scores and hard negative scores.

```python
from splade.losses.pairwise import InBatchPairwiseNLL

loss_fn = InBatchPairwiseNLL()
in_batch_scores = torch.randn(8, 8)  # batch_size x batch_size
neg_scores = torch.randn(8, 1)  # hard negatives
loss = loss_fn({"pos_score": in_batch_scores, "neg_score": neg_scores})
```

--------------------------------

### Compute Metrics with SPLADE

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Calculate Mean Reciprocal Rank (MRR) and Normalized Discounted Cumulative Gain (NDCG) using the provided run and qrel data. Ensure the 'splade' library is installed and imported.

```python
from splade.evaluate import mrr_k, evaluate

mrr = mrr_k(run, qrel, k=10)
print(f"MRR@10: {mrr}")

# Other metrics
ndcg = evaluate(run, qrel, metric="ndcg", agg=True)
recall_100 = evaluate(run, qrel, metric="recall", agg=True, select="100")
print(f"NDCG: {ndcg}")
print(f"Recall@100: {recall_100}")
```

--------------------------------

### Initialize SparsityRatio Regularizer

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/losses.md

Instantiates the SparsityRatio regularizer, which calculates the fraction of zero elements. Requires the output dimension (vocabulary size) to be specified.

```python
from splade.losses.regularization import SparsityRatio

reg = SparsityRatio(output_dim=30522)
sparcity = reg(batch_repr)
```

--------------------------------

### Initialize SparseIndexing

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/evaluation.md

Instantiate the SparseIndexing class to prepare for creating an inverted index. Ensure the model and configuration, including 'index_dir', are provided.

```python
class SparseIndexing(Evaluator):
    def __init__(self, model, config, compute_stats=False, dim_voc=None, is_query=False, 
                 force_new=True, **kwargs)
```

--------------------------------

### Build Custom SPLADE Index

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Guides through building a custom inverted index using SPLADE. This involves loading a model, preparing a document collection, encoding documents in batches, and adding them to the index. The index and document IDs are then saved.

```python
from splade.models.transformer_rep import Splade
from splade.indexing.inverted_index import IndexDictOfArray
from splade.datasets.dataloaders import CollectionDataLoader
from splade.datasets.datasets import CollectionDatasetPreLoad
import torch

# Load model
model = Splade("bert-base-uncased").cuda()

# Load documents
dataset = CollectionDatasetPreLoad("data/docs", id_style="row_id")
loader = CollectionDataLoader(
    dataset=dataset,
    tokenizer_type="bert-base-uncased",
    max_length=512,
    batch_size=64
)

# Create index
index = IndexDictOfArray("my_index", force_new=True, dim_voc=30522)
doc_ids = []

# Encode and index
model.eval()
with torch.no_grad():
    for batch in loader:
        inputs = {k: v.cuda() for k, v in batch.items() if k != "id"}
        outputs = model(d_kwargs=inputs)
        doc_repr = outputs["d_rep"]
        
        # Convert to COO format
        row, col = torch.nonzero(doc_repr, as_tuple=True)
        data = doc_repr[row, col]
        
        # Add to index
        batch_ids = batch["id"].tolist()
        doc_ids.extend(batch_ids)
        index.add_batch_document(
            row.cpu().numpy(),
            col.cpu().numpy(),
            data.cpu().numpy()
        )

# Save index
import pickle
index.save(dim=30522)
pickle.dump(doc_ids, open("my_index/doc_ids.pkl", "wb"))
```

--------------------------------

### Initializing Regularizers with Valid Names in SPLADE

Source: https://github.com/naver/splade/blob/main/_autodocs/errors.md

Demonstrates the correct usage of 'init_regularizer' by providing valid regularizer names. Using an invalid name will result in a NotImplementedError.

```python
# Incorrect
reg = init_regularizer("invalid")  # Raises NotImplementedError

# Correct
reg = init_regularizer("L1")
reg = init_regularizer("FLOPS")
reg = init_regularizer("sparsity_ratio", output_dim=30522)
```

--------------------------------

### Train a SPLADE Model

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Initiate the training process for a SPLADE model. Specify directories for checkpoints, index, and output. This command loads configuration, initializes the model, prepares data loaders, trains the model, and saves checkpoints.

```bash
python -m splade.train \
  config.checkpoint_dir=experiments/checkpoint \
  config.index_dir=experiments/index \
  config.out_dir=experiments/out
```

--------------------------------

### Run Indexing with Custom Configuration

Source: https://github.com/naver/splade/blob/main/README.md

This command initiates the indexing process for SPLADE. It allows overriding specific configuration parameters directly from the command line.

```bash
python3 -m splade.index
```

--------------------------------

### Run All-in-One SPLADE Workflow

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Execute training, indexing, and retrieval in a single command. This simplifies the SPLADE workflow by chaining the core operations.

```bash
python -m splade.all \
  config.checkpoint_dir=experiments/checkpoint \
  config.index_dir=experiments/index \
  config.out_dir=experiments/out
```

--------------------------------

### Set SPLADE Configuration Name

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Specify the configuration file to use for the experiment. This is an alternative to providing the full path.

```bash
export SPLADE_CONFIG_NAME="config_splade++_ensembledistil.yaml"
```

--------------------------------

### Download and extract indexes and queries

Source: https://github.com/naver/splade/blob/main/efficient_splade_pisa/README.md

Download the Pisa index and query data archive and extract its contents to the current directory.

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

```

--------------------------------

### Train SPLADE Model with Configuration

Source: https://github.com/naver/splade/blob/main/_autodocs/README.md

Train a SPLADE model using a specified configuration file and output directories for checkpoints, index, and results. Ensure the SPLADE_CONFIG_NAME environment variable is set.

```bash
export SPLADE_CONFIG_NAME="config_splade++_ensembledistil.yaml"
python -m splade.train \
  config.checkpoint_dir=experiments/checkpoint \
  config.index_dir=experiments/index \
  config.out_dir=experiments/out
```

--------------------------------

### Create Anserini Readable Files

Source: https://github.com/naver/splade/blob/main/README.md

This command prepares data in a format readable by Anserini, including JSON collection and TSV queries. It requires specifying the model and index directory.

```bash
python3 -m splade.create_anserini \
  init_dict.model_type_or_dir=naver/splade-cocondenser-ensembledistil \
  config.pretrained_no_yamlconfig=true \
  config.index_dir=experiments/pre-trained/index \
  +quantization_factor_document=100 \
  +quantization_factor_query=100
```

--------------------------------

### Complete Indexing Workflow

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/indexing.md

Demonstrates a full indexing process using a Splade model, collection data loader, and IndexDictOfArray. It encodes documents, converts them to COO format, adds them to the index, and saves the final index and document IDs.

```python
import torch
import numpy as np
from splade.models.transformer_rep import Splade
from splade.datasets.dataloaders import CollectionDataLoader
from splade.datasets.datasets import CollectionDatasetPreLoad
from splade.indexing.inverted_index import IndexDictOfArray
import pickle

# Initialize model
model = Splade("bert-base-uncased", agg="max")
model.cuda()

# Load collection
dataset = CollectionDatasetPreLoad("data/documents", id_style="row_id")
loader = CollectionDataLoader(
    dataset=dataset,
    tokenizer_type="bert-base-uncased",
    max_length=512,
    batch_size=64,
    shuffle=False
)

# Create index
index = IndexDictOfArray("index_dir", force_new=True, dim_voc=model.output_dim)
doc_ids = []

# Encode documents and add to index
model.eval()
with torch.no_grad():
    for batch in loader:
        inputs = {k: v.cuda() for k, v in batch.items() if k != "id"}
        outputs = model(d_kwargs=inputs)
        doc_repr = outputs["d_rep"]  # (batch_size, vocab_size)
        
        # Convert to COO format
        row, col = torch.nonzero(doc_repr, as_tuple=True)
        data = doc_repr[row, col]
        
        # Add to index
        batch_ids = batch["id"].tolist()
        doc_ids.extend(batch_ids)
        index.add_batch_document(
            row.cpu().numpy(),
            col.cpu().numpy(),
            data.cpu().numpy()
        )

# Save index
index.save(dim=model.output_dim)
pickle.dump(doc_ids, open("index_dir/doc_ids.pkl", "wb"))
```

--------------------------------

### Download and Extract Pre-trained Models

Source: https://github.com/naver/splade/blob/main/conf/efficient_splade/README.md

Download the middle trained PLMs from the provided URL and extract them. This is a prerequisite for training using the provided configurations.

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

--------------------------------

### Valid Match Parameters for Siamese Model Initialization

Source: https://github.com/naver/splade/blob/main/_autodocs/errors.md

Shows the correct 'match' parameters ('dot_product' or 'cosine_sim') for initializing a Siamese model to avoid an AssertionError.

```python
model = Siamese("bert-base-uncased", match="dot_product")
model = Siamese("bert-base-uncased", match="cosine_sim")
```

--------------------------------

### Training Parameters Configuration

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Sets core training parameters including the number of iterations, learning rate, warmup steps, weight decay, and early stopping criteria. Also defines metrics for monitoring and checkpoint selection.

```yaml
nb_iterations: 400000
lr: 2e-5
warmup_steps: 10000
weight_decay: 0.01
early_stopping: "loss"             # Metric to monitor: "loss" or metric name
patience: 3
monitoring_ckpt: "loss"            # Metric for checkpoint selection
seed: 42
```

--------------------------------

### Initialize and Use RegWeightScheduler

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/losses.md

Demonstrates how to use the RegWeightScheduler to gradually increase the regularization weight over training steps. The scheduler implements a quadratic increase schedule.

```python
from splade.losses.regularization import RegWeightScheduler

scheduler = RegWeightScheduler(lambda_=0.1, T=10000)
for step in range(10000):
    lambda_t = scheduler.step()
    loss = main_loss + lambda_t * regularizer(batch_repr)
```

--------------------------------

### Initialize MsMarcoHardNegatives Dataset

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/datasets.md

Initializes the MsMarcoHardNegatives dataset. Requires paths to the dataset, document collection, query collection, and relevance judgments.

```python
class MsMarcoHardNegatives(Dataset):
    def __init__(self, dataset_path, document_dir, query_dir, qrels_path)
    def __len__(self) -> int
    def __getitem__(self, idx) -> Tuple[str, str, List[str], List[str]]
```

--------------------------------

### Initialize L1 Regularizer

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/losses.md

Instantiates the L1 sparsity regularizer. Use this to encourage sparse representations in your model.

```python
from splade.losses.regularization import L1

reg = L1()
batch_repr = torch.randn(8, 30522)
l1_loss = reg(batch_repr)
```

--------------------------------

### Train SPLADE Model (With Titles)

Source: https://github.com/naver/splade/blob/main/main_config/two_msmarco/README.md

Train the SPLADE model using a configuration that includes titles. Requires NUMGPU to be set to the number of GPUs available.

```bash
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
```

--------------------------------

### Run Retrieval with Custom Configuration

Source: https://github.com/naver/splade/blob/main/README.md

This command initiates the retrieval process for SPLADE. It allows overriding specific configuration parameters directly from the command line.

```bash
python3 -m splade.retrieve
```

--------------------------------

### Validate Configuration with OmegaConf

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Load and validate your experiment configuration using OmegaConf. This utility function helps in initializing the configuration and identifying any potential issues.

```python
from omegaconf import OmegaConf
from splade.utils.utils import get_initialize_config

# Load and validate config
exp_dict = OmegaConf.create({...})
exp_dict, config, init_dict, _ = get_initialize_config(exp_dict)
print(OmegaConf.to_yaml(config))
```

--------------------------------

### Train Toy Splade Model (Default)

Source: https://github.com/naver/splade/blob/main/conf/README.md

Train a toy Splade model using default configurations. Specify the checkpoint directory for the output.

```python
python -m src.train config.checkpoint_dir=<MY_RES_DIR>
```

--------------------------------

### Loss Functions and Regularizers

Source: https://github.com/naver/splade/blob/main/_autodocs/FILE_INDEX.txt

Documentation for various loss functions and regularizers used in SPLADE training, including PairwiseNLL, InBatchPairwiseNLL, and regularization parameters.

```APIDOC
## Loss Functions and Regularizers

This section describes the loss functions and regularization techniques available in SPLADE.

### Loss Functions

- **PairwiseNLL**: Negative Log Likelihood loss for pairwise comparisons.
- **InBatchPairwiseNLL**: In-batch variant of PairwiseNLL.
- **PairwiseBPR**: Bayesian Personalized Ranking loss.
- **DistilMarginMSE**: Mean Squared Error loss with a distillation margin.
- **DistilKLLoss**: Kullback-Leibler divergence loss for distillation.
- **BCEWithLogitsLoss**: Binary Cross-Entropy loss with logits.

### Regularizers

- **L1**: L1 regularization.
- **L0**: L0 regularization.
- **FLOPS**: Floating Point Operations regularization.
- **RegWeightScheduler**: Scheduler for regularization weights.
- **SparsityRatio**: Regularizer for sparsity.

### Usage

Refer to `api-reference/losses.md` for detailed information on each loss function and regularizer, including their parameters and usage.
```

--------------------------------

### Run Training with Custom Regularizer Parameters

Source: https://github.com/naver/splade/blob/main/README.md

This command runs the SPLADE training process, allowing external modification of regularizer parameters like lambda_q and lambda_d.

```bash
python3 -m splade.all config.regularizer.FLOPS.lambda_q=0.06 config.regularizer.FLOPS.lambda_d=0.02
```

--------------------------------

### Training Configuration

Source: https://github.com/naver/splade/blob/main/_autodocs/types.md

Parameters for training SPLADE models, such as the number of iterations, learning rate, warmup steps, weight decay, and early stopping criteria.

```python
{
    "nb_iterations": int,           # Total training iterations
    "lr": float,                    # Learning rate
    "warmup_steps": int,            # Linear warmup steps
    "weight_decay": float,          # L2 regularization
    "early_stopping": str,          # Validation metric for early stopping
    "patience": int,                # Patience for early stopping
    "monitoring_ckpt": str,         # Metric to monitor for checkpointing
}
```

--------------------------------

### PyTorch Model Training and Inference

Source: https://github.com/naver/splade/blob/main/_autodocs/types.md

Demonstrates standard PyTorch patterns for setting a model to training or inference mode, performing forward passes, and backpropagation.

```python
# Training mode
model.train()
output = model(**input_dict)
loss.backward()

# Inference mode
model.eval()
with torch.no_grad():
    output = model(**input_dict)
```

--------------------------------

### Build Sparse Inverted Index

Source: https://github.com/naver/splade/blob/main/_autodocs/README.md

Initialize a sparse inverted index using a pre-trained model checkpoint. This command sets up the index directory for efficient retrieval.

```bash
python -m splade.index \
  init_dict.model_type_or_dir=experiments/checkpoint/model \
  config.index_dir=experiments/index
```

--------------------------------

### Train SPLADE Model (Without Titles)

Source: https://github.com/naver/splade/blob/main/main_config/two_msmarco/README.md

Train the SPLADE model using the default configuration, excluding titles. Requires NUMGPU to be set to the number of GPUs available.

```bash
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
```

--------------------------------

### Load SPLADE Model

Source: https://github.com/naver/splade/blob/main/_autodocs/README.md

Instantiate and prepare a SPLADE model for evaluation. Ensure the model name is valid.

```python
from splade.models.transformer_rep import Splade

model = Splade("naver/splade-cocondenser-ensembledistil")
model.eval()
```

--------------------------------

### Run All Pruning Experiments

Source: https://github.com/naver/splade/blob/main/pruning/README.md

Execute this script to perform all pruning experiments for a specified model. The script handles folder creation and calls the individual pruning script.

```bash
bash run_all.sh MODELNAME 1
```

--------------------------------

### Train Splade Model Using Full Config Path

Source: https://github.com/naver/splade/blob/main/conf/README.md

Train a Splade model by providing the full path to a saved configuration file using the SPLADE_CONFIG_FULLPATH environment variable.

```bash
SPLADE_CONFIG_FULLPATH=/path/to/my/exp/config_splade.yaml python -m src.train
```

--------------------------------

### Initialize SPLADE Model

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Instantiate the standard SPLADE model, which includes both query and document encoders with expansion. This model uses a 'max' aggregation strategy.

```python
from splade.models.transformer_rep import Splade

model = Splade("bert-base-uncased", agg="max")
```

--------------------------------

### Download and Extract Data

Source: https://github.com/naver/splade/blob/main/pruning/README.md

Use these commands to download the necessary data archive and extract its contents for the pruning experiments.

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

--------------------------------

### Create and Populate Sparse Index

Source: https://github.com/naver/splade/blob/main/_autodocs/README.md

Initialize a sparse inverted index and add document representations in batches. Requires `torch`. Ensure `force_new=True` is used for a fresh index.

```python
from splade.indexing.inverted_index import IndexDictOfArray

index = IndexDictOfArray("index_dir", force_new=True, dim_voc=30522)

# Add documents in batches
for batch in loader:
    row, col = torch.nonzero(doc_repr, as_tuple=True)
    data = doc_repr[row, col]
    index.add_batch_document(row.numpy(), col.numpy(), data.numpy())

index.save()
```

--------------------------------

### Initialize DistilPairsDatasetPreLoad

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/datasets.md

Initialize DistilPairsDatasetPreLoad for training with teacher scores, suitable for distillation. The data directory must contain a raw.tsv file with query, positive document, negative document, and their respective teacher scores.

```python
from splade.datasets.datasets import DistilPairsDatasetPreLoad

dataset = DistilPairsDatasetPreLoad("data/train_distil")
query, pos, neg, s_pos, s_neg = dataset[0]
```

--------------------------------

### Pre-trained Model Indexing Configuration

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Configures the environment and command-line arguments for indexing a pre-trained model from HuggingFace. Sets `config.pretrained_no_yamlconfig=true` to load the model directly.

```bash
export SPLADE_CONFIG_NAME="config_splade++_cocondenser_ensembledistil"
python -m splade.index \
  init_dict.model_type_or_dir=naver/splade-cocondenser-ensembledistil \
  config.pretrained_no_yamlconfig=true \
  config.index_dir=experiments/pre-trained/index
```

--------------------------------

### Root Configuration Defaults

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Defines the default configuration subsets for training, data, model, indexing, retrieval, and FLOPS estimation. Specifies required output directories and precision settings.

```yaml
defaults:
  - train/config: splade_toy        # Training config subset
  - train/data: toy                 # Data paths subset
  - train/model: splade             # Model parameters subset
  - index: toy                      # Indexing config
  - retrieve_evaluate: toy          # Retrieval/evaluation config
  - flops: toy                       # FLOPS estimation config

config:
  checkpoint_dir: ???               # Required: where to save checkpoints
  index_dir: ???                    # Required: where to save/load index
  out_dir: ???                      # Required: where to save results
  fp16: false                       # Use float16 mixed precision
```

--------------------------------

### Initialize SpladeLexical Model

Source: https://github.com/naver/splade/blob/main/_autodocs/api-reference/transformer_rep.md

Instantiate the SpladeLexical model, which enhances sparse representations with lexical term matching. Configure the lexical type and aggregation method during initialization.

```python
class SpladeLexical(SiameseBase):
    def __init__(self, model_type_or_dir, model_type_or_dir_q=None, freeze_d_model=False, 
                 lexical_type="query", agg="sum", fp16=False)
```

--------------------------------

### Load SPLADE Model and Encode Document

Source: https://github.com/naver/splade/blob/main/_autodocs/QUICKSTART.md

Demonstrates how to load a trained SPLADE model and encode a document for retrieval. Requires PyTorch and Hugging Face Transformers. Ensure the model and tokenizer are correctly loaded and the document is tokenized.

```python
import torch
from splade.models.transformer_rep import Splade

# Load model
model = Splade("path/to/checkpoint/model", agg="max")
model.eval()

# Encode document
doc_text = "example document"
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
tokens = tokenizer(doc_text, return_tensors="pt", padding="longest", 
                   truncation=True, max_length=512)
tokens = {k: v.cuda() for k, v in tokens.items()}

with torch.no_grad():
    output = model(d_kwargs=tokens)
    doc_repr = output["d_rep"]  # sparse representation
```

--------------------------------

### Set Full Path to SPLADE Configuration

Source: https://github.com/naver/splade/blob/main/_autodocs/configuration.md

Provide the absolute path to the desired configuration file. This overrides SPLADE_CONFIG_NAME if both are set.

```bash
export SPLADE_CONFIG_FULLPATH="/path/to/experiments/checkpoint_dir/config.yaml"
```

--------------------------------

### Load SPLADE model and tokenizer

Source: https://github.com/naver/splade/blob/main/inference_splade.ipynb

Initializes the SPLADE model and tokenizer from the specified Hugging Face model directory. The model is set to evaluation mode.

```python
# loading model and tokenizer

model = Splade(model_type_or_dir, agg="max")
model.eval()
tokenizer = AutoTokenizer.from_pretrained(model_type_or_dir)
reverse_voc = {v: k for k, v in tokenizer.vocab.items()}
```