langchain/docs/docs/integrations/retrievers/cohere-reranker.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fc0db1bc",
   "metadata": {},
   "source": [
    "# Cohere reranker\n",
    "\n",
    ">[Cohere](https://cohere.ai/about) is a Canadian startup that provides natural language processing models that help companies improve human-machine interactions.\n",
    "\n",
    "This notebook shows how to use [Cohere's rerank endpoint](https://docs.cohere.com/docs/reranking) in a retriever. This builds on top of ideas in the [ContextualCompressionRetriever](/docs/modules/data_connection/retrievers/contextual_compression/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f5973bb-7897-4340-a8ce-c3365ee73b2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet  cohere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b37bd138-4f3c-4d2c-bc4b-be705ce27a09",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%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": 13,
   "id": "c47b0b26-6d51-4beb-aedb-ad09740a9a2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get a new token: https://dashboard.cohere.ai/\n",
    "\n",
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"COHERE_API_KEY\"] = getpass.getpass(\"Cohere API Key:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6fa3d916",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "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": "f5057385",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Set up the base vector store retriever\n",
    "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": "code",
   "execution_count": 15,
   "id": "b7648612",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document 1:\n",
      "\n",
      "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
      "\n",
      "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 2:\n",
      "\n",
      "We cannot let this happen. \n",
      "\n",
      "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
      "\n",
      "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 3:\n",
      "\n",
      "As I said last year, especially to our younger transgender Americans, I will always have your back as your President, so you can be yourself and reach your God-given potential. \n",
      "\n",
      "While it often appears that we never agree, that isn’t true. I signed 80 bipartisan bills into law last year. From preventing government shutdowns to protecting Asian-Americans from still-too-common hate crimes to reforming military justice.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 4:\n",
      "\n",
      "I spoke with their families and told them that we are forever in debt for their sacrifice, and we will carry on their mission to restore the trust and safety every community deserves. \n",
      "\n",
      "I’ve worked on these issues a long time. \n",
      "\n",
      "I know what works: Investing in crime prevention and community police officers who’ll walk the beat, who’ll know the neighborhood, and who can restore trust and safety. \n",
      "\n",
      "So let’s not abandon our streets. Or choose between safety and equal justice.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 5:\n",
      "\n",
      "He met the Ukrainian people. \n",
      "\n",
      "From President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world. \n",
      "\n",
      "Groups of citizens blocking tanks with their bodies. Everyone from students to retirees teachers turned soldiers defending their homeland. \n",
      "\n",
      "In this struggle as President Zelenskyy said in his speech to the European Parliament “Light will win over darkness.” The Ukrainian Ambassador to the United States is here tonight.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 6:\n",
      "\n",
      "So let’s not abandon our streets. Or choose between safety and equal justice. \n",
      "\n",
      "Let’s come together to protect our communities, restore trust, and hold law enforcement accountable. \n",
      "\n",
      "That’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 7:\n",
      "\n",
      "But that trickle-down theory led to weaker economic growth, lower wages, bigger deficits, and the widest gap between those at the top and everyone else in nearly a century. \n",
      "\n",
      "Vice President Harris and I ran for office with a new economic vision for America. \n",
      "\n",
      "Invest in America. Educate Americans. Grow the workforce. Build the economy from the bottom up  \n",
      "and the middle out, not from the top down.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 8:\n",
      "\n",
      "A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n",
      "\n",
      "And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 9:\n",
      "\n",
      "The widow of Sergeant First Class Heath Robinson.  \n",
      "\n",
      "He was born a soldier. Army National Guard. Combat medic in Kosovo and Iraq. \n",
      "\n",
      "Stationed near Baghdad, just yards from burn pits the size of football fields. \n",
      "\n",
      "Heath’s widow Danielle is here with us tonight. They loved going to Ohio State football games. He loved building Legos with their daughter. \n",
      "\n",
      "But cancer from prolonged exposure to burn pits ravaged Heath’s lungs and body. \n",
      "\n",
      "Danielle says Heath was a fighter to the very end.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 10:\n",
      "\n",
      "As I’ve told Xi Jinping, it is never a good bet to bet against the American people. \n",
      "\n",
      "We’ll create good jobs for millions of Americans, modernizing roads, airports, ports, and waterways all across America. \n",
      "\n",
      "And we’ll do it all to withstand the devastating effects of the climate crisis and promote environmental justice.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 11:\n",
      "\n",
      "As Ohio Senator Sherrod Brown says, “It’s time to bury the label “Rust Belt.” \n",
      "\n",
      "It’s time. \n",
      "\n",
      "But with all the bright spots in our economy, record job growth and higher wages, too many families are struggling to keep up with the bills.  \n",
      "\n",
      "Inflation is robbing them of the gains they might otherwise feel. \n",
      "\n",
      "I get it. That’s why my top priority is getting prices under control.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 12:\n",
      "\n",
      "This was a bipartisan effort, and I want to thank the members of both parties who worked to make it happen. \n",
      "\n",
      "We’re done talking about infrastructure weeks. \n",
      "\n",
      "We’re going to have an infrastructure decade. \n",
      "\n",
      "It is going to transform America and put us on a path to win the economic competition of the 21st Century that we face with the rest of the world—particularly with China.  \n",
      "\n",
      "As I’ve told Xi Jinping, it is never a good bet to bet against the American people.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 13:\n",
      "\n",
      "He will never extinguish their love of freedom. He will never weaken the resolve of the free world. \n",
      "\n",
      "We meet tonight in an America that has lived through two of the hardest years this nation has ever faced. \n",
      "\n",
      "The pandemic has been punishing. \n",
      "\n",
      "And so many families are living paycheck to paycheck, struggling to keep up with the rising cost of food, gas, housing, and so much more. \n",
      "\n",
      "I understand.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 14:\n",
      "\n",
      "I understand. \n",
      "\n",
      "I remember when my Dad had to leave our home in Scranton, Pennsylvania to find work. I grew up in a family where if the price of food went up, you felt it. \n",
      "\n",
      "That’s why one of the first things I did as President was fight to pass the American Rescue Plan.  \n",
      "\n",
      "Because people were hurting. We needed to act, and we did. \n",
      "\n",
      "Few pieces of legislation have done more in a critical moment in our history to lift us out of crisis.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 15:\n",
      "\n",
      "My administration is providing assistance with job training and housing, and now helping lower-income veterans get VA care debt-free.  \n",
      "\n",
      "Our troops in Iraq and Afghanistan faced many dangers. \n",
      "\n",
      "One was stationed at bases and breathing in toxic smoke from “burn pits” that incinerated wastes of war—medical and hazard material, jet fuel, and more. \n",
      "\n",
      "When they came home, many of the world’s fittest and best trained warriors were never the same. \n",
      "\n",
      "Headaches. Numbness. Dizziness.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 16:\n",
      "\n",
      "Danielle says Heath was a fighter to the very end. \n",
      "\n",
      "He didn’t know how to stop fighting, and neither did she. \n",
      "\n",
      "Through her pain she found purpose to demand we do better. \n",
      "\n",
      "Tonight, Danielle—we are. \n",
      "\n",
      "The VA is pioneering new ways of linking toxic exposures to diseases, already helping more veterans get benefits. \n",
      "\n",
      "And tonight, I’m announcing we’re expanding eligibility to veterans suffering from nine respiratory cancers.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 17:\n",
      "\n",
      "Cancer is the #2 cause of death in America–second only to heart disease. \n",
      "\n",
      "Last month, I announced our plan to supercharge  \n",
      "the Cancer Moonshot that President Obama asked me to lead six years ago. \n",
      "\n",
      "Our goal is to cut the cancer death rate by at least 50% over the next 25 years, turn more cancers from death sentences into treatable diseases.  \n",
      "\n",
      "More support for patients and families. \n",
      "\n",
      "To get there, I call on Congress to fund ARPA-H, the Advanced Research Projects Agency for Health.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 18:\n",
      "\n",
      "My plan to fight inflation will lower your costs and lower the deficit. \n",
      "\n",
      "17 Nobel laureates in economics say my plan will ease long-term inflationary pressures. Top business leaders and most Americans support my plan. And here’s the plan: \n",
      "\n",
      "First – cut the cost of prescription drugs. Just look at insulin. One in ten Americans has diabetes. In Virginia, I met a 13-year-old boy named Joshua Davis.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 19:\n",
      "\n",
      "Let’s pass the Paycheck Fairness Act and paid leave.  \n",
      "\n",
      "Raise the minimum wage to $15 an hour and extend the Child Tax Credit, so no one has to raise a family in poverty. \n",
      "\n",
      "Let’s increase Pell Grants and increase our historic support of HBCUs, and invest in what Jill—our First Lady who teaches full-time—calls America’s best-kept secret: community colleges. \n",
      "\n",
      "And let’s pass the PRO Act when a majority of workers want to form a union—they shouldn’t be stopped.\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 20:\n",
      "\n",
      "Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans.  \n",
      "\n",
      "Last year COVID-19 kept us apart. This year we are finally together again. \n",
      "\n",
      "Tonight, we meet as Democrats Republicans and Independents. But most importantly as Americans. \n",
      "\n",
      "With a duty to one another to the American people to the Constitution. \n",
      "\n",
      "And with an unwavering resolve that freedom will always triumph over tyranny.\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_community.embeddings import CohereEmbeddings\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "documents = TextLoader(\"../../modules/state_of_the_union.txt\").load()\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "retriever = FAISS.from_documents(texts, CohereEmbeddings()).as_retriever(\n",
    "    search_kwargs={\"k\": 20}\n",
    ")\n",
    "\n",
    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
    "docs = retriever.invoke(query)\n",
    "pretty_print_docs(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f230c065",
   "metadata": {},
   "source": [
    "## Doing reranking with CohereRerank\n",
    "Now let's wrap our base retriever with a `ContextualCompressionRetriever`. We'll add an `CohereRerank`, uses the Cohere rerank endpoint to rerank the returned results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b83dfedb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.contextual_compression import ContextualCompressionRetriever\n",
    "from langchain_cohere import CohereRerank\n",
    "from langchain_community.llms import Cohere\n",
    "\n",
    "llm = Cohere(temperature=0)\n",
    "compressor = CohereRerank()\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor, base_retriever=retriever\n",
    ")\n",
    "\n",
    "compressed_docs = compression_retriever.invoke(\n",
    "    \"What did the president say about Ketanji Jackson Brown\"\n",
    ")\n",
    "pretty_print_docs(compressed_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70727c2f",
   "metadata": {},
   "source": [
    "You can of course use this retriever within a QA pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "367dafe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ae697ca4",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = RetrievalQA.from_chain_type(\n",
    "    llm=Cohere(temperature=0), retriever=compression_retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "46ee62fc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'query': 'What did the president say about Ketanji Brown Jackson',\n",
       " 'result': \" The president speaks highly of Ketanji Brown Jackson, stating that she is one of the nation's top legal minds, and will continue the legacy of excellence of Justice Breyer. The president also mentions that he worked with her family and that she comes from a family of public school educators and police officers. Since her nomination, she has received support from various groups, including the Fraternal Order of Police and judges from both major political parties. \\n\\nWould you like me to extract another sentence from the provided text? \"}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain({\"query\": query})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}