Citation Prediction

A citation recommendation system that allows users to find relevant citations for their paper drafts. The tool is backed by Semantic Scholar's OpenCorpus dataset (released on 2017-02-21).

This repository contains code to train and evaluate citation recommendation models.

Data

You can download the associated OpenCorpus dataset here.

For models and data associated with experiments from the paper, please refer below.

Installation

1. Download and install Conda

2. [Optional] Install direnv

3. Install GPU Drivers

citation-prediction uses the tensorflow framework to run its neural network models. It was tested on NVIDIA GPUs (GeForce GTX 1080 and Tesla V100-SXM2). To use GPUs, install CUDA and cudnn compatible with your OS and GPU. The version of tensorflow used in this repo is 1.12, and requires CUDA 9.0.

1. Run:

   `./env.sh`
   

2. Activate Conda Environment

Citation-prediction Evaluation

This section details how to run the end-to-end system using pre-trained models and evaluate performance of citation-prediction for each dataset. If you successfully executed the previous steps, trained models should already be available.

• Open Corpus

  ATTENTION: Experimenting with the open corpus dataset needs a system with at least 100G of RAM.
  bash python citation_prediction/scripts/evaluate.py --dataset_type oc --candidate_selector_type ann --split test --paper_embedder_dir data/open_corpus/models/paper_embedder/ --num_candidates 5 --ranker_type neural --citation_ranker_dir data/open_corpus/models/citation_ranker/ --n_eval 20000

• Pubmed ```bash python citationprediction/scripts/evaluate.py --datasettype pubmed --candidateselectortype ann --split test --paperembedderdir data/comparison/pubmed/models/paperembedder/ --numcandidates 10 --rankertype neural --citationrankerdir data/comparison/pubmed/models/citationranker/

```

• DBLP bash python citation_prediction/scripts/evaluate.py --dataset_type dblp --candidate_selector_type ann --split test --paper_embedder_dir data/comparison/dblp/models/paper_embedder/ --num_candidates 10 --ranker_type neural --citation_ranker_dir data/comparison/dblp/models/citation_ranker/

BM25 Baseline

• Open Corpus bash python citation_prediction/scripts/evaluate.py --dataset_type oc --candidate_selector_type bm25 --split test --ranker_type none --num_candidates 5

• Pubmed bash python citation_prediction/scripts/evaluate.py --dataset_type pubmed --candidate_selector_type bm25 --split test --ranker_type none --num_candidates 100

• DBLP bash python citation_prediction/scripts/evaluate.py --dataset_type dblp --candidate_selector_type bm25 --split test --ranker_type none --num_candidates 50

Train.py

The main script to train and tune hyperparameters for various models is train.py. Usage:

bash python train.py [options]

• General Parameters:

  • --mode (Required): The mode to run the train.py script in. Possible values: train or hyperopt. The train mode will train a single model and save to a given location. The hyperopt mode will run hyperparamter-optimization and return the best found model.
  • --dataset_type: Dataset to use. Possible values: dblp (default), pubmed or oc
  • --model_name: Possible values: paper_embedder (default) or citation_ranker

• Parameters specific to Hyperparameter Optimization

  • --max_evals_initial: No. of models to train in the first phase. Our hyperparameter optimization method runs in two steps. In the first step, a large number of models are run for a few epochs and the best performing max_evals_secondary models are run for more number of epochs in the second phase.
  • --max_evals_secondary: No. of models to train in the second phase. Best max_evals_secondary models from Phase 1 are trained for a longer time
  • --total_samples_initial: No. of samples to train first phase models on
  • --total_samples_secondary: No. of samples to train second phase models on
  • --models_dir_base: Base directory to store hyperopt results in
  • --n_eval: No. of validation examples to evaluate a trained model
  • --run_identifier: A string to identify the experiment
  • --version: Version string to be appended to the directory used to store model in

• Parameters specific to Training a single model

  • hyperopts_results_pkl: Path to the .pkl file generated by the hyperopt mode
  • options_json: Optional json file containing all options required to train a model

  Refer to the ModelOptions class for more options.

Training

We use the hyperopt package to tune hyperparameters. Here we describe how to run hyperopt for both the paper embedder model and the citation ranker and train the citation_prediction system.

• Hyperopt for Paper Embedder Model for DBLP

  bash python citation_prediction/scripts/train.py --mode hyperopt --dataset_type dblp --n_eval 500 --model_name paper_embedder --models_dir_base data/hyperopts/dblp/ --version 1 &> data/hyperopts/dblp/dblp.paper_embedder.hyperopt.log

  Execution Time: ~19 hours (For 25 "initial" and 5 "secondary" trial models)

• Paper Embedder Model for DBLP

  Create an empty directory: bash mkdir data/comparison/dblp/models/trained/

  bash python citation_prediction/scripts/train.py --mode train --dataset_type dblp --n_eval 500 --model_name paper_embedder --hyperopts_results_pkl data/hyperopts/dblp/citation_prediction_hyperopt_paper_embedder_dblp_2018-XX-XX_1/hyperopt_results.pickle --models_dir_base data/comparison/dblp/models/trained/ &> data/comparison/dblp/models/dblp.paper_embedder.trained.log This should produce a trained paperembedder model in the `--modelsdirbase data/comparison/dblp/models/trained/paperembedder/` directory.

  Execution time: ~5 hours

  • Evaluating the Paper Embedder for DBLP bash python citation_prediction/script/evaluate.py --dataset_type dblp --candidate_selector_type ann --split test --paper_embedder_dir data/comparison/dblp/models/trained/paper_embedder/ --num_candidates 10 --ranker_type none

• Hyperopt for Citation Ranker Model bash python citation_prediction/scripts/train.py --mode hyperopt --dataset_type dblp --models_ann_dir data/comparison/dblp/models/trained/paper_embedder/ --n_eval 500 --model_name citation_ranker --models_dir_base data/hyperopts/dblp/ --version 1 &> data/hyperopts/dblp/dblp.citation_ranker.hyperopt.log

• Citation Ranker Model for DBLP bash python citation_prediction/scripts/train.py --mode train --dataset_type dblp --hyperopts_results_pkl data/hyperopts/dblp/citation_prediction_hyperopt_citation_ranker_dblp_2018-XX-XX_1/hyperopt_results.pickle --n_eval 500 --model_name citation_ranker --models_ann_dir data/comparison/dblp/models/trained/paper_embedder/ --models_dir data/comparison/dblp/models/trained/citation_ranker/ --version 1 &> data/comparison/dblp/models/trained/dblp.citation_ranker.trained.log

Change the --dataset_type parameter to train for the other datasets. Hyperopt on the OC dataset takes a long time. We suggest setting the parameters manually or by using values described in our paper and training the two models as in steps 2 and 4 above.