allRank is a PyTorch-based framework for training neural Learning-to-Rank (LTR) models, featuring implementations of:
- common pointwise, pairwise and listwise loss functions
- fully connected and Transformer-like scoring functions
- commonly used evaluation metrics like Normalized Discounted Cumulative Gain (NDCG) and Mean Reciprocal Rank (MRR)
allRank provides an easy and flexible way to experiment with various LTR neural network models and loss functions. It is easy to add a custom loss, and to configure the model and the training procedure. We hope that allRank will facilitate both research in neural LTR and its industrial applications.
- ListNet (For a binary and graded relevance)
- ListMLE
- RankNet
- Ordinal loss
- LambdaRank
- LambdaLoss
- ApproxNDCG
- RMSE
To help you get started, we provide a run_example.sh script which generates dummy ranking data in libsvm format and trains
a Transformer model on the data using provided example config.json config file. Once you run the script, the dummy data can be found in dummy_data directory
and the results of the experiment in test_run directory. To run the example, Docker is required.
To train your own model, configure your experiment in config.json file and run
python allrank/main.py --config_file_name allrank/config.json --run_id <the_name_of_your_experiment> --output <the_place_to_save_results>
All the hyperparameters of the training procedure: i.e. model defintion, data location, loss and metrics used, training hyperparametrs etc. are controlled
by the config.json file. We provide a template file config_template.json where supported attributes, their meaning and possible values are explained.
Note that following MSLR-WEB30K convention, your libsvm file with training data should be named train.txt. You can specify the name of the validation dataset
(eg. valid or test) in the config.
To experiment with your own custom loss, you need to implement a function that takes two tensors (model prediction and ground truth) as input
and put it in the losses package, making sure it is exposed on a package level.
To use it in training, simply pass the name (and args, if your loss method has some hyperparameters) of your function in the correct place in the config file:
"loss": {
"name": "yourLoss",
"args": {
"arg1": val1,
"arg2: val2
}
}
You should run scripts/ci.sh to verify that code passes style guidelines and unit tests.
Apache 2 License