examples/rag_with_quantized_llm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6b1aeaea-43d9-476b-8fce-5c4b55707874",
   "metadata": {},
   "source": [
    "# Run Quantized LLMs on CPUs\n",
    "\n",
    "In this tutorial, we will demo how to build a RAG pipeline running on CPUs with [optimum-intel](https://github.com/huggingface/optimum-intel), using an LLM.\n",
    "\n",
    "We will use a pipeline that will:\n",
    "\n",
    "- Quantize an LLM\n",
    "- Fetch relevant documents for our question.\n",
    "- Rerank the documents for better performance.\n",
    "- Run the LLM on CPU to answer the question.\n",
    "\n",
    "For more information about optimum-intel, we refer to the [original repository](https://github.com/huggingface/optimum-intel)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a12b4ef-6626-4c2f-af4b-8023f93678b6",
   "metadata": {},
   "source": [
    "First, we export the model to an ONNX file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d92d7e09-c09e-484a-a569-f91339f9875f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'facebook/opt-iml-max-1.3b'\n",
    "\n",
    "converted_model_path = f\"/tmp/{model_name.replace('/','_')}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "96a05063-37da-49dd-ae65-9662bd90c4fd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "69932ca4bf3e4f14b45aeb6982eb91c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/596 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Framework not specified. Using pt to export the model.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7988c897270d4be3a31bbec056f79b34",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/2.63G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b488cabeb7a04591b27c3431d01be0c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/682 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e6ab8c78ec9d46aebfa97284bfe2c896",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.json:   0%|          | 0.00/899k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9721d31688e54472b5cc5eca5d04a646",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76b5a1da69cf4ab58e94ad326111c0a2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/221 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using the export variant default. Available variants are:\n",
      "    - default: The default ONNX variant.\n",
      "Using framework PyTorch: 2.2.1+cu121\n",
      "Overriding 1 configuration item(s)\n",
      "\t- use_cache -> True\n",
      "  if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):\n",
      "Saving external data to one file...\n",
      "Post-processing the exported models...\n",
      "Deduplicating shared (tied) weights...\n",
      "Found different candidate ONNX initializers (likely duplicate) for the tied weights:\n",
      "\tlm_head.weight: {'onnx::MatMul_5598'}\n",
      "\tmodel.decoder.embed_tokens.weight: {'model.decoder.embed_tokens.weight'}\n",
      "Removing duplicate initializer onnx::MatMul_5598...\n"
     ]
    }
   ],
   "source": [
    "from optimum.onnxruntime import ORTModelForCausalLM\n",
    "\n",
    "model = ORTModelForCausalLM.from_pretrained(model_name, export=True, trust_remote_code=True)\n",
    "model.save_pretrained(converted_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cea3cd81-5080-4796-9f4b-d9c6f6172913",
   "metadata": {},
   "source": [
    "Then, we load the exported model back, and quantize it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6326e749-49df-4d16-8be1-fb5f3af4b603",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Creating dynamic quantizer: QOperator (mode: IntegerOps, schema: u8/u8, channel-wise: True)\n",
      "Quantizing model...\n",
      "Saving quantized model at: /tmp/facebook_opt-iml-max-1.3b/quantized (external data format: False)\n",
      "Configuration saved in /tmp/facebook_opt-iml-max-1.3b/quantized/ort_config.json\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PosixPath('/tmp/facebook_opt-iml-max-1.3b/quantized')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from optimum.onnxruntime.configuration import AutoQuantizationConfig\n",
    "from optimum.onnxruntime import ORTQuantizer\n",
    "import os\n",
    "\n",
    "model = ORTModelForCausalLM.from_pretrained(converted_model_path)\n",
    "\n",
    "qconfig = AutoQuantizationConfig.avx2(is_static=False)\n",
    "quantizer = ORTQuantizer.from_pretrained(model)\n",
    "\n",
    "quantizer.quantize(save_dir=os.path.join(converted_model_path, 'quantized'), quantization_config=qconfig)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fb51e32-c066-418c-a021-73a3e3f13f1f",
   "metadata": {},
   "source": [
    "Now that our model is quantized, we can create a RAG pipeline with it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c25d1abb-1e57-4290-8a6e-43aa20ceca17",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from haystack import Pipeline, Document\n",
    "from haystack.document_stores.in_memory import InMemoryDocumentStore\n",
    "from haystack.components.retrievers.in_memory import InMemoryBM25Retriever\n",
    "from haystack.components.builders.answer_builder import AnswerBuilder\n",
    "from haystack.components.builders.prompt_builder import PromptBuilder\n",
    "from haystack.components.rankers import TransformersSimilarityRanker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d1cfe9c-7861-4ed8-836e-689a612faf15",
   "metadata": {},
   "source": [
    "We start from a collection of paragraphs from Wikipedia, for the retrieval phase:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "23442f1b-ede3-4ef8-9768-d0e598315cc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "document_collection = [{'id': '11457596',\n",
    "  'text': 'Quest\", the \"Ultima\" series, \"EverQuest\", the \"Warcraft\" series, and the \"Elder Scrolls\" series of games as well as video games set in Middle-earth itself. Research also suggests that some consumers of fantasy games derive their motivation from trying to create an epic fantasy narrative which is influenced by \"The Lord of the Rings\". In 1965, songwriter Donald Swann, who was best known for his collaboration with Michael Flanders as Flanders & Swann, set six poems from \"The Lord of the Rings\" and one from \"The Adventures of Tom Bombadil\" (\"Errantry\") to music. When Swann met with Tolkien to play the',\n",
    "  'title': 'The Lord of the Rings'},\n",
    " {'id': '11457582',\n",
    "  'text': 'helped \"The Lord of the Rings\" become immensely popular in the United States in the 1960s. The book has remained so ever since, ranking as one of the most popular works of fiction of the twentieth century, judged by both sales and reader surveys. In the 2003 \"Big Read\" survey conducted in Britain by the BBC, \"The Lord of the Rings\" was found to be the \"Nation\\'s best-loved book\". In similar 2004 polls both Germany and Australia also found \"The Lord of the Rings\" to be their favourite book. In a 1999 poll of Amazon.com customers, \"The Lord of the',\n",
    "  'title': 'The Lord of the Rings'},\n",
    " {'id': '11457540',\n",
    "  'text': 'of Tolkien\\'s works is such that the use of the words \"Tolkienian\" and \"Tolkienesque\" has been recorded in the \"Oxford English Dictionary\". The enduring popularity of \"The Lord of the Rings\" has led to numerous references in popular culture, the founding of many societies by fans of Tolkien\\'s works, and the publication of many books about Tolkien and his works. \"The Lord of the Rings\" has inspired, and continues to inspire, artwork, music, films and television, video games, board games, and subsequent literature. Award-winning adaptations of \"The Lord of the Rings\" have been made for radio, theatre, and film. In',\n",
    "  'title': 'The Lord of the Rings'},\n",
    " {'id': '11457587',\n",
    "  'text': 'has been read as fitting the model of Joseph Campbell\\'s \"monomyth\". \"The Lord of the Rings\" has been adapted for film, radio and stage. The book has been adapted for radio four times. In 1955 and 1956, the BBC broadcast \"The Lord of the Rings\", a 13-part radio adaptation of the story. In the 1960s radio station WBAI produced a short radio adaptation. A 1979 dramatization of \"The Lord of the Rings\" was broadcast in the United States and subsequently issued on tape and CD. In 1981, the BBC broadcast \"The Lord of the Rings\", a new dramatization in 26',\n",
    "  'title': 'The Lord of the Rings'},\n",
    " {'id': '11457592',\n",
    "  'text': '\"The Lord of the Rings\", was released on the internet in May 2009 and has been covered in major media. \"Born of Hope\", written by Paula DiSante, directed by Kate Madison, and released in December 2009, is a fan film based upon the appendices of \"The Lord of the Rings\". In November 2017, Amazon acquired the global television rights to \"The Lord of the Rings\", committing to a multi-season television series. The series will not be a direct adaptation of the books, but will instead introduce new stories that are set before \"The Fellowship of the Ring\". Amazon said the',\n",
    "  'title': 'The Lord of the Rings'},\n",
    " {'id': '7733817',\n",
    "  'text': 'The Lord of the Rings Online The Lord of the Rings Online: Shadows of Angmar is a massive multiplayer online role-playing game (MMORPG) for Microsoft Windows and OS X set in a fantasy universe based upon J. R. R. Tolkien\\'s Middle-earth writings, taking place during the time period of \"The Lord of the Rings\". It launched in North America, Australia, Japan, and Europe in 2007. Originally subscription-based, it is free-to-play, with a paid VIP subscription available that provides players various perks.  The game\\'s environment is based on \"The Lord of the Rings\" and \"The Hobbit\". However, Turbine does not',\n",
    "  'title': 'The Lord of the Rings Online'},\n",
    " {'id': '22198847',\n",
    "  'text': 'of \"The Lord of the Rings\", including Ian McKellen, Andy Serkis, Hugo Weaving, Elijah Wood, Ian Holm, Christopher Lee, Cate Blanchett and Orlando Bloom who reprised their roles. Although the \"Hobbit\" films were even more commercially successful than \"The Lord of the Rings\", they received mixed reviews from critics. Numerous video games were released to supplement the film series. They include: \",\" Pinball, \"\", \"\", , \"\", \"\", \"\", \"\", \"The Lord of the Rings Online\", \"\", \"\", \"\", \"Lego The Lord of the Rings\", \"Guardians of Middle-earth\", \"\", and \"\".',\n",
    "  'title': 'The Lord of the Rings (film series)'},\n",
    " {'id': '24071573',\n",
    "  'text': 'Lord of the Rings (musical) The Lord of the Rings is the most prominent of several theatre adaptations of J. R. R. Tolkien\\'s epic high fantasy novel of the same name, with music by A. R. Rahman, Christopher Nightingale and the band Värttinä, and book and lyrics by Matthew Warchus and Shaun McKenna. Set in the world of Middle-earth, \"The Lord of the Rings\" tells the tale of a humble hobbit who is asked to play the hero and undertake a treacherous mission to destroy an evil, magic ring without being seduced by its power. The show was first performed',\n",
    "  'title': 'Lord of the Rings (musical)'},\n",
    " {'id': '11457536',\n",
    "  'text': 'The Lord of the Rings The Lord of the Rings is an epic high fantasy novel written by English author and scholar J. R. R. Tolkien. The story began as a sequel to Tolkien\\'s 1937 fantasy novel \"The Hobbit\", but eventually developed into a much larger work. Written in stages between 1937 and 1949, \"The Lord of the Rings\" is one of the best-selling novels ever written, with over 150 million copies sold. The title of the novel refers to the story\\'s main antagonist, the Dark Lord Sauron, who had in an earlier age created the One Ring to rule',\n",
    "  'title': 'The Lord of the Rings'},\n",
    " {'id': '13304003',\n",
    "  'text': \"The Lord of the Rings (disambiguation) The Lord of the Rings is a fantasy novel by J. R. R. Tolkien. The title refers to Sauron, the story's main antagonist. The Lord of the Rings may also refer to:\",\n",
    "  'title': 'The Lord of the Rings (disambiguation)'}]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d381c2c-a3ca-4bb8-b003-481d6aed5148",
   "metadata": {},
   "source": [
    "We then create an InMemoryDocumentStore document store, to store all the documents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d8f3e6df-3efa-41c7-a096-11f90b4b9c44",
   "metadata": {},
   "outputs": [],
   "source": [
    "store = InMemoryDocumentStore()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "99700ae4-fd3f-4d16-ad7c-8cd746c9999d",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = [Document(id=item[\"id\"], content=item[\"text\"], meta={\"title\": item[\"title\"]}) for item in document_collection]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a2d901a6-a350-45a9-b69e-ab44a4dad9f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "store.write_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beb6fb8f-eb34-421c-9ceb-32f8df00ec6b",
   "metadata": {},
   "source": [
    "Next, we create a simple BM25 retriever on top of our store, and an additional reranker component to improve the ranking of the documents used for answering the question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "702ce4e4-6a26-473f-9804-91136ddb60e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = InMemoryBM25Retriever(document_store=store)\n",
    " \n",
    "ranker = TransformersSimilarityRanker()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5334ce5-48ff-4f09-804e-45bf96e8c2b2",
   "metadata": {},
   "source": [
    "Now that we have created the retrieval components, we move to the LLM usage.\n",
    "\n",
    "We create a document template that contains a placeholder for the retrieved documents to be inserted:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "343eae7e-9c83-4d0e-9b60-706c573466be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build a RAG pipeline\n",
    "prompt_template = \"\"\"\n",
    "Given these documents, answer the question.\n",
    "Documents:\n",
    "{% for doc in documents %}\n",
    "    {{ doc.content }}\n",
    "{% endfor %}\n",
    "Question: {{query}}\n",
    "Answer:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a868eb36-c972-4343-aa12-0d5bee9fb24b",
   "metadata": {},
   "source": [
    "Next, we use the ```OpenVINOGenerator```, to load the LLM as a quantized ONNX model on our CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5706b86d-4393-466e-a578-d773683bc626",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastrag.generators.ort import ORTGenerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "37016502-d9f3-4084-8e8f-aac90f7a0831",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "44d8cb98-ad91-4f37-b5f8-ee190ec81735",
   "metadata": {},
   "outputs": [],
   "source": [
    "converted_model_path = \"/tmp/facebook_opt-iml-max-1.3b\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2d3c2e4c-6871-4d50-8aec-0d93627575f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = ORTGenerator(\n",
    "    model=os.path.join(converted_model_path, 'quantized'),\n",
    "    task=\"text-generation\",\n",
    "    generation_kwargs={\n",
    "        \"max_new_tokens\": 100,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e9a00ea-e59b-4492-b175-eaef76f33c07",
   "metadata": {},
   "source": [
    "With the model and the prompt template now ready, we create a PromptNode to unify both modules:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e03a9423-fcaa-4c2a-b677-f32d79f642ce",
   "metadata": {},
   "source": [
    "Our components are now ready. We can now create a pipeline, to connect all of them together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7d4a517e-9fa6-407f-9527-e372e1623c5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pipeline created\n"
     ]
    }
   ],
   "source": [
    "pipe = Pipeline()\n",
    "\n",
    "pipe.add_component(\"retriever\", retriever)\n",
    "pipe.add_component(\"ranker\", ranker)\n",
    "pipe.add_component(\"prompt_builder\", PromptBuilder(template=prompt_template))\n",
    "pipe.add_component(\"llm\", generator)\n",
    "\n",
    "pipe.connect(\"retriever.documents\", \"ranker.documents\")\n",
    "pipe.connect(\"ranker\", \"prompt_builder.documents\")\n",
    "pipe.connect(\"prompt_builder\", \"llm\")\n",
    "print(\"Pipeline created\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae3804a9-a067-4033-bded-3a76c7294be1",
   "metadata": {},
   "source": [
    "Finally, lets ask it a question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3c0206f4-adf6-4c0d-bd7b-be5fa9851f39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bde189bee4184e4e8ba156874b5f8754",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Ranking by BM25...:   0%|          | 0/10 [00:00<?, ? docs/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query = \"Who is the main villan in Lord of the Rings?\"\n",
    "answer_result = pipe.run({\n",
    "    \"prompt_builder\": {\n",
    "        \"query\": query\n",
    "    },\n",
    "    \"retriever\": {\n",
    "        \"query\": query\n",
    "    },\n",
    "    \"ranker\": {\n",
    "        \"query\": query,\n",
    "        \"top_k\": 1\n",
    "    }\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c0cfab76-495d-471d-99a5-5aef4ac31e28",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' Sauron'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "answer_result[\"llm\"][\"replies\"][0]"
   ]
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}