minimal_bert_cat_usage_example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "minimal_bert_cat_usage_example.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3.7.1 64-bit ('deep_learning': conda)",
      "metadata": {
        "interpreter": {
          "hash": "c850e6f77ff7c9cbece5364f8526ec42dd183cf59251b1cfd7b71a0467b242c1"
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ng63tDwZSSm5"
      },
      "source": [
        "# Using Our Margin-MSE trained BERT_Cat Checkpoint\n",
        "\n",
        "We provide a fully retrieval trained (with Margin-MSE using a 3 teacher Bert_Cat Ensemble on MSMARCO-Passage) DistilBert-based instance on the HuggingFace model hub here: https://huggingface.co/sebastian-hofstaetter/distilbert-cat-margin_mse-T2-msmarco\n",
        "\n",
        "This instance can be used to **re-rank a candidate set** or **directly for a vector index based dense retrieval**. The architecure is a 6-layer DistilBERT, with an additional single linear layer at the end. \n",
        "\n",
        "If you want to know more about our simple, yet effective knowledge distillation method for efficient information retrieval models for a variety of student architectures, check out our paper: https://arxiv.org/abs/2010.02666 🎉\n",
        "\n",
        "This notebook gives you a minimal usage example of downloading our BERT_Cat checkpoint to encode concatenated passages and queries to create a score of their relevance. \n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "Let's get started by installing the awesome *transformers* library from HuggingFace:\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r2WyNOE2R2rW"
      },
      "source": [
        "pip install transformers"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YqkWDa_jWu7c"
      },
      "source": [
        "The next step is to download our checkpoint and initialize the tokenizer and models:\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oTYEtziISSDl"
      },
      "source": [
        "from transformers import AutoTokenizer,AutoModel, PreTrainedModel,PretrainedConfig\n",
        "from typing import Dict\n",
        "import torch\n",
        "\n",
        "class BERT_Cat_Config(PretrainedConfig):\n",
        "    model_type = \"BERT_Cat\"\n",
        "    bert_model: str\n",
        "    trainable: bool = True\n",
        "\n",
        "class BERT_Cat(PreTrainedModel):\n",
        "    \"\"\"\n",
        "    The vanilla/mono BERT concatenated (we lovingly refer to as BERT_Cat) architecture \n",
        "    -> requires input concatenation before model, so that batched input is possible\n",
        "    \"\"\"\n",
        "    config_class = BERT_Cat_Config\n",
        "    base_model_prefix = \"bert_model\"\n",
        "\n",
        "    def __init__(self,\n",
        "                 cfg) -> None:\n",
        "        super().__init__(cfg)\n",
        "        \n",
        "        self.bert_model = AutoModel.from_pretrained(cfg.bert_model)\n",
        "\n",
        "        for p in self.bert_model.parameters():\n",
        "            p.requires_grad = cfg.trainable\n",
        "\n",
        "        self._classification_layer = torch.nn.Linear(self.bert_model.config.hidden_size, 1)\n",
        "\n",
        "    def forward(self,\n",
        "                query_n_doc_sequence):\n",
        "\n",
        "        vecs = self.bert_model(**query_n_doc_sequence)[0][:,0,:] # assuming a distilbert model here\n",
        "        score = self._classification_layer(vecs)\n",
        "        return score\n",
        "\n",
        "#\n",
        "# init the model & tokenizer (using the distilbert tokenizer)\n",
        "#\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased\") # honestly not sure if that is the best way to go, but it works :)\n",
        "model = BERT_Cat.from_pretrained(\"sebastian-hofstaetter/distilbert-cat-margin_mse-T2-msmarco\")"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EOGT8YQQX1Ot"
      },
      "source": [
        "Now we are ready to use the model to encode two sample passage and query pairs:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Rzt9Ix9UYMLy",
        "outputId": "529e338e-b4e7-4251-cf9b-4363ac8a3ed8"
      },
      "source": [
        "# our relevant example (with the query)\n",
        "passage1_input = tokenizer(\"what is the transformers library\",\"We are very happy to show you the 🤗 Transformers library for pre-trained language models. We are helping the community work together towards the goal of advancing NLP 🔥.\",return_tensors=\"pt\")\n",
        "# a non-relevant example (with the query)\n",
        "passage2_input = tokenizer(\"what is the transformers library\",\"Hmm I don't like this new movie about transformers that i got from my local library. Those transformers are robots?\",return_tensors=\"pt\")\n",
        "\n",
        "#print(\"Passage 1 Tokenized:\",passage1_input)\n",
        "#print(\"Passage 2 Tokenized:\",passage2_input)\n",
        "\n",
        "# note how we call the bert model for pairs, can not be changed (look for colbert or bert_dot for independent forward calls)\n",
        "score_for_p1 = model.forward(passage1_input).squeeze(0)\n",
        "score_for_p2 = model.forward(passage2_input).squeeze(0)\n",
        "\n",
        "print(\"---\")\n",
        "print(\"Score passage 1 <-> query: \",float(score_for_p1))\n",
        "print(\"Score passage 2 <-> query: \",float(score_for_p2))"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "---\n",
            "Score passage 1 <-> query:  5.899686336517334\n",
            "Score passage 2 <-> query:  2.2803378105163574\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_1bY5qB9b-AI"
      },
      "source": [
        "As we see the model gives the first passage a higher score than the second - these scores would now be used to generate a list (if we run this comparison on all passages in our candidate set).\n",
        "\n",
        "- If you want to look at more complex usages and training code we have a library for that: https://github.com/sebastian-hofstaetter/transformer-kernel-ranking 👏\n",
        "\n",
        "- If you use our model checkpoint please cite our work as:\n",
        "\n",
        "    ```\n",
        "@misc{hofstaetter2020_crossarchitecture_kd,\n",
        "      title={Improving Efficient Neural Ranking Models with Cross-Architecture Knowledge Distillation}, \n",
        "      author={Sebastian Hofst{\\\"a}tter and Sophia Althammer and Michael Schr{\\\"o}der and Mete Sertkan and Allan Hanbury},\n",
        "      year={2020},\n",
        "      eprint={2010.02666},\n",
        "      archivePrefix={arXiv},\n",
        "      primaryClass={cs.IR}\n",
        "}\n",
        "    ```\n",
        "\n",
        "Thank You 😊 If you have any questions feel free to reach out to Sebastian via mail (email in the paper). \n"
      ]
    }
  ]
}