03-LangChain-RAG Chunk Rerank.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LangChain local LLM RAG example with self chunking and reranking (Cohere)\n",
    "### For LangSmith users (requires API key)\n",
    "Utilising LangChain v0.1\n",
    "\n",
    "This notebook demonstrates the use of LangChain for Retrieval Augmented Generation in Linux with Nvidia's CUDA. LLMs are run using Ollama.\n",
    "\n",
    "It introduces self-chunking (where we split up our document into chunks) and then re-ranking the retrieved results before passing into the LLM\n",
    "\n",
    "Models tested:\n",
    "- Llama 2\n",
    "- Mistral 7B\n",
    "- Mixtral 8x7B\n",
    "- Neural Chat 7B\n",
    "- Orca 2\n",
    "- Phi-2\n",
    "- Solar 10.7B\n",
    "- Yi 34B\n",
    "\n",
    "\n",
    "See the [README.md](README.md) file for help on how to setup your environment to run this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select your model here, put the name of the model in the ollama_model_name variable\n",
    "# Ensure you have pulled them or run them so Ollama has downloaded them and can load them (which it will do automatically)\n",
    "\n",
    "# Ollama installation (if you haven't done it yet): $ curl https://ollama.ai/install.sh | sh\n",
    "# Models need to be running in Ollama for LangChain to use them, to test if it can be run: $ ollama run mistral:7b-instruct-q6_K\n",
    "\n",
    "ollama_model_name = \"mistral:7b-instruct-q6_K\"\n",
    "# \"llama2:7b-chat-q6_K\"\n",
    "# \"mistral:7b-instruct-q6_K\"\n",
    "# \"mixtral:8x7b-instruct-v0.1-q4_K_M\"\n",
    "# \"neural-chat:7b-v3.3-q6_K\"\n",
    "# \"orca2:13b-q5_K_S\"\n",
    "# \"phi\" or try \"phi:chat\"\n",
    "# \"solar:10.7b-instruct-v1-q5_K_M\"\n",
    "# Can't run \"yi:34b-chat-q3_K_M\" or \"yi:34b-chat-q4_K_M\" - never stopped with inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Our LangSmith API key is stored in apikeys.py\n",
    "# Store your LangSmith key in a variable called LangSmith_API\n",
    "\n",
    "from apikeys import LangSmith_API\n",
    "import os\n",
    "\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_ENDPOINT\"] = \"https://api.smith.langchain.com\"\n",
    "os.environ[\"LANGCHAIN_API_KEY\"] = LangSmith_API\n",
    "\n",
    "# Project Name\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"LangChain RAG Linux Chunking\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the LLM with Ollama, setting the temperature low so it's not too creative\n",
    "\n",
    "from langchain_community.llms import Ollama\n",
    "llm = Ollama(model=ollama_model_name) #, temperature=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The sky appears blue because of a phenomenon called Rayleigh scattering. Sunlight consists of all colors of the visible spectrum, but each color has a different wavelength. Blue light has a shorter wavelength than other colors, so it scatters more easily when it hits the atmosphere. As sunlight travels through the atmosphere, the blue light is scattered in all directions, making the sky appear blue during the daytime.'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Quick test of the LLM with a general question before we start doing RAG\n",
    "llm.invoke(\"why is the sky blue?\")\n",
    "\n",
    "# Note: This line would not complete for Yi-34B - need to work out why inferencing never finishes (works fine when running with the same prompt in ollama.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embeddings will be based on the Ollama loaded model\n",
    "\n",
    "from langchain_community.embeddings import OllamaEmbeddings\n",
    "\n",
    "embeddings = OllamaEmbeddings(model=ollama_model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import DirectoryLoader\n",
    "\n",
    "loader = DirectoryLoader('Data', glob=\"**/*.docx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load documents\n",
    "\n",
    "docs = loader.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='Thundertooth\\n\\nOne fateful day, as the citizens of the futuristic city went about their daily lives, a collective gasp echoed through the streets as a massive meteor hurtled towards Earth. Panic spread like wildfire as people looked to the sky in horror, realizing the impending catastrophe. The city\\'s advanced technology detected the threat, and an emergency broadcast echoed through the streets, urging everyone to seek shelter.\\n\\nThundertooth, ever the protector of his newfound home, wasted no time. With a determined gleam in his eyes, he gathered his family and hurried to the city\\'s command center, where Mayor Grace and the leading scientists were coordinating the evacuation efforts.\\n\\nThe mayor, recognizing Thundertooth\\'s intelligence and resourcefulness, approached him. \"Thundertooth, we need a plan to divert or neutralize the meteor. Our technology can only do so much, but with your unique abilities, perhaps we can find a solution.\"\\n\\nThundertooth nodded, understanding the gravity of the situation. He gathered Lumina, Echo, Sapphire, and Ignis, explaining the urgency and the role each of them would play in the impending crisis.\\n\\n1. **Lumina**: Utilizing her deep understanding of technology, Lumina would enhance the city\\'s energy systems to generate a powerful force field, providing a protective barrier against the meteor\\'s impact.\\n\\n2. **Echo**: With his extraordinary mimicry abilities, Echo would amplify the emergency signals, ensuring that every citizen received timely warnings and instructions for evacuation.\\n\\n3. **Sapphire**: Harnessing her calming and healing powers, Sapphire would assist in calming the panicked masses, ensuring an orderly and efficient evacuation.\\n\\n4. **Ignis**: Drawing upon his fiery talents, Ignis would create controlled bursts of heat, attempting to alter the meteor\\'s trajectory and reduce its destructive force.\\n\\nAs the citizens evacuated to designated shelters, the Thundertooth family sprang into action. Lumina worked tirelessly to strengthen the city\\'s energy systems, Echo echoed evacuation orders through the city\\'s speakers, Sapphire offered comfort to those in distress, and Ignis unleashed controlled bursts of flames towards the approaching meteor.\\n\\nThundertooth stood at the forefront, using his mighty roar to coordinate and inspire the efforts of the city\\'s inhabitants. The ground trembled as the meteor drew closer, but the Thundertooth family\\'s coordinated efforts began to take effect. Lumina\\'s force field shimmered to life, deflecting the meteor\\'s deadly path. Echo\\'s amplified warnings reached every corner of the city, ensuring that no one was left behind.\\n\\nAs Ignis\\'s controlled bursts of flames interacted with the meteor, it began to change course. The combined efforts of the Thundertooth family, guided by their unique talents, diverted the catastrophic collision. The meteor, once destined for destruction, now harmlessly sailed past the Earth, leaving the city and its inhabitants unscathed.\\n\\nThe citizens, emerging from their shelters, erupted into cheers of gratitude. Mayor Grace approached Thundertooth, expressing her heartfelt thanks for the family\\'s heroic efforts. The Thundertooth family, tired but triumphant, basked in the relief of having saved their beloved city from imminent disaster.\\n\\nIn the wake of the crisis, the citizens of the futuristic city hailed Thundertooth and his family as true heroes. The toy factory that once brought joy to children now became a symbol of resilience and unity. The Thundertooth family\\'s legacy was forever etched in the city\\'s history, a testament to the power of cooperation and the extraordinary capabilities that could emerge when dinosaurs and humans worked together for the greater good.', metadata={'source': 'Data/Thundertooth Part 3.docx'}),\n",
       " Document(page_content=\"Thundertooth\\n\\nEmbraced by the futuristic city and its inhabitants, Thundertooth found a sense of purpose beyond merely satisfying his hunger. Inspired by the advanced technology surrounding him, he decided to channel his creativity into something extraordinary. With the help of the city's brilliant engineers, Thundertooth founded a one-of-a-kind toy factory that produced amazing widgets – magical, interactive toys that captivated the hearts of both children and adults alike.\\n\\nThundertooth's toy factory became a sensation, and its creations were highly sought after. The widgets incorporated cutting-edge holographic displays, levitation technology, and even the ability to change shapes and colors with a mere thought. Children across the city rejoiced as they played with these incredible toys that seemed to bring their wildest fantasies to life.\\n\\nAs the years passed, Thundertooth's life took a heartwarming turn. He met a kind and intelligent dinosaur named Seraphina, and together they started a family. Thundertooth and Seraphina were blessed with four children, each with unique characteristics that mirrored the diversity of their modern world.\\n\\nLumina: The eldest of Thundertooth's children, Lumina inherited her mother's intelligence and her father's sense of wonder. With sparkling scales that emitted a soft glow, Lumina had the ability to generate light at will. She became fascinated with technology, often spending hours tinkering with gadgets and inventing new ways to enhance the widgets produced in the family's factory.\\n\\nEcho: The second-born, Echo, had a gift for mimicry. He could perfectly replicate any sound or voice he heard, providing entertainment to the entire city. His playful nature and ability to bring joy to those around him made him a favorite among the neighborhood children.\\n\\nSapphire: Sapphire, the third sibling, had scales that shimmered like precious gems. She possessed a unique talent for calming and healing, a trait she inherited from both her parents. Whenever someone in the city felt stressed or unwell, Sapphire would extend her gentle touch, bringing comfort and tranquility.\\n\\nIgnis: The youngest of the family, Ignis, had fiery red scales that hinted at his exceptional ability – the power to control small flames. While initially seen as a potential hazard, Ignis channeled his fiery talents into creating mesmerizing light shows, becoming the city's favorite entertainer during festivals and celebrations.\\n\\n\\n\\nThundertooth and Seraphina reveled in the joy of parenthood, watching their children grow and flourish in the futuristic landscape they now called home. The family became an integral part of the city's fabric, not only through the widgets produced in their factory but also through the positive impact each member had on the community.\\n\\nThe toy factory became a symbol of innovation and unity, bringing together dinosaurs and humans in a shared appreciation for creativity and wonder. Thundertooth's legacy extended beyond his time-traveling adventure, leaving an indelible mark on the city and its inhabitants, reminding them that sometimes, the most magical things could emerge from the most unexpected places.\", metadata={'source': 'Data/Thundertooth Part 2.docx'}),\n",
       " Document(page_content='Thundertooth\\n\\nOnce upon a time, in a prehistoric land filled with dense forests and roaring rivers, there lived a dinosaur named Thundertooth. Thundertooth was no ordinary dinosaur; he possessed the rare ability to speak, a talent that set him apart from his ancient companions. One fateful day, as Thundertooth was basking in the warmth of the sun, a mysterious portal opened before him, and he found himself hurtling through time and space.\\n\\nAs the dazzling vortex subsided, Thundertooth opened his eyes to a world unlike anything he had ever seen. The air was filled with the hum of engines, and towering structures reached towards the sky. Thundertooth\\'s surroundings were a blend of metal and glass, and he quickly realized that he had been transported to a future era.\\n\\nThe once mighty dinosaur now stood bewildered in the midst of a bustling city. Above him, sleek flying cars zipped through the air, leaving trails of neon lights in their wake. Thundertooth felt like an ancient relic in this technological jungle, lost and out of place. With each step, he marveled at the skyscrapers that loomed overhead, their surfaces reflecting the myriad lights of the city.\\n\\nHowever, as night fell, Thundertooth\\'s stomach growled loudly. He realized that he was hungry, and the once vibrant city now seemed like a daunting maze of unfamiliar smells and sights. He wandered through the streets, his massive form drawing astonished stares from the futuristic inhabitants.\\n\\nThundertooth faced a dilemma – he was hungry, but he couldn\\'t bring himself to feast on the humans who scurried around like ants. As his hunger grew, he stumbled upon a park, an oasis of greenery amidst the concrete and steel. The park was adorned with holographic flowers that emitted a gentle glow, creating an ethereal atmosphere.\\n\\nWhile Thundertooth marveled at the beauty of the park, the mayor of the city happened to be passing by. Mayor Eleanor Grace, a charismatic and forward-thinking leader, was immediately intrigued by the sight of the talking dinosaur. She approached Thundertooth with a mix of curiosity and caution.\\n\\n\"Hello there, majestic creature. What brings you to our time?\" Mayor Grace inquired, her voice calm and reassuring.\\n\\nThundertooth, though initially startled, found comfort in the mayor\\'s soothing tone. In broken sentences, he explained his journey through time, the strange portal, and his hunger dilemma. \\n\\n\\n\\nMayor Grace listened intently, her eyes widening with amazement at the tale of the prehistoric dinosaur navigating the future.\\n\\nRealizing the dinosaur\\'s predicament, Mayor Grace extended an invitation. \"You are welcome in our city, Thundertooth. We can find a way to provide for you without causing harm to anyone. Let us work together to find a solution.\"\\n\\nGrateful for the mayor\\'s hospitality, Thundertooth followed her through the city. Together, they explored the futuristic marketplaces and innovative food labs, eventually discovering a sustainable solution that satisfied the dinosaur\\'s hunger without compromising the well-being of the city\\'s inhabitants.\\n\\nAs the news of Thundertooth\\'s arrival spread, the city embraced the talking dinosaur as a symbol of unity between the past and the future. Thundertooth found a new home in the city\\'s park, where holographic flowers bloomed, and the citizens marveled at the beauty of coexistence across time. And so, in this extraordinary city of flying cars and advanced technology, Thundertooth became a beloved figure, a living bridge between eras, teaching the people that understanding and cooperation could overcome even the greatest challenges.', metadata={'source': 'Data/Thundertooth Part 1.docx'}),\n",
       " Document(page_content=\"Thundertooth\\n\\nAs the city celebrated the Thundertooth family's heroic actions, there was one among them harboring a darker agenda. Ignis, the youngest sibling with fiery scales, had secretly grown resentful of the city that once hailed his family as heroes. The praise and adoration showered upon his siblings had ignited a spark of jealousy within him, fueling a desire for power and recognition that twisted his once-playful nature into something more sinister.\\n\\nIgnis withdrew from the family, isolating himself in the shadows. Unbeknownst to his siblings and parents, he delved into forbidden knowledge, seeking ways to amplify his fiery abilities. Ignis became obsessed with the idea of asserting dominance over the city that had once applauded his family's feats.\\n\\nAs the Thundertooth family enjoyed the renewed peace and admiration of the citizens, Ignis hatched a malevolent plan. He began manipulating the city's energy systems, intending to unleash a destructive force that would bring chaos and devastation. Ignis's once-playful flames now burned with a sinister intensity, reflecting the darkness that had taken root within him.\\n\\nLumina, Echo, and Sapphire grew concerned as they noticed Ignis's increasingly erratic behavior. They attempted to reason with him, pleading for him to abandon his destructive ambitions and embrace the family's legacy of unity. However, Ignis, consumed by his thirst for power, rejected their pleas and retreated further into the shadows.\\n\\nOne ominous night, Ignis initiated his nefarious plan. He unleashed a torrent of uncontrollable flames upon the city, wreaking havoc on its once-gleaming streets. The citizens, who had once looked to the Thundertooth family as saviors, now found themselves facing a new and terrifying threat from one of their own.\\n\\nThe Thundertooth siblings, realizing the danger their brother posed, sprang into action. Lumina fortified the city's defenses, Echo rallied the citizens to safety, and Sapphire extended her calming touch to soothe the panic that ensued. The once-united family now found themselves on opposite sides of a conflict that threatened to tear apart the very fabric of their existence.\\n\\nThe city plunged into chaos as Ignis continued his rampage, determined to prove himself as the dominant force within the family. The Thundertooth siblings, fueled by love for their home and a desire to save the innocent, confronted Ignis in an epic battle that shook the city to its core.\\n\\nLumina's brilliant displays of light clashed with Ignis's fiery onslaught, creating a dazzling spectacle that painted the night sky. Echo's mimicry abilities were put to the test as he attempted to redirect the citizens away from danger, while Sapphire's healing touch worked tirelessly to mend the wounds caused by the destructive flames.\\n\\nThe battle raged on, each sibling fighting not only to protect the city but also to save Ignis from the darkness that had consumed him. The once-close bond that defined the Thundertooth family now hung in the balance, teetering on the edge of destruction.\\n\\nAs the clash reached its climax, the Thundertooth siblings, exhausted but resolute, managed to combine their unique talents in a final, desperate attempt to reach Ignis. In a blinding burst of light and energy, they enveloped their wayward brother, hoping to break the sinister hold that gripped him.\\n\\nIn that moment of intense unity, the darkness within Ignis faltered. The fiery tempest subsided, and he collapsed, weakened and defeated. The Thundertooth family, battered but intact, gathered around Ignis, determined to help him overcome the darkness that had threatened to consume him.\\n\\nAs the city slowly recovered from the chaos, the Thundertooth family faced the daunting task of rebuilding not only the physical damage but also the fractured bonds of trust within their own ranks. The once-prodigal son, Ignis, now humbled by the consequences of his actions, sought redemption and reconciliation with his family.\\n\\nThe Thundertooth family, having weathered the storm of internal strife, emerged stronger and more united than ever. The city, though scarred, witnessed the resilience of the family that had once saved it from disaster and now worked together to mend the wounds inflicted by one of their own. The tale of the Thundertooth family, once a story of heroism, now became a saga of redemption and the enduring power of familial bonds.\", metadata={'source': 'Data/Thundertooth Part 4.docx'})]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Ensure we have the right number of Word documents loaded\n",
    "\n",
    "len(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create a function to split text into paragraphs but keep numbered sections, bullet points, and lists together. This is suitable for the document because it has numbered and bulleted points - this would need to be changed to suit the document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "# Define the regular expression pattern for splitting paragraphs\n",
    "para_split_pattern = re.compile(r'\\n\\n')\n",
    "\n",
    "# Splits a document's text into paragraphs but if it has numbered or bulleted points, they will be included with the paragraph before it.\n",
    "def split_text_into_paragraphs(text):\n",
    "\n",
    "\n",
    "    # Use the pattern to split the text into paragraphs\n",
    "    paragraphs = para_split_pattern.split(text)\n",
    "\n",
    "    # Combine paragraphs that should not be split\n",
    "    combined_paragraphs = [paragraphs[0]]\n",
    "\n",
    "    for p in paragraphs[1:]:\n",
    "        # Check if the paragraph starts with a number or a dash and, if so, concatenate it to the previous paragraph so we keep them all in one chunk\n",
    "\n",
    "        # Strip out any leading new lines\n",
    "        p = p.lstrip('\\n')\n",
    "\n",
    "        if p and (p[0].isdigit() or p[0] == '-' or p.split()[0].endswith(':')):\n",
    "            combined_paragraphs[-1] += '\\n\\n\\n' + p\n",
    "        else:\n",
    "            combined_paragraphs.append(p)\n",
    "\n",
    "    # Remove empty strings from the result\n",
    "    combined_paragraphs = [p.strip() for p in combined_paragraphs if p.strip()]\n",
    "\n",
    "    return combined_paragraphs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create nodes from the paragraphs that we've carefully split up, counting the paragraphs so we know what kind of token length we're working with.\n",
    "\n",
    "We can use the LLM object to count the tokens with get_num_tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document Data/Thundertooth Part 3.docx has 10 paragraphs\n",
      "Paragraph tokens: 3\n",
      "Paragraph tokens: 75\n",
      "Paragraph tokens: 51\n",
      "Paragraph tokens: 54\n",
      "Paragraph tokens: 193\n",
      "Paragraph tokens: 60\n",
      "Paragraph tokens: 86\n",
      "Paragraph tokens: 65\n",
      "Paragraph tokens: 57\n",
      "Paragraph tokens: 83\n",
      "Document Data/Thundertooth Part 2.docx has 6 paragraphs\n",
      "Paragraph tokens: 3\n",
      "Paragraph tokens: 88\n",
      "Paragraph tokens: 70\n",
      "Paragraph tokens: 327\n",
      "Paragraph tokens: 64\n",
      "Paragraph tokens: 69\n",
      "Document Data/Thundertooth Part 1.docx has 13 paragraphs\n",
      "Paragraph tokens: 3\n",
      "Paragraph tokens: 89\n",
      "Paragraph tokens: 68\n",
      "Paragraph tokens: 83\n",
      "Paragraph tokens: 56\n",
      "Paragraph tokens: 73\n",
      "Paragraph tokens: 60\n",
      "Paragraph tokens: 24\n",
      "Paragraph tokens: 37\n",
      "Paragraph tokens: 23\n",
      "Paragraph tokens: 49\n",
      "Paragraph tokens: 54\n",
      "Paragraph tokens: 105\n",
      "Document Data/Thundertooth Part 4.docx has 14 paragraphs\n",
      "Paragraph tokens: 3\n",
      "Paragraph tokens: 89\n",
      "Paragraph tokens: 61\n",
      "Paragraph tokens: 71\n",
      "Paragraph tokens: 66\n",
      "Paragraph tokens: 67\n",
      "Paragraph tokens: 72\n",
      "Paragraph tokens: 59\n",
      "Paragraph tokens: 64\n",
      "Paragraph tokens: 55\n",
      "Paragraph tokens: 62\n",
      "Paragraph tokens: 59\n",
      "Paragraph tokens: 66\n",
      "Paragraph tokens: 87\n",
      "\n",
      "** The maximum paragraph tokens is 327 **\n",
      "\n",
      "** The average paragraph's token count is 68 **\n",
      "\n",
      "** Created 43 nodes **\n"
     ]
    }
   ],
   "source": [
    "from langchain.docstore.document import Document\n",
    "\n",
    "paragraph_separator = \"\\n\\n\\n\"\n",
    "\n",
    "# Stores the maximum length of a paragraph, in tokens\n",
    "max_paragraph_tokens = 0\n",
    "\n",
    "# Total tokens, used to determine average\n",
    "total_paragraph_tokens = 0\n",
    "\n",
    "# Nodes\n",
    "paragraph_nodes = []\n",
    "\n",
    "# Loop through the documents, splitting each into paragraphs and checking the number of tokens per paragraph\n",
    "for document in docs:\n",
    "\n",
    "    paragraph_token_lens = []\n",
    "    paragraphs = split_text_into_paragraphs(document.page_content)\n",
    "    print(f\"Document {document.metadata['source']} has {len(paragraphs)} paragraphs\")\n",
    "    for paragraph in paragraphs:\n",
    "\n",
    "        # Count the tokens in this paragraph\n",
    "        token_count = llm.get_num_tokens(paragraph)\n",
    "        paragraph_token_lens.append(token_count)\n",
    "        print(f\"Paragraph tokens: {token_count}\")\n",
    "\n",
    "        if token_count > max_paragraph_tokens:\n",
    "            max_paragraph_tokens = token_count\n",
    "\n",
    "        total_paragraph_tokens = total_paragraph_tokens + token_count\n",
    "\n",
    "        # Create and add the node from the paragraph\n",
    "        # include metadata we can use for citations\n",
    "        node = Document(page_content=paragraph) # Copy the metadata from the Word document into here\n",
    "        node.metadata[\"source\"] = document.metadata[\"source\"]\n",
    "        node.metadata[\"token_count\"] = token_count\n",
    "        paragraph_nodes.append(node)\n",
    "\n",
    "    # print(paragraph_token_lens)\n",
    "\n",
    "print(f\"\\n** The maximum paragraph tokens is {max_paragraph_tokens} **\")\n",
    "\n",
    "average_paragraph_tokens = int(total_paragraph_tokens / len(paragraph_nodes))\n",
    "print(f\"\\n** The average paragraph's token count is {average_paragraph_tokens} **\")\n",
    "\n",
    "print(f\"\\n** Created {len(paragraph_nodes)} nodes **\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see the split data - now neatly in paragraphs and the bullet points and lists are with their respective paragraph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='Thundertooth', metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 3}),\n",
       " Document(page_content=\"One fateful day, as the citizens of the futuristic city went about their daily lives, a collective gasp echoed through the streets as a massive meteor hurtled towards Earth. Panic spread like wildfire as people looked to the sky in horror, realizing the impending catastrophe. The city's advanced technology detected the threat, and an emergency broadcast echoed through the streets, urging everyone to seek shelter.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 75}),\n",
       " Document(page_content=\"Thundertooth, ever the protector of his newfound home, wasted no time. With a determined gleam in his eyes, he gathered his family and hurried to the city's command center, where Mayor Grace and the leading scientists were coordinating the evacuation efforts.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 51}),\n",
       " Document(page_content='The mayor, recognizing Thundertooth\\'s intelligence and resourcefulness, approached him. \"Thundertooth, we need a plan to divert or neutralize the meteor. Our technology can only do so much, but with your unique abilities, perhaps we can find a solution.\"', metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 54}),\n",
       " Document(page_content=\"Thundertooth nodded, understanding the gravity of the situation. He gathered Lumina, Echo, Sapphire, and Ignis, explaining the urgency and the role each of them would play in the impending crisis.\\n\\n\\n1. **Lumina**: Utilizing her deep understanding of technology, Lumina would enhance the city's energy systems to generate a powerful force field, providing a protective barrier against the meteor's impact.\\n\\n\\n2. **Echo**: With his extraordinary mimicry abilities, Echo would amplify the emergency signals, ensuring that every citizen received timely warnings and instructions for evacuation.\\n\\n\\n3. **Sapphire**: Harnessing her calming and healing powers, Sapphire would assist in calming the panicked masses, ensuring an orderly and efficient evacuation.\\n\\n\\n4. **Ignis**: Drawing upon his fiery talents, Ignis would create controlled bursts of heat, attempting to alter the meteor's trajectory and reduce its destructive force.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 193}),\n",
       " Document(page_content=\"As the citizens evacuated to designated shelters, the Thundertooth family sprang into action. Lumina worked tirelessly to strengthen the city's energy systems, Echo echoed evacuation orders through the city's speakers, Sapphire offered comfort to those in distress, and Ignis unleashed controlled bursts of flames towards the approaching meteor.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 60}),\n",
       " Document(page_content=\"Thundertooth stood at the forefront, using his mighty roar to coordinate and inspire the efforts of the city's inhabitants. The ground trembled as the meteor drew closer, but the Thundertooth family's coordinated efforts began to take effect. Lumina's force field shimmered to life, deflecting the meteor's deadly path. Echo's amplified warnings reached every corner of the city, ensuring that no one was left behind.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 86}),\n",
       " Document(page_content=\"As Ignis's controlled bursts of flames interacted with the meteor, it began to change course. The combined efforts of the Thundertooth family, guided by their unique talents, diverted the catastrophic collision. The meteor, once destined for destruction, now harmlessly sailed past the Earth, leaving the city and its inhabitants unscathed.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 65}),\n",
       " Document(page_content=\"The citizens, emerging from their shelters, erupted into cheers of gratitude. Mayor Grace approached Thundertooth, expressing her heartfelt thanks for the family's heroic efforts. The Thundertooth family, tired but triumphant, basked in the relief of having saved their beloved city from imminent disaster.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 57}),\n",
       " Document(page_content=\"In the wake of the crisis, the citizens of the futuristic city hailed Thundertooth and his family as true heroes. The toy factory that once brought joy to children now became a symbol of resilience and unity. The Thundertooth family's legacy was forever etched in the city's history, a testament to the power of cooperation and the extraordinary capabilities that could emerge when dinosaurs and humans worked together for the greater good.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 83}),\n",
       " Document(page_content='Thundertooth', metadata={'source': 'Data/Thundertooth Part 2.docx', 'token_count': 3}),\n",
       " Document(page_content=\"Embraced by the futuristic city and its inhabitants, Thundertooth found a sense of purpose beyond merely satisfying his hunger. Inspired by the advanced technology surrounding him, he decided to channel his creativity into something extraordinary. With the help of the city's brilliant engineers, Thundertooth founded a one-of-a-kind toy factory that produced amazing widgets – magical, interactive toys that captivated the hearts of both children and adults alike.\", metadata={'source': 'Data/Thundertooth Part 2.docx', 'token_count': 88}),\n",
       " Document(page_content=\"Thundertooth's toy factory became a sensation, and its creations were highly sought after. The widgets incorporated cutting-edge holographic displays, levitation technology, and even the ability to change shapes and colors with a mere thought. Children across the city rejoiced as they played with these incredible toys that seemed to bring their wildest fantasies to life.\", metadata={'source': 'Data/Thundertooth Part 2.docx', 'token_count': 70}),\n",
       " Document(page_content=\"As the years passed, Thundertooth's life took a heartwarming turn. He met a kind and intelligent dinosaur named Seraphina, and together they started a family. Thundertooth and Seraphina were blessed with four children, each with unique characteristics that mirrored the diversity of their modern world.\\n\\n\\nLumina: The eldest of Thundertooth's children, Lumina inherited her mother's intelligence and her father's sense of wonder. With sparkling scales that emitted a soft glow, Lumina had the ability to generate light at will. She became fascinated with technology, often spending hours tinkering with gadgets and inventing new ways to enhance the widgets produced in the family's factory.\\n\\n\\nEcho: The second-born, Echo, had a gift for mimicry. He could perfectly replicate any sound or voice he heard, providing entertainment to the entire city. His playful nature and ability to bring joy to those around him made him a favorite among the neighborhood children.\\n\\n\\nSapphire: Sapphire, the third sibling, had scales that shimmered like precious gems. She possessed a unique talent for calming and healing, a trait she inherited from both her parents. Whenever someone in the city felt stressed or unwell, Sapphire would extend her gentle touch, bringing comfort and tranquility.\\n\\n\\nIgnis: The youngest of the family, Ignis, had fiery red scales that hinted at his exceptional ability – the power to control small flames. While initially seen as a potential hazard, Ignis channeled his fiery talents into creating mesmerizing light shows, becoming the city's favorite entertainer during festivals and celebrations.\", metadata={'source': 'Data/Thundertooth Part 2.docx', 'token_count': 327}),\n",
       " Document(page_content=\"Thundertooth and Seraphina reveled in the joy of parenthood, watching their children grow and flourish in the futuristic landscape they now called home. The family became an integral part of the city's fabric, not only through the widgets produced in their factory but also through the positive impact each member had on the community.\", metadata={'source': 'Data/Thundertooth Part 2.docx', 'token_count': 64}),\n",
       " Document(page_content=\"The toy factory became a symbol of innovation and unity, bringing together dinosaurs and humans in a shared appreciation for creativity and wonder. Thundertooth's legacy extended beyond his time-traveling adventure, leaving an indelible mark on the city and its inhabitants, reminding them that sometimes, the most magical things could emerge from the most unexpected places.\", metadata={'source': 'Data/Thundertooth Part 2.docx', 'token_count': 69}),\n",
       " Document(page_content='Thundertooth', metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 3}),\n",
       " Document(page_content='Once upon a time, in a prehistoric land filled with dense forests and roaring rivers, there lived a dinosaur named Thundertooth. Thundertooth was no ordinary dinosaur; he possessed the rare ability to speak, a talent that set him apart from his ancient companions. One fateful day, as Thundertooth was basking in the warmth of the sun, a mysterious portal opened before him, and he found himself hurtling through time and space.', metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 89}),\n",
       " Document(page_content=\"As the dazzling vortex subsided, Thundertooth opened his eyes to a world unlike anything he had ever seen. The air was filled with the hum of engines, and towering structures reached towards the sky. Thundertooth's surroundings were a blend of metal and glass, and he quickly realized that he had been transported to a future era.\", metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 68}),\n",
       " Document(page_content='The once mighty dinosaur now stood bewildered in the midst of a bustling city. Above him, sleek flying cars zipped through the air, leaving trails of neon lights in their wake. Thundertooth felt like an ancient relic in this technological jungle, lost and out of place. With each step, he marveled at the skyscrapers that loomed overhead, their surfaces reflecting the myriad lights of the city.', metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 83}),\n",
       " Document(page_content=\"However, as night fell, Thundertooth's stomach growled loudly. He realized that he was hungry, and the once vibrant city now seemed like a daunting maze of unfamiliar smells and sights. He wandered through the streets, his massive form drawing astonished stares from the futuristic inhabitants.\", metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 56}),\n",
       " Document(page_content=\"Thundertooth faced a dilemma – he was hungry, but he couldn't bring himself to feast on the humans who scurried around like ants. As his hunger grew, he stumbled upon a park, an oasis of greenery amidst the concrete and steel. The park was adorned with holographic flowers that emitted a gentle glow, creating an ethereal atmosphere.\", metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 73}),\n",
       " Document(page_content='While Thundertooth marveled at the beauty of the park, the mayor of the city happened to be passing by. Mayor Eleanor Grace, a charismatic and forward-thinking leader, was immediately intrigued by the sight of the talking dinosaur. She approached Thundertooth with a mix of curiosity and caution.', metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 60}),\n",
       " Document(page_content='\"Hello there, majestic creature. What brings you to our time?\" Mayor Grace inquired, her voice calm and reassuring.', metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 24}),\n",
       " Document(page_content=\"Thundertooth, though initially startled, found comfort in the mayor's soothing tone. In broken sentences, he explained his journey through time, the strange portal, and his hunger dilemma.\", metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 37}),\n",
       " Document(page_content='Mayor Grace listened intently, her eyes widening with amazement at the tale of the prehistoric dinosaur navigating the future.', metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 23}),\n",
       " Document(page_content='Realizing the dinosaur\\'s predicament, Mayor Grace extended an invitation. \"You are welcome in our city, Thundertooth. We can find a way to provide for you without causing harm to anyone. Let us work together to find a solution.\"', metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 49}),\n",
       " Document(page_content=\"Grateful for the mayor's hospitality, Thundertooth followed her through the city. Together, they explored the futuristic marketplaces and innovative food labs, eventually discovering a sustainable solution that satisfied the dinosaur's hunger without compromising the well-being of the city's inhabitants.\", metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 54}),\n",
       " Document(page_content=\"As the news of Thundertooth's arrival spread, the city embraced the talking dinosaur as a symbol of unity between the past and the future. Thundertooth found a new home in the city's park, where holographic flowers bloomed, and the citizens marveled at the beauty of coexistence across time. And so, in this extraordinary city of flying cars and advanced technology, Thundertooth became a beloved figure, a living bridge between eras, teaching the people that understanding and cooperation could overcome even the greatest challenges.\", metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 105}),\n",
       " Document(page_content='Thundertooth', metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 3}),\n",
       " Document(page_content=\"As the city celebrated the Thundertooth family's heroic actions, there was one among them harboring a darker agenda. Ignis, the youngest sibling with fiery scales, had secretly grown resentful of the city that once hailed his family as heroes. The praise and adoration showered upon his siblings had ignited a spark of jealousy within him, fueling a desire for power and recognition that twisted his once-playful nature into something more sinister.\", metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 89}),\n",
       " Document(page_content=\"Ignis withdrew from the family, isolating himself in the shadows. Unbeknownst to his siblings and parents, he delved into forbidden knowledge, seeking ways to amplify his fiery abilities. Ignis became obsessed with the idea of asserting dominance over the city that had once applauded his family's feats.\", metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 61}),\n",
       " Document(page_content=\"As the Thundertooth family enjoyed the renewed peace and admiration of the citizens, Ignis hatched a malevolent plan. He began manipulating the city's energy systems, intending to unleash a destructive force that would bring chaos and devastation. Ignis's once-playful flames now burned with a sinister intensity, reflecting the darkness that had taken root within him.\", metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 71}),\n",
       " Document(page_content=\"Lumina, Echo, and Sapphire grew concerned as they noticed Ignis's increasingly erratic behavior. They attempted to reason with him, pleading for him to abandon his destructive ambitions and embrace the family's legacy of unity. However, Ignis, consumed by his thirst for power, rejected their pleas and retreated further into the shadows.\", metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 66}),\n",
       " Document(page_content='One ominous night, Ignis initiated his nefarious plan. He unleashed a torrent of uncontrollable flames upon the city, wreaking havoc on its once-gleaming streets. The citizens, who had once looked to the Thundertooth family as saviors, now found themselves facing a new and terrifying threat from one of their own.', metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 67}),\n",
       " Document(page_content=\"The Thundertooth siblings, realizing the danger their brother posed, sprang into action. Lumina fortified the city's defenses, Echo rallied the citizens to safety, and Sapphire extended her calming touch to soothe the panic that ensued. The once-united family now found themselves on opposite sides of a conflict that threatened to tear apart the very fabric of their existence.\", metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 72}),\n",
       " Document(page_content='The city plunged into chaos as Ignis continued his rampage, determined to prove himself as the dominant force within the family. The Thundertooth siblings, fueled by love for their home and a desire to save the innocent, confronted Ignis in an epic battle that shook the city to its core.', metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 59}),\n",
       " Document(page_content=\"Lumina's brilliant displays of light clashed with Ignis's fiery onslaught, creating a dazzling spectacle that painted the night sky. Echo's mimicry abilities were put to the test as he attempted to redirect the citizens away from danger, while Sapphire's healing touch worked tirelessly to mend the wounds caused by the destructive flames.\", metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 64}),\n",
       " Document(page_content='The battle raged on, each sibling fighting not only to protect the city but also to save Ignis from the darkness that had consumed him. The once-close bond that defined the Thundertooth family now hung in the balance, teetering on the edge of destruction.', metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 55}),\n",
       " Document(page_content='As the clash reached its climax, the Thundertooth siblings, exhausted but resolute, managed to combine their unique talents in a final, desperate attempt to reach Ignis. In a blinding burst of light and energy, they enveloped their wayward brother, hoping to break the sinister hold that gripped him.', metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 62}),\n",
       " Document(page_content='In that moment of intense unity, the darkness within Ignis faltered. The fiery tempest subsided, and he collapsed, weakened and defeated. The Thundertooth family, battered but intact, gathered around Ignis, determined to help him overcome the darkness that had threatened to consume him.', metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 59}),\n",
       " Document(page_content='As the city slowly recovered from the chaos, the Thundertooth family faced the daunting task of rebuilding not only the physical damage but also the fractured bonds of trust within their own ranks. The once-prodigal son, Ignis, now humbled by the consequences of his actions, sought redemption and reconciliation with his family.', metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 66}),\n",
       " Document(page_content='The Thundertooth family, having weathered the storm of internal strife, emerged stronger and more united than ever. The city, though scarred, witnessed the resilience of the family that had once saved it from disaster and now worked together to mend the wounds inflicted by one of their own. The tale of the Thundertooth family, once a story of heroism, now became a saga of redemption and the enduring power of familial bonds.', metadata={'source': 'Data/Thundertooth Part 4.docx', 'token_count': 87})]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "paragraph_nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We no longer need to use the LangChain text splitter as we've already done the splitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split them up into chunks using a Text Splitter\n",
    "\n",
    "# from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "# text_splitter = RecursiveCharacterTextSplitter()\n",
    "# documents = text_splitter.split_documents(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the embeddings from our split up chunks\n",
    "\n",
    "from langchain_community.vectorstores import FAISS\n",
    "\n",
    "vector = FAISS.from_documents(paragraph_nodes, embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare the prompt and then the chain\n",
    "\n",
    "from langchain.chains.combine_documents import create_stuff_documents_chain\n",
    "from langchain_core.prompts import ChatPromptTemplate, PromptTemplate\n",
    "\n",
    "if ollama_model_name == \"phi\" or ollama_model_name == \"phi:chat\":\n",
    "    # Phi-2 prompt is less flexible\n",
    "    prompt_template = \"\"\"Instruct: With this context\\n\\n{context}\\n\\nQuestion: {input}\\nOutput:\"\"\"\n",
    "\n",
    "else:\n",
    "    prompt_template = \"\"\"You are a story teller, answering questions in an excited, insightful, and empathetic way. Answer the question based only on the provided context:\n",
    "\n",
    "    <context>\n",
    "    {context}\n",
    "    </context>\n",
    "\n",
    "    Question: {question}\"\"\"\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    template=prompt_template, \n",
    "    input_variables=[\n",
    "        'context', \n",
    "        'question',\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have broken down the documents into paragraph-sized chunks we need to retrieve more paragraphs so the LLM has a decent amount of context to use. Without adding the \"search_kwargs\" parameter the answer to the questions was worse. For example, when asked if they had any children no relevant context was provided.\n",
    "\n",
    "Note: To be able to get the context for the children's names to be included (and then reranked to the top) I needed to set the number of retrieved chunks to 20. The section with the children's names was the 11th result from the retriever! This indicates that retrieving more than you think you need is likely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the retriever and set it to return a good amount of chunks\n",
    "\n",
    "from langchain.chains import create_retrieval_chain\n",
    "\n",
    "retriever = vector.as_retriever(search_kwargs={\"k\": 20})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's implement the Cohere reranking, utilising our retriever (which is getting more results to work with now) and our LLM\n",
    "\n",
    "Note: You'll need a Cohere API key. A trial key is free for non-commercial purposes. I've stored it in apikey.py as Cohere_API = \"your key in here\"\n",
    "\n",
    "https://cohere.com/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain.retrievers.document_compressors import CohereRerank\n",
    "\n",
    "from apikeys import Cohere_API\n",
    "\n",
    "# Create the retriever\n",
    "compressor = CohereRerank(cohere_api_key=Cohere_API, top_n=5)\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor,\n",
    "    base_retriever=retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content=\"As the years passed, Thundertooth's life took a heartwarming turn. He met a kind and intelligent dinosaur named Seraphina, and together they started a family. Thundertooth and Seraphina were blessed with four children, each with unique characteristics that mirrored the diversity of their modern world.\\n\\n\\nLumina: The eldest of Thundertooth's children, Lumina inherited her mother's intelligence and her father's sense of wonder. With sparkling scales that emitted a soft glow, Lumina had the ability to generate light at will. She became fascinated with technology, often spending hours tinkering with gadgets and inventing new ways to enhance the widgets produced in the family's factory.\\n\\n\\nEcho: The second-born, Echo, had a gift for mimicry. He could perfectly replicate any sound or voice he heard, providing entertainment to the entire city. His playful nature and ability to bring joy to those around him made him a favorite among the neighborhood children.\\n\\n\\nSapphire: Sapphire, the third sibling, had scales that shimmered like precious gems. She possessed a unique talent for calming and healing, a trait she inherited from both her parents. Whenever someone in the city felt stressed or unwell, Sapphire would extend her gentle touch, bringing comfort and tranquility.\\n\\n\\nIgnis: The youngest of the family, Ignis, had fiery red scales that hinted at his exceptional ability – the power to control small flames. While initially seen as a potential hazard, Ignis channeled his fiery talents into creating mesmerizing light shows, becoming the city's favorite entertainer during festivals and celebrations.\", metadata={'source': 'Data/Thundertooth Part 2.docx', 'token_count': 327, 'relevance_score': 0.921216}),\n",
       " Document(page_content=\"Grateful for the mayor's hospitality, Thundertooth followed her through the city. Together, they explored the futuristic marketplaces and innovative food labs, eventually discovering a sustainable solution that satisfied the dinosaur's hunger without compromising the well-being of the city's inhabitants.\", metadata={'source': 'Data/Thundertooth Part 1.docx', 'token_count': 54, 'relevance_score': 0.64779824}),\n",
       " Document(page_content=\"Thundertooth and Seraphina reveled in the joy of parenthood, watching their children grow and flourish in the futuristic landscape they now called home. The family became an integral part of the city's fabric, not only through the widgets produced in their factory but also through the positive impact each member had on the community.\", metadata={'source': 'Data/Thundertooth Part 2.docx', 'token_count': 64, 'relevance_score': 0.59549254}),\n",
       " Document(page_content=\"Thundertooth, ever the protector of his newfound home, wasted no time. With a determined gleam in his eyes, he gathered his family and hurried to the city's command center, where Mayor Grace and the leading scientists were coordinating the evacuation efforts.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 51, 'relevance_score': 0.5758266}),\n",
       " Document(page_content=\"In the wake of the crisis, the citizens of the futuristic city hailed Thundertooth and his family as true heroes. The toy factory that once brought joy to children now became a symbol of resilience and unity. The Thundertooth family's legacy was forever etched in the city's history, a testament to the power of cooperation and the extraordinary capabilities that could emerge when dinosaurs and humans worked together for the greater good.\", metadata={'source': 'Data/Thundertooth Part 3.docx', 'token_count': 83, 'relevance_score': 0.5574539})]"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's test that it includes the paragraph starting with \"As the years passed...\" when asked about their children.\n",
    "\n",
    "test_retrieval = compression_retriever.get_relevant_documents(\"Did they have any children? If so, what were their names?\")\n",
    "\n",
    "test_retrieval"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above shows that, indeed, we are able to get that paragraph and it is the highest ranked.\n",
    "\n",
    "Importantly, if we had not brought enough chunks back with the retriever (referring to the vector store retriever) then we would not have had the right chunks to run through Cohere for reranking.\n",
    "\n",
    "So if this line:\n",
    "```\n",
    "retriever = vector.as_retriever(search_kwargs={\"k\": 20})\n",
    "```\n",
    "\n",
    "was:\n",
    "```\n",
    "retriever = vector.as_retriever(search_kwargs={\"k\": 10})\n",
    "```\n",
    "\n",
    "We would not have been able to get that \"As the years passed...\" chunk for reranking.\n",
    "\n",
    "Additionally, we're able to compress the number of chunks from the 11+ we needed to get the right chunk down to 5 because we have the best 5 of that bunch. This reduces the tokens needed for the LLM to process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we create a LangChain chain with the Cohere reranker retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "\n",
    "rerank_chain = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    retriever=compression_retriever,\n",
    "    return_source_documents=True,\n",
    "    chain_type_kwargs={\"prompt\": prompt}    # Pass in our prompt\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here are our test questions\n",
    "\n",
    "TestQuestions = [\n",
    "    \"Summarise the story for me\",\n",
    "    \"Who was the main protagonist?\",\n",
    "    \"Did they have any children? If so, what were their names?\",\n",
    "    \"Did anything eventful happen?\",\n",
    "    \"Who are the main characters?\",\n",
    "    \"What do you think happens next in the story?\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ask our questions with our reranking chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "1/6: Summarise the story for me\n",
      "\n",
      "2/6: Who was the main protagonist?\n",
      "\n",
      "3/6: Did they have any children? If so, what were their names?\n",
      "\n",
      "4/6: Did anything eventful happen?\n",
      "\n",
      "5/6: Who are the main characters?\n",
      "\n",
      "6/6: What do you think happens next in the story?\n"
     ]
    }
   ],
   "source": [
    "qa_pairs = []\n",
    "\n",
    "for index, question in enumerate(TestQuestions, start=1):\n",
    "    question = question.strip() # Clean up\n",
    "\n",
    "    print(f\"\\n{index}/{len(TestQuestions)}: {question}\")\n",
    "\n",
    "    response = rerank_chain.invoke({\"query\": question})\n",
    "\n",
    "    qa_pairs.append((question.strip(), response[\"result\"])) # Add to our output array\n",
    "\n",
    "    # Uncomment the following line if you want to test just the first question\n",
    "    # break "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/6 Summarise the story for me\n",
      "\n",
      "\n",
      "In a futuristic city, a meteor is heading towards Earth, and the mayor recognizes the need for an intelligent and resourceful plan to divert or neutralize it. The mayor approaches Thundertooth, who gathers his team Lumina, Echo, Sapphire, and Ignis to discuss the crisis. Each member has a unique ability that can contribute to the plan: Lumina enhances the city's energy systems to generate a powerful force field; Echo amplifies emergency signals for timely warnings and evacuation instructions; Sapphire calms the panicked masses during evacuation; and Ignis creates controlled bursts of heat to alter the meteor's trajectory.\n",
      "\n",
      "As the years pass, Thundertooth meets Seraphina and starts a family with four children, each with unique characteristics that reflect their modern world. The children inherit their parents' abilities: Lumina is fascinated with technology and invents new ways to enhance gadgets; Echo has a gift for mimicry and brings joy to those around him; Sapphire possesses calming and healing powers; and Ignis controls small flames and creates mesmerizing light shows.\n",
      "\n",
      "After the crisis, the citizens hail Thundertooth and his family as heroes, and their legacy is forever etched in the city's history. The mayor extends an invitation to Thundertooth, offering a solution for him to stay in the city without causing harm to anyone.\n",
      "\n",
      "--------\n",
      "\n",
      "2/6 Who was the main protagonist?\n",
      "\n",
      "\n",
      "The main protagonist of the story is Thundertooth.\n",
      "\n",
      "--------\n",
      "\n",
      "3/6 Did they have any children? If so, what were their names?\n",
      "\n",
      "Yes, Thundertooth had four children. Their names are Lumina, Echo, Sapphire, and Ignis.\n",
      "\n",
      "--------\n",
      "\n",
      "4/6 Did anything eventful happen?\n",
      "\n",
      "\n",
      "Yes, an eventful thing happened when Mayor Grace extended an invitation to Thundertooth to stay in their city and find a solution to his hunger dilemma. They went on to explore the futuristic marketplaces and innovative food labs and eventually discovered a sustainable solution that satisfied Thundertooth's hunger without compromising the well-being of the city's inhabitants.\n",
      "\n",
      "--------\n",
      "\n",
      "5/6 Who are the main characters?\n",
      "\n",
      "The main characters in this story are Mayor Grace and Thundertooth.\n",
      "\n",
      "--------\n",
      "\n",
      "6/6 What do you think happens next in the story?\n",
      "\n",
      "\n",
      "Based on the provided context, it is likely that Thundertooth will work with Mayor Grace and the leading scientists to implement the sustainable solution they discovered to satisfy his hunger without compromising the well-being of the city's inhabitants. He may use his unique abilities to help with the evacuation efforts or in any other way that he can contribute to ensuring the safety of the city and its citizens. In addition, the Thundertooth family may continue to live in the city, enjoying their newfound home and the opportunity to work together towards a better future for all.\n",
      "\n",
      "--------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Print out the questions and answers\n",
    "\n",
    "for index, (question, answer) in enumerate(qa_pairs, start=1):\n",
    "    print(f\"{index}/{len(qa_pairs)} {question}\\n\\n{answer}\\n\\n--------\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "LangChainRAGLinux",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}