tutorials/26_Hybrid_Retrieval.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kTas9ZQ7lXP7"
   },
   "source": [
    "# Tutorial: Creating a Hybrid Retrieval Pipeline\n",
    "\n",
    "> This tutorial is based on Haystack 1.x. If you're using Haystack 2.x and would like to follow the updated version of this tutorial, check out [Creating a Hybrid Pipeline](https://haystack.deepset.ai/tutorials/33_hybrid_retrieval). \n",
    ">\n",
    "> For more information on Haystack 2.0, read the [Haystack 2.0 announcement](https://haystack.deepset.ai/blog/haystack-2-release).\n",
    "\n",
    "- **Level**: Intermediate\n",
    "- **Time to complete**: 15 minutes\n",
    "- **Nodes Used**: `EmbeddingRetriever`, `BM25Retriever`, `JoinDocuments`, `SentenceTransformersRanker` and `InMemoryDocumentStore`\n",
    "- **Goal**: After completing this tutorial, you will have learned about creating your first hybrid retrieval and when it's useful."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0hw_zoKolXQL"
   },
   "source": [
    "## Overview\n",
    "\n",
    "\n",
    "**Hybrid Retrieval** merges dense and sparse vectors together to deliver the best of both search methods. Generally speaking, dense vectors excel at understanding the context of the query, whereas sparse vectors excel at keyword matches.\n",
    "\n",
    "There are many cases when a simple sparse retrieval like BM25 performs better than a dense retrieval (for example in a specific domain like healthcare) because a dense encoder model needs to be trained on data. For more details about Hybrid Retrieval, check out [Blog Post: Hybrid Document Retrieval](https://haystack.deepset.ai/blog/hybrid-retrieval)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ITs3WTT5lXQT"
   },
   "source": [
    "## Preparing the Colab Environment\n",
    "\n",
    "- [Enable GPU Runtime in Colab](https://docs.haystack.deepset.ai/docs/enabling-gpu-acceleration#enabling-the-gpu-in-colab)\n",
    "- [Set logging level to INFO](https://docs.haystack.deepset.ai/docs/log-level)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2g9fhjxDlXQb"
   },
   "source": [
    "## Installing Haystack\n",
    "\n",
    "To start, let's install the latest release of Haystack with `pip`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L40ZxZW8lXQh"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "pip install --upgrade pip\n",
    "pip install \"datasets>=2.6.1\"\n",
    "pip install farm-haystack[inference]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CJBcPNbBlXQq"
   },
   "source": [
    "### Enabling Telemetry\n",
    "\n",
    "Knowing you're using this tutorial helps us decide where to invest our efforts to build a better product but you can always opt out by commenting the following line. See [Telemetry](https://docs.haystack.deepset.ai/docs/telemetry) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lUbTGVo4lXQv"
   },
   "outputs": [],
   "source": [
    "from haystack.telemetry import tutorial_running\n",
    "\n",
    "tutorial_running(26)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5HLBUYOplXQ1"
   },
   "source": [
    "## Creating a Hybrid Retrieval Pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "usdANiAGlXQ9"
   },
   "source": [
    "### 1) Initialize the DocumentStore and Clean Documents\n",
    "\n",
    "\n",
    "You'll start creating a hybrid pipeline by initializing a DocumentStore and preprocessing documents before storing them in the DocumentStore.\n",
    "\n",
    "You will use the PubMed Abstracts as Documents. There are a lot of datasets from PubMed on Hugging Face Hub; you will use [anakin87/medrag-pubmed-chunk](https://huggingface.co/datasets/anakin87/medrag-pubmed-chunk) in this tutorial.\n",
    "\n",
    "Initialize `InMemoryDocumentStore` and don't forget to set `use_bm25=True` and the dimension of your embeddings in `embedding_dim`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cLbh-UtelXRL"
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from haystack.document_stores import InMemoryDocumentStore\n",
    "\n",
    "dataset = load_dataset(\"anakin87/medrag-pubmed-chunk\", split=\"train\")\n",
    "\n",
    "document_store = InMemoryDocumentStore(use_bm25=True, embedding_dim=384)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WgxFjbGgdQla"
   },
   "source": [
    "You can create your document list with a simple for loop.\n",
    "The data has 3 features:\n",
    "* *pmid*\n",
    "* *title*\n",
    "* *content*: the abstract\n",
    "* *contents*: abstract + title\n",
    "\n",
    "For searching, you will use the *contents* feature. The other features will be stored as metadata, and you will use them to have a **pretty print** of the search results.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RvrG_QzirSsq"
   },
   "outputs": [],
   "source": [
    "from haystack.schema import Document\n",
    "\n",
    "docs = []\n",
    "for doc in dataset:\n",
    "    docs.append(\n",
    "        Document(content=doc[\"contents\"], meta={\"title\": doc[\"title\"], \"abstract\": doc[\"content\"], \"pmid\": doc[\"id\"]})\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tNJkztzWaWzZ"
   },
   "source": [
    "The PreProcessor class is designed to help you clean and split text into sensible units.\n",
    "\n",
    "> To learn about the preprocessing step, check out [Tutorial: Preprocessing Your Documents](https://haystack.deepset.ai/tutorials/08_preprocessing).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RrCCmLvGqhYw"
   },
   "outputs": [],
   "source": [
    "from haystack.nodes import PreProcessor\n",
    "\n",
    "preprocessor = PreProcessor(\n",
    "    clean_empty_lines=True,\n",
    "    clean_whitespace=True,\n",
    "    clean_header_footer=True,\n",
    "    split_by=\"word\",\n",
    "    split_length=512,\n",
    "    split_overlap=32,\n",
    "    split_respect_sentence_boundary=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8PzBU_jnsBTZ"
   },
   "outputs": [],
   "source": [
    "docs_to_index = preprocessor.process(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ii9x0gr9lXRT"
   },
   "source": [
    "### 2) Initialize the Retrievers\n",
    "\n",
    "Initialize a sparse retriever using [BM25](https://docs.haystack.deepset.ai/docs/retriever#bm25-recommended) and a dense retriever using a [sentence-transformers model](https://docs.haystack.deepset.ai/docs/retriever#embedding-retrieval-recommended)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rXHbHru0lXRY"
   },
   "outputs": [],
   "source": [
    "from haystack.nodes import EmbeddingRetriever, BM25Retriever\n",
    "\n",
    "sparse_retriever = BM25Retriever(document_store=document_store)\n",
    "dense_retriever = EmbeddingRetriever(\n",
    "    document_store=document_store,\n",
    "    embedding_model=\"sentence-transformers/all-MiniLM-L6-v2\",\n",
    "    use_gpu=True,\n",
    "    scale_score=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cx8307ZglXRd"
   },
   "source": [
    "### 3) Write Documents and Update Embeddings\n",
    "\n",
    "Write documents to the DocumentStore, first by deleting any remaining documents and then calling `write_documents()`. The `update_embeddings()` method uses the given retriever to create an embedding for each document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7S-QdaDYlXRg"
   },
   "outputs": [],
   "source": [
    "document_store.delete_documents()\n",
    "document_store.write_documents(docs_to_index)\n",
    "document_store.update_embeddings(retriever=dense_retriever)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_gugk_k2lXRi"
   },
   "source": [
    "### 4) Initialize JoinDocuments and Ranker\n",
    "\n",
    "While exploring hybrid search, we needed a way to combine the results of BM25 and dense vector search into a single ranked list. It may not be obvious how to combine them:\n",
    "\n",
    "* Different retrievers use incompatible score types, like BM25 and cosine similarity.\n",
    "* Documents may come from single or multiple sources at the same time. There should be a way to deal with duplicates in the final ranking.\n",
    "\n",
    "The merging and ranking of the documents from different retrievers is an open problem, however, Haystack offers several methods in [`JoinDocuments`](https://docs.haystack.deepset.ai/docs/join_documents). Here, you will use the simplest, `concatenate`, and pass the task to the ranker.\n",
    "\n",
    "Use a [re-ranker based on a cross-encoder](https://docs.haystack.deepset.ai/docs/ranker#sentencetransformersranker) that scores the relevancy of all candidates for the given search query.\n",
    "For more information about the `Ranker`, check the Haystack [docs](https://docs.haystack.deepset.ai/docs/ranker)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d_RiKspTlXRl"
   },
   "outputs": [],
   "source": [
    "from haystack.nodes import JoinDocuments, SentenceTransformersRanker\n",
    "\n",
    "join_documents = JoinDocuments(join_mode=\"concatenate\")\n",
    "rerank = SentenceTransformersRanker(model_name_or_path=\"cross-encoder/ms-marco-MiniLM-L-6-v2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PexSrsBLlXRp"
   },
   "source": [
    "### 5) Create the Hybrid Retrieval Pipeline\n",
    "\n",
    "With a Haystack `Pipeline`, you can connect your building blocks into a search pipeline. Under the hood, `Pipelines` are Directed Acyclic Graphs (DAGs) that you can easily customize for your own use cases.\n",
    "You can learn more about Pipelines in the [docs](https://docs.haystack.deepset.ai/docs/pipelines)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i0XLbnAXlXRt"
   },
   "outputs": [],
   "source": [
    "from haystack.pipelines import Pipeline\n",
    "\n",
    "pipeline = Pipeline()\n",
    "pipeline.add_node(component=sparse_retriever, name=\"SparseRetriever\", inputs=[\"Query\"])\n",
    "pipeline.add_node(component=dense_retriever, name=\"DenseRetriever\", inputs=[\"Query\"])\n",
    "pipeline.add_node(component=join_documents, name=\"JoinDocuments\", inputs=[\"SparseRetriever\", \"DenseRetriever\"])\n",
    "pipeline.add_node(component=rerank, name=\"ReRanker\", inputs=[\"JoinDocuments\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V3bsFkHuhHn4"
   },
   "source": [
    "### Generating a Pipeline Diagram\n",
    "\n",
    "With any Pipeline, whether prebuilt or custom constructed, you can save a diagram showing how all the components are connected. For example, the hybrid pipeline should look like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oCIMtwmThQG4"
   },
   "outputs": [],
   "source": [
    "# Uncomment the following to generate the images\n",
    "# !apt install libgraphviz-dev\n",
    "# !pip install pygraphviz\n",
    "\n",
    "# pipeline.draw(\"pipeline_hybrid.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sTTVLUJylXRx"
   },
   "source": [
    "## Trying Out the Hybrid Pipeline\n",
    "\n",
    "Search an article with Hybrid Retrieval. If you want to see all the steps, enable `debug=True` in `JoinDocuments`'s `params`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p-5WbeBulXR0"
   },
   "outputs": [],
   "source": [
    "prediction = pipeline.run(\n",
    "    query=\"apnea in infants\",\n",
    "    params={\n",
    "        \"SparseRetriever\": {\"top_k\": 10},\n",
    "        \"DenseRetriever\": {\"top_k\": 10},\n",
    "        \"JoinDocuments\": {\"top_k_join\": 15},  # comment for debug\n",
    "        # \"JoinDocuments\": {\"top_k_join\": 15, \"debug\":True}, #uncomment for debug\n",
    "        \"ReRanker\": {\"top_k\": 5},\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WvPv1cJ6gbBJ"
   },
   "source": [
    "Create a function to print a kind of *search page*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "raL_z_sByDoQ"
   },
   "outputs": [],
   "source": [
    "def pretty_print_results(prediction):\n",
    "    for doc in prediction[\"documents\"]:\n",
    "        print(doc.meta[\"title\"], \"\\t\", doc.score)\n",
    "        print(doc.meta[\"abstract\"])\n",
    "        print(\"\\n\", \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mSUiizGNytwX"
   },
   "outputs": [],
   "source": [
    "pretty_print_results(prediction)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "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.11.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}