A simple RAG demo from AWS re:Invent BOA308 2023

Spells_Embed.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3, json, faiss \n",
    "import numpy as np\n",
    "from jinja2 import Template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrock_runtime_client = boto3.client(\"bedrock-runtime\", region_name=\"us-west-2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Spells\n",
    "- These spells represent a repository of information.  In real systems these would represent a large collection of documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spells = [\n",
    "    \"To turn a frog into a horse, use three upward swishes of the wand while saying neigh.\",\n",
    "    \"To make an apple pie appear, tap the mixing bowl twice while chanting yummy yummy.\",\n",
    "    \"To become invisible, spin in a circle three times while whispering now you see me, no your don't.\",\n",
    "    \"To fly like a bird, flap your arms wildly while shouting tweet tweet flap flap flap.\",\n",
    "    \"To turn your enemy into a newt, point at them and say newt newt you are a newt.\",\n",
    "    \"To heal a papercut, kiss your finger gently and hum a soothing melody.\",\n",
    "    \"To unlock a door, trace a key shape in the air while turning an imaginary key.'\",\n",
    "    \"To turn off the rain, stretch your arms to the sky and yell no more rain today.\",\n",
    "    \"To make a flower bloom, softly blow on it while murmuring grow grow little flower grow.\",\n",
    "    \"To light a candle, snap your fingers twice over the wick and say flame on!\",\n",
    "    \"To get started with any AWS service quickly, open the console and use Amazon Q.\",\n",
    "    \"To become a fish, puff out your cheeks and say bloop bloop bloop I'm a fish.\",\n",
    "    \"To bake a cake, wave your hands over the mixing bowl chanting cake cake cake.\",\n",
    "    \"To turn back time, spin clockwise while counting backward from ten to one.\",\n",
    "    \"To become a giant, stand tall and say fee fi fo fum I'm a giant on the run.\",\n",
    "    \"To talk to animals, place two fingers to your temple and whisper what do you say today?\",\n",
    "    \"To travel through time, close your eyes and count down from five back.\",\n",
    "    \"To make it snow, wiggle your fingers while saying let it snow let it snow let is snow.\",\n",
    "    \"To turn off the lights, clap twice and firmly say lights out everyone lights out.\",\n",
    "    \"To read someone's mind, place your hands on their head and say I know your thoughts.\",\n",
    "    \"To levitate an object, keep your eyes on it and repeat up up up you go!\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create a vector database (Magic Book Shelf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def embed(text):\n",
    "    \"\"\"\n",
    "    Embeds a text (or spell) using the Amazon titan-embed-text-v1 model.\n",
    "    \"\"\"\n",
    "\n",
    "    kwargs = {\n",
    "        \"modelId\": \"amazon.titan-embed-text-v1\",\n",
    "        \"contentType\": \"application/json\",\n",
    "        \"accept\": \"*/*\",\n",
    "        \"body\": json.dumps({\n",
    "            \"inputText\": text\n",
    "        })\n",
    "    }\n",
    "\n",
    "    resp = bedrock_runtime_client.invoke_model(**kwargs)\n",
    "\n",
    "    resp_body = json.loads(resp.get('body').read())\n",
    "    return resp_body.get('embedding')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an array to store the embeddings.\n",
    "spell_embeddings = np.array([]).reshape(0, 1536)\n",
    "\n",
    "# Embed each fact.\n",
    "for spell in spells:\n",
    "    embeddings = embed(spell)\n",
    "    spell_embeddings = np.append(spell_embeddings, np.array(embeddings).reshape(1, -1), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a vector store (this is just an in memory vector store for this sample)\n",
    "magic_book_shelf = faiss.IndexFlatL2(1536)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add out vectors to the vector store. And print the number of vectors in the vector store.\n",
    "magic_book_shelf.add(spell_embeddings)\n",
    "print(magic_book_shelf.ntotal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ask a magical question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set a question to ask.  We also need to get an embedded version of the question. \n",
    "query = \"How do I become a fish?\"\n",
    "embedded_query = embed(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using the vector store, find facts that that are the most similar to our question. Find 4 facts.\n",
    "k = 4                 \n",
    "D, I = magic_book_shelf.search(np.array([embedded_query]), k) \n",
    "print(I) # <- The indexes of the relevant facts\n",
    "print(D) # <- The distances of the relevant facts (or how similar they are)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"Given the spells provided in the following list of spells between the <spell> tags,\n",
    "find the answer to the question written between the <question> tags.\n",
    "\n",
    "<spells>\n",
    "{%- for spell in spells %}\n",
    " - '{{spell}}' {% endfor %} \n",
    "</spells>\n",
    "\n",
    "<question>{{question}}</question>\n",
    "\n",
    "Provide an answer in full including parts of the question, so the answer can be understood alone.\n",
    "If the spells are not relevant to the question being asked, you must not provide any other opinion and you \n",
    "must respond with \"Given the information I have access to, I am unable to answer that question at this time.\"\n",
    "\n",
    "Just provide the answer, nothing else.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use Jinja to fill out our prompt template, adding in all the facts.\n",
    "\n",
    "data = {\n",
    "    'spells': [spells[index] for index in I[0]], \n",
    "    'question': query\n",
    "}\n",
    "\n",
    "template = Template(prompt_template)\n",
    "prompt = template.render(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Send the prompt to a LLM to figure out the answer to our question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kwargs = {\n",
    "    \"modelId\": \"amazon.titan-text-express-v1\",\n",
    "    \"contentType\": \"application/json\",\n",
    "    \"accept\": \"*/*\",\n",
    "    \"body\": json.dumps({\n",
    "        \"inputText\": prompt,\n",
    "        \"textGenerationConfig\": {\n",
    "            \"maxTokenCount\": 4096,\n",
    "            \"stopSequences\": [],\n",
    "            \"temperature\": 0,\n",
    "            \"topP\": 1\n",
    "        }\n",
    "    })\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = bedrock_runtime_client.invoke_model(**kwargs)\n",
    "response_body = json.loads(response.get('body').read())\n",
    "\n",
    "generation = response_body['results'][0]['outputText']\n",
    "print(generation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now try:\n",
    "query = \"How can I get started with any AWS service quickly?\"\n",
    "\n",
    "query = \"How do I become a fish?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python-dev",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}