{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c21dd5cd-0e42-419d-8e82-468c93b6139d",
   "metadata": {},
   "source": [
    "# Extraction des données et création de la base de donnée pour les autres notebooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f5436125-4dff-4238-ae98-9863f694111b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mdelvallez/conda_env/raghn/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.documents import Document\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_huggingface import HuggingFaceEmbeddings\n",
    "from datasets import load_dataset, Dataset\n",
    "from tqdm import tqdm\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "74b11ffb-831e-4ad2-8fbb-d2a81a095a9b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---Load Dataset & Embedder---\n"
     ]
    }
   ],
   "source": [
    "print(\"---Load Dataset & Embedder---\")\n",
    "db_size = 1000 # 100000\n",
    "batch_size = 50 # 5000\n",
    "dataset = load_dataset(\"miracl/miracl-corpus\", \"fr\", split=\"train\", trust_remote_code=True, streaming=True)\n",
    "model_name = \"intfloat/e5-small\"\n",
    "model_kwargs = {'device': 'cpu'}\n",
    "encode_kwargs = {'normalize_embeddings': True}\n",
    "hf = HuggingFaceEmbeddings(\n",
    "    model_name=model_name,\n",
    "    model_kwargs=model_kwargs,\n",
    "    encode_kwargs=encode_kwargs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1bf81c29-2414-40af-8058-24a21b90a837",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transfer(start, end):\n",
    "    data = list(dataset.skip(start).take(end-start))\n",
    "    docs = []\n",
    "    ids = []\n",
    "    for idx in range(end-start):\n",
    "        doc = Document(page_content=data[idx]['text'],\n",
    "                    metadata={\n",
    "                        'title': data[idx]['title']\n",
    "                    })\n",
    "        docs.append(doc)\n",
    "        ids.append(data[idx]['docid'])\n",
    "    return docs, ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "28cbddb2-a127-49e0-b6f9-502d0542d3be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---Docs extraction---\n"
     ]
    }
   ],
   "source": [
    "print(\"---Docs extraction---\")\n",
    "docs, ids = transfer(0, batch_size)\n",
    "persist_directory = \"corpus_db\"\n",
    "\n",
    "vectordb = Chroma.from_documents(documents=docs, embedding=hf, ids = ids, persist_directory=persist_directory) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "133425c3-ce1e-4ebd-9bcc-3e701b7c71b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_process(batch_size):\n",
    "    for i in tqdm(range(batch_size, db_size, batch_size)):\n",
    "        docs, ids = transfer(i, i+batch_size)\n",
    "        vectordb.add_documents(documents = docs, ids = ids)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b006d35-e769-418d-bede-b045d7fcb201",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"---Batching Documents---\")\n",
    "\n",
    "batch_process(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cfcd5b6-3b07-4cc8-b2ac-f095e0e9acca",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"---Titles extraction---\")\n",
    "\n",
    "metadata = vectordb.get(include=['metadatas'], limit=20, offset=0)['metadatas']\n",
    "\n",
    "title_set = set([doc['title'] for doc in metadata])\n",
    "\n",
    "offset = len(title_set)\n",
    "limit = batch_size\n",
    "while len(title_set) == limit:\n",
    "    metadata = vectordb.get(include=['metadatas'], limit=limit, offset=offset)['metadatas']\n",
    "    title_set.update([doc['title'] for doc in metadata])\n",
    "    # mise à jour des bornes d'extraction pour le batch suivant\n",
    "    offset += len(title_set)\n",
    "\n",
    "\n",
    "docs = []\n",
    "for title in tqdm(title_set):\n",
    "    doc = Document(page_content=title)\n",
    "    docs.append(doc)\n",
    "persist_directory = \"title_db\"\n",
    "titledb = Chroma.from_documents(documents=docs, embedding=hf, persist_directory=persist_directory)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "raghn",
   "language": "python",
   "name": "raghn"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
