Ecosyste.ms: Awesome

An open API service indexing awesome lists of open source software.

Awesome Lists | Featured Topics | Projects

https://github.com/allegro/allrank

allRank is a framework for training learning-to-rank neural models based on PyTorch.
https://github.com/allegro/allrank

click-model deep-learning information-retrieval learning-to-rank machine-learning ndcg python pytorch ranking transformer

Last synced: about 6 hours ago
JSON representation

allRank is a framework for training learning-to-rank neural models based on PyTorch.

Awesome Lists containing this project

README 

        
# allRank : Learning to Rank in PyTorch



## About



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)

* click-models for experiments on simulated click-through data



### Motivation



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.



## Features



### Implemented loss functions:

 1. ListNet (for binary and graded relevance)

 2. ListMLE

 3. RankNet

 4. Ordinal loss

 5. LambdaRank

 6. LambdaLoss

 7. ApproxNDCG

 8. RMSE

 9. NeuralNDCG (introduced in https://arxiv.org/pdf/2102.07831)



### Getting started guide



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.



### Getting the right architecture version (GPU vs CPU-only)



Since torch binaries are different for GPU and CPU and GPU version doesn't work on CPU - one must select & build appropriate docker image version.



To do so pass `gpu` or `cpu` as `arch_version` build-arg in 



```docker build --build-arg arch_version=${ARCH_VERSION}```



When calling `run_example.sh` you can select the proper version by a first cmd line argument e.g. 



```run_example.sh gpu ...```



with `cpu` being the default if not specified.



### Configuring your model & training



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



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. Results will be saved under the path ```/results/```

 

Google Cloud Storage is supported in allRank as a place for data and job results.



### Implementing custom loss functions



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

    }

  }

```



### Applying click-model



To apply a click model you need to first have an allRank model trained.

Next, run:



```python allrank/rank_and_click.py --input-model-path  --roles  --config_file_name allrank/config.json --run_id  --job_dir ``` 



The model will be used to rank all slates from the dataset specified in config. Next - a click model configured in config will be applied and the resulting click-through dataset will be written under ```/results/``` in a libSVM format.

The path to the results directory may then be used as an input for another allRank model training.



## Continuous integration



You should run `scripts/ci.sh` to verify that code passes style guidelines and unit tests.



## Research



This framework was developed to support the research project [Context-Aware Learning to Rank with Self-Attention](https://arxiv.org/abs/2005.10084). If you use allRank in your research, please cite:

```

@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 you use the NeuralNDCG loss function, please cite the corresponding work, [NeuralNDCG: Direct Optimisation of a Ranking Metric via Differentiable Relaxation of Sorting](https://arxiv.org/abs/2102.07831):

```

@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



Apache 2 License