notebooks/open_mds.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Open-domain MDS Experiments\n",
    "\n",
    "This notebook organizes the analysis of the open-domain MDS experiments.\n",
    "\n",
    "__Note__: if you are running this notebook in colab, uncomment and run the following cell to install the project and its dependencies\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install \"git+https://github.com/allenai/open-mds.git\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following cells to import the required packages and load some helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from typing import Dict, List, Optional, Tuple\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import scipy.stats as stats\n",
    "import seaborn as sns\n",
    "from datasets import load_dataset\n",
    "\n",
    "from open_mds.common import util\n",
    "\n",
    "\n",
    "# Threshold under which to reject the null hypothesis\n",
    "THRESHOLD = 0.01\n",
    "\n",
    "# Controls the max number of studies to consider for MS2 and Cochrane.\n",
    "# Following https://aclanthology.org/2021.emnlp-main.594/, take the first 25 articles.\n",
    "MAX_INCLUDED_STUDIES = 25\n",
    "\n",
    "# Use the same styling for all plots & figures in the paper\n",
    "sns.set_theme(context=\"paper\", style=\"ticks\", palette=\"tab10\", font_scale=1.8)\n",
    "\n",
    "# Display all rows/columns in DataFrame\n",
    "pd.set_option(\"display.max_columns\", None)\n",
    "pd.set_option(\"display.max_rows\", None)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate Summarization in the Open-domain Setting\n",
    "\n",
    "Here we load the results from the document retrieval experiements to produce a table comparing baseline summarization performance to performance when the input document set is retrieved."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Point the variable `data_dir` to the location of a directory that contains the results of running the [`run_summarization.py`](../scripts/run_summarization.py) script for one or more models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = \"../output/results/\"\n",
    "# Make sure the directory exists and contains the expected subdirectories\n",
    "!ls $data_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {\n",
    "    \"dataset\": [],\n",
    "    \"model\": [],\n",
    "    \"retriever\": [],\n",
    "    \"top_k_strategy\": [],\n",
    "    \"metric\": [],\n",
    "    \"baseline\": [],\n",
    "    \"retrieval\": [],\n",
    "    \"difference\": [],\n",
    "    \"significant\": [],\n",
    "}\n",
    "\n",
    "for subdir in Path(data_dir).iterdir():\n",
    "    # Some datasets have blind test splits, and so we evaluate on the validation set\n",
    "    # HuggingFace assigns a different prefix to the keys in the output json, so set that here\n",
    "    metric_key_prefix = \"eval\" if subdir.name in {\"ms2\", \"cochrane\"} else \"predict\"\n",
    "\n",
    "    # The metrics we want to record results for\n",
    "    metric_columns = [\n",
    "        f\"{metric_key_prefix}_rouge_avg_fmeasure\",\n",
    "        f\"{metric_key_prefix}_bertscore_f1\",\n",
    "    ]\n",
    "\n",
    "    # Load the results as dataframes\n",
    "    baseline_df, retrieval_df = util.load_results_dicts(\n",
    "        data_dir=subdir,\n",
    "        metric_columns=metric_columns,\n",
    "        metric_key_prefix=metric_key_prefix,\n",
    "        # Only retain data that pertains to the retrieval experiments\n",
    "        load_perturbation_results=False,\n",
    "        load_training_results=False,\n",
    "    )\n",
    "\n",
    "    for model_name_or_path in retrieval_df.model_name_or_path.unique():\n",
    "        for retriever in [\"sparse\"]:\n",
    "            for top_k_strat in [\"max\", \"mean\", \"oracle\"]:\n",
    "                for metric in metric_columns:\n",
    "                    # Isolate the results from one experiment\n",
    "                    experiment_df = retrieval_df[retrieval_df.model_name_or_path == model_name_or_path][\n",
    "                        retrieval_df[f\"{metric_key_prefix}_retriever\"] == retriever\n",
    "                    ][retrieval_df[f\"{metric_key_prefix}_top_k_strategy\"] == top_k_strat]\n",
    "\n",
    "                    baseline_scores = baseline_df[baseline_df.model_name_or_path == model_name_or_path][metric]\n",
    "                    retrieval_scores = experiment_df[metric]\n",
    "                    retrieval_scores_delta = experiment_df[f\"{metric}_delta\"]\n",
    "\n",
    "                    # Sanity check that we are comparing the same number of samples\n",
    "                    assert len(baseline_scores) == len(retrieval_scores) == len(retrieval_scores_delta)\n",
    "\n",
    "                    # Report any significant differences\n",
    "                    _, pvalue = stats.ttest_rel(baseline_scores, retrieval_scores)\n",
    "\n",
    "                    # Collect the results we are interested in\n",
    "                    metric_key = metric.removeprefix(f\"{metric_key_prefix}_\")\n",
    "                    results[\"dataset\"].append(subdir.name)\n",
    "                    results[\"model\"].append(model_name_or_path)\n",
    "                    results[\"retriever\"].append(retriever)\n",
    "                    results[\"top_k_strategy\"].append(top_k_strat)\n",
    "                    results[\"metric\"].append(metric_key)\n",
    "                    results[\"baseline\"].append(round(baseline_scores.mean(), 2))\n",
    "                    results[\"retrieval\"].append(round(retrieval_scores.mean(), 2))\n",
    "                    results[\"difference\"].append(round(retrieval_scores_delta.mean(), 2))\n",
    "                    results[\"significant\"].append(pvalue < THRESHOLD)\n",
    "\n",
    "results_df = pd.DataFrame(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_df.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may wish to subset the results dataframe by dataset, retriever, and top-k strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_df[results_df.dataset == \"wcep\"][results_df.retriever == \"sparse\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Counting Retrieval Errors\n",
    "\n",
    "Here we tally the errors made by each retriever for each dataset and top-k strategy and plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _get_error_stats(\n",
    "    *, ground_truth_inputs: List[str], retrieved_inputs: List[str], doc_sep_token: Optional[str] = None\n",
    ") -> Tuple[Dict[str, int], Dict[str, List[int]]]:\n",
    "    \"\"\"Given a list of retrieved documents IDs and ground truth documents IDs, return a dictionary that\n",
    "    contains the counts of each type of error.\n",
    "    \"\"\"\n",
    "    error_stats = {\"addition\": 0, \"deletion\": 0, \"replacement\": 0}\n",
    "    retrieved_docs_ids = []\n",
    "\n",
    "    for ground_truth_docs, retrieved_docs in zip(ground_truth_inputs, retrieved_inputs):\n",
    "        # Get the individual documents\n",
    "        if doc_sep_token:\n",
    "            ground_truth_docs = util.split_docs(ground_truth_docs, doc_sep_token=doc_sep_token)\n",
    "            retrieved_docs = util.split_docs(retrieved_docs, doc_sep_token=doc_sep_token)\n",
    "\n",
    "        # Shouldn't be necessary, but strip whitespace and lowercase the strings for most robust equality checks\n",
    "        ground_truth_docs = [util.sanitize_text(doc, lowercase=True) for doc in ground_truth_docs]\n",
    "        retrieved_docs = [util.sanitize_text(doc, lowercase=True) for doc in retrieved_docs]\n",
    "\n",
    "        retrieved_docs_ids.append([doc for doc in retrieved_docs])\n",
    "\n",
    "        # Count up the number of additions and deletions\n",
    "        additions: int = sum(True for doc in retrieved_docs if doc not in ground_truth_docs)\n",
    "        deletions: int = sum(True for doc in ground_truth_docs if doc not in retrieved_docs)\n",
    "        replacements = 0\n",
    "\n",
    "        # Count all cases of 1 addition + 1 deletion as a single replacement error.\n",
    "        if additions and deletions:\n",
    "            replacements = min(additions, deletions)\n",
    "            additions -= replacements\n",
    "            deletions -= replacements\n",
    "\n",
    "        error_stats[\"addition\"] += additions\n",
    "        error_stats[\"deletion\"] += deletions\n",
    "        error_stats[\"replacement\"] += replacements\n",
    "\n",
    "    return error_stats, retrieved_docs_ids"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Collect the error statistics. Note that this can take several minutes and will cache several GBs worth of datasets to `~/.cache/huggingface/datasets`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {\"dataset\": [], \"retriever\": [], \"strategy\": [], \"error_count\": [], \"error_type\": []}\n",
    "# Also collect IDs of the retrieved docs to compare sparse and dense retrievers\n",
    "doc_ids = {}\n",
    "\n",
    "for dataset_name in [\"multinews\", \"wcep\", \"multixscience\", \"ms2\", \"cochrane\"]:\n",
    "    for retriever in [\"sparse\", \"dense\"]:\n",
    "        for strategy in [\"max\", \"mean\", \"oracle\"]:\n",
    "\n",
    "            if dataset_name not in doc_ids:\n",
    "                doc_ids[dataset_name] = {\"sparse\": {}, \"dense\": {}}\n",
    "\n",
    "            # In some cases we replaced the validation set with retrieved results\n",
    "            split = \"validation\" if dataset_name in {\"ms2\", \"cochrane\"} else \"test\"\n",
    "            retrieved_dataset = load_dataset(f\"allenai/{dataset_name}_{retriever}_{strategy}\")[split]\n",
    "\n",
    "            if dataset_name in {\"multinews\", \"wcep\"}:\n",
    "                doc_sep_token = (\n",
    "                    util.DOC_SEP_TOKENS[\"multi_news\"]\n",
    "                    if dataset_name == \"multinews\"\n",
    "                    else util.DOC_SEP_TOKENS[\"ccdv/WCEP-10\"]\n",
    "                )\n",
    "                ground_truth_dataset = load_dataset(\n",
    "                    \"multi_news\" if dataset_name == \"multinews\" else \"ccdv/WCEP-10\", split=split\n",
    "                )\n",
    "                error_stats, retrieved_docs_ids = _get_error_stats(\n",
    "                    ground_truth_inputs=ground_truth_dataset[\"document\"],\n",
    "                    retrieved_inputs=retrieved_dataset[\"document\"],\n",
    "                    doc_sep_token=doc_sep_token,\n",
    "                )\n",
    "            elif dataset_name == \"multixscience\":\n",
    "                ground_truth_dataset = load_dataset(\"multi_x_science_sum\", split=split)\n",
    "                error_stats, retrieved_docs_ids = _get_error_stats(\n",
    "                    ground_truth_inputs=[example[\"abstract\"] for example in ground_truth_dataset[\"ref_abstract\"]],\n",
    "                    retrieved_inputs=[example[\"abstract\"] for example in retrieved_dataset[\"ref_abstract\"]],\n",
    "                )\n",
    "            elif dataset_name in {\"ms2\", \"cochrane\"}:\n",
    "                ground_truth_dataset = load_dataset(\"allenai/mslr2022\", name=dataset_name, split=split)\n",
    "                # Following https://aclanthology.org/2021.emnlp-main.594/, take the first 25 articles.\n",
    "                error_stats, retrieved_docs_ids = _get_error_stats(\n",
    "                    ground_truth_inputs=[example[\"pmid\"][:_MAX_INCLUDED_STUDIES] for example in ground_truth_dataset],\n",
    "                    retrieved_inputs=[example[\"pmid\"][:_MAX_INCLUDED_STUDIES] for example in retrieved_dataset],\n",
    "                )\n",
    "            else:\n",
    "                raise ValueError(f\"Unrecognized dataset_name: {dataset_name}.\")\n",
    "\n",
    "            doc_ids[dataset_name][retriever][strategy] = retrieved_docs_ids\n",
    "\n",
    "            # Give each dataset a nicely formatted name for plotting\n",
    "            if dataset_name == \"multinews\":\n",
    "                nice_dataset_name = \"Multi-News\"\n",
    "            elif dataset_name == \"wcep\":\n",
    "                nice_dataset_name = \"WCEP-10\"\n",
    "            elif dataset_name == \"multixscience\":\n",
    "                nice_dataset_name = \"Multi-XScience\"\n",
    "            elif dataset_name == \"ms2\":\n",
    "                nice_dataset_name = \"MS2\"\n",
    "            elif dataset_name == \"cochrane\":\n",
    "                nice_dataset_name = \"Cochrane\"\n",
    "\n",
    "            # Collect the error stats for each dataset in a way amendable to plotting\n",
    "            results[\"dataset\"].extend([nice_dataset_name] * len(error_stats))\n",
    "            results[\"retriever\"].extend([retriever] * len(error_stats))\n",
    "            results[\"strategy\"].extend([strategy] * len(error_stats))\n",
    "            results[\"error_count\"].extend(\n",
    "                [error_stats[\"addition\"], error_stats[\"deletion\"], error_stats[\"replacement\"]]\n",
    "            )\n",
    "            results[\"error_type\"].extend([\"addition\", \"deletion\", \"replacement\"])\n",
    "\n",
    "results_df = pd.DataFrame(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = results_df.copy()\n",
    "# Display error counts in the 100s\n",
    "df.error_count = df.error_count / 100\n",
    "\n",
    "# Setup the grid\n",
    "g = sns.FacetGrid(\n",
    "    df,\n",
    "    row=\"retriever\",\n",
    "    col=\"dataset\",\n",
    "    sharex=True,\n",
    "    sharey=False,\n",
    "    row_order=[\"sparse\", \"dense\"],\n",
    "    col_order=[\"Multi-News\", \"WCEP-10\", \"Multi-XScience\", \"MS2\", \"Cochrane\"],\n",
    "    margin_titles=True,\n",
    ")\n",
    "\n",
    "# Plot the barplots\n",
    "_ = g.map_dataframe(\n",
    "    sns.barplot,\n",
    "    x=\"strategy\",\n",
    "    y=\"error_count\",\n",
    "    hue=\"error_type\",\n",
    "    order=[\"max\", \"mean\", \"oracle\"],\n",
    "    palette=\"tab10\",\n",
    ")\n",
    "\n",
    "# Setup a legend\n",
    "_ = g.add_legend(loc=\"lower center\", bbox_to_anchor=(0.4, 1.0), frameon=False, ncol=5, columnspacing=0.8)\n",
    "# Setup global axis titles\n",
    "_ = g.set_axis_labels(\"\", \"\")\n",
    "_ = g.fig.supylabel(\"Absolute Error Count (100s)\", x=0.025, horizontalalignment=\"left\", verticalalignment=\"center\")\n",
    "# Change the default subplot title format, see: https://wckdouglas.github.io/2016/12/seaborn_annoying_title\n",
    "_ = g.set_titles(row_template=r\"{row_name}\", col_template=r\"{col_name}\")\n",
    "# Rotate the x-axis labels\n",
    "_ = [plt.setp(ax.get_xticklabels(), rotation=45) for ax in g.axes.flat]\n",
    "\n",
    "# Save the figure\n",
    "Path(\"../output/plots\").mkdir(parents=True, exist_ok=True)\n",
    "plt.savefig(f\"../output/plots/retrieval_errors.svg\", facecolor=\"white\", bbox_inches=\"tight\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 ('open-mds-KCQg3FFx-py3.9')",
   "language": "python",
   "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"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "1de8a9ac52ca6a8e43f03f98931f021acf0adbba0952f200363bc18db767aa0e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}