examples/knowledge_graph_construction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Knowledge Graph Construction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we will demonstrate how to build an open-domain QA pipeline using the unique components from fastRAG. \n",
    "\n",
    "We will use a `TfidfRetriever` retriever, a `Fusion-in-Decoder` model to generate answers given the retrieved evidence, and a knowledge graph creator `KGCreator`, that will peform the following:\n",
    "\n",
    "* Extract the entities in the relevant passages.\n",
    "* Assign each valid pair of entities an relation, fitting from the context in the text. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Document Store\n",
    "\n",
    "First, we have a list of 20 passages about former president Barack Obama, fetched from a Wikipedia Elastic Index. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "metadata": {},
   "outputs": [],
   "source": [
    "from haystack.document_stores import InMemoryDocumentStore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {},
   "outputs": [],
   "source": [
    "passages = [\n",
    "   {\n",
    "      \"title\":\"Barack Obama (disambiguation)\",\n",
    "      \"text\":\"Barack Obama (disambiguation) Barack Obama (born 1964) is an American attorney and politician who served as the 44th President of the United States from 2009 to 2017.  Barack Obama may also refer to :\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama\",\n",
    "      \"text\":\"Barack Obama Barack Hussein Obama II (; born August 4, 1961) is an American attorney and politician who served as the 44th president of the United States from 2009 to 2017. A member of the Democratic Party, he was the first African American to be elected to the presidency. He previously served as a U.S. senator from Illinois from 2005 to 2008 and an Illinois state senator from 1997 to 2004. Obama was born in Honolulu, Hawaii. After graduating from Columbia University in 1983, he worked as a community organizer in Chicago. In 1988, he enrolled in Harvard Law School,\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama Sr.\",\n",
    "      \"text\":\"Barack Obama Sr. Barack Hussein Obama Sr. (; 18 June 1936 – 24 November 1982) was a Kenyan senior governmental economist and the father of Barack Obama, the 44th President of the United States. He is a central figure of his son\\\\'s memoir, \\\"Dreams from My Father\\\" (1995). Obama married in 1954 and had two children with his first wife, Kezia. He was selected for a special program to attend college in the United States and studied at the University of Hawaii. There, Obama met Stanley Ann Dunham, whom he married in 1961, and with whom he had a son,\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama in comics\",\n",
    "      \"text\":\"the face. Barack Obama is the subject of graphic novel \\\"Barack Hussein Obama\\\" by Steven Weissman. In this, President Obama and his cast of characters (Secretary Clinton, VP Joe Biden, his family) experience life in a parallel universe. Barack Obama has also appeared in Archie Comics Veronica #199, and Archie #616 and #617. President Obama was in the Flashpoint Storyline of DC comics of 2011. He discusses the earth members of The Green Lantern Corp with Amanda Waller.\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Family of Barack Obama\",\n",
    "      \"text\":\"Family of Barack Obama The family of Barack Obama, the 44th President of the United States, and his wife Michelle Obama is made up of people of Kenyan (Luo), African-American, and Old Stock American (including originally English, Scots-Irish, Welsh, German, and Swiss) ancestry. Their immediate family was the First Family of the United States from 2009 to 2017. The Obamas are the first First Family of African-American descent. Michelle LaVaughn Robinson Obama (born January 17, 1964) is an American lawyer, university administrator, and writer who served as the First Lady of the United States from 2009 to 2017. She is\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Bibliography of Barack Obama\",\n",
    "      \"text\":\"Bibliography of Barack Obama This bibliography of Barack Obama is a list of written and published works, both books and films, about Barack Obama, 44th President of the United States.\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"President Barack Obama (painting)\",\n",
    "      \"text\":\"President Barack Obama (painting) President Barack Obama is a 2018 portrait of Barack Obama by the artist Kehinde Wiley for the National Portrait Gallery. In October 2017, it was announced that Wiley had been chosen by Barack Obama to paint an official portrait of the former president to appear in Smithsonian\\\\'s National Portrait Gallery \\\"America\\\\'s Presidents\\\" exhibition. The painting depicts Obama sitting in a chair seemingly floating among foliage. The foliage is described by the author as \\\"chrysanthemums (the official flower of Chicago), jasmine (symbolic of Hawaii where the president spent most of his childhood) and African blue lilies (alluding\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama presidential campaign\",\n",
    "      \"text\":\"Barack Obama presidential campaign Barack Obama, the 44th President of the United States, has successfully run for president twice: Barack Obama presidential campaign may refer to:\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama: The Story\",\n",
    "      \"text\":\"Barack Obama: The Story Barack Obama: The Story is a book written by David Maraniss on the life of United States President Barack Obama. The biography was published on June 19, 2012.\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama Day\",\n",
    "      \"text\":\"Barack Obama Day Barack Obama Day refers to two days of recognition in the United States in honor of Barack Obama, who served as the 44th President of the United States from 2009 to 2017. The State of Illinois celebrates the commemorative holiday every August 4, which is Obama's birthday, beginning in 2018. Obama was a member of the Illinois Senate from 1997 to 2004 and represented the state in the United States Senate from 2005 to 2008 before becoming president. Similar to other commemorative holidays, it is not a legal state holiday, meaning workplaces are not closed on the\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama \\\"Joker\\\" poster\",\n",
    "      \"text\":\"Barack Obama &quot;Joker&quot; poster The Barack Obama \\\"Joker\\\" poster is a digitally manipulated image of United States President Barack Obama, designed by Firas Alkhateeb in January 2009, that was adopted by some critics of the Obama administration and described as the most famous anti-Obama image. The image portrays Obama as comic book supervillain the Joker, based on the portrayal by Heath Ledger in \\\"The Dark Knight\\\" (2008). Alkhateeb has said the image was not intended to make a political statement. He uploaded the image to the photo-sharing website Flickr, from where it was downloaded by an unknown individual who added\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama in comics\",\n",
    "      \"text\":\"who proclaims that he is endorsing him for president. The issue sold out four print runs. A month later, the comic was followed up by \\\"Presidential Material: Barack Obama\\\" by Jeff Mariotte and in November 2008 with \\\"Obama: The Comic Book\\\" by Rod Espinosa. In November 2008, two things led to an explosion in popularity of the Obama comic book character. One of Obama\\\\'s advisers gave an interview to journalist Jon Swaine of \\\"The Daily Telegraph\\\" titled, \\\"Barack Obama: The 50 facts you might not know.\\\" In the interview, it emerged that Obama collects \\\"Spider-Man and Conan the Barbarian.\\\" Then\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Family of Barack Obama\",\n",
    "      \"text\":\"Akumu Obama. She is the sole surviving full sibling of Barack Obama Sr. Sarah Onyango Obama was the third wife of Obama's paternal grandfather. She is known for short as Sarah Obama; she is sometimes referred to as Sarah Ogwel, Sarah Hussein Obama, or Sarah Anyango Obama. She lives in Nyang'oma Kogelo village, 30 miles west of western Kenya's main town, Kisumu, on the edge of Lake Victoria. (She should not be confused with her stepdaughter of the same name, Sarah Obama, a daughter of Onyango's second wife Akumu.) Although she is not a blood relation, Barack Obama calls her\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama\",\n",
    "      \"text\":\"11th consecutive year, although Dwight D. Eisenhower was selected most admired in twelve non-consecutive years. Obama was born on August 4, 1961, at Kapiolani Medical Center for Women and Children in Honolulu, Hawaii. He is the only president who was born outside of the contiguous 48 states. He was born to a white mother and a black father. His mother, Ann Dunham (1942–1995), was born in Wichita, Kansas; she was mostly of English descent, with some German, Irish (3.13%), Scottish, Swiss, and Welsh ancestry. His father, Barack Obama Sr. (1936–1982), was a Luo Kenyan from Nyang'oma Kogelo. Obama's parents met\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama\",\n",
    "      \"text\":\"since Democratic President Jimmy Carter. By contrast, the federal prison population increased significantly under presidents Ronald Reagan, George H. W. Bush, Bill Clinton, and George W. Bush. Obama left office in January 2017 with a 60% approval rating. A 2017 C-SPAN \\\"Presidential Historians Survey\\\" ranked Obama as the 12th-best US president. The Barack Obama Presidential Center is Obama\\\\'s planned presidential library. It will be hosted by the University of Chicago and located in Jackson Park on the South Side of Chicago.\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama Plaza\",\n",
    "      \"text\":\"Barack Obama Plaza Barack Obama Plaza (Moneygall services), is an off line service area, at Junction 23 of the M7 on the outskirts of the village of Moneygall in Counties Tipperary, Ireland. The plaza was opened for business in June 2014 and is accessed using the existing junction 23 slip roads. It is named after US president Barack Obama, whose 3rd great grandfather lived nearby. The ancestor was reportedly Falmouth Kearney, who emigrated to the US in 1850. The service station cost 7 million euros to construct, and opened in 2014. The Plaza is owned and operated by Supermacs Ireland\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama Academy\",\n",
    "      \"text\":\"Barack Obama Academy Barack Obama Academy is a small alternative middle school in Oakland, California. It is part of the Oakland Unified School District. It became notable as the first middle school in the United States to be officially named or renamed after US President Barack Obama in March 2009. The middle school, which opened in 2007, was formerly known as the Alternative Learning Community. The name change was prompted by the school's students. As of 2011 it had 24 students, most of whom were low income African Americans.\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Inauguration of Barack Obama\",\n",
    "      \"text\":\"Inauguration of Barack Obama Inauguration of Barack Obama may refer to:\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Barack Obama Day\",\n",
    "      \"text\":\"Bill 55, which designated August 4 as Barack Obama Day but did not make it an official state holiday. The bill passed both houses of the Illinois General Assembly with no votes against, and was signed into law by Illinois Governor Bruce Rauner on August 4, 2017. The bill amended the State Commemorative Dates Act to include a new section: Barack Obama Day. August 4th of each year is designated as Barack Obama Day, to be observed throughout the State as a day set apart to honor the 44th President of the United States of America who began his career\"\n",
    "   },\n",
    "   {\n",
    "      \"title\":\"Protests against Barack Obama\",\n",
    "      \"text\":\"Protests against Barack Obama Protests against Barack Obama occurred throughout the United States during Barack Obama's 2008 presidential campaign and during Obama's presidency. During the 2008 presidential election, particularly in the lead up to November 4, election day, numerous incidents against Obama were documented.\"\n",
    "   }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "metadata": {},
   "outputs": [],
   "source": [
    "document_store = InMemoryDocumentStore()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "metadata": {},
   "outputs": [],
   "source": [
    "from haystack.schema import Document\n",
    "\n",
    "documents = []\n",
    "for i, passage in enumerate(passages):\n",
    "    documents.append(Document(content=passage[\"text\"], meta={\"title\": passage[\"title\"]}, id=i))\n",
    "\n",
    "document_store.write_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize the pipeline components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the components we are going to use in our pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 226,
   "metadata": {},
   "outputs": [],
   "source": [
    "from haystack.nodes import TfidfRetriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = TfidfRetriever(top_k=20, document_store=document_store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/store/.local/lib/python3.8/site-packages/transformers/models/t5/tokenization_t5_fast.py:156: FutureWarning: This tokenizer was incorrectly instantiated with a model max length of 512 which will be corrected in Transformers v5.\n",
      "For now, this behavior is kept to avoid breaking backwards compatibility when padding/encoding with `truncation is True`.\n",
      "- Be aware that you SHOULD NOT rely on t5-base automatically truncating your input to 512 when padding/encoding.\n",
      "- If you want to encode/pad to sequences longer than 512 you can either instantiate this tokenizer with `model_max_length` or pass `max_length` when encoding/padding.\n",
      "- To avoid this warning, please instantiate this tokenizer with `model_max_length` set to your preferred value.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from fastrag.readers import FiDReader\n",
    "\n",
    "fid_model_path = None  ## change this to the local FID model\n",
    "assert fid_model_path is not None, \"Please change fid_model_path to the path of your trained FiD model\"\n",
    "\n",
    "reader = FiDReader(\n",
    "    input_converter_tokenizer_max_len= 256,\n",
    "    model_name_or_path=fid_model_path, \n",
    "    num_beams=1, \n",
    "    min_length=2, \n",
    "    max_length=100, \n",
    "    use_gpu=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the KGCreator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastrag.kg_creators.luke import LukeKGCreator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to /store/nltk_data...\n",
      "[nltk_data]   Package punkt is already up-to-date!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[38;5;2m✔ Download and installation successful\u001b[0m\n",
      "You can now load the package via spacy.load('en_core_web_sm')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at studio-ousia/luke-large-finetuned-tacred were not used when initializing LukeForEntityPairClassification: ['luke.embeddings.position_ids']\n",
      "- This IS expected if you are initializing LukeForEntityPairClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing LukeForEntityPairClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "kg_creator = LukeKGCreator(\n",
    "    model_name=\"studio-ousia/luke-large-finetuned-tacred\",\n",
    "    use_gpu=False,\n",
    "    batch_size=4,\n",
    "    max_length=256,\n",
    "    spacy_package=\"en_core_web_sm\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 213,
   "metadata": {},
   "outputs": [],
   "source": [
    "from haystack import Pipeline\n",
    "\n",
    "p = Pipeline()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add the components in the right order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "metadata": {},
   "outputs": [],
   "source": [
    "p.add_node(component=retriever, name=\"Retriever\", inputs=[\"Query\"])\n",
    "p.add_node(component=reader, name=\"Reader\", inputs=[\"Retriever\"])\n",
    "p.add_node(component=kg_creator, name=\"KGCreator\", inputs=[\"Reader\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run a query through the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:00<00:00, 26.36it/s]\n",
      "100%|██████████| 42/42 [00:02<00:00, 20.95it/s]\n"
     ]
    }
   ],
   "source": [
    "res = p.run(query=\"Who is Barack Obama?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display the answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 218,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'44th President'"
      ]
     },
     "execution_count": 218,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res['answers'][0].answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Show the KG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we can parse the relations, and display an interactive knowledge graph, showing all the relations in a single visualization.\n",
    "On the display, we use the pyvis package, with networkx for the graph representation itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 219,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastrag.ui.kg import edge_type_to_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "metadata": {},
   "outputs": [],
   "source": [
    "relations = res['relations']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 223,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "import pandas as pd\n",
    "from pyvis.network import Network\n",
    "\n",
    "def get_knowledge_graph_ui(relations):\n",
    "    df = pd.DataFrame(relations, columns=[\"source\", \"target\", \"edge\"]).drop_duplicates(\n",
    "        subset=[\"source\", \"target\"]\n",
    "    )\n",
    "\n",
    "    g_custom = nx.Graph()\n",
    "    value_counts = pd.value_counts(list(df[\"source\"]) + list(df[\"target\"]))\n",
    "    value_counts = value_counts[value_counts > 1]\n",
    "\n",
    "    for node_name, node_count in value_counts.iteritems():\n",
    "        g_custom.add_node(node_name, size=node_count + 5)\n",
    "\n",
    "    value_counts_df = pd.DataFrame(value_counts).reset_index()[[\"index\"]]\n",
    "    df_min = df.join(\n",
    "        value_counts_df.rename(columns={\"index\": \"source\"}).set_index(\"source\"),\n",
    "        on=\"source\",\n",
    "        how=\"inner\",\n",
    "    ).join(\n",
    "        value_counts_df.rename(columns={\"index\": \"target\"}).set_index(\"target\"),\n",
    "        on=\"target\",\n",
    "        how=\"inner\",\n",
    "    )\n",
    "\n",
    "    for _, df_row in df_min.iterrows():\n",
    "        g_custom.add_edge(\n",
    "            df_row[\"source\"], df_row[\"target\"], label=edge_type_to_label[df_row[\"edge\"]]\n",
    "        )\n",
    "\n",
    "    nt = Network(height='800px', width=\"75%\", notebook=True)\n",
    "    nt.from_nx(g_custom, show_edge_weights=True)\n",
    "    nt.repulsion(node_distance=200, spring_length=100)\n",
    "    return nt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 225,
   "metadata": {},
   "outputs": [],
   "source": [
    "nt = get_knowledge_graph_ui(relations)\n",
    "nt.show_buttons()\n",
    "nt.show(\"nx.html\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "01262dede09baa68616418263efd26d33bafc508f82c218954990f624836b45d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}