notebooks/unisim-gemma-text_rag_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# On-Device RAG powered AI agent demo\n",
    "\n",
    "In this notebook we are going to explore how you can build your own on-device LLM-powered AI agent that leverages  RAG (Retrieval-Augmented Generation) to correctly answer questions about the characters of The Wizarding World of Harry Potter. \n",
    "\n",
    "To do this, we are going to combine [Ollama](https://github.com/ollama/ollama) as our local inference engine, [Gemma](https://ai.google.dev/gemma) as our local LLM, our newly released [RETSim](https://arxiv.org/abs/2311.17264) ultra-fast near-duplicate text embeddings, and [USearch](https://github.com/unum-cloud/usearch) for efficient indexing and retrieval. \n",
    "\n",
    "For those who want a more detailed write-up, you can read it at [Wingardium Trivia-osa! On-Device Sorting Hatbot Powered by Gemma, Ollama, USearch, and RETSim](https://elie.net/blog/ai/wingardium-trivia-osa-on-device-sorting-hatbot-powered-by-gemma-ollama-usearch-and-retsim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First things first, we are installing the packages we need for Ollama to run Gemma locally, and UniSim to index data with RETSim and retrieve it with USearch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# installing dependencies\n",
    "!pip install -U tqdm Iprogress unisim ollama tabulate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing libraries\n",
    "import json\n",
    "from tabulate import tabulate\n",
    "from tqdm.auto import tqdm\n",
    "import ollama\n",
    "import unisim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-Flight checks\n",
    "Quickly testing that ollama, Gemma and Unisim are all setup and working well including downloading the latest Gemma version if needed.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gemma 9B loaded\n"
     ]
    }
   ],
   "source": [
    "# Making sure Gemma is installed with Ollama otherwise installing it\n",
    "MODEL = 'gemma'\n",
    "try:\n",
    "    ollama.show(MODEL)\n",
    "except Exception as e:\n",
    "    print(f\"can't find {MODEL}: {e} installing it\")\n",
    "    ollama.pull(MODEL)\n",
    "info = ollama.show(MODEL)\n",
    "print(f\"{MODEL.capitalize()} {info['details']['parameter_size']} loaded\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Hello Elie! 👋 It's lovely to hear from you. How can I help you today? 😊\""
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# small wrapper function to make generation easier and check it all works\n",
    "# we use generate as we are going for a RAG style system so streaming is not useful\n",
    "def generate(prompt: str) -> str:\n",
    "    res = ollama.generate(model=MODEL, prompt=prompt)\n",
    "    if res['done']:\n",
    "        return res['response']\n",
    "    else:\n",
    "        raise Exception(f\"Generation failed: {res['done_reason']}\")\n",
    "\n",
    "generate(\"Hello Gemma it is Elie\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Loading model]\n",
      "|-model_id: text/retsim/v1\n",
      "|-model path: /Users/elieb/git/unisim/unisim/embedder/models/text/retsim/v1.onnx\n",
      "|-input: ['unk__340', 512, 24], tensor(float)\n",
      "|-output: ['unk__341', 256], tensor(float)\n",
      "INFO: UniSim is storing a copy of the indexed data\n",
      "INFO: If you are using large data corpus, consider disabling this behavior using store_data=False\n",
      "INFO: Accelerator is not available, using CPU\n",
      "Similarity 0.9258707761764526 - TextSim works as intended\n"
     ]
    }
   ],
   "source": [
    "# initializizing TextSim for near duplicate text similarity\n",
    "VERBOSE = True  # interactive demo so we want to see what happens\n",
    "txtsim = unisim.TextSim(verbose=True)\n",
    "# check it works as intended\n",
    "sim_value = txtsim.similarity(\"Gemma\", \"Gemmaa\")\n",
    "if sim_value > 0.9:\n",
    "    print(f\"Similarity {sim_value} - TextSim works as intended\")\n",
    "else:\n",
    "    print(f\"Similarity {sim_value} - Something is very wrong with TextSim\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing model without retrivial augmentations\n",
    "\n",
    "Before building RAG let’s evaluate how much Gemma knows about the Wizarding world by asking a few questions with increasing difficulty and let’s throw some typos into the mix to also see how it affects the model performance. I added type next to the question to express what type of test it is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "questions = [\n",
    "             {\"q\":'Which School is Harry Potter part of?', 'type': 'basic fact'},\n",
    "             {\"q\": 'Who is ermionne?', 'type': 'typo'},\n",
    "             {\"q\": 'What is Aberforth job?', 'type': 'harder fact'},\n",
    "             {\"q\": \"what is dubldore job?\", 'type': 'harder fact and typo'},\n",
    "             {\"q\": 'Which school is  Nympadora from?', 'type': 'hard fact'},\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Direct Generation Answers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating direct answers: 100%|██████████| 5/5 [00:28<00:00,  5.74s/it]\n"
     ]
    }
   ],
   "source": [
    "# Let’s run those through Gemma via Ollama and see what type of answer we get.\n",
    "for q in tqdm(questions, desc=\"Generating direct answers\"):\n",
    "    q['direct'] = generate(q['q'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Questions Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[answers without retrival]\n",
      "\n",
      "Q:Which School is Harry Potter part of?? (type: basic fact)\n",
      "Direct answer: Hogwarts School of Witchcraft and Wizardry is the school that Harry Potter attends...\n",
      "\n",
      "Q:Who is ermionne?? (type: typo)\n",
      "Direct answer: Ermionne is a French fashion designer known for her colorful and playful designs, primarily focused ..\n",
      "\n",
      "Q:What is Aberforth job?? (type: harder fact)\n",
      "Direct answer: Aberforth is a fictional character in the Harry Potter series of books and films. He does not have a..\n",
      "\n",
      "Q:what is dubldore job?? (type: harder fact and typo)\n",
      "Direct answer: **Dublador** is a voice actor who provides voices for characters in animated films, television shows..\n",
      "\n",
      "Q:Which school is  Nympadora from?? (type: hard fact)\n",
      "Direct answer: Nympadora is a character from the book series \"Harry Potter\" and did not attend any school. She is a..\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"[answers without retrieval]\\n\")\n",
    "for q in questions:\n",
    "    a =  q['direct'][:100].replace('\\n', ' ')\n",
    "    print(f\"Q:{q['q']}? (type: {q['type']})\")\n",
    "    print(f\"Direct answer: {a}..\")\n",
    "    print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Overall results are not great so let's use UniSim and RetSim to index Harry Potter Data and get better results using the RAG pattern\n",
    "\n",
    "**Note**: Some of those mistakes can probably be reduced and the answers improved by using a better prompt. Feel free to experiment with prompt tuning. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Indexing Harry Potter characters data \n",
    "\n",
    "The first step to build our RAG pipeline to help the LLM with additional context is to load the data, compute the embeddings and index them. We are simply indexing the characters name using RETSim embedding and will return the data associated with it during the retrieval process to help the model.\n",
    "\n",
    "The data used is from the Kaggle [Characters in Harry Potter Books dataset](https://www.kaggle.com/datasets/zez000/characters-in-harry-potter-books)\n",
    "\n",
    "Each character has its name and a few fields. Our game plan is to use \n",
    "unisim/textsim to perform typo-resilient name lookup and retrieve the relevant \n",
    "fields to help Gemma answer about the characters\n",
    "\n",
    "\n",
    "| Field               | Description    | Retrieval Strategy |\n",
    "| :------------------ | :------------- | :------------------|\n",
    "| Name                |  char name     | unisim embedding   |\n",
    "| Descr               |  Char info     | Retrieved          |\n",
    "| Link                |  link to wiki  | Retrieved          |\n",
    "| Gender              |  Char gender   | Retrieved          |\n",
    "| Species/Race        |                | Retrieved          |\n",
    "| Blood School        |  Magic school  | Retrieved          |\n",
    "| Profession          |                | Retrieved          |\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1350 characters loaded from harry_potter_characters.json\n"
     ]
    }
   ],
   "source": [
    "raw_data = json.loads(open('data/harry_potter_characters.json').read())\n",
    "CHARACTERS_INFO = {}  # we are deduping the data using the name as key\n",
    "for d in raw_data:\n",
    "    name = d['Name'].lower().strip()\n",
    "    CHARACTERS_INFO[name] = d\n",
    "print(f'{len(CHARACTERS_INFO)} characters loaded from harry_potter_characters.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['mrs. abbott',\n",
       " 'hannah abbott',\n",
       " 'abel treetops',\n",
       " 'euan abercrombie',\n",
       " 'aberforth dumbledore',\n",
       " 'abernathy',\n",
       " 'abraham peasegood',\n",
       " 'abraham potter',\n",
       " 'abraxas malfoy',\n",
       " 'achilles tolliver']"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# indexing data with TextSim\n",
    "txtsim.reset_index()  # clean up in case we run this cell multiple times\n",
    "idx = txtsim.add(list(CHARACTERS_INFO.keys()))\n",
    "txtsim.indexed_data[:10]  # display what we added"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s quickly write `lookup()` function to make it easier to query the index \n",
    "and test it on one of my favorite characters, Newt Scamander, but with a typo in his name.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query 0: \"new scamramber\"\n",
      "Most similar matches:\n",
      "\n",
      "  idx  is_match      similarity  text\n",
      "-----  ----------  ------------  -----------------------\n",
      " 1005  False               0.81  newt scamander\n",
      " 1006  False               0.71  newt scamander's mother\n",
      " 1172  False               0.63  sam\n",
      "\n",
      "[best lookup result]\n",
      "name: Newt Scamander / School: Hogwarts - Hufflepuff / Profession: Magizoologist\n",
      "Description: Newton “Newt” Scamander is a famous magizoologist and author of Fantastic Beasts and Where To Find Them (PS5) as well as a number of other books. Now retired, Scamander lives in Dorset with his wife Porpentina (FB). He received the Order of Merlin, second… \n"
     ]
    }
   ],
   "source": [
    "# writing a small lookup function wrapper and testing it\n",
    "def lookup(name: str, k: int = 3, threshold: float = 0.9, verbose: bool = False) -> dict:\n",
    "    data = []\n",
    "    name = [name.lower().strip()]\n",
    "    lkp = txtsim.search(name, similarity_threshold=threshold, k=k)\n",
    "    # visualize results for each query using .visualize\n",
    "    res = lkp.results[0]\n",
    "    if verbose:\n",
    "        txtsim.visualize(res)\n",
    "    for m in res.matches:\n",
    "        if m.is_match:\n",
    "            data.append(CHARACTERS_INFO[m.data])\n",
    "\n",
    "    # no match? then let's use the first best result\n",
    "    if not len(data):\n",
    "        data.append(CHARACTERS_INFO[res.matches[0].data])\n",
    "    return data\n",
    "\n",
    "r = lookup(\"New Scamramber\", verbose=True)   # verbose to show all the matches\n",
    "print('')\n",
    "print('[best lookup result]')\n",
    "print(f\"name: {r[0]['Name']} / School: {r[0]['School']} / Profession: {r[0]['Profession']}\")\n",
    "print(f\"Description: {r[0]['Descr']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a RAG with Gemma and UniSim\n",
    "\n",
    "The RAG implementation is going to be in four steps: \n",
    "1. Ask Gemma what is the name of the character so we can look it up. Given that we have access to a powerful LLM, using it to extract the named entity is I think the simplest and more robust way to do so\n",
    "\n",
    "2. Retrieve the nearest match info from our UniSim index\n",
    "\n",
    "3. Replace the name in the user query with the looked up name to fix the typo, which is very important and often overlooked, and then inject in the query the information we retrieve\n",
    "\n",
    "4. Answer the user’s question and impress them with our extensive knowledge of the Wizarding world of Harry Potter!\n",
    "\n",
    "This translates to this simple code, with helper functions defined earlier and available in the colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rag(prompt: str, k: int = 5, threshold: float = 0.9, verbose: bool = False) -> str:\n",
    "    # normalizing the prompt\n",
    "    prompt = prompt.lower().strip()\n",
    "\n",
    "    # ask Gemma who is the character\n",
    "    char_prompt = f\"In the following sentence: '{prompt}' who is the subject? reply only with name.\"\n",
    "    if verbose:\n",
    "        print(f\"Char prompt: {char_prompt}\")\n",
    "    character = generate(char_prompt)\n",
    "    if verbose:\n",
    "        print(f\"Character: '{character}'\")\n",
    "\n",
    "    # lookup the character\n",
    "    data = lookup(character, k=k, threshold=threshold, verbose=verbose)\n",
    "\n",
    "    # augmented prompt\n",
    "    # replace the name in the prompt with the one in the rag\n",
    "    prompt = prompt.replace(character.lower().strip(), data[0]['Name'].lower().strip())\n",
    "\n",
    "    aug_prompt = f\"Using the following data: {data} answer the following question: '{prompt}'. Don't mention your sources - just the answer.\"\n",
    "\n",
    "    if verbose:\n",
    "        print(f\"Augmented prompt: {aug_prompt}\")\n",
    "    response = generate(aug_prompt)\n",
    "\n",
    "    return response\n",
    "# rag(questions[-1]['q'], verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RAG answers vs direct answer generation\n",
    "Lets see our RAG in action and compare it to the directly generated answers we got before.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating rag answers: 100%|██████████| 5/5 [00:28<00:00,  5.69s/it]\n"
     ]
    }
   ],
   "source": [
    "# generate direct answers using Gemma\n",
    "for q in tqdm(questions, desc=\"Generating rag answers\"):\n",
    "    q['direct'] = generate(q['q'])  # redoing in case the questoins were changed\n",
    "    q['rag'] = rag(q['q'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Here are the answers we got from Gemma with retrieval and how they compare to the direct generation one.\n",
    " Overall we see that the RAG powered answers are more precise (Harry is correctly from Hogwarts - Gryffindor),\n",
    " have the right context (E.g Albus Dumbledor job) and answers very precise questions (Nympadora Tonks school)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Answers with RAG vs without]\n",
      "\n",
      "Q:Which School is Harry Potter part of?? (type: basic fact)\n",
      "Direct answer: Hogwarts School of Witchcraft and Wizardry is the school that Harry Potter attends...\n",
      "RAG anwer: Harry Potter is part of Hogwarts - Gryffindor.\n",
      "\n",
      "Q:Who is ermionne?? (type: typo)\n",
      "Direct answer: Ermionne is a fictional character from Greek mythology, most famously known for her role in the stor..\n",
      "RAG anwer: Hermione Granger is a resourceful, principled, and brilliant witch known for her academic prowess an\n",
      "\n",
      "Q:What is Aberforth job?? (type: harder fact)\n",
      "Direct answer: Aberforth is a fictional character in the Harry Potter series of books and films. He does not have a..\n",
      "RAG anwer: Aberforth was a barman.\n",
      "\n",
      "Q:what is dubldore job?? (type: harder fact and typo)\n",
      "Direct answer: **Dublador** is a voice-over artist who provides voices for characters, narration, or other elements..\n",
      "RAG anwer: Headmaster at Hogwarts School.\n",
      "\n",
      "Q:Which school is  Nympadora from?? (type: hard fact)\n",
      "Direct answer: Nympadora is a character from the Harry Potter series of books and films, and does not attend any sc..\n",
      "RAG anwer: Hogwarts - Hufflepuff\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"[Answers with RAG vs without]\\n\")\n",
    "for q in questions:\n",
    "    a =  q['direct'][:100].replace('\\n', ' ')\n",
    "    print(f\"Q:{q['q']}? (type: {q['type']})\")\n",
    "    print(f\"Direct answer: {a}..\")\n",
    "    r =  q['rag'][:100].replace('\\n', ' ')\n",
    "    print(f'RAG answer: {r}')\n",
    "    print(\"\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}