02_kg_construction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "01e3ee7f-4236-421c-b382-6215c8fcd3e5",
   "metadata": {},
   "source": [
    "## Knowledge Graph (KG) Construction from Unstructured Text\n",
    "\n",
    "> Slides link: https://docs.google.com/presentation/d/1ta1Gw004GIBky70Bm1-fYULBIKBBb02hKFg1l0k81vY/edit#slide=id.p\n",
    "\n",
    "**Useful Related Links:**\n",
    "- https://www.boundaryml.com/blog/structured-output-from-llms\n",
    "- https://huggingface.co/spaces/urchade/gliner_multiv2.1\n",
    "- https://huggingface.co/urchade/gliner_large-v2\n",
    "- https://github.com/hitz-zentroa/GoLLIE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "680f6d25-4188-42bc-a0b0-c1bfed6508a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install gliner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b4bfeefc-d3fa-400f-9258-3c7f41986beb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/mehdiallahyari/environments/llama_index_env/lib/python3.11/site-packages/transformers/convert_slow_tokenizer.py:560: UserWarning: The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option which is not implemented in the fast tokenizers. In practice this means that the fast version of the tokenizer can produce unknown tokens whereas the sentencepiece version would have converted these unknown tokens into a sequence of byte tokens matching the original piece of text.\n",
      "  warnings.warn(\n",
      "/Users/mehdiallahyari/environments/llama_index_env/lib/python3.11/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from gliner import GLiNER\n",
    "\n",
    "def merge_entities(entities):\n",
    "    if not entities:\n",
    "        return []\n",
    "    merged = []\n",
    "    current = entities[0]\n",
    "    for next_entity in entities[1:]:\n",
    "        if next_entity['label'] == current['label'] and (next_entity['start'] == current['end'] + 1 or next_entity['start'] == current['end']):\n",
    "            current['text'] = text[current['start']: next_entity['end']].strip()\n",
    "            current['end'] = next_entity['end']\n",
    "        else:\n",
    "            merged.append(current)\n",
    "            current = next_entity\n",
    "    # Append the last entity\n",
    "    merged.append(current)\n",
    "    return merged\n",
    "\n",
    "\n",
    "# model = GLiNER.from_pretrained(\"numind/NuNerZero\")\n",
    "model = GLiNER.from_pretrained(\"numind/NuZero_token\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f13628b2-d993-47ca-97ac-afbe36b0504f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fiat => organization\n",
      "Chrysler => company\n",
      "U.S. => location\n",
      "Italian => location\n",
      "Europe => location\n",
      "January 1 => date\n",
      "Chrysler => company\n",
      "Chief Executive => position\n",
      "Sergio Marchionne => person\n"
     ]
    }
   ],
   "source": [
    "# NuZero requires labels to be lower-cased!\n",
    "labels = [\"location\",\"date\",\"person\",\"event\", \"company\", \"organization\", \"position\"]\n",
    "labels = [l.lower() for l in labels]\n",
    "\n",
    "text = \"\"\"Fiat has completed its buyout of Chrysler, making the U.S. business a wholly-owned subsidiary of the Italian\n",
    "carmaker as it gears up to use their combined resources to turn around its loss-making operations in\n",
    "Europe. The company announced on January 1 that it had struck a $4.35 billion deal - cheaper than analysts\n",
    "had expected - to gain full control of Chrysler, ending more than a year of tense talks that had obstructed Chief Executive Sergio Marchionne's efforts to create the\n",
    "world's seventh-largest auto maker.\"\"\"\n",
    "\n",
    "entities = model.predict_entities(text, labels, threshold=0.4)\n",
    "\n",
    "entities = merge_entities(entities)\n",
    "\n",
    "for entity in entities:\n",
    "    print(entity[\"text\"], \"=>\", entity[\"label\"])\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "id": "2cfd02a7-25fb-4ebc-88e7-8be746d9b9ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install wikipedia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8b9383cf-b300-4de1-8cc1-a591d063cec7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wikipedia\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e09b77b9-35c2-41c3-862e-7fb3f2269b39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Thomas Jeffrey Hanks (born July 9, 1956) is an American actor and filmmaker. Known for both his comedic and dramatic roles, he is one of the most popular and recognizable film stars worldwide, and is regarded as an American cultural icon. Hanks's films have grossed more than $4.9 billion in North America and more than $9.96 billion worldwide, making him the fourth-highest-grossing actor in North America. He has received numerous honors including the AFI Life Achievement Award in 2002, the Kennedy Center Honor in 2014, the Presidential Medal of Freedom and the French Legion of Honor both in 2016, as well as the Golden Globe Cecil B. DeMille Award in 2020.\\nHanks made his breakthrough with leading roles in a series of comedy films that received positive media attention, such as Splash (1984), The Money Pit (1986), Big (1988) and A League of Their Own (1992). He won two consecutive Academy Awards for Best Actor for starring as a gay lawyer suffering from AIDS in Philadelphia (1993) and the\""
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "page = wikipedia.page(title=\"Tom Hanks\", auto_suggest=False)\n",
    "page.content[:1000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b776a48a-4770-4b25-99aa-93d2e56cbc20",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_text_splitters import RecursiveCharacterTextSplitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "55b2fcbb-db9f-499a-b121-dea10c4e80e3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "97"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=200,\n",
    "    chunk_overlap=20,\n",
    "    separators=[\"\\n\\n\", \"\\n\"]\n",
    ")\n",
    "\n",
    "chunks = text_splitter.split_text(page.content)\n",
    "len(chunks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "52a76790-c035-457e-8471-d1776e93662e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "662"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(chunks[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "06d9ff94-ad5e-45e2-8311-6c77031df428",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = [\"award\", \"location\", \"organization\", \"person\", \"movie\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "82f9e3b2-3f16-4122-87c2-f671e3362a95",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 97/97 [00:31<00:00,  3.11it/s]\n"
     ]
    }
   ],
   "source": [
    "chunks_entities = []\n",
    "entity_list = []\n",
    "duplicates = set()\n",
    "for text in tqdm(chunks):\n",
    "    entities = model.predict_entities(text, labels, threshold=0.7)\n",
    "    entities = merge_entities(entities)\n",
    "    chunk_entities = set()\n",
    "    for entity in entities:\n",
    "        # print(entity[\"text\"], \"=>\", entity[\"label\"])\n",
    "        chunk_entities.add(entity[\"text\"])\n",
    "        if entity[\"text\"] in duplicates:\n",
    "            continue\n",
    "        duplicates.add(entity[\"text\"])\n",
    "        entity_list.append((entity[\"text\"], \"=>\", entity[\"label\"]))\n",
    "\n",
    "    chunks_entities.append(list(chunk_entities))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b5b19ef9-4d40-4bb0-a2ee-7a578c32d723",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['Thomas Jeffrey Hanks', 'North America', 'July', ','], ['Ron', 'Hanks']]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chunks_entities[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9bf68b8e-426a-46ba-8f85-09ece82da015",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nHaving grown up in the Bay Area, Hanks says that some of his first movie memories were seeing movies in the Alameda Theatre in Alameda, California. Hanks studied theater at Chabot College in Hayward, California, and transferred to California State University, Sacramento after two years. During a 2001 interview with sportscaster Bob Costas, Hanks was asked whether he would rather have an Oscar or a Heisman Trophy. He replied that he would rather win a Heisman by playing halfback for the California Golden Bears. He told New York magazine in 1986, \"Acting classes looked like the best place for a guy who liked to make a lot of noise and be rather flamboyant. I spent a lot of time going to plays. I wouldn\\'t take dates with me. I\\'d just drive to a theater, buy myself a ticket, sit in the seat and read the program, and then get into the play completely. I spent a lot of time like that, seeing Brecht, Tennessee Williams, Ibsen, and all that.\"'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chunks[9]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "7fa1f470-ea7f-40c4-9bd1-ed5b12576bfe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('Thomas Jeffrey Hanks', '=>', 'person'),\n",
       " ('July', '=>', 'date'),\n",
       " (',', '=>', 'date'),\n",
       " ('North America', '=>', 'location')]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "entity_list[:4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "c6633498-3fa7-4073-a55c-b1336709210b",
   "metadata": {},
   "outputs": [],
   "source": [
    "locs = []\n",
    "orgs = []\n",
    "persons = []\n",
    "awards = []\n",
    "movies = []\n",
    "for e in entity_list:\n",
    "    s,p, o = e\n",
    "    if o == 'person':\n",
    "        persons.append(s.lower())\n",
    "    elif o == 'organization':\n",
    "        orgs.append(s.lower())\n",
    "    elif o == 'location':\n",
    "        locs.append(s.lower())\n",
    "    elif o == 'award':\n",
    "        awards.append(s.lower())\n",
    "    elif o == 'movie':\n",
    "        movies.append(s.lower())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "2e530fbe-6756-455c-9c61-7f2921da1840",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "88"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(movies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "25e9263c-1fa7-446f-a14a-74b15a32df07",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['north america',\n",
       " 'philadelphia',\n",
       " 'broadway',\n",
       " 'concord',\n",
       " 'california',\n",
       " 'red bluff',\n",
       " 'oakland',\n",
       " 'bay area',\n",
       " 'alameda',\n",
       " 'hayward',\n",
       " 'cleveland',\n",
       " 'ohio',\n",
       " 'new york city',\n",
       " 'los angeles',\n",
       " 'us',\n",
       " 'hollywood',\n",
       " 'wall street',\n",
       " 'moon',\n",
       " 'france',\n",
       " 'u.s.',\n",
       " 'texas',\n",
       " 'soviet union',\n",
       " 'neighborhood',\n",
       " 'studio 8h',\n",
       " 'queensland',\n",
       " 'australia',\n",
       " 'new orleans',\n",
       " 'greece',\n",
       " 'ketchum',\n",
       " 'idaho',\n",
       " 'las vegas',\n",
       " 'schöneck',\n",
       " 'hesse',\n",
       " 'germany',\n",
       " 'united states',\n",
       " 'kentucky',\n",
       " 'mati',\n",
       " 'athens',\n",
       " 'white house',\n",
       " 'new york',\n",
       " 'rock and roll hall of fame',\n",
       " 'pittsburgh',\n",
       " 'worldwide',\n",
       " 'asteroid 12818 tomhanks',\n",
       " 'world',\n",
       " 'london',\n",
       " 'secaucus',\n",
       " 'new jersey',\n",
       " 'boston',\n",
       " 'edina',\n",
       " 'minnesota']"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "locs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "id": "92cbf349-c8dd-4d51-8a99-8a9284f8c2b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install SPARQLWrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "109f9d33-f2ea-4d68-aef4-e2cd6e614df3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "http://dbpedia.org/resource/Category:Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://www.w3.org/2004/02/skos/core#Concept\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://www.w3.org/2002/07/owl#Thing\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://xmlns.com/foaf/0.1/Person\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/ontology/Person\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://www.ontologydesignpatterns.org/ont/dul/DUL.owl#NaturalPerson\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://www.wikidata.org/entity/Q19088\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://www.wikidata.org/entity/Q215627\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://www.wikidata.org/entity/Q5\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://www.wikidata.org/entity/Q729\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanFilmDirectors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanFilmProducers\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanMaleFilmActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanMaleTelevisionActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanMaleVoiceActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanPeople\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanPeopleOfBritishDescent\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanPeopleOfEnglishDescent\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanPeopleOfGermanDescent\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanPeopleOfPortugueseDescent\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanScreenwriters\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanTelevisionActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanTelevisionActresses\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatAmericanVoiceActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/ontology/Animal\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/ontology/Eukaryote\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/ontology/Species\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatCaliforniaStateUniversity,SacramentoAlumni\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatChabotCollegeAlumni\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://schema.org/Person\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatLGBTRightsActivistsFromTheUnitedStates\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatLivingPeople\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatWritersFromLosAngeles,California\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatPeopleFromCalifornia\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatPeopleFromConcord,California\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatPeopleFromOakland,California\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatSpaceAdvocates\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Actor109765278\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Actress109767700\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Advocate109774783\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Alumnus109786338\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/CausalAgent100007347\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Communicator109610660\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Creator109614315\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Disputant109615465\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Emigrant110051975\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Entertainer109616922\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/FilmDirector110088200\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/FilmMaker110088390\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Intellectual109621545\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/LivingThing100004258\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Migrant110314952\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Militant110315837\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Object100002684\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Organism100004475\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Performer110415638\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Person100007846\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/PhysicalEntity100001930\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Producer110480018\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Reformer110515194\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Writer110794014\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/YagoLegalActor\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/YagoLegalActorGeo\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Scholar110557854\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Screenwriter110564400\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Scriptwriter110564905\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Traveler109629752\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Whole100003553\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Wikicat20th-centuryActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Wikicat20th-centuryAmericanActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Wikicat20th-centuryAmericanMaleActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Wikicat21st-centuryActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Wikicat21st-centuryAmericanActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/Wikicat21st-centuryAmericanMaleActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatActors\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatActorsFromCalifornia\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatActressesFromSacramento,California\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatFilmDirectorsFromCalifornia\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatFilmProducers\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/1999/02/22-rdf-syntax-ns#type http://dbpedia.org/class/yago/WikicatPortugueseEmigrantsToTheUnitedStates\n",
      "http://dbpedia.org/resource/Category:Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label توم هانكس\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Τομ Χανκς\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label 톰 행크스\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label トム・ハンクス\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n",
      "http://dbpedia.org/resource/Tom_Hanks http://www.w3.org/2000/01/rdf-schema#label Tom Hanks\n"
     ]
    }
   ],
   "source": [
    "from SPARQLWrapper import SPARQLWrapper, JSON\n",
    "\n",
    "# Define the SPARQL endpoint\n",
    "sparql = SPARQLWrapper(\"http://dbpedia.org/sparql\")\n",
    "\n",
    "# Define the SPARQL query\n",
    "query = \"\"\"\n",
    "PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>\n",
    "\n",
    "SELECT ?subject ?predicate ?object\n",
    "WHERE {\n",
    "  { \n",
    "    ?subject ?predicate ?object .\n",
    "    ?subject rdfs:label \"Tom Hanks\"@en .\n",
    "  }\n",
    "  UNION\n",
    "  {\n",
    "    ?subject ?predicate ?object .\n",
    "    ?subject rdfs:label \"Killing Lincoln\"@en .\n",
    "  }\n",
    "}\n",
    "LIMIT 100\n",
    "\"\"\"\n",
    "\n",
    "# Set the query\n",
    "sparql.setQuery(query)\n",
    "\n",
    "# Set the return format to JSON\n",
    "sparql.setReturnFormat(JSON)\n",
    "\n",
    "# Execute the query and convert the result to a Python dictionary\n",
    "results = sparql.query().convert()\n",
    "\n",
    "# Process and print the results\n",
    "for result in results[\"results\"][\"bindings\"]:\n",
    "    print(f\"{result['subject']['value']} {result['predicate']['value']} {result['object']['value']}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "aabd82a8-c60d-42d2-9ab5-d1a214684cc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"\"\"PREFIX dbo: <http://dbpedia.org/ontology/>\n",
    "PREFIX dbp: <http://dbpedia.org/property/>\n",
    "PREFIX dbr: <http://dbpedia.org/resource/>\n",
    "PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>\n",
    "\n",
    "SELECT ?relation ?object\n",
    "WHERE {\n",
    "  {\n",
    "    dbr:Tom_Hanks ?relation ?object .\n",
    "    ?object rdfs:label \"Killing Lincoln\"@en .\n",
    "  }\n",
    "  UNION\n",
    "  {\n",
    "    dbr:Killing_Lincoln ?relation ?object .\n",
    "    ?object rdfs:label \"Tom Hanks\"@en .\n",
    "  }\n",
    "}\n",
    "\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "6341e220-5aba-4711-a569-c544b47a39d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "http://dbpedia.org/ontology/wikiPageWikiLink http://dbpedia.org/resource/Tom_Hanks\n"
     ]
    }
   ],
   "source": [
    "# Set the query\n",
    "sparql.setQuery(query)\n",
    "\n",
    "# Set the return format to JSON\n",
    "sparql.setReturnFormat(JSON)\n",
    "\n",
    "# Execute the query and convert the result to a Python dictionary\n",
    "results = sparql.query().convert()\n",
    "\n",
    "# Process and print the results\n",
    "for result in results[\"results\"][\"bindings\"]:\n",
    "    print(f\"{result['relation']['value']} {result['object']['value']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a9f9c669-4eb2-4eaa-b186-7fd8cacd4691",
   "metadata": {},
   "outputs": [],
   "source": [
    "from litellm import completion\n",
    "from typing import List\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "bbb4ba07-5cf5-4cf4-9d28-3273818162dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_entities(ent_list:List[str]) -> str:\n",
    "    return \"\\n\\n\".join([e for e in ent_list])\n",
    "\n",
    "# print(format_entities(chunks_entities[9]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "c42e47ad-f0fe-4fec-8651-adb0a2daf084",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'subject': 'Hanks', 'relationship': 'launched', 'object': 'Playtone'},\n",
       " {'subject': 'Playtone',\n",
       "  'relationship': 'has',\n",
       "  'object': 'exclusive television development deal with HBO'},\n",
       " {'subject': 'Hanks',\n",
       "  'relationship': 'won',\n",
       "  'object': 'seven Primetime Emmy Awards'},\n",
       " {'subject': 'Hanks',\n",
       "  'relationship': 'won',\n",
       "  'object': 'Emmy Awards for his work as a producer of various limited series and television movies'},\n",
       " {'subject': 'Hanks',\n",
       "  'relationship': 'won',\n",
       "  'object': 'Emmy Awards for From the Earth to the Moon, Band of Brothers, John Adams, The Pacific, Game Change, and Olive Kitteridge'},\n",
       " {'subject': 'Hanks',\n",
       "  'relationship': 'made',\n",
       "  'object': \"Broadway debut in Nora Ephron's Lucky Guy\"},\n",
       " {'subject': 'Hanks',\n",
       "  'relationship': 'earned',\n",
       "  'object': 'Tony Award for Best Actor in a Play nomination'}]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "system_message = \"\"\"Extract all the relationships between the following entities ONLY based on the given context. \n",
    "Return a list of JSON objects. For example:\n",
    "\n",
    "<Examples>\n",
    "    [{{\"subject\": \"John\", \"relationship\": \"lives in\", \"object\": \"US\"}},\n",
    "    {{\"subject\": \"Eifel towel\", \"relationship\": \"is located in\", \"object\": \"Paris\"}},\n",
    "    {{\"subject\": \"Hayao Miyazaki\", \"relationship\": \"is\", \"object\": \"Japanese animator\"}}]\n",
    "</Examples>\n",
    "\n",
    "- ONLY return triples and nothing else. None of 'subject', 'relationship' and 'object' can be empty.\n",
    "\n",
    "Entities: \\n\\n{entities}\n",
    "\n",
    "\"\"\"\n",
    "# example for a particular set of entities\n",
    "i = 3\n",
    "\n",
    "ents = format_entities(chunks_entities[i])\n",
    "text = chunks[i]\n",
    "\n",
    "user_message = \"Context: {text}\\n\\nTriples:\"\n",
    "response = completion(\n",
    "  # api_key=OPENAI_API_KEY,\n",
    "  # model=\"ollama/adrienbrault/nous-hermes2pro-llama3-8b:q8_0\",\n",
    "  # model=\"ollama/llama3\",\n",
    "  model=\"gpt-3.5-turbo\",\n",
    "  messages=[{\"content\": system_message.format(entities=ents),\"role\": \"system\"}, {\"content\": user_message.format(text=text),\"role\": \"user\"}],\n",
    "  max_tokens=1000,\n",
    "  format = \"json\"\n",
    ")\n",
    "\n",
    "triples = json.loads(response.choices[0].message.content)\n",
    "triples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 366,
   "id": "1c72d9d1-e61f-49da-ac2d-9c9f9cf157e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|███████████▍                                                                                                                              | 8/97 [00:51<06:27,  4.36s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 7, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 11%|███████████████▌                                                                                                                         | 11/97 [01:13<08:50,  6.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 10, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|████████████████▉                                                                                                                        | 12/97 [01:14<06:26,  4.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 11, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██████████████████████████▊                                                                                                              | 19/97 [01:55<06:12,  4.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 18, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 21%|████████████████████████████▏                                                                                                            | 20/97 [01:56<04:51,  3.79s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 19, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 32%|███████████████████████████████████████████▊                                                                                             | 31/97 [03:09<05:49,  5.29s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 30, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|████████████████████████████████████████████████████████████▋                                                                            | 43/97 [04:23<04:18,  4.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 42, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 54%|█████████████████████████████████████████████████████████████████████████▍                                                               | 52/97 [05:16<03:20,  4.45s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 51, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 61%|███████████████████████████████████████████████████████████████████████████████████▎                                                     | 59/97 [05:46<02:43,  4.29s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 58, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 82%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████▉                        | 80/97 [07:34<01:14,  4.41s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 79, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 92%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋           | 89/97 [08:13<00:31,  3.97s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 88, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 97%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊    | 94/97 [08:39<00:12,  4.29s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error for chunk 93, Expecting value: line 1 column 1 (char 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 97/97 [08:52<00:00,  5.49s/it]\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "errors = []\n",
    "all_triples = []\n",
    "for i in tqdm(range(len(chunks_entities))):\n",
    "    try:\n",
    "        ents = format_entities(chunks_entities[i])\n",
    "        text = chunks[i]\n",
    "        \n",
    "        user_message = \"Context: {text}\\n\\nTriples:\"\n",
    "        response = completion(\n",
    "            # api_key=OPENAI_API_KEY,\n",
    "            # model=\"ollama/adrienbrault/nous-hermes2pro-llama3-8b:q8_0\",\n",
    "            # model=\"ollama/llama3\",\n",
    "            model=\"gpt-3.5-turbo\",\n",
    "            messages=[{\"content\": system_message.format(entities=ents),\"role\": \"system\"}, {\"content\": user_message.format(text=text),\"role\": \"user\"}],\n",
    "            max_tokens=1000,\n",
    "            format=\"json\"\n",
    "        )\n",
    "        triples = json.loads(response.choices[0].message.content)\n",
    "        all_triples.append(triples)\n",
    "        time.sleep(3)\n",
    "    except Exception as e:\n",
    "        print(f\"Error for chunk {i}, {e}\")\n",
    "        errors.append(response.choices[0].message.content)\n",
    "        all_triples.append([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 372,
   "id": "f789430f-5c18-4581-9443-b63902f0cecf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# output_file = \"triples.json\"\n",
    "# json_data = json.dumps(all_triples, indent=4)\n",
    "# with open(output_file, \"w\") as file:\n",
    "#     file.write(json_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "0f1fe0b0-0d17-4e76-9ba5-1ee7eba5aac1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'subject': 'Thomas Jeffrey Hanks',\n",
       "  'relationship': 'has received',\n",
       "  'object': 'AFI Life Achievement Award'},\n",
       " {'subject': 'Thomas Jeffrey Hanks',\n",
       "  'relationship': 'has received',\n",
       "  'object': 'Kennedy Center Honor'},\n",
       " {'subject': 'Thomas Jeffrey Hanks',\n",
       "  'relationship': 'has received',\n",
       "  'object': 'Presidential Medal of Freedom'},\n",
       " {'subject': 'Thomas Jeffrey Hanks',\n",
       "  'relationship': 'has received',\n",
       "  'object': 'French Legion of Honor'},\n",
       " {'subject': 'Thomas Jeffrey Hanks',\n",
       "  'relationship': 'has received',\n",
       "  'object': 'Golden Globe Cecil B. DeMille Award'},\n",
       " {'subject': 'Thomas Jeffrey Hanks',\n",
       "  'relationship': 'has grossed',\n",
       "  'object': '$4.9 billion in North America'},\n",
       " {'subject': 'Thomas Jeffrey Hanks',\n",
       "  'relationship': 'has grossed',\n",
       "  'object': '$9.96 billion worldwide'},\n",
       " {'subject': 'Thomas Jeffrey Hanks',\n",
       "  'relationship': 'is',\n",
       "  'object': 'American actor'},\n",
       " {'subject': 'Thomas Jeffrey Hanks',\n",
       "  'relationship': 'is',\n",
       "  'object': 'filmmaker'},\n",
       " {'subject': 'North America',\n",
       "  'relationship': 'has grossed',\n",
       "  'object': '$4.9 billion'},\n",
       " {'subject': 'North America',\n",
       "  'relationship': 'has grossed',\n",
       "  'object': 'fourth-highest-grossing actor'},\n",
       " {'subject': 'North America',\n",
       "  'relationship': 'has grossed',\n",
       "  'object': '$9.96 billion worldwide'}]"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_file = \"triples.json\"\n",
    "with open(input_file, \"r\") as file:\n",
    "    all_triples = json.load(file)\n",
    "\n",
    "all_triples[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "becb4e9f-de8e-42a8-892c-8708c5371266",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install pyvis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "2da12aa9-6512-4fd7-abd0-64aef7b471ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_color(n: str) -> str:\n",
    "    type_to_color = {\n",
    "        \"person\": \"#6495ED\",\n",
    "        \"location\": \"#3CB371\",\n",
    "        \"award\": \"#F4A460\",\n",
    "        \"organization\": \"#CD5C5C\",\n",
    "        \"movie\": \"#6A5ACD\"\n",
    "    }\n",
    "    if n.lower() in persons:\n",
    "        return type_to_color[\"person\"]\n",
    "    if n.lower() in locs:\n",
    "        return type_to_color[\"location\"]\n",
    "    if n.lower() in awards:\n",
    "        return type_to_color[\"award\"]\n",
    "    if n.lower() in orgs:\n",
    "        return type_to_color[\"organization\"]\n",
    "    if n.lower() in movies:\n",
    "        return type_to_color[\"movie\"]\n",
    "    return \"red\"\n",
    "    \n",
    "\n",
    "def get_size(n: str) -> int:\n",
    "    type_to_size = {\n",
    "        \"person\": 50,\n",
    "        \"location\": 30,\n",
    "        \"award\": 20,\n",
    "        \"organization\": 10,\n",
    "        \"movie\": 40\n",
    "    }\n",
    "    if n.lower() in persons:\n",
    "        return type_to_size[\"person\"]\n",
    "    if n.lower() in locs:\n",
    "        return type_to_size[\"location\"]\n",
    "    if n.lower() in awards:\n",
    "        return type_to_size[\"award\"]\n",
    "    if n.lower() in orgs:\n",
    "        return type_to_size[\"organization\"]\n",
    "    if n.lower() in movies:\n",
    "        return type_to_size[\"movie\"]\n",
    "    \n",
    "    return 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "a49c6510-a31b-4655-8d0c-ca67bb93f313",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error in item: {'subject': 'William Dodd', 'relationship': 'is an American diplomat'}\n"
     ]
    }
   ],
   "source": [
    "from pyvis.network import Network\n",
    "import networkx as nx\n",
    "\n",
    "G = nx.Graph()\n",
    "\n",
    "for items in all_triples:\n",
    "    for item in items:\n",
    "        try:\n",
    "            node_1 = item[\"subject\"]\n",
    "            node_2 = item[\"object\"]\n",
    "            G.add_node(node_1, title=node_1, color=get_color(node_1), size=get_size(node_1), label=node_1)\n",
    "            G.add_node(node_2, title=node_2, color=get_color(node_2), size=get_size(node_2), label=node_2)\n",
    "            G.add_edge(node_1, node_2, title=item[\"relationship\"], weight=4)\n",
    "        except Exception as e:\n",
    "            print(f\"Error in item: {item}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "b491a903-673e-477b-9b8f-3a60b7cdcc6a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "graph.html\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"1000\"\n",
       "            height=\"800\"\n",
       "            src=\"graph.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x32a7c6d10>"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nt =  Network(height=\"750px\", width=\"100%\")\n",
    "# nt =  Network(height=\"750px\", width=\"100%\", bgcolor=\"#222222\", font_color=\"white\")\n",
    "\n",
    "nt.from_nx(G)\n",
    "# nt.toggle_physics(True)\n",
    "nt.force_atlas_2based(central_gravity=0.015, gravity=-31)\n",
    "nt.show(\"graph.html\", notebook=False)\n",
    "# Generate the HTML\n",
    "# html = nt.generate_html()\n",
    "\n",
    "# # Write the HTML to a file\n",
    "# with open(\"graph.html\", \"w\") as file:\n",
    "#     file.write(html)\n",
    "\n",
    "# # Display the graph in a Jupyter Notebook\n",
    "from IPython.display import IFrame\n",
    "IFrame(\"graph.html\", width=1000, height=800)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ab5b143-6bc6-4db7-9a41-59c677e9c2e0",
   "metadata": {},
   "source": [
    "![image](kg_graph.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb7b6c67-f118-4a37-aac8-2b2331cae44f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}