langchain/docs/docs/integrations/document_transformers/jina_rerank.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f6ff09ab-c736-4a18-a717-563b4e29d22d",
   "metadata": {},
   "source": [
    "# Jina Reranker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1288789a-4c30-4fc3-90c7-dd1741a2550b",
   "metadata": {},
   "source": [
    "This notebook shows how to use Jina Reranker for document compression and retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0e4d52e-3968-4f8b-9865-a886f27e5feb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -qU langchain langchain-openai langchain-community langchain-text-splitters langchainhub\n",
    "\n",
    "%pip install --upgrade --quiet  faiss\n",
    "\n",
    "# OR  (depending on Python version)\n",
    "\n",
    "%pip install --upgrade --quiet  faiss_cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1fc07a6-8e01-4aa5-8ed4-ca2b0bfca70c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function for printing docs\n",
    "\n",
    "\n",
    "def pretty_print_docs(docs):\n",
    "    print(\n",
    "        f\"\\n{'-' * 100}\\n\".join(\n",
    "            [f\"Document {i+1}:\\n\\n\" + d.page_content for i, d in enumerate(docs)]\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8ec4823-fdc1-4339-8a25-da598a1e2a4c",
   "metadata": {},
   "source": [
    "## Set up the base vector store retriever"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9db25269-e798-496f-8fb9-2bb280735118",
   "metadata": {},
   "source": [
    "Let's start by initializing a simple vector store retriever and storing the 2023 State of the Union speech (in chunks). We can set up the retriever to retrieve a high number (20) of docs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce01a2b5-d7f4-4902-9156-9a3a86704f40",
   "metadata": {},
   "source": [
    "##### Set the Jina and OpenAI API keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6692d5c5-c84a-4d42-8dd8-5ce90ff56d20",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()\n",
    "os.environ[\"JINA_API_KEY\"] = getpass.getpass()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "981159af-fa3c-4f75-adb4-1a4de1950f2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_community.embeddings import JinaEmbeddings\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "documents = TextLoader(\n",
    "    \"../../how_to/state_of_the_union.txt\",\n",
    ").load()\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "\n",
    "embedding = JinaEmbeddings(model_name=\"jina-embeddings-v2-base-en\")\n",
    "retriever = FAISS.from_documents(texts, embedding).as_retriever(search_kwargs={\"k\": 20})\n",
    "\n",
    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
    "docs = retriever.get_relevant_documents(query)\n",
    "pretty_print_docs(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5a514b7-027a-4dd4-9cfc-63fb4d50aa66",
   "metadata": {},
   "source": [
    "## Doing reranking with JinaRerank"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdd9e0ca-d728-42cb-88ad-459fb8a56b33",
   "metadata": {},
   "source": [
    "Now let's wrap our base retriever with a ContextualCompressionRetriever, using Jina Reranker as a compressor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3000019e-cc0d-4365-91d0-72247ee4d624",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain_community.document_compressors import JinaRerank\n",
    "\n",
    "compressor = JinaRerank()\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor, base_retriever=retriever\n",
    ")\n",
    "\n",
    "compressed_docs = compression_retriever.get_relevant_documents(\n",
    "    \"What did the president say about Ketanji Jackson Brown\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f314f74c-48a9-4243-8d3c-2b7f820e1e40",
   "metadata": {},
   "outputs": [],
   "source": [
    "pretty_print_docs(compressed_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87164f04-194b-4138-8d94-f179f6f34a31",
   "metadata": {},
   "source": [
    "## QA reranking with Jina Reranker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2b4ab60b-5a26-4cfb-9b58-3dc2d83b772b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m System Message \u001b[0m================================\n",
      "\n",
      "Answer any use questions based solely on the context below:\n",
      "\n",
      "<context>\n",
      "\u001b[33;1m\u001b[1;3m{context}\u001b[0m\n",
      "</context>\n",
      "\n",
      "=============================\u001b[1m Messages Placeholder \u001b[0m=============================\n",
      "\n",
      "\u001b[33;1m\u001b[1;3m{chat_history}\u001b[0m\n",
      "\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "\u001b[33;1m\u001b[1;3m{input}\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from langchain import hub\n",
    "from langchain.chains import create_retrieval_chain\n",
    "from langchain.chains.combine_documents import create_stuff_documents_chain\n",
    "\n",
    "retrieval_qa_chat_prompt = hub.pull(\"langchain-ai/retrieval-qa-chat\")\n",
    "retrieval_qa_chat_prompt.pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72af3eb3-b644-4b5f-bf5f-f1dc43c96882",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "combine_docs_chain = create_stuff_documents_chain(llm, retrieval_qa_chat_prompt)\n",
    "chain = create_retrieval_chain(compression_retriever, combine_docs_chain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "126401a7-c545-4de0-92dc-e9bc1001a6ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain.invoke({\"input\": query})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "poetry-venv-2",
   "language": "python",
   "name": "poetry-venv-2"
  },
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}