Pacama95/chatbot_agent

1{2 "cells": [3  {4   "cell_type": "markdown",5   "id": "1f4bf6d8",6   "metadata": {},7   "source": [8    "# RAG: Retrieval Augmented Generation"9   ]10  },11  {12   "cell_type": "markdown",13   "id": "f23d7c55",14   "metadata": {},15   "source": [16    "## Define the embedding model"17   ]18  },19  {20   "cell_type": "code",21   "execution_count": 2,22   "id": "cdde09a8",23   "metadata": {},24   "outputs": [],25   "source": [26    "## RAG\n",27    "\n",28    "from langchain.embeddings.base import Embeddings\n",29    "\n",30    "from sentence_transformers import SentenceTransformer\n",31    "from typing import List\n",32    "\n",33    "# Embedding model\n",34    "# Create a class for SentenceTransformers compatibility with Chroma\n",35    "class CustomEmbeddings(Embeddings):\n",36    "    def __init__(self, model_name: str):\n",37    "        self.model = SentenceTransformer(model_name)\n",38    "\n",39    "    def embed_documents(self, documents: List[str]) -> List[List[float]]:\n",40    "        return [self.model.encode(d).tolist() for d in documents]\n",41    "\n",42    "    def embed_query(self, query: str) -> List[float]:\n",43    "        return self.model.encode([query])[0].tolist()\n",44    "    \n",45    "# Create the custom embedding function\n",46    "embedding_model = CustomEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\")"47   ]48  },49  {50   "cell_type": "markdown",51   "id": "cb288efc",52   "metadata": {},53   "source": [54    "## Load Vector DB with documents"55   ]56  },57  {58   "cell_type": "code",59   "execution_count": null,60   "id": "72b41214",61   "metadata": {},62   "outputs": [],63   "source": [64    "## Load some custom data...\n",65    "from langchain_chroma import Chroma\n",66    "from langchain.docstore.document import Document\n",67    "\n",68    "import re\n",69    "\n",70    "def read_messages_from_file(file_path):\n",71    "    # Pattern to match lines like: [27/7/18, 21:54:49] Name: Message\n",72    "    pattern = re.compile(r'^\\[\\d{1,2}/\\d{1,2}/\\d{2}, \\d{1,2}:\\d{2}:\\d{2}\\] (.*?): (.*)$')\n",73    "    lrm_char = '\\u200e'\n",74    "    \n",75    "    documents = []\n",76    "    \n",77    "    with open(file_path, 'r', encoding='utf-8') as file:\n",78    "        for line in file:\n",79    "            if lrm_char in line:\n",80    "                continue  # Skip lines with U+200E\n",81    "            line = line.strip()\n",82    "            match = pattern.match(line)\n",83    "            if match:\n",84    "                name, message = match.groups()\n",85    "                documents.append(Document(\n",86    "                    page_content=message,\n",87    "                    metadata={\"name\": name}\n",88    "                ))\n",89    "    \n",90    "    return documents\n",91    "\n",92    "documents = read_messages_from_file('group_chat/_chat.txt')\n"93   ]94  },95  {96   "cell_type": "code",97   "execution_count": null,98   "id": "590d2b57",99   "metadata": {},100   "outputs": [],101   "source": [102    "#Index\n",103    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",104    "from tqdm import tqdm\n",105    "\n",106    "# Splitter\n",107    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",108    "    chunk_size=100,\n",109    "    chunk_overlap=50\n",110    ")\n",111    "\n",112    "# Make splits\n",113    "splits = text_splitter.split_documents(documents)\n",114    "\n",115    "print(splits[:5])\n",116    "print(f\"Total docs: {len(splits)}\")\n",117    "\n",118    "def embed_and_store_in_batches(docs, embedding_model, persist_directory, batch_size=100):\n",119    "    # Initialize the Vector DB with an initial batch\n",120    "    vectorstore = Chroma.from_documents(\n",121    "        documents=docs[:1],\n",122    "        embedding=embedding_model,\n",123    "        persist_directory=persist_directory\n",124    "    )\n",125    "\n",126    "    for i in tqdm(range(1, len(docs), batch_size), desc=\"Processing batches\"):\n",127    "        batch = docs[i:i+batch_size]\n",128    "        try:\n",129    "            vectorstore.add_documents(batch)\n",130    "        except Exception as e:\n",131    "            print(f\"Error en batch {i}-{i+batch_size}: {e}\")\n",132    "    \n",133    "    print(\"✅ Embeddings generated and stored.\")\n",134    "    return vectorstore\n",135    "\n",136    "vectorstore = embed_and_store_in_batches(\n",137    "    docs=splits,\n",138    "    embedding_model=embedding_model,\n",139    "    persist_directory=\"./chroma_db\",\n",140    "    batch_size=100\n",141    ")\n",142    "\n",143    "query = \"Feria\"\n",144    "results = vectorstore.similarity_search(query)\n",145    "\n",146    "print(results)"147   ]148  },149  {150   "cell_type": "markdown",151   "id": "effaf3bb",152   "metadata": {},153   "source": [154    "### Start Chroma DB from an existing DB"155   ]156  },157  {158   "cell_type": "code",159   "execution_count": null,160   "id": "95bf953c",161   "metadata": {},162   "outputs": [],163   "source": [164    "# Testing\n",165    "from langchain_chroma import Chroma\n",166    "\n",167    "vectordb = Chroma(persist_directory=\"./chroma_db\", embedding_function=embedding_model)\n",168    "\n",169    "retriever = vectordb.as_retriever()\n",170    "\n",171    "vectordb.similarity_search(\"Hi\", k = 10)"172   ]173  }174 ],175 "metadata": {176  "kernelspec": {177   "display_name": "venv",178   "language": "python",179   "name": "python3"180  },181  "language_info": {182   "codemirror_mode": {183    "name": "ipython",184    "version": 3185   },186   "file_extension": ".py",187   "mimetype": "text/x-python",188   "name": "python",189   "nbconvert_exporter": "python",190   "pygments_lexer": "ipython3",191   "version": "3.13.3"192  }193 },194 "nbformat": 4,195 "nbformat_minor": 5196}197