examples/notebooks/ltr.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F5Ng-_HyW5LP"
      },
      "source": [
        "# Terrier Learning to Rank Examples\n",
        "\n",
        "This notebook demonstrates the use of Pyterrier in a learning-to-rank fashion.\n",
        "\n",
        "## Preparation\n",
        "\n",
        "Lets install pyterrier, as usual."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eypl7XPrkifV"
      },
      "outputs": [],
      "source": [
        "%pip install -q python-terrier"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hPK5k4g2kkKo",
        "outputId": "67632f17-c2e4-4229-b7dc-d5671ec19ea6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading terrier-assemblies 5.x-SNAPSHOT jar-with-dependencies to /Users/craigm/.pyterrier...\n",
            "Done\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "PyTerrier 0.10.0 has loaded Terrier 5.9-SNAPSHOT (built by jitpack on 2024-04-22 17:11) and terrier-helper 0.0.8\n",
            "\n",
            "No etc/terrier.properties, using terrier.default.properties for bootstrap configuration.\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import pyterrier as pt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M5BmNjqoXGow"
      },
      "source": [
        "## Load Files and Index\n",
        "\n",
        "Again, lets focus on the small Vaswani test collection. Its easily accessible via the dataset API."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "1MCH20mGB8EG"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading vaswani topics to /Users/craigm/.pyterrier/corpora/vaswani/query-text.trec\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "query-text.trec: 10.7kiB [00:00, 2.76MiB/s]                  \n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading vaswani qrels to /Users/craigm/.pyterrier/corpora/vaswani/qrels\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "qrels: 24.3kiB [00:00, 9.58MiB/s]                  \n"
          ]
        }
      ],
      "source": [
        "dataset = pt.datasets.get_dataset(\"vaswani\")\n",
        "\n",
        "indexref = dataset.get_index()\n",
        "topics = dataset.get_topics()\n",
        "qrels = dataset.get_qrels()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c8hUuA_KKPUH"
      },
      "source": [
        "## Multi-stage Retrieval\n",
        "\n",
        "In this experiment, we will be re-ranking the results obtaind from a BM25 ranking, by adding more features. Will then pass these for re-ranking by a regression technique, such as Random Forests.\n",
        "\n",
        "Conceptually, this pipeline has three stages:\n",
        "1. PL2 ranking\n",
        "2. Re-rank by each of the feaures (\"TF_IDF\" and \"PL2\")\n",
        "3. Apply the RandomForests\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "QEjmsD3ya8Pc"
      },
      "outputs": [],
      "source": [
        "#this ranker will make the candidate set of documents for each query\n",
        "BM25 = pt.terrier.Retriever(indexref, wmodel=\"BM25\")\n",
        "\n",
        "#these rankers we will use to re-rank the BM25 results\n",
        "TF_IDF =  pt.terrier.Retriever(indexref, wmodel=\"TF_IDF\")\n",
        "PL2 =  pt.terrier.Retriever(indexref, wmodel=\"PL2\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T07YF3-ULGsG"
      },
      "source": [
        "OK, so how do we combine these?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "vTLh6SrCLGM0"
      },
      "outputs": [],
      "source": [
        "pipe = BM25 >> (TF_IDF ** PL2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l7M4cUxCLMTo"
      },
      "source": [
        "Here, we are using two Pyterrer operators:\n",
        " - `>>` means \"then\", and takes the output documents of BM25 and puts them into the next stage. This means that TF_IDF and PL2 are ONLY applied on the documents that BM25 has identified.\n",
        " - `**` means feature-union - which makes each ranker into a feature in the `features` column of the results.\n",
        "\n",
        "Lets give a look at the output to see what it gives:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 112
        },
        "id": "DYNOf_TwLp0Z",
        "outputId": "66f01ca8-43da-4c89-a11e-f36a73d2ac44"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>qid</th>\n",
              "      <th>docid</th>\n",
              "      <th>docno</th>\n",
              "      <th>rank</th>\n",
              "      <th>score</th>\n",
              "      <th>query</th>\n",
              "      <th>features</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>10702</td>\n",
              "      <td>10703</td>\n",
              "      <td>0</td>\n",
              "      <td>13.472012</td>\n",
              "      <td>chemical</td>\n",
              "      <td>[7.38109017620895, 6.9992254918907575]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1</td>\n",
              "      <td>1055</td>\n",
              "      <td>1056</td>\n",
              "      <td>1</td>\n",
              "      <td>12.517082</td>\n",
              "      <td>chemical</td>\n",
              "      <td>[6.857899681644975, 6.358419229871986]</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "  qid  docid  docno  rank      score     query  \\\n",
              "0   1  10702  10703     0  13.472012  chemical   \n",
              "1   1   1055   1056     1  12.517082  chemical   \n",
              "\n",
              "                                 features  \n",
              "0  [7.38109017620895, 6.9992254918907575]  \n",
              "1  [6.857899681644975, 6.358419229871986]  "
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "pipe.search(\"chemical\").head(2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-ZMvd3qjLkrs"
      },
      "source": [
        "See, we now have a \"features\" column with numbers representing the TF_IDF and PL2 feature scores."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ye6ZpcZaMBjT"
      },
      "source": [
        "*A note about efficiency*: doing retrieval, then re-ranking the documents again can be slow. For this reason, Terrier has a FeaturesRetriever. Lets try this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 112
        },
        "id": "5gCHuDiJMNJZ",
        "outputId": "4e6ec41d-1c1b-4b7e-d318-de6ed0ef0883"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>qid</th>\n",
              "      <th>query</th>\n",
              "      <th>docid</th>\n",
              "      <th>rank</th>\n",
              "      <th>features</th>\n",
              "      <th>docno</th>\n",
              "      <th>score</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>chemical</td>\n",
              "      <td>10702</td>\n",
              "      <td>0</td>\n",
              "      <td>[7.38109017620895, 6.9992254918907575]</td>\n",
              "      <td>10703</td>\n",
              "      <td>13.472012</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1</td>\n",
              "      <td>chemical</td>\n",
              "      <td>1055</td>\n",
              "      <td>1</td>\n",
              "      <td>[6.857899681644975, 6.358419229871986]</td>\n",
              "      <td>1056</td>\n",
              "      <td>12.517082</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "  qid     query  docid  rank                                features  docno  \\\n",
              "0   1  chemical  10702     0  [7.38109017620895, 6.9992254918907575]  10703   \n",
              "1   1  chemical   1055     1  [6.857899681644975, 6.358419229871986]   1056   \n",
              "\n",
              "       score  \n",
              "0  13.472012  \n",
              "1  12.517082  "
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "fbr = pt.terrier.FeaturesRetriever(indexref, wmodel=\"BM25\", features=[\"WMODEL:TF_IDF\", \"WMODEL:PL2\"])\n",
        "#lets look at the top 2 results\n",
        "(fbr %2).search(\"chemical\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fo567qmCMZ41"
      },
      "source": [
        "However, this kind of optimisation is common in Pyterrier, so Pyterrier actually supports automatic pipeline optimisation, using the `.compile()` function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 130
        },
        "id": "jmrnqg9YMpl2",
        "outputId": "a1fd9210-5cb2-4d3f-9a4a-8045502092b4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Applying 8 rules\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>qid</th>\n",
              "      <th>docid</th>\n",
              "      <th>docno</th>\n",
              "      <th>rank</th>\n",
              "      <th>score</th>\n",
              "      <th>query</th>\n",
              "      <th>features</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>10702</td>\n",
              "      <td>10703</td>\n",
              "      <td>0</td>\n",
              "      <td>13.472012</td>\n",
              "      <td>chemical</td>\n",
              "      <td>[7.38109017620895, 6.9992254918907575]</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1</td>\n",
              "      <td>1055</td>\n",
              "      <td>1056</td>\n",
              "      <td>1</td>\n",
              "      <td>12.517082</td>\n",
              "      <td>chemical</td>\n",
              "      <td>[6.857899681644975, 6.358419229871986]</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "  qid  docid  docno  rank      score     query  \\\n",
              "0   1  10702  10703     0  13.472012  chemical   \n",
              "1   1   1055   1056     1  12.517082  chemical   \n",
              "\n",
              "                                 features  \n",
              "0  [7.38109017620895, 6.9992254918907575]  \n",
              "1  [6.857899681644975, 6.358419229871986]  "
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "pipe_fast = pipe.compile()\n",
        "(pipe_fast %2).search(\"chemical\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "siS6M5t_hugs"
      },
      "source": [
        "Finally, often we want our initial retrieval score to be a feature also. We can do this in one of two ways:\n",
        " - by adding a `SAMPLE` feature to FeaturesRetriever\n",
        " - or in the original feature-union definition, including an identity Transformer\n",
        "\n",
        "In doing so, the BM25 score (13.47 andf 12.51) are now copied in as the first position of the features column."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 112
        },
        "id": "iXxeKfPXhuPA",
        "outputId": "1fd3bf99-ec04-4572-cc6d-625d50cd1529"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>qid</th>\n",
              "      <th>query</th>\n",
              "      <th>docid</th>\n",
              "      <th>rank</th>\n",
              "      <th>features</th>\n",
              "      <th>docno</th>\n",
              "      <th>score</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>chemical</td>\n",
              "      <td>10702</td>\n",
              "      <td>0</td>\n",
              "      <td>[13.472012496423268, 7.38109017620895, 6.99922...</td>\n",
              "      <td>10703</td>\n",
              "      <td>13.472012</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1</td>\n",
              "      <td>chemical</td>\n",
              "      <td>1055</td>\n",
              "      <td>1</td>\n",
              "      <td>[12.517081895047532, 6.857899681644975, 6.3584...</td>\n",
              "      <td>1056</td>\n",
              "      <td>12.517082</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "  qid     query  docid  rank  \\\n",
              "0   1  chemical  10702     0   \n",
              "1   1  chemical   1055     1   \n",
              "\n",
              "                                            features  docno      score  \n",
              "0  [13.472012496423268, 7.38109017620895, 6.99922...  10703  13.472012  \n",
              "1  [12.517081895047532, 6.857899681644975, 6.3584...   1056  12.517082  "
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "fbr3f = pt.terrier.FeaturesRetriever(indexref, wmodel=\"BM25\", features=[\"SAMPLE\", \"WMODEL:TF_IDF\", \"WMODEL:PL2\"])\n",
        "(fbr3f %2).search(\"chemical\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>qid</th>\n",
              "      <th>docid</th>\n",
              "      <th>docno</th>\n",
              "      <th>rank</th>\n",
              "      <th>score</th>\n",
              "      <th>query</th>\n",
              "      <th>features</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>10702</td>\n",
              "      <td>10703</td>\n",
              "      <td>0</td>\n",
              "      <td>13.472012</td>\n",
              "      <td>chemical</td>\n",
              "      <td>[13.472012496423268, 7.38109017620895, 6.99922...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1</td>\n",
              "      <td>1055</td>\n",
              "      <td>1056</td>\n",
              "      <td>1</td>\n",
              "      <td>12.517082</td>\n",
              "      <td>chemical</td>\n",
              "      <td>[12.517081895047532, 6.857899681644975, 6.3584...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "  qid  docid  docno  rank      score     query  \\\n",
              "0   1  10702  10703     0  13.472012  chemical   \n",
              "1   1   1055   1056     1  12.517082  chemical   \n",
              "\n",
              "                                            features  \n",
              "0  [13.472012496423268, 7.38109017620895, 6.99922...  \n",
              "1  [12.517081895047532, 6.857899681644975, 6.3584...  "
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "pipe3f = BM25 >> (pt.Transformer.identity() ** TF_IDF ** PL2)\n",
        "(pipe3f %2).search(\"chemical\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As you can see, the results of both pipelines are identical."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R47HlFoMYAhi"
      },
      "source": [
        "# Learning models and re-ranking\n",
        "\n",
        "Ok, lets get onto the actual machine learning. We can use standard Python ML techniques. We will demonstrate a few here, including from sci-kit learn and xgBoost.\n",
        "\n",
        "In each case, the pattern is the same:\n",
        " - Create a transformer that does the re-ranking\n",
        " - Call the `fit()` method on the created object with the training topics (and validation topics as necessary)\n",
        " - Evaluate the results with the Experiment function by using the test topics\n",
        "\n",
        " Firstly, lets separate our topics into train/validation/test."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "e7r10lR3DvzM"
      },
      "outputs": [],
      "source": [
        "train_topics, valid_topics, test_topics = np.split(topics, [int(.6*len(topics)), int(.8*len(topics))])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3PYw_jasN6Vk"
      },
      "source": [
        "## sci-kit learn RandomForestRegressor\n",
        "\n",
        "Our first learning-to-rank will be done using sci-kit learn's [RandomForestRegressor](https://scikit-learn.org/stable/modules/generated/sklearn.ensemble.RandomForestRegressor.html).\n",
        "\n",
        "We use `pt.ltr.apply_learned_model()`, which returns a PyTerrier Transformer that passes the document features as \"X\" features to RandomForest. To learn the model (called fitting) the RandomForest, we invoke the `fit()` method - on the entire pipeline, specifying the queries (topics) and relevance assessment (qrels). The latter are used for the \"Y\" labels for the RandomForest fitting.\n",
        "\n",
        "NB: due to their bootstrap nature, Random Forests do not overfit, so we do not provide validation data to `fit()`.\n",
        "\n",
        "On the other hand, we could use any regression learner from sklearn, and adjust its parameters ourselves.\n",
        "\n",
        "Finally, we Experiment() on the test data to compare performances."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 112
        },
        "id": "YTI_ax4K19nl",
        "outputId": "4973c7b5-14fd-4034-b5cc-6557a5156485"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>name</th>\n",
              "      <th>map</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>PL2 Baseline</td>\n",
              "      <td>0.206031</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>LTR Baseline</td>\n",
              "      <td>0.144662</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "           name       map\n",
              "0  PL2 Baseline  0.206031\n",
              "1  LTR Baseline  0.144662"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from sklearn.ensemble import RandomForestRegressor\n",
        "\n",
        "BaselineLTR = fbr3f >> pt.ltr.apply_learned_model(RandomForestRegressor(n_estimators=400))\n",
        "BaselineLTR.fit(train_topics, qrels)\n",
        "\n",
        "results = pt.Experiment([PL2, BaselineLTR], test_topics, qrels, [\"map\"], names=[\"PL2 Baseline\", \"LTR Baseline\"])\n",
        "results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Here, the RandomForest pipeline wasnt very good. LambdaMART is normally a bit better. Lets try that next..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iGw58PCuumuT"
      },
      "source": [
        "## XgBoost Pipeline\n",
        "\n",
        "We now demonstrate the use of a LambdaMART implementation from [xgBoost](https://xgboost.readthedocs.io/en/latest/). Again, PyTerrier provides a Transformer object from `pt.ltr.apply_learned_model()`, this time passing `form='ltr'` as kwarg.\n",
        "\n",
        "This takes in the constrcutor the actual xgBoost model that you want to train. We took the xgBoost configuration from [their example code](https://github.com/dmlc/xgboost/blob/master/demo/rank/rank.py).\n",
        "\n",
        "Call the `fit()` method on the full pipeline with the training *and validation* topics.\n",
        "\n",
        "The same pipeline can also be used with [LightGBM](https://github.com/microsoft/LightGBM).\n",
        "\n",
        "Evaluate the results with the Experiment function by using the test topics."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "nM0r8EgFuGtQ"
      },
      "outputs": [],
      "source": [
        "import xgboost as xgb\n",
        "params = {'objective': 'rank:ndcg',\n",
        "          'learning_rate': 0.1,\n",
        "          'gamma': 1.0, \n",
        "          'min_child_weight': 0.1,\n",
        "          'max_depth': 6,\n",
        "          'random_state': 42\n",
        "         }\n",
        "\n",
        "BaseLTR_LM = fbr3f >> pt.ltr.apply_learned_model(xgb.sklearn.XGBRanker(**params), form='ltr')\n",
        "BaseLTR_LM.fit(train_topics, qrels, valid_topics, qrels)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HVXoNhzSP-k2"
      },
      "source": [
        "And evaluate the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 112
        },
        "id": "Dn56DKZMTQ_m",
        "outputId": "133260ca-e979-4006-9120-5339682331e0"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>name</th>\n",
              "      <th>map</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>PL2 Baseline</td>\n",
              "      <td>0.206031</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>LambdaMART</td>\n",
              "      <td>0.210969</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>"
            ],
            "text/plain": [
              "           name       map\n",
              "0  PL2 Baseline  0.206031\n",
              "1    LambdaMART  0.210969"
            ]
          },
          "execution_count": 20,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "allresultsLM = pt.Experiment([PL2, BaseLTR_LM],\n",
        "                                test_topics,\n",
        "                                qrels, [\"map\"],\n",
        "                                names=[\"PL2 Baseline\", \"LambdaMART\"])\n",
        "allresultsLM"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Excellent, event on this small dataset, adding a few more features and LambdaMART can enhance effectiveness!"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}