03-RAG-local-llm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4c6d44f7-b10f-4d06-92e2-e889445eed75",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "<img src=\"images/ragna-logo.png\" width=\"200px\" align=\"right\"/>\n",
    "\n",
    "# Use Local LLM with Ragna\n",
    "\n",
    "<hr>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e1610ef-6fa7-438e-a76c-c89cfb8b59ec",
   "metadata": {},
   "source": [
    "## Create a new Ragna assistant"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6e2c8b6-ab62-472f-98ce-24bf63a7fc45",
   "metadata": {},
   "source": [
    "To use a local LLM in Ragna, we have to subclass the `ragna.core.Assistant` abstract base class. The only abstract method that we have to overwrite is [`.answer()`](https://ragna.chat/en/stable/references/python-api/#ragna.core.Assistant.answer). It gets passed the `prompt` of the user as well as the `sources` retrieved from the source storage. In there we combine these two parts of information into one large prompt for the LLM, start the generation, and `yield` back the individual chunks.\n",
    "\n",
    "<details>\n",
    "<summary> <b>Expand to read <code>local_llm.py</code> → </b></summary>\n",
    "\n",
    "```python\n",
    "from pathlib import Path\n",
    "from typing import Iterator\n",
    "\n",
    "from ragna.core import Assistant, PackageRequirement, Source\n",
    "\n",
    "class Llama38BInstruct(Assistant):\n",
    "    @classmethod\n",
    "    def display_name(cls):\n",
    "        return \"turboderp/Llama-3-8B-Instruct-exl2\"\n",
    "\n",
    "    @classmethod\n",
    "    def requirements(cls):\n",
    "        return [\n",
    "            PackageRequirement(\"torch\"),\n",
    "            PackageRequirement(\"exllamav2\"),\n",
    "        ]\n",
    "\n",
    "    @classmethod\n",
    "    def is_available(cls):\n",
    "        requirements_available = super().is_available()\n",
    "        if not requirements_available:\n",
    "            return False\n",
    "\n",
    "        import torch\n",
    "\n",
    "        return torch.cuda.is_available()\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        from exllamav2 import (\n",
    "            ExLlamaV2,\n",
    "            ExLlamaV2Cache,\n",
    "            ExLlamaV2Config,\n",
    "            ExLlamaV2Tokenizer,\n",
    "        )\n",
    "        from exllamav2.generator import ExLlamaV2Sampler, ExLlamaV2StreamingGenerator\n",
    "\n",
    "        config = ExLlamaV2Config()\n",
    "        config.model_dir = str(Path.home() / \"shared/pycon/models\" / self.display_name())\n",
    "        config.prepare()\n",
    "\n",
    "        self.tokenizer = ExLlamaV2Tokenizer(config)\n",
    "\n",
    "        model = ExLlamaV2(config)\n",
    "        cache = ExLlamaV2Cache(model, lazy=True)\n",
    "        model.load_autosplit(cache)\n",
    "        self.generator = ExLlamaV2StreamingGenerator(model, cache, self.tokenizer)\n",
    "        self.generator.set_stop_conditions({self.tokenizer.eos_token_id, 78191})\n",
    "\n",
    "        self.settings = ExLlamaV2Sampler.Settings()\n",
    "        self.settings.temperature = 0.0\n",
    "\n",
    "    def _make_prompt(self, prompt: str, sources: list[Source]) -> str:\n",
    "        return \"\\n\".join(\n",
    "            [\n",
    "                f\"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\",\n",
    "                f\"\",\n",
    "                f\"Answer the question based only on the following context:\",\n",
    "                *[source.content for source in sources],\n",
    "                f\"<|eot_id|><|start_header_id|>user<|end_header_id|>\",\n",
    "                f\"\",\n",
    "                f\"{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\",\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def answer(\n",
    "        self, prompt: str, sources: list[Source], *, max_new_tokens: int = 256\n",
    "    ) -> Iterator[str]:\n",
    "        input_ids = self.tokenizer.encode(\n",
    "            self._make_prompt(prompt, sources), add_bos=False\n",
    "        )\n",
    "\n",
    "        self.generator.begin_stream_ex(input_ids, self.settings)\n",
    "\n",
    "        for _ in range(max_new_tokens):\n",
    "            result = self.generator.stream_ex()\n",
    "            if result[\"eos\"]:\n",
    "                break\n",
    "            yield result[\"chunk\"]\n",
    "```\n",
    "\n",
    "</details>\n",
    "\n",
    "<br>\n",
    "\n",
    "☝🏼 **Note:** Apart from the prompt generation in `._make_prompt()`, the code is actually model agnostic and can be used with any LLM with `exl2` weights."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef967823-bc69-4a06-8e72-9b8713fbaf58",
   "metadata": {},
   "source": [
    "## Use the assistant\n",
    "\n",
    "You can directly import the class and start using Llama3-8B."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65b07a54-12b0-46b8-8bc0-0531978d179a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from local_llm import Llama38BInstruct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d79a6139-9bb3-4d0d-b338-e9584927b8fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "Llama38BInstruct.display_name()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24974543-918e-471c-95b7-3af6d4770adc",
   "metadata": {},
   "outputs": [],
   "source": [
    "Llama38BInstruct.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3873ab13-6f23-476f-89c0-50549169a97d",
   "metadata": {},
   "source": [
    "Let's share the Python Software Foundation's annual reports, and ask questions about the PSF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84700998-f92e-4207-aa68-77bdd0c8f82a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ragna import Rag, source_storages\n",
    "\n",
    "documents = [\n",
    "    \"files/psf-report-2021.pdf\",\n",
    "    \"files/psf-report-2022.pdf\",\n",
    "    \"files/psf-report-2023.pdf\",\n",
    "]\n",
    "\n",
    "chat = Rag().chat(\n",
    "    documents=documents,\n",
    "    source_storage=source_storages.Chroma,\n",
    "    assistant=Llama38BInstruct,\n",
    ")\n",
    "\n",
    "await chat.prepare();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a719c026-9da5-4d87-9234-7a08bf0171f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "message = await chat.answer(\"What does the PSF do?\", stream=True)\n",
    "\n",
    "async for chunk in message:\n",
    "    print(chunk, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6323cfb8-b26a-45a6-a809-bd8b63ab4a82",
   "metadata": {},
   "source": [
    "Learn more about custom components: https://ragna.chat/en/latest/generated/tutorials/gallery_custom_components/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe59b269-902c-42af-aa80-082d99b1f6f6",
   "metadata": {},
   "source": [
    "## Your turn ✨\n",
    "\n",
    "Query the reports with GPT4 & Llama3 to find out PSF's net income over the years. We'll come back to this at the end of the tutorial. :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4af5e264-ce3b-4ad3-83bb-15f209b52248",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat = rag.chat(\n",
    "    ...\n",
    ")\n",
    "\n",
    "await ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bf387fd-e34d-4411-9b47-288d7db85aef",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = ...\n",
    "\n",
    "print(f\"\\nLLM Response: \\n\\n{answer.content}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c7a2e43-6d6a-46b9-a87d-9b00901d27f9",
   "metadata": {},
   "source": [
    "<hr>\n",
    "\n",
    "_❗️ **Warning:** Make sure to stop the Jupyter Kernel (in the JupyterLab Menu Bar, click on \"Kernel\" -> \"Shut down Kernel\") before proceeding to prevent the \"insufficient VRAM\" error._\n",
    "\n",
    "<br>\n",
    "\n",
    "**✨ Next: [Building Chat Applications with Panel](04-panel-intro.ipynb) →**\n",
    "\n",
    "\n",
    "💬 _Wish to continue discussions after the tutorial? Contact the presenters: [@pavithraes](https://github.com/pavithraes), [@dharhas](https://github.com/dharhas), [@ahuang11](https://github.com/ahuang11)_\n",
    "\n",
    "<hr>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09a172d1-e920-4af1-930b-1143b58fd2d5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pycon-pycon-ragna",
   "language": "python",
   "name": "conda-env-pycon-pycon-ragna-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}