RAG_with_LMstudio_GH.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6c1003c3-43b3-45bb-9874-423b76f2195b",
   "metadata": {},
   "source": [
    "RAG with LM studio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de844bb-4b76-4cd7-9576-ec756594a441",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "import os\n",
    "\n",
    "def read_pdf(path):\n",
    "    r_splitter = RecursiveCharacterTextSplitter(\n",
    "        chunk_size=2000,\n",
    "        chunk_overlap=450,\n",
    "        separators=[\"\\n\\n\", \"\\n\", \".\", \" \", \"\"]\n",
    "    )\n",
    "\n",
    "    loaders = []\n",
    "    docs = []\n",
    "\n",
    "    for filename in os.listdir(path):\n",
    "        if filename.endswith(\".pdf\"):\n",
    "            file_path = os.path.join(path, filename)\n",
    "            loader = PyPDFLoader(file_path)\n",
    "            loaders.append(loader)\n",
    "            docs.extend(loader.load())  # Load documents immediately\n",
    "\n",
    "    r_splits = r_splitter.split_documents(docs)\n",
    "\n",
    "    return r_splits\n",
    "\n",
    "path = r\"sample\"\n",
    "\n",
    "r_splits = read_pdf(path)\n",
    "print(len(r_splits))\n",
    "print(len(r_splits))\n",
    "page = r_splits[0]\n",
    "page.metadata\n",
    "print(page.page_content[0:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2f40797-2003-409e-bdb3-a61454fe057b",
   "metadata": {},
   "outputs": [],
   "source": [
    "page = r_splits[5]\n",
    "print(page.metadata)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12aa3217-47bf-495a-9b62-c6d1633052f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import Chroma\n",
    "from langchain.docstore.document import Document\n",
    "from openai import OpenAI\n",
    "from typing import List\n",
    "import os\n",
    "import shutil\n",
    "import tempfile\n",
    "import time\n",
    "\n",
    "\n",
    "\n",
    "client = OpenAI(base_url=\"http://localhost:1238/v1\", api_key=\"lm-studio\")\n",
    "\n",
    "\n",
    "\n",
    "class CustomEmbedding2:\n",
    "    def __init__(self):\n",
    "        self.embeddings = []  # List to store the embeddings\n",
    "\n",
    "    def embed_documents(self, texts: List[str]) -> List[List[float]]:\n",
    "        embeddings = [get_embedding(text) for text in texts]\n",
    "        self.embeddings = embeddings  # Store the embeddings in the `embeddings` attribute\n",
    "        return embeddings\n",
    "\n",
    "    def embed_query(self, text: str) -> List[float]:\n",
    "        embedding = get_embedding(text)\n",
    "        self.embeddings = [embedding]\n",
    "        return embedding\n",
    "\n",
    "    def get_embeddings(self) -> List[List[float]]:\n",
    "        return self.embeddings\n",
    "\n",
    "def get_embedding(text, model=\"TheBloke/nomic-embed-text\"):\n",
    "    text = text.replace(\"\\n\", \" \")\n",
    "    return client.embeddings.create(input=[text], model=model).data[0].embedding\n",
    "\n",
    "\n",
    "def create_vectordb(embedding, r_splits, new_vectordb=True, user_path=None, temp_dir=None):\n",
    "    cwd = os.getcwd()  # Get the current working directory\n",
    "\n",
    "    if new_vectordb:\n",
    "        if user_path is None:\n",
    "            if temp_dir is None:\n",
    "                # Create a temporary directory for persisting the vector database\n",
    "                temp_dir = tempfile.mkdtemp()\n",
    "                persist_directory = os.path.join(cwd, temp_dir, 'chroma_db')\n",
    "            else:\n",
    "                persist_directory = os.path.join(cwd, temp_dir, 'chroma_db')\n",
    "        else:\n",
    "            persist_directory = os.path.join(cwd, user_path, 'chroma_db')\n",
    "\n",
    "        # Remove the existing persist directory (if any)\n",
    "        if os.path.exists(persist_directory):\n",
    "            shutil.rmtree(persist_directory)\n",
    "\n",
    "        print(\"Creating a new vector database...\")\n",
    "        start_time = time.time()\n",
    "        vectordb = Chroma.from_documents(\n",
    "            documents=r_splits,\n",
    "            embedding=embedding,\n",
    "            persist_directory=persist_directory\n",
    "        )\n",
    "        end_time = time.time()\n",
    "        print(f\"New vector database created in {end_time - start_time:.2f} seconds. Directory: {persist_directory}\")\n",
    "\n",
    "    else:\n",
    "        if user_path is None:\n",
    "            print(\"Please provide a valid user path to load the existing vector database.\")\n",
    "            return None\n",
    "        else:\n",
    "            persist_directory = os.path.join(cwd, user_path, 'chroma_db')\n",
    "            if not os.path.exists(persist_directory):\n",
    "                print(f\"Vector database not found in the specified path: {persist_directory}\")\n",
    "                return None\n",
    "\n",
    "            print(\"Loading existing vector database...\")\n",
    "            start_time = time.time()\n",
    "            vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding)\n",
    "            end_time = time.time()\n",
    "            print(f\"Existing vector database loaded in {end_time - start_time:.2f} seconds. Directory: {persist_directory}\")\n",
    "\n",
    "    return vectordb, persist_directory\n",
    "\n",
    "# Default values\n",
    "new_vectordb = True\n",
    "user_path = os.path.join(os.getcwd(), \"docs\", \"sample\")\n",
    "temp_dir = os.path.join(os.getcwd(), \"docs\", \"temp\")\n",
    "\n",
    "\n",
    "# Extract the text from the Document objects\n",
    "texts = [doc.page_content for doc in r_splits]\n",
    "# Create the custom embedding object\n",
    "embedding = CustomEmbedding2()\n",
    "# Optionally, you can override the default values here\n",
    "# For example:\n",
    "# new_vectordb = False\n",
    "# user_path = \"/path/to/your/directory\"\n",
    "\n",
    "vectordb, temp_dir = create_vectordb(embedding, r_splits, new_vectordb, user_path, temp_dir)\n",
    "print(\"\\nfinished\", temp_dir, len(vectordb))\n",
    "print(vectordb._collection.count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef1dd2b6-5f9f-43fb-acd7-faf9c237081a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(vectordb._collection.count())\n",
    "query= \"Raise your hand if the text on the small screens is legible\"\n",
    "most_similar = vectordb.similarity_search(query, k=5)#, filter={\"source\":\"docs/cs229_lectures/MachineLearning-Lecture03.pdf\"})\n",
    "print(most_similar)\n",
    "# Check the size of the vector database\n",
    "print(f\"Vector database size: {len(vectordb)}\")\n",
    "print(f\" database size: {len(r_splits)}\")\n",
    "for doc in most_similar:\n",
    "    print(doc.metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a93f1a93-bad9-4885-849f-7c0938f4e24e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.vectorstores import Chroma\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.prompts import PromptTemplate\n",
    "import textwrap\n",
    "\n",
    "# ... (previous code) ...\n",
    "\n",
    "# Create a RAG model\n",
    "llm = OpenAI(base_url=\"http://localhost:1238/v1\", api_key=\"lm-studio\")\n",
    "retriever = vectordb.as_retriever()\n",
    "\n",
    "# Define a custom prompt template\n",
    "template = \"\"\"\n",
    "Use the following context from the vector database and the previous conversation to answer the query. Incorporate your own knowledge and reasoning as an AI assistant:\n",
    "\n",
    "Previous Conversation: {history}\n",
    "\n",
    "Context: {context}\n",
    "\n",
    "Query: {query}\n",
    "\n",
    "Answer:\n",
    "\"\"\"\n",
    "\n",
    "prompt_template = PromptTemplate(input_variables=[\"history\", \"context\", \"query\"], template=template)\n",
    "\n",
    "# Create a custom chain\n",
    "qa_chain = LLMChain(llm=llm, prompt=prompt_template)\n",
    "\n",
    "history = []\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "373b8648-4e8e-4b7e-a23c-90f808038503",
   "metadata": {},
   "outputs": [],
   "source": [
    "while True:\n",
    "    query = input(\"> \")\n",
    "    context = retriever.get_relevant_documents(query)\n",
    "    result_str = qa_chain.run(\n",
    "        history=\"\\n\".join([f\"{msg['role']}: {msg['content']}\" for msg in history]),\n",
    "        query=query,\n",
    "        max_tokens=4096,\n",
    "        context=\"\".join([doc.page_content for doc in context])\n",
    "    )\n",
    "\n",
    "    # Print the query\n",
    "    #print(\"\\nQuery:\")\n",
    "    #print(f\"{query}\")\n",
    "\n",
    "    # Print the context\n",
    "    #print(\"\\nContext:\")\n",
    "    #for doc in context:\n",
    "      #  print(doc.page_content)\n",
    "      #  print(\"-\" * 80)\n",
    "\n",
    "    # Print the final answer\n",
    "    final_answer_start = result_str.find(\"Final Answer:\")\n",
    "    if final_answer_start != -1:\n",
    "        final_answer = result_str[final_answer_start + len(\"Final Answer:\"):].strip()\n",
    "        print(\"\\nFinal Answer:\")\n",
    "        print(textwrap.fill(final_answer, width=80))\n",
    "    else:\n",
    "        print(\"\\nAnswer:\")\n",
    "        print(textwrap.fill(result_str, width=80))\n",
    "\n",
    "    # Update the history\n",
    "    history.append({\"role\": \"user\", \"content\": query})\n",
    "    history.append({\"role\": \"assistant\", \"content\": result_str})"
   ]
  }
 ],
 "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.8.19"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}