Skip to content

yflyl613/Tiny-NewsRec

Repository files navigation

Tiny-NewsRec

The source code and data for our paper "Tiny-NewsRec: Effective and Efficient PLM-based News Recommendation" in EMNLP 2022.

Requirements

  • PyTorch == 1.6.0
  • TensorFlow == 1.15.0
  • horovod == 0.19.5
  • transformers == 3.0.2

Getting Started

  • Prepare Data

    You can download and unzip the public MIND dataset with the following command:

    # Under Tiny-NewsRec/
    mkdir MIND && mkdir log_all && mkdir model_all
    cd MIND
    wget https://mind201910small.blob.core.windows.net/release/MINDlarge_train.zip
    wget https://mind201910small.blob.core.windows.net/release/MINDlarge_dev.zip
    wget https://mind201910small.blob.core.windows.net/release/MINDlarge_test.zip
    unzip MINDlarge_train.zip -d MINDlarge_train
    unzip MINDlarge_dev.zip -d MINDlarge_dev
    unzip MINDlarge_test.zip -d MINDlarge_test
    cd ../
    

    Then, you should run python split_file.py under Tiny-NewsRec/ to prepare the training data. Set N in line 13 of split_file.py as the number of available GPUs. This script will construct the training samples and split them into N files for multi-GPU training.

  • Run Experiments

    Our Tiny-NewsRec method contains the following 4 steps:

    • Step 1

      Run the notebook Domain-specific_Post-train.ipynb to domain-specifically post-train the PLM-based news encoder. This will generate a checkpoint named DP_12_layer_{step}.pt every args.T steps under Tiny-NewsRec/. Then you need to set the variable ckpt_paths as the paths to the last $M$ checkpoints and run the rest cells. For each checkpoint, it will generate two .pkl files named teacher_title_emb_{idx}.pkl and teacher_body_emb_{idx}.pkl which are used for the post-training stage knowledge distillation in our method.

    • Step 2

      Run the notebook Post-train_KD.ipynb to run the post-training stage knowledge distillation in our method. Modify args.num_hidden_layers to change the number of Transformer layers in the student model. This will generate a checkpoint of the student model under Tiny-NewsRec/.

    • Step 3

      Use the script Tiny-NewsRec/PLM-NR/demo.sh to finetune the M teacher models post-trained in Step 1 with the news recommendation task. Remember to set use_pretrain_model as True and set pretrain_model_path as the path to one of these teacher models respectively.

    • Step 4

      Run bash demo.sh get_teacher_emb under Tiny-NewsRec/Tiny-NewsRec to generate the news embeddings of the $M$ teacher model finetuned in Step 3 for the finetuning stage knowledge distillation in our method. Set teacher_ckpts as the path to these teacher models (separate by space).

      Use the script Tiny-NewsRec/Tiny-NewsRec/demo.sh to run the finetuning stage knowledge distillation in our method. Modify the value of num_student_layers to change the number of Transformer layers in the student model and set bert_trainable_layers to the indexes of its last two layers (start from 0). Set use_pretrain_model as True and set pretrain_model_path as the path to the checkpoint of the student model generated in Step 2. Then you can start training with bash demo.sh train.

About

[EMNLP 2022] Official Pytorch implementation for "Tiny-NewsRec: Efficient and Effective PLM-based News Recommendation"

Topics

Resources

Stars

Watchers

Forks