minimal_prettr_usage_example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "minimal_prettr_usage_example.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3.8.10 64-bit ('matchmaker': conda)"
    },
    "interpreter": {
      "hash": "3ba72a6f94069afe0ef6de271f70e2a1476bd2023db51a918bb2eedc65fba870"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ng63tDwZSSm5"
      },
      "source": [
        "# Using Our Margin-MSE trained PreTTR 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/prettr-distilbert-split_at_3-margin_mse-T2-msmarco\n",
        "\n",
        "This instance can be used to **re-rank a candidate set**. The architecture is a 6-layer DistilBERT, split at layer 3,with an additional single linear layer at the end for scoring the cls token. \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 PreTTR checkpoint to encode passages and queries. \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 DistilBertModel,AutoTokenizer\n",
        "from transformers.models.distilbert.modeling_distilbert import *\n",
        "import math\n",
        "import torch\n",
        "from torch import nn as nn\n",
        "\n",
        "class PreTTRConfig(DistilBertConfig):\n",
        "    join_layer_idx = 3\n",
        "\n",
        "class PreTTR(DistilBertModel):\n",
        "    '''\n",
        "    PreTTR changes the distilbert model from huggingface to be able to split query and document until a set layer,\n",
        "    we skipped compression present in the original\n",
        "\n",
        "    from: Efficient Document Re-Ranking for Transformers by Precomputing Term Representations\n",
        "          MacAvaney, et al. https://arxiv.org/abs/2004.14255\n",
        "    '''\n",
        "    config_class = PreTTRConfig\n",
        "\n",
        "    def __init__(self, config):\n",
        "        super().__init__(config)\n",
        "        self.transformer = SplitTransformer(config)  # Encoder, we override the classes, but the names stay the same -> so it gets properly initialized\n",
        "        self.embeddings = PosOffsetEmbeddings(config)  # Embeddings\n",
        "        self._classification_layer = torch.nn.Linear(self.config.hidden_size, 1, bias=False)\n",
        "\n",
        "        self.join_layer_idx = config.join_layer_idx\n",
        "\n",
        "    def forward(\n",
        "            self,\n",
        "            query,\n",
        "            document,\n",
        "            use_fp16: bool = False) -> torch.Tensor:\n",
        "\n",
        "        with torch.cuda.amp.autocast(enabled=use_fp16):\n",
        "\n",
        "            query_input_ids = query[\"input_ids\"]\n",
        "            query_attention_mask = query[\"attention_mask\"]\n",
        "\n",
        "            document_input_ids = document[\"input_ids\"][:, 1:]\n",
        "            document_attention_mask = document[\"attention_mask\"][:, 1:]\n",
        "\n",
        "            query_embs = self.embeddings(query_input_ids)  # (bs, seq_length, dim)\n",
        "            document_embs = self.embeddings(document_input_ids, query_input_ids.shape[-1])  # (bs, seq_length, dim)\n",
        "\n",
        "            tfmr_output = self.transformer(\n",
        "                query_embs=query_embs,\n",
        "                query_mask=query_attention_mask,\n",
        "                doc_embs=document_embs,\n",
        "                doc_mask=document_attention_mask,\n",
        "                join_layer_idx=self.join_layer_idx\n",
        "            )\n",
        "            hidden_state = tfmr_output[0]\n",
        "\n",
        "            score = self._classification_layer(hidden_state[:, 0, :]).squeeze()\n",
        "\n",
        "            return score\n",
        "\n",
        "\n",
        "class PosOffsetEmbeddings(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.word_embeddings = nn.Embedding(config.vocab_size, config.dim, padding_idx=config.pad_token_id)\n",
        "        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.dim)\n",
        "        if config.sinusoidal_pos_embds:\n",
        "            create_sinusoidal_embeddings(\n",
        "                n_pos=config.max_position_embeddings, dim=config.dim, out=self.position_embeddings.weight\n",
        "            )\n",
        "\n",
        "        self.LayerNorm = nn.LayerNorm(config.dim, eps=1e-12)\n",
        "        self.dropout = nn.Dropout(config.dropout)\n",
        "\n",
        "    def forward(self, input_ids, pos_offset=0):\n",
        "        \"\"\"\n",
        "        Parameters\n",
        "        ----------\n",
        "        input_ids: torch.tensor(bs, max_seq_length)\n",
        "            The token ids to embed.\n",
        "\n",
        "        Outputs\n",
        "        -------\n",
        "        embeddings: torch.tensor(bs, max_seq_length, dim)\n",
        "            The embedded tokens (plus position embeddings, no token_type embeddings)\n",
        "        \"\"\"\n",
        "        seq_length = input_ids.size(1)\n",
        "        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)  # (max_seq_length)\n",
        "        position_ids = position_ids.unsqueeze(0).expand_as(input_ids) + pos_offset  # (bs, max_seq_length)\n",
        "\n",
        "        word_embeddings = self.word_embeddings(input_ids)  # (bs, max_seq_length, dim)\n",
        "        position_embeddings = self.position_embeddings(position_ids)  # (bs, max_seq_length, dim)\n",
        "\n",
        "        embeddings = word_embeddings + position_embeddings  # (bs, max_seq_length, dim)\n",
        "        embeddings = self.LayerNorm(embeddings)  # (bs, max_seq_length, dim)\n",
        "        embeddings = self.dropout(embeddings)  # (bs, max_seq_length, dim)\n",
        "        return embeddings\n",
        "\n",
        "\n",
        "class SplitTransformer(nn.Module):\n",
        "    def __init__(self, config):\n",
        "        super().__init__()\n",
        "        self.n_layers = config.n_layers\n",
        "\n",
        "        layer = TransformerBlock(config)\n",
        "        self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.n_layers)])\n",
        "\n",
        "    def forward(self, query_embs, query_mask, doc_embs, doc_mask, join_layer_idx, output_attentions=False, output_hidden_states=False):\n",
        "        \"\"\"\n",
        "        Parameters\n",
        "        ----------\n",
        "        x: torch.tensor(bs, seq_length, dim)\n",
        "            Input sequence embedded.\n",
        "        attn_mask: torch.tensor(bs, seq_length)\n",
        "            Attention mask on the sequence.\n",
        "\n",
        "        Outputs\n",
        "        -------\n",
        "        hidden_state: torch.tensor(bs, seq_length, dim)\n",
        "            Sequence of hiddens states in the last (top) layer\n",
        "        all_hidden_states: Tuple[torch.tensor(bs, seq_length, dim)]\n",
        "            Tuple of length n_layers with the hidden states from each layer.\n",
        "            Optional: only if output_hidden_states=True\n",
        "        all_attentions: Tuple[torch.tensor(bs, n_heads, seq_length, seq_length)]\n",
        "            Tuple of length n_layers with the attention weights from each layer\n",
        "            Optional: only if output_attentions=True\n",
        "        \"\"\"\n",
        "        all_hidden_states = ()\n",
        "        all_attentions = ()\n",
        "\n",
        "        #\n",
        "        # query / doc sep.\n",
        "        #\n",
        "        hidden_state_q = query_embs\n",
        "        hidden_state_d = doc_embs\n",
        "        for layer_module in self.layer[:join_layer_idx]:\n",
        "\n",
        "            layer_outputs_q = layer_module(\n",
        "                x=hidden_state_q, attn_mask=query_mask, head_mask=None, output_attentions=output_attentions\n",
        "            )\n",
        "            hidden_state_q = layer_outputs_q[-1]\n",
        "\n",
        "            layer_outputs_d = layer_module(\n",
        "                x=hidden_state_d, attn_mask=doc_mask, head_mask=None, output_attentions=output_attentions\n",
        "            )\n",
        "            hidden_state_d = layer_outputs_d[-1]\n",
        "\n",
        "        #\n",
        "        # combine\n",
        "        #\n",
        "        x = torch.cat([hidden_state_q, hidden_state_d], dim=1)\n",
        "        attn_mask = torch.cat([query_mask, doc_mask], dim=1)\n",
        "\n",
        "        #\n",
        "        # combined\n",
        "        #\n",
        "        hidden_state = x\n",
        "        for layer_module in self.layer[join_layer_idx:]:\n",
        "            layer_outputs = layer_module(\n",
        "                x=hidden_state, attn_mask=attn_mask, head_mask=None, output_attentions=output_attentions\n",
        "            )\n",
        "            hidden_state = layer_outputs[-1]\n",
        "\n",
        "        # Add last layer\n",
        "        if output_hidden_states:\n",
        "            all_hidden_states = all_hidden_states + (hidden_state,)\n",
        "\n",
        "        outputs = (hidden_state,)\n",
        "        if output_hidden_states:\n",
        "            outputs = outputs + (all_hidden_states,)\n",
        "        if output_attentions:\n",
        "            outputs = outputs + (all_attentions,)\n",
        "        return outputs  # last-layer hidden state, (all hidden states), (all attentions)\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 = PreTTR.from_pretrained(\"sebastian-hofstaetter/prettr-distilbert-split_at_3-margin_mse-T2-msmarco\")"
      ],
      "execution_count": 1,
      "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\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\")\n",
        "# a non-relevant example\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\")\n",
        "\n",
        "# the user query -> which should give us a better score for the first passage\n",
        "query_input = tokenizer(\"what is the transformers library\")\n",
        "# adding the mask augmentation, we used 8 as the fixed number for training regardless of batch-size\n",
        "# it has a somewhat (although not huge) positive impact on effectiveness, we hypothesize that might be due to the increased\n",
        "# capacity of the query encoding, not so much because of the [MASK] pre-training, but who knows :)\n",
        "query_input.input_ids += [103] * 8 # [MASK]\n",
        "query_input.attention_mask += [1] * 8\n",
        "query_input[\"input_ids\"] = torch.LongTensor(query_input.input_ids).unsqueeze(0)\n",
        "query_input[\"attention_mask\"] = torch.LongTensor(query_input.attention_mask).unsqueeze(0)\n",
        "\n",
        "#print(\"Passage 1 Tokenized:\",passage1_input)\n",
        "#print(\"Passage 2 Tokenized:\",passage2_input)\n",
        "#print(\"Query Tokenized:\",query_input)\n",
        "\n",
        "# note how we call the bert model for pairs, can be changed to: forward_representation and forward_aggregation\n",
        "score_for_p1 = model.forward(query_input,passage1_input).squeeze(0)\n",
        "score_for_p2 = model.forward(query_input,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": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "---\nScore passage 1 <-> query:  5.672976493835449\nScore passage 2 <-> query:  1.37656831741333\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_1bY5qB9b-AI"
      },
      "source": [
        "As we see the model gives the first passage a higher score than the second - yeah 🏃‍♂️\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"
      ]
    }
  ]
}