allRank

allRank: Learning to Rank in PyTorch

About allRank

allRank is a specialized framework based on PyTorch, designed to facilitate the training of neural Learning-to-Rank (LTR) models. It offers implementations for various ranking loss functions, ranging from pointwise and pairwise to listwise. The framework also provides different scoring functions, such as fully connected and Transformer-like models, while incorporating popular evaluation metrics like Normalized Discounted Cumulative Gain (NDCG) and Mean Reciprocal Rank (MRR). For those looking to experiment with simulated click-through data, allRank includes click-models as well.

Motivation Behind allRank

The primary aim of allRank is to make it simple and flexible to explore different neural LTR models and their corresponding loss functions. Users can effortlessly add custom loss functions and fine-tune the model configuration and training procedures. allRank is intended to aid research in neural LTR and help integrate these approaches into practical industrial applications.

Key Features

Implemented Loss Functions:

• ListNet: Suitable for both binary and graded relevance.
• ListMLE
• RankNet
• Ordinal Loss
• LambdaRank
• LambdaLoss
• Approximate NDCG
• Root Mean Square Error (RMSE)
• NeuralNDCG: A recent addition that introduces direct optimization of a ranking metric through differentiable relaxation of sorting.

Getting Started Guide

To make the onboarding process smoother, allRank provides a run_example.sh script. This script helps generate synthetic ranking data in libsvm format and trains a Transformer model using a pre-configured config.json file. Running this script requires Docker, and after execution, users can access the sample data in the dummy_data directory and the experiment results in the test_run directory.

Choosing the Right Architecture Version (GPU vs CPU-only)

allRank has separate binaries for GPU and CPU, requiring users to select the appropriate architecture version for Docker image. A compatible image version can be selected during the Docker build process by specifying either gpu or cpu as the arch_version argument. While executing run_example.sh, specify the architecture as a command-line argument like so:

run_example.sh gpu ...

If no specific version is specified, cpu is used by default.

Configuring the Model & Training

Users interested in training a custom model can modify the config.json file to configure experiments. Running the following command with your specifications will start the training process:

python allrank/main.py --config_file_name allrank/config.json --run_id <the_name_of_your_experiment> --job_dir <the_place_to_save_results>

The config.json file governs all aspects of training, including model definition, data paths, chosen losses and metrics, and other training hyperparameters. A template file, config_template.json, is provided to guide users in this process. For consistency, training data should be named train.txt, following the MSLR-WEB30K convention. Results will be saved under the <job_dir>/results/<run_id> path.

Custom Loss Function Implementation

allRank supports custom loss functions. To introduce a new loss function, developers should implement it to accept two tensors—model prediction and ground truth—and include it in the losses package. To use a custom loss in training, specify the function's name and its arguments in the config.json file:

"loss": {
  "name": "yourLoss",
  "args": {
    "arg1": "val1",
    "arg2": "val2"
  }
}

Utilizing Click-Models

To make use of a click model, begin by training an allRank model. Subsequently, execute:

python allrank/rank_and_click.py --input-model-path <path_to_the_model_weights_file> --roles <comma_separated_list_of_ds_roles_to_process e.g. train,valid> --config_file_name allrank/config.json --run_id <the_name_of_your_experiment> --job_dir <the_place_to_save_results>

The model will apply its ranking capabilities across all datasets defined in the configuration, and a click model will process these and generate a click-through dataset in libSVM format. This dataset can further be used for subsequent allRank trainings.

Continuous Integration

To ensure code quality and correctness, execute scripts/ci.sh, which validates style guidelines and conducts unit tests.

Research Context

allRank supports the research initiative titled "Context-Aware Learning to Rank with Self-Attention." If allRank aids your research, kindly reference:

@article{Pobrotyn2020ContextAwareLT,
  title={Context-Aware Learning to Rank with Self-Attention},
  author={Przemyslaw Pobrotyn and Tomasz Bartczak and Mikolaj Synowiec and Radoslaw Bialobrzeski and Jaroslaw Bojar},
  journal={ArXiv},
  year={2020},
  volume={abs/2005.10084}
}

Additionally, if employing the NeuralNDCG loss function, please cite the following:

@article{Pobrotyn2021NeuralNDCG,
  title={NeuralNDCG: Direct Optimisation of a Ranking Metric via Differentiable Relaxation of Sorting},
  author={Przemyslaw Pobrotyn and Radoslaw Bialobrzeski},
  journal={ArXiv},
  year={2021},
  volume={abs/2102.07831}
}

License

allRank is licensed under the Apache 2 License, granting users the freedom to use and modify the framework under specified conditions.