minimal_bert_dot_usage_example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "minimal_bert_dot_usage_example.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ng63tDwZSSm5"
      },
      "source": [
        "# Using Our Margin-MSE trained Bert_Dot (or BERT Dense Retrieval) Checkpoint\r\n",
        "\r\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-dot-margin_mse-T2-msmarco\r\n",
        "\r\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, without architecture additions or modifications (we only change the weights during training) - to receive a query/passage representation we pool the CLS vector. \r\n",
        "\r\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 🎉\r\n",
        "\r\n",
        "This notebook gives you a minimal usage example of downloading our Bert_Dot checkpoint to encode passages and queries to create a dot-product based score of their relevance. \r\n",
        "\r\n",
        "\r\n",
        "\r\n",
        "---\r\n",
        "\r\n",
        "\r\n",
        "Let's get started by installing the awesome *transformers* library from HuggingFace:\r\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\r\n",
        "\r\n",
        "# you can switch the model to the original \"distilbert-base-uncased\" to see that the usage example then breaks and the score ordering is reversed :O\r\n",
        "#pre_trained_model_name = \"distilbert-base-uncased\"\r\n",
        "pre_trained_model_name = \"sebastian-hofstaetter/distilbert-dot-margin_mse-T2-msmarco\"\r\n",
        "\r\n",
        "tokenizer = AutoTokenizer.from_pretrained(pre_trained_model_name) \r\n",
        "bert_model = AutoModel.from_pretrained(pre_trained_model_name)"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EOGT8YQQX1Ot"
      },
      "source": [
        "Now we are ready to use the model to encode two sample passages and a query:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Rzt9Ix9UYMLy",
        "outputId": "529e338e-b4e7-4251-cf9b-4363ac8a3ed8"
      },
      "source": [
        "# our relevant example\r\n",
        "passage1_input = tokenizer(\"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\")\r\n",
        "# a non-relevant example\r\n",
        "passage2_input = tokenizer(\"Hmm I don't like this new movie about transformers that i got from my local library. Those transformers are robots?\",return_tensors=\"pt\")\r\n",
        "# the user query -> which should give us a better score for the first passage\r\n",
        "query_input = tokenizer(\"what is the transformers library\",return_tensors=\"pt\")\r\n",
        "\r\n",
        "print(\"Passage 1 Tokenized:\",passage1_input)\r\n",
        "print(\"Passage 2 Tokenized:\",passage2_input)\r\n",
        "print(\"Query Tokenized:\",query_input)\r\n",
        "\r\n",
        "# note how we call the bert model independently between passages and query :)\r\n",
        "# [0][:,0,:] pools (or selects) the CLS vector from the full output\r\n",
        "passage1_encoded = bert_model(**passage1_input)[0][:,0,:].squeeze(0)\r\n",
        "passage2_encoded = bert_model(**passage2_input)[0][:,0,:].squeeze(0)\r\n",
        "query_encoded    = bert_model(**query_input)[0][:,0,:].squeeze(0)\r\n",
        "\r\n",
        "print(\"---\")\r\n",
        "print(\"Passage Encoded Shape:\",passage1_encoded.shape)\r\n",
        "print(\"Query Encoded Shape:\",query_encoded.shape)"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Passage 1 Tokenized: {'input_ids': tensor([[  101,  2057,  2024,  2200,  3407,  2000,  2265,  2017,  1996,   100,\n",
            "         19081,  3075,  2005,  3653,  1011,  4738,  2653,  4275,  1012,  2057,\n",
            "          2024,  5094,  1996,  2451,  2147,  2362,  2875,  1996,  3125,  1997,\n",
            "         10787, 17953,  2361,   100,  1012,   102]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
            "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}\n",
            "Passage 2 Tokenized: {'input_ids': tensor([[  101, 17012,  1045,  2123,  1005,  1056,  2066,  2023,  2047,  3185,\n",
            "          2055, 19081,  2008,  1045,  2288,  2013,  2026,  2334,  3075,  1012,\n",
            "          2216, 19081,  2024, 13507,  1029,   102]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
            "         1, 1]])}\n",
            "Query Tokenized: {'input_ids': tensor([[  101,  2054,  2003,  1996, 19081,  3075,   102]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])}\n",
            "---\n",
            "Passage Encoded Shape: torch.Size([768])\n",
            "Query Encoded Shape: torch.Size([768])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V_96RCg7Y1cP"
      },
      "source": [
        "Now that we have our encoded vectors, we can generate the score with a simple dot product! \r\n",
        "\r\n",
        "(This can be offloaded to a vector indexing library like Faiss)\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RzDL1qKDalIR",
        "outputId": "ee6271b2-0da4-4717-8c7e-2730f84475fe"
      },
      "source": [
        "score_for_p1 = query_encoded.dot(passage1_encoded)\r\n",
        "print(\"Score passage 1 <-> query: \",float(score_for_p1))\r\n",
        "\r\n",
        "score_for_p2 = query_encoded.dot(passage2_encoded)\r\n",
        "print(\"Score passage 2 <-> query: \",float(score_for_p2))"
      ],
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Score passage 1 <-> query:  108.82856750488281\n",
            "Score passage 2 <-> query:  99.5865249633789\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "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 collection or candidate set). The scores are in the 100+ range (as we create a dot-product of 768 dimensional vectors, which naturally gives a larger score)\n",
        "\n",
        "*As a fun exercise you can swap the pre-trained model to the initial distilbert checkpoint and see that the example doesn't work anymore*\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"
      ]
    }
  ]
}