Turorial_Local_Citation_Recommendation.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/nianlonggu/Local-Citation-Recommendation/blob/main/Turorial_Local_Citation_Recommendation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e6Ou4o3ywuzY"
      },
      "source": [
        "# Local-Citation-Recommendation\n",
        "Code for ECIR 2022 paper [Local Citation Recommendation with Hierarchical-Attention Text Encoder and SciBERT-based Reranking](https://link.springer.com/chapter/10.1007/978-3-030-99736-6_19)\n",
        "\n",
        "# Update 07-11-2022\n",
        "1. Cleaned the code for training and testing the prefetching system, to make it easier to read and to run.\n",
        "2. Simplify the information in config file, now there is only one global configuration file for prefetching and it is more readable.\n",
        "3. Optimize the GPU usage, now the system can be trained using a single GPU.\n",
        "4. Introduced the structure of the dataset and showed how to build your custom dataset and train a citation recommendation system on that.\n",
        "5. Provided a step-by-step tutorial on google colab, illustrating the whole process of training and testing of the entire prefetching and reranking system.\n",
        "\n",
        "# Hardware Requirement\n",
        "1. OS: Ubuntu 20.04 or 18.04\n",
        "2. 1 or more GPUs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wppGf6Hjw_zM"
      },
      "source": [
        "# Install Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AK2JER0TJrI3",
        "outputId": "52ae8fe3-c4f0-4109-bde9-10bb1995a909"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'Local-Citation-Recommendation'...\n",
            "remote: Enumerating objects: 302, done.\u001b[K\n",
            "remote: Counting objects: 100% (302/302), done.\u001b[K\n",
            "remote: Compressing objects: 100% (224/224), done.\u001b[K\n",
            "remote: Total 302 (delta 123), reused 139 (delta 46), pack-reused 0\u001b[K\n",
            "Receiving objects: 100% (302/302), 20.41 MiB | 10.19 MiB/s, done.\n",
            "Resolving deltas: 100% (123/123), done.\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: cupy-cuda11x in /usr/local/lib/python3.7/dist-packages (11.0.0)\n",
            "Requirement already satisfied: fastrlock>=0.5 in /usr/local/lib/python3.7/dist-packages (from cupy-cuda11x) (0.8.1)\n",
            "Requirement already satisfied: numpy<1.26,>=1.20 in /usr/local/lib/python3.7/dist-packages (from cupy-cuda11x) (1.21.6)\n"
          ]
        }
      ],
      "source": [
        "!rm -r *\n",
        "!git clone https://github.com/nianlonggu/Local-Citation-Recommendation.git\n",
        "\n",
        "!cd Local-Citation-Recommendation; pip install -r requirements.txt\n",
        "!pip install cupy-cuda12x  ## suppose CUDA version 12.x\n",
        "import torch\n",
        "torch.__version__\n",
        "import nltk\n",
        "nltk.download('omw-1.4', quiet=True)\n",
        "nltk.download('stopwords',quiet=True)\n",
        "nltk.download('wordnet',quiet=True)\n",
        "import os\n",
        "os.chdir(\"Local-Citation-Recommendation\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "53L_-ZYpxRDx"
      },
      "source": [
        "# Download Glove Embedding \n",
        "For simplicity, we refer **MAIN** as the main folder of the repo."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8jtK74cDhBw-",
        "outputId": "7049f5f2-28f5-4000-deaa-a8135fe602bc"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1T2R1H8UstSILH_JprUaPNY0fasxD2Txr\n",
            "To: /content/Local-Citation-Recommendation/model.zip\n",
            "100% 295M/295M [00:00<00:00, 300MB/s]\n",
            "Archive:  model.zip\n",
            "   creating: model/\n",
            "   creating: model/glove/\n",
            "  inflating: model/glove/unigram_embeddings_200dim.pkl  \n",
            "  inflating: model/glove/vocabulary_200dim.pkl  \n"
          ]
        }
      ],
      "source": [
        "!gdown  https://drive.google.com/uc?id=1T2R1H8UstSILH_JprUaPNY0fasxD2Txr; unzip model.zip; rm model.zip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HlsPhzDMxp3F"
      },
      "source": [
        "# Prepare Dataset\n",
        "\n",
        "## Option 1: Build your custom dataset \n",
        "**This github repo contains a \"pseudo\" custom dataset that is actually ACL-200**\n",
        "\n",
        "The custom dataset will contain 5 components: contexts, papers, training/validation/test sets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "yzA4jnytxrSp"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "contexts = json.load(open(\"data/custom/contexts.json\"))\n",
        "papers = json.load(open(\"data/custom/papers.json\"))\n",
        "train_set = json.load(open(\"data/custom/train.json\"))\n",
        "val_set = json.load(open(\"data/custom/val.json\"))\n",
        "test_set = json.load(open(\"data/custom/test.json\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wcWkeEVqxuNf"
      },
      "source": [
        "### contexts contain the local contexts that cite a paper."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5CdRkGqYx2Wb",
        "outputId": "b715d60d-086c-4686-d737-69ff4b3258bc"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'masked_text': ' retaining all stopwords. These measures have been shown to correlate best with human judgments in general, but among the automatic measures, ROUGE-1 and ROUGE-2 also correlate best with the Pyramid ( TARGETCIT ; OTHERCIT) and Responsiveness manual metrics (OTHERCIT). Moreover, ROUGE-1 has been shown to best reflect human-automatic summary comparisons (OTHERCIT). For single concept systems, the results are s',\n",
              " 'context_id': 'P15-2138_N04-1019_0',\n",
              " 'citing_id': 'P15-2138',\n",
              " 'refid': 'N04-1019'}"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ],
      "source": [
        "for key in contexts:\n",
        "    break\n",
        "contexts[key]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UNNS3qSwxwc0"
      },
      "source": [
        "### papers contain the papers database, each paper has title and abstract.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6COKl-CUx7Nz",
        "outputId": "8f1d2687-42b6-43b1-e8d9-7cac8800818e"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'title': 'Shared task system description:\\nFrustratingly hard compositionality prediction',\n",
              " 'abstract': 'We considered a wide range of features for the DiSCo 2011 shared task about compositionality prediction for word pairs, including COALS-based endocentricity scores, compositionality scores based on distributional clusters, statistics about wordnet-induced paraphrases, hyphenation, and the likelihood of long translation equivalents in other languages. Many of the features we considered correlated significantly with human compositionality scores, but in support vector regression experiments we obtained the best results using only COALS-based endocentricity scores. Our system was nevertheless the best performing system in the shared task, and average error reductions over a simple baseline in cross-validation were 13.7% for English and 50.1% for German.'}"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ],
      "source": [
        "for key in papers:\n",
        "    break\n",
        "papers[key]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "etc8i422xzre"
      },
      "source": [
        "### train/val/test set contain the context_id (used for get the local context information and cited and citing papers information)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rciarXCSx-Tx",
        "outputId": "3900dcda-3a2b-4007-acc8-b0890b908999"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'context_id': 'P12-1066_D10-1120_1', 'positive_ids': ['D10-1120']}"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ],
      "source": [
        "train_set[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-94hgS2_yE66"
      },
      "source": [
        "positive_ids means the paper that is actually cited by the context. In this experiment the positive_ids always has one paper.\n",
        "\n",
        "You can create you own dataset with the same structure, and then train the citation recommendation system."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fjpeovp4yJHK"
      },
      "source": [
        "## Option 2: Download Processed Dataset\n",
        "You can also download our **processed dataset** (and **pretrained models**) from [Google Drive](https://drive.google.com/drive/folders/1QwQuJsBOGEESFTgl-7wWbqcig7vJNlQ2?usp=sharing)\n",
        "\n",
        "(There can be some additional information in the processed dataset other than what have been displayed in the examples above. They are irrelevant information.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V2cEBGLHyYh1"
      },
      "source": [
        "# Prefetching Part\n",
        "In the following experiment, we use the \"custom\" dataset as an example. This dataset is the same as the ACL dataset. If you have created you dataset, you need to modify the config file at\n",
        "\n",
        "MAIN/src.prefetch/config/YOUR_DATASET_NAME/global_config.cfg\n",
        "## Training\n",
        "\n",
        "This can take around 2h for this custom dataset (the same as ACL-200)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1SY5U-HVyI5r",
        "outputId": "c8a7477f-af27-4dd9-a5e6-2fa02a5be90b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  3% 133/5000 [01:11<43:44,  1.85it/s]\n",
            "Traceback (most recent call last):\n",
            "  File \"train.py\", line 193, in <module>\n",
            "    loss = train_iteration(batch)\n",
            "  File \"train.py\", line 36, in train_iteration\n",
            "    return loss.item()\n",
            "KeyboardInterrupt\n",
            "Traceback (most recent call last):\n",
            "  File \"run.py\", line 183, in <module>\n",
            "    ))), shell = True\n",
            "  File \"/usr/lib/python3.7/subprocess.py\", line 490, in run\n",
            "    stdout, stderr = process.communicate(input, timeout=timeout)\n",
            "  File \"/usr/lib/python3.7/subprocess.py\", line 956, in communicate\n",
            "    self.wait()\n",
            "  File \"/usr/lib/python3.7/subprocess.py\", line 1019, in wait\n",
            "    return self._wait(timeout=timeout)\n",
            "  File \"/usr/lib/python3.7/subprocess.py\", line 1653, in _wait\n",
            "    (pid, sts) = self._try_wait(0)\n",
            "  File \"/usr/lib/python3.7/subprocess.py\", line 1611, in _try_wait\n",
            "    (pid, sts) = os.waitpid(self.pid, wait_flags)\n",
            "KeyboardInterrupt\n"
          ]
        }
      ],
      "source": [
        "!cd src/prefetch; python run.py -config_file_path config/custom/global_config.cfg -mode train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zIvdmqGI0T26"
      },
      "source": [
        "This code will automatically handle the loop of  training -> updating paper embeddings -> updating prefetched candidates for contructing triplets -> resuming training.\n",
        "\n",
        "\n",
        "In this case, the model checkpoint will be stored in the folder \"MAIN/model/prefetch/custom/\"; <br>\n",
        "              the paper embeddings are stored in \"MAIN/embedding/prefetch/custom/\"; <br>\n",
        "              the training log files are stored in \"MAIN/log/prefetch/custom/\"; \n",
        "              \n",
        "The file MAIN/log/prefetch/custom/validate_NN.log contains the validation performance of each checkpoint during training. With this information, we can pick up the best-performance model for testing. \n",
        "\n",
        "You can specify where to store the checkpoint, log files and other parammeters by modifying the config/custom/global_config.cfg configuration file.\n",
        "\n",
        "Note: **Before testing, after determining the best checkpoint, removing all the other checkpoints. If there are multiple checkpoints in MAIN/model/prefetch/custom/, the model will use the latest checkpoint by default.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P_2XtXxM0a_d"
      },
      "source": [
        "## Testing\n",
        "To test the performance of the prefetching model, we need to first use the model checkpoint to compute the embedding for each paper in the database. This paper embedding is the index of the paper database, which is then used to perform nearest neighbor search. Then the next step is the test the prefetching performance.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BY3bcxCj0iGJ"
      },
      "source": [
        "Here I download the pretrained ACL model only for demonstration. If you are training the model using your custom data, you can wait until the training ends and select the best checkpoint based on the validation performance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "ZEs_ofWyyI7E",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "fa048fee-ea1c-417f-d35f-a48b0e8d8e8d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=13J5mtRg6t3Lcsn6fCdLJ1pZhtXnpREEq\n",
            "To: /content/Local-Citation-Recommendation/model/prefetch/custom/model_batch_35000.pt\n",
            "100% 337M/337M [00:01<00:00, 199MB/s]\n"
          ]
        }
      ],
      "source": [
        "try:\n",
        "    os.system(\"rm -r model/prefetch/custom\")\n",
        "    os.makedirs(\"model/prefetch/custom\")\n",
        "except:\n",
        "    pass\n",
        "!cd model/prefetch/custom; gdown  https://drive.google.com/uc?id=13J5mtRg6t3Lcsn6fCdLJ1pZhtXnpREEq;"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kww9s0-Z2d7a"
      },
      "source": [
        "Step 1: Compute the embeddings of all papers in the database, using the trained text encoder (HAtten)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "Bw2t6OSD2Ohy",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c2b46b84-7de9-4ef7-b730-81853852bd90"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "100% 19776/19776 [00:43<00:00, 451.54it/s]\n"
          ]
        }
      ],
      "source": [
        "!cd src/prefetch; python run.py -config_file_path config/custom/global_config.cfg -mode compute_paper_embedding"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Step 2: Test the prefetching performance on the test set"
      ],
      "metadata": {
        "id": "PDfO1ruCuLEx"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "zGywipQo2D0m",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ddbee2f1-c4e6-4b1a-8cab-5f59d29eeced"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "embedding loaded\n",
            "100% 9585/9585 [00:47<00:00, 200.26it/s]\n",
            "100% 8/8 [00:01<00:00,  4.48it/s]\n",
            "{'ckpt_name': '../../model/prefetch/custom/model_batch_35000.pt', 'recall': {10: 0.2811684924360981, 20: 0.3732916014606155, 50: 0.5021387584767867, 100: 0.6032342201356286, 200: 0.7003651538862806, 500: 0.8028169014084507, 1000: 0.8700052164840897, 2000: 0.9236306729264476}}\n",
            "Finished!\n"
          ]
        }
      ],
      "source": [
        "!cd src/prefetch; python run.py -config_file_path config/custom/global_config.cfg -mode test"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Reranking Part"
      ],
      "metadata": {
        "id": "YZOGpWivvHaJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training\n",
        "In order to train the reranker, we need to first create the training dataset. More specifically, for each query in the training set, we first use the trained HAtten prefetcher to prefetch a list of (2000) candidates. Then within the 2000 prefetched candidates we can construct triplets to train the reranker.\n"
      ],
      "metadata": {
        "id": "Lxb5LLyFvcSW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "At this stage, we should have trained the prefetching model. We need to 1) compute paper embeddings for prefetching; 2) use the prefetching model to obtain prefetched candidates for each training example; 3) use the training examples with prefetched candidates to fine-tune SciBERT reranker. "
      ],
      "metadata": {
        "id": "NgnfXYf5vioq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!cd src/prefetch; python run.py -config_file_path config/custom/global_config.cfg -mode compute_paper_embedding"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ewmk83WtvJ22",
        "outputId": "4d6c7a2b-c7cb-4fa0-cb6b-18be08fd7ad1"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "100% 19776/19776 [00:43<00:00, 451.58it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cd src/prefetch; python run.py -config_file_path config/custom/global_config.cfg -mode get_training_examples_with_prefetched_ids_for_reranking"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "26YhaS2nvJ4Y",
        "outputId": "41ce4b1c-e97e-483f-af86-bc7d720d598e"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "embedding loaded\n",
            "100% 10000/10000 [00:49<00:00, 200.75it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cd src/prefetch; python run.py -config_file_path config/custom/global_config.cfg -mode get_val_examples_with_prefetched_ids_for_reranking"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kgXC802ovJ6N",
        "outputId": "257de256-dfbe-475a-ac4b-c1d6996fa3cb"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "embedding loaded\n",
            "100% 1000/1000 [00:05<00:00, 185.86it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "After get the prefetched ids for training and validstion set, we can start training the reranker: (This can take 2.5 h to finish one epoch on this custom dataset)"
      ],
      "metadata": {
        "id": "AZBxCfc-2L-Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!cd src/rerank; python train.py -config_file_path  config/custom/scibert/training_NN_prefetch.config"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rZGfD8D92LfP",
        "outputId": "fa49f44b-5abd-4ab5-8035-565c29d5df01"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading: 100% 385/385 [00:00<00:00, 367kB/s]\n",
            "Downloading: 100% 228k/228k [00:00<00:00, 245kB/s]\n",
            "Downloading: 100% 442M/442M [00:06<00:00, 70.1MB/s]\n",
            "Some weights of the model checkpoint at allenai/scibert_scivocab_uncased were not used when initializing BertModel: ['cls.predictions.transform.dense.weight', 'cls.predictions.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.weight']\n",
            "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "/usr/local/lib/python3.7/dist-packages/transformers/optimization.py:310: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
            "  FutureWarning,\n",
            "  0% 0/10000 [00:00<?, ?it/s]/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:1960: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
            "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n",
            "  2% 216/10000 [03:07<2:21:32,  1.15it/s]\n",
            "Traceback (most recent call last):\n",
            "  File \"train.py\", line 161, in <module>\n",
            "    loss = train_iteration( batch )\n",
            "  File \"train.py\", line 37, in train_iteration\n",
            "    loss.backward()\n",
            "  File \"/usr/local/lib/python3.7/dist-packages/torch/_tensor.py\", line 396, in backward\n",
            "    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)\n",
            "  File \"/usr/local/lib/python3.7/dist-packages/torch/autograd/__init__.py\", line 175, in backward\n",
            "    allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass\n",
            "KeyboardInterrupt\n",
            "^C\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Use The HAtten Prefetcher and the SciBERT Reranker in Python Code"
      ],
      "metadata": {
        "id": "7HFXRm5p3tC2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Before run this code, we should have trained the prefetching and the reranking models and know the path to the checkpoint of the saved model. (Here I download the pretrained model on ACL-200 for demonstration. If you train using your custom data, skip this and put the trained models' checkpoint to the corresponding model folder.)"
      ],
      "metadata": {
        "id": "gsDs4C94339-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "try:\n",
        "    os.system(\"rm -r model/prefetch/custom\")\n",
        "    os.makedirs(\"model/prefetch/custom\")\n",
        "except:\n",
        "    pass\n",
        "!cd model/prefetch/custom; gdown  https://drive.google.com/uc?id=13J5mtRg6t3Lcsn6fCdLJ1pZhtXnpREEq;\n",
        "\n",
        "try:\n",
        "    os.system(\"rm -r model/rerank/custom\")\n",
        "    os.makedirs(\"model/rerank/custom/scibert/NN_prefetch\")\n",
        "except:\n",
        "    pass\n",
        "!cd model/rerank/custom/scibert/NN_prefetch; gdown  https://drive.google.com/uc?id=1DmSw6HR2W4fbUKp24K1_TREPhlLiQuEb;"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "44o7PBnP2Lho",
        "outputId": "7c57fbb1-bb6f-46ab-9988-fbbc2d828fb9"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=13J5mtRg6t3Lcsn6fCdLJ1pZhtXnpREEq\n",
            "To: /content/Local-Citation-Recommendation/model/prefetch/custom/model_batch_35000.pt\n",
            "100% 337M/337M [00:01<00:00, 222MB/s]\n",
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1DmSw6HR2W4fbUKp24K1_TREPhlLiQuEb\n",
            "To: /content/Local-Citation-Recommendation/model/rerank/custom/scibert/NN_prefetch/model_batch_91170.pt\n",
            "100% 1.32G/1.32G [00:28<00:00, 46.1MB/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load models"
      ],
      "metadata": {
        "id": "KTpuaoyf7HQ1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from citation_recommender import * \n",
        "prefetcher = Prefetcher( \n",
        "       model_path=\"model/prefetch/custom/model_batch_35000.pt\",\n",
        "       embedding_path=\"embedding/prefetch/custom/paper_embedding.pkl\", ## make sure the papers embeddings have been computed\n",
        "       gpu_list= [0,] \n",
        ")\n",
        "reranker = Reranker( model_path = \"model/rerank/custom/scibert/NN_prefetch/model_batch_91170.pt\", \n",
        "                     gpu_list = [0,] )"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x56aHUSJ7KlA",
        "outputId": "0cdd9d8b-f227-4718-beea-3b3ecd774879"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "embedding loaded\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of the model checkpoint at allenai/scibert_scivocab_uncased were not used when initializing BertModel: ['cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.predictions.decoder.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.bias']\n",
            "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Get paper recommendations"
      ],
      "metadata": {
        "id": "r-mTXfeA7be4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then we can construct a query, and use the query to find n most relevant papers. The query is a dictionary containing 3 keys: \"citing_title\",\"citing_abstract\" and \"local_context\". We can get some real example from the test set, or we can contruct a simple query as follows:\n",
        "\n",
        "(## You can specify any other values, e.g., 100, 1000 or 2000. Note that the reranking time is proportional to the number of candidates to rerank.) "
      ],
      "metadata": {
        "id": "Vj357UWI6cNg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "idx = 100\n",
        "context_info = contexts[test_set[idx][\"context_id\"]]\n",
        "citing_id = context_info[\"citing_id\"]\n",
        "refid = context_info[\"refid\"]  ## The ground-truth cited paper\n",
        "\n",
        "local_context = context_info[\"masked_text\"]\n",
        "citing_paper = papers[citing_id]\n",
        "citing_title = citing_paper[\"title\"]\n",
        "citing_abstract = citing_paper[\"abstract\"]"
      ],
      "metadata": {
        "id": "cTrU2bbF87-B"
      },
      "execution_count": 73,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "citing_title, citing_abstract, local_context, refid"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HxHPm-bC9WqL",
        "outputId": "39b95001-aeb9-483d-c8c1-3b4c21c96560"
      },
      "execution_count": 74,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "('Multiple System Combination for Transliteration',\n",
              " 'We report the results of our experiments in the context of the NEWS 2015 Shared Task on Transliteration. We focus on methods of combining multiple base systems, and leveraging transliterations from multiple languages. We show error reductions over the best base system of up to 10% when using supplemental transliterations, and up to 20% when using system combination. We also discuss the quality of the shared task datasets.',\n",
              " ' score of the last prediction in the list. Our development experiments indicated that this method of combination was more accurate than a simpler method that uses only the prediction ranks. 4.2 RERANK TARGETCIT propose a reranking approach to transliteration to leverage supplemental representations, such as phonetic transcriptions and transliterations from other languages. The reranker utilizes many features',\n",
              " 'N12-1044')"
            ]
          },
          "metadata": {},
          "execution_count": 74
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "candi_list = prefetcher.get_top_n(\n",
        "  {\n",
        "      \"citing_title\":citing_title,\n",
        "      \"citing_abstract\":citing_abstract,\n",
        "      \"local_context\":local_context\n",
        "  }, 500\n",
        ")\n",
        "print(candi_list[:10])\n",
        "\n",
        "for pos, cadi_id in enumerate(candi_list):\n",
        "    if cadi_id == refid:\n",
        "        print(\"The truely cited paper's id %s appears in position: %d among the prefetched ids.\"%( refid, pos ))\n",
        "        break\n",
        "if refid not in candi_list:\n",
        "    print(\"The truely cited paper's id %s is not included in the prefetched ids\"%( refid ))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yEtPzDpl556o",
        "outputId": "30501342-e876-4b5d-a649-d1f355094b3d"
      },
      "execution_count": 75,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['W15-3911', 'W09-3506', 'W10-2406', 'P04-1021', 'W10-2401', 'W12-4402', 'W09-3519', 'W09-3510', 'W11-3201', 'W12-4407']\n",
            "The truely cited paper's id N12-1044 appears in position: 49 among the prefetched ids.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Then we can rerank the prefetched candidates"
      ],
      "metadata": {
        "id": "T9RG8bdw7qEZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "candidate_list = [  {\"paper_id\": pid,\n",
        "                     \"title\":papers[pid].get(\"title\",\"\"),\n",
        "                     \"abstract\":papers[pid].get(\"abstract\",\"\")}\n",
        "                            for pid in candi_list ] \n",
        "# start reranking\n",
        "reranked_candidate_list = reranker.rerank( citing_title,citing_abstract,local_context, candidate_list )\n",
        "reranked_candidate_ids = [item[\"paper_id\"] for item in reranked_candidate_list]"
      ],
      "metadata": {
        "id": "sv0_ECwG7ggf"
      },
      "execution_count": 76,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for pos, cadi_id in enumerate(reranked_candidate_ids):\n",
        "    if cadi_id == refid:\n",
        "        print(\"The truely cited paper's id %s appears in position: %d among the reranked ids.\"%( refid, pos ))\n",
        "        break\n",
        "if refid not in reranked_candidate_ids:\n",
        "    print(\"The truely cited paper's id %s is not included in the reranked ids\"%( refid ))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MBsofFuA7giN",
        "outputId": "265d6ef2-9c7b-437d-93ec-5a1747917676"
      },
      "execution_count": 77,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The truely cited paper's id N12-1044 appears in position: 2 among the reranked ids.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Evaluation of the whole prefetching and reranking pipeline\n",
        "\n",
        "\n",
        "We use HAtten to prefetch 100 candidates and then rerank them and we record the Recall@10 in the final recommendations  (We test this on 100 test examples only for demonstration)"
      ],
      "metadata": {
        "id": "Mkn0tcWVBMEj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "hit_list = []\n",
        "top_K = 10\n",
        "\n",
        "for idx in tqdm(range(100)):\n",
        "\n",
        "    context_info = contexts[test_set[idx][\"context_id\"]]\n",
        "    citing_id = context_info[\"citing_id\"]\n",
        "    refid = context_info[\"refid\"]  ## The ground-truth cited paper\n",
        "\n",
        "    local_context = context_info[\"masked_text\"]\n",
        "    citing_paper = papers[citing_id]\n",
        "    citing_title = citing_paper[\"title\"]\n",
        "    citing_abstract = citing_paper[\"abstract\"]\n",
        "\n",
        "    candi_list = prefetcher.get_top_n(\n",
        "        {\n",
        "            \"citing_title\":citing_title,\n",
        "            \"citing_abstract\":citing_abstract,\n",
        "            \"local_context\":local_context\n",
        "        }, 100  ## 100 candidates \n",
        "    )\n",
        "\n",
        "    candidate_list = [  {\"paper_id\": pid,\n",
        "                     \"title\":papers[pid].get(\"title\",\"\"),\n",
        "                     \"abstract\":papers[pid].get(\"abstract\",\"\")}\n",
        "                            for pid in candi_list ] \n",
        "    # start reranking\n",
        "    reranked_candidate_list = reranker.rerank( citing_title,citing_abstract,local_context, candidate_list )\n",
        "    reranked_candidate_ids = [item[\"paper_id\"] for item in reranked_candidate_list]\n",
        "\n",
        "    hit_list.append( refid in reranked_candidate_ids[:top_K])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Qxfqf47n7gka",
        "outputId": "d8ecf2fa-1300-4d8d-f6d7-191eb9a453f1"
      },
      "execution_count": 78,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r  0%|          | 0/100 [00:00<?, ?it/s]/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:1960: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
            "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n",
            "100%|██████████| 100/100 [05:13<00:00,  3.14s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The average recall@10:0.4500\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"The average recall@%d: %.4f\"%( top_K, np.mean(hit_list)))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tbgBOFnLD0Jm",
        "outputId": "cca927cb-2432-4cf8-8fb5-06ee12cc4917"
      },
      "execution_count": 80,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The average recall@10: 0.4500\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This value is close to the results on ACL-200 in Table 4 in the paper, where we tested using full test set."
      ],
      "metadata": {
        "id": "e_gbho5MCwlM"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# References\n",
        "When using our code or models for your application, please cite the following paper:\n",
        "\n",
        "```\n",
        "@InProceedings{10.1007/978-3-030-99736-6_19,\n",
        "author=\"Gu, Nianlong\n",
        "and Gao, Yingqiang\n",
        "and Hahnloser, Richard H. R.\",\n",
        "editor=\"Hagen, Matthias\n",
        "and Verberne, Suzan\n",
        "and Macdonald, Craig\n",
        "and Seifert, Christin\n",
        "and Balog, Krisztian\n",
        "and N{\\o}rv{\\aa}g, Kjetil\n",
        "and Setty, Vinay\",\n",
        "title=\"Local Citation Recommendation with Hierarchical-Attention Text Encoder and SciBERT-Based Reranking\",\n",
        "booktitle=\"Advances in Information Retrieval\",\n",
        "year=\"2022\",\n",
        "publisher=\"Springer International Publishing\",\n",
        "address=\"Cham\",\n",
        "pages=\"274--288\",\n",
        "abstract=\"The goal of local citation recommendation is to recommend a missing reference from the local citation context and optionally also from the global context. To balance the tradeoff between speed and accuracy of citation recommendation in the context of a large-scale paper database, a viable approach is to first prefetch a limited number of relevant documents using efficient ranking methods and then to perform a fine-grained reranking using more sophisticated models. In that vein, BM25 has been found to be a tough-to-beat approach to prefetching, which is why recent work has focused mainly on the reranking step. Even so, we explore prefetching with nearest neighbor search among text embeddings constructed by a hierarchical attention network. When coupled with a SciBERT reranker fine-tuned on local citation recommendation tasks, our hierarchical Attention encoder (HAtten) achieves high prefetch recall for a given number of candidates to be reranked. Consequently, our reranker requires fewer prefetch candidates to rerank, yet still achieves state-of-the-art performance on various local citation recommendation datasets such as ACL-200, FullTextPeerRead, RefSeer, and arXiv.\",\n",
        "isbn=\"978-3-030-99736-6\"\n",
        "}\n",
        "```"
      ],
      "metadata": {
        "id": "isXIJvhUDL7e"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "e6Ou4o3ywuzY",
        "fjpeovp4yJHK"
      ],
      "provenance": [],
      "authorship_tag": "ABX9TyPjzV5P1HqKtfRQMn4HhGmj",
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}