LLM_RAG_CHAINLIT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyMR2MECsNzmJ4SoJvWJB7hj",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/phongloihong/RAG_LLM_BASIC/blob/main/LLM_RAG_CHAINLIT.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Piqk-Zd-PANe",
        "outputId": "a01ccb68-5bf0-4c56-fc06-6f7deb0e4ece"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Object \"install\" is unknown, try \"ip help\".\n",
            "\u001b[K\u001b[?25h/tools/node/bin/lt -> /tools/node/lib/node_modules/localtunnel/bin/lt.js\n",
            "+ localtunnel@2.0.2\n",
            "updated 1 package in 0.818s\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "chainlit 1.1.304 requires numpy<2.0,>=1.26; python_version >= \"3.9\", but you have numpy 1.24.4 which is incompatible.\n",
            "pandas-stubs 2.0.3.230814 requires numpy>=1.25.0; python_version >= \"3.9\", but you have numpy 1.24.4 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "!pip install -q transformers==4.41.2\n",
        "!pip install -q bitsandbytes==0.43.1\n",
        "!ip install -q accelerate==0.31.0\n",
        "!pip install -q langchain==0.2.5\n",
        "!pip install -q langchainhub==0.1.20\n",
        "!pip install -q langchain-chroma==0.1.1\n",
        "!pip install -q langchain-community==0.2.5\n",
        "!pip install -q langchain-openai==0.1.9\n",
        "!pip install -q langchain_huggingface==0.0.3\n",
        "!pip install -q chainlit==1.1.304\n",
        "!pip install -q python-dotenv==1.0.1\n",
        "!pip install -q pypdf==4.2.0\n",
        "!npm install -g localtunnel\n",
        "!pip install -q numpy==1.24.4\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Init on_chat_start function"
      ],
      "metadata": {
        "id": "EgkiAjjUSdkT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile app.py\n",
        "import chainlit as cl\n",
        "import torch\n",
        "\n",
        "from chainlit.types import AskFileResponse\n",
        "\n",
        "\n",
        "from transformers import BitsAndBytesConfig\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline\n",
        "from langchain_huggingface.llms import HuggingFacePipeline\n",
        "\n",
        "from langchain.memory import ConversationBufferMemory\n",
        "from langchain_community.chat_message_histories import ChatMessageHistory\n",
        "from langchain.chains import ConversationalRetrievalChain\n",
        "\n",
        "from langchain_huggingface import HuggingFaceEmbeddings\n",
        "from langchain_chroma import Chroma\n",
        "from langchain_community.document_loaders import PyPDFLoader, TextLoader\n",
        "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
        "from langchain_core.runnables import RunnablePassthrough\n",
        "from langchain_core.output_parsers import StrOutputParser\n",
        "from langchain import hub\n",
        "\n",
        "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)\n",
        "embedding = HuggingFaceEmbeddings()\n",
        "\n",
        "def process_file(file: AskFileResponse):\n",
        "  if file.type == \"text/plain\":\n",
        "    Loader = TextLoader\n",
        "  elif file.type == \"application/pdf\":\n",
        "    Loader = PyPDFLoader\n",
        "\n",
        "  loader = Loader(file.path)\n",
        "  documents = loader.load()\n",
        "  docs = text_splitter.split_documents(documents)\n",
        "  for i, doc in enumerate(docs):\n",
        "    doc.metadata[\"source\"] = f\"source_{i}\"\n",
        "\n",
        "  return docs\n",
        "\n",
        "def get_vector_db(file: AskFileResponse):\n",
        "  docs = process_file(file)\n",
        "  cl.user_session.set(\"docs\", docs)\n",
        "  vector_db = Chroma.from_documents(documents=docs, embedding=embedding)\n",
        "  return vector_db\n",
        "\n",
        "def get_huggingface_llm(model_name: str = \"lmsys/vicuna-7b-v1.5\", max_new_token: int = 512):\n",
        "  nf4_config = BitsAndBytesConfig(\n",
        "      load_in_4bit=True,\n",
        "      bnb_4bit_quant_type=\"nf4\",\n",
        "      bnb_4bit_use_double_quant=True,\n",
        "      bnb_4bit_compute_dtype=torch.bfloat16\n",
        "  )\n",
        "\n",
        "  model = AutoModelForCausalLM.from_pretrained(\n",
        "      model_name,\n",
        "      quantization_config=nf4_config,\n",
        "      low_cpu_mem_usage=True\n",
        "  )\n",
        "\n",
        "  tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "  model_pipeline = pipeline(\n",
        "      \"text-generation\",\n",
        "      model=model,\n",
        "      tokenizer=tokenizer,\n",
        "      max_new_tokens=max_new_token,\n",
        "      pad_token_id=tokenizer.eos_token_id,\n",
        "      device_map=\"auto\"\n",
        "  )\n",
        "\n",
        "  llm = HuggingFacePipeline(\n",
        "      pipeline=model_pipeline,\n",
        "  )\n",
        "  return llm\n",
        "\n",
        "LLM = get_huggingface_llm()\n",
        "\n",
        "\n",
        "welcome_message = \"\"\"Welcome to the PDF QA! To get started:\n",
        "1. Upload a PDF or text file\n",
        "2. Ask a question about the file\n",
        "\"\"\"\n",
        "\n",
        "@cl.on_chat_start\n",
        "async def on_chat_start():\n",
        "  files = None\n",
        "  while files is None:\n",
        "    files = await cl.AskFileMessage(\n",
        "        content=welcome_message,\n",
        "        accept=[\"text/plain\", \"application/pdf\"],\n",
        "        max_size_mb=20,\n",
        "        timeout=180,\n",
        "    ).send()\n",
        "  file = files[0]\n",
        "\n",
        "  msg = cl.Message(content=f\"Processing {file.name}...\", disable_feedback=True)\n",
        "  await msg.send()\n",
        "\n",
        "  vector_db = await cl.make_async(get_vector_db)(file)\n",
        "\n",
        "  message_history = ChatMessageHistory()\n",
        "  memory = ConversationBufferMemory(\n",
        "      memory_key=\"chat_history\",\n",
        "      output_key=\"answer\",\n",
        "      chat_memory=message_history,\n",
        "      return_messages=True,\n",
        "  )\n",
        "  retriever = vector_db.as_retriever(search_type=\"mmr\", search_kwargs={\"k\": 3})\n",
        "  chain = ConversationalRetrievalChain.from_llm(\n",
        "      llm=LLM,\n",
        "      chain_type=\"stuff\",\n",
        "      retriever=retriever,\n",
        "      memory=memory,\n",
        "      return_source_documents=True\n",
        "  )\n",
        "\n",
        "  msg.content = f\"'{file.name}' processed. You can now ask questions !\"\n",
        "  await msg.update()\n",
        "\n",
        "  cl.user_session.set(\"chain\", chain)\n",
        "\n",
        "@cl.on_message\n",
        "async def on_message(message: cl.Message):\n",
        "  chain = cl.user_session.get(\"chain\")\n",
        "  cb = cl.AsyncLangchainCallbackHandler()\n",
        "  res = await chain.ainvoke(message.content, callbacks=[cb])\n",
        "  answer = res[\"answer\"]\n",
        "  source_documents = res[\"source_documents\"]\n",
        "  text_elements = []\n",
        "\n",
        "  if source_documents:\n",
        "    for source_idx, source_doc in enumerate(source_documents):\n",
        "      source_name = f\"source_{source_idx}\"\n",
        "      text_elements.append(\n",
        "          cl.Text(content=source_doc.page_content, name=source_name)\n",
        "      )\n",
        "\n",
        "    source_names = [text_el.name for text_el in text_elements]\n",
        "\n",
        "    if source_names:\n",
        "      answer += f\"\\nSources: {', '.join(source_names)}\"\n",
        "    else:\n",
        "      answer += \"\\nNo sources found\"\n",
        "\n",
        "  await cl.Message(content=answer, elements=text_elements).send()\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "F-Q9392nShDn",
        "outputId": "d169efa6-6bd6-4ad9-a302-a6bbc548d22e"
      },
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Overwriting app.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Init on_message function"
      ],
      "metadata": {
        "id": "VWtBzWjuTx_j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!chainlit run app.py &>/content/logs.txt &"
      ],
      "metadata": {
        "id": "AXbSL1IdU2wE"
      },
      "execution_count": 44,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import urllib"
      ],
      "metadata": {
        "id": "auSzLiieU-vD"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print (\"Password/Enpoint IP for localtunnel is:\",urllib.request.urlopen('https://ipv4.icanhazip.com').read().decode('utf8').strip (\"\\n\"))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_aPy4yaWVATE",
        "outputId": "dd0002d0-9f36-4def-803f-8a92ad410d64"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Password/Enpoint IP for localtunnel is: 34.83.126.189\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!lt --port 8000 ---subdomain aivn-simple-rag"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Y3xmcSZiVP8u",
        "outputId": "9c8b4bf9-e001-4a0b-ce3d-7b6f555c18ce"
      },
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "your url is: https://aivn-simple-rag.loca.lt\n",
            "^C\n"
          ]
        }
      ]
    }
  ]
}