embed_wikipedia_articles.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pipeline to Save Embeddings for Wikipedia Articles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = 'false'\n",
    "\n",
    "import numpy as np\n",
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "import transformers\n",
    "from datasets import load_dataset\n",
    "from sentence_transformers import SentenceTransformer, util\n",
    "\n",
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters \n",
    "Change these if running on your machine. Can set some to download from HF Hub."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_model = '/home/stefanwebb/models/llms/multi-qa-MiniLM-L6-cos-v1'\n",
    "embeddings_path = '/home/stefanwebb/embeddings/wikimedia/wikipedia/20231101.en'\n",
    "\n",
    "dataset_path = '/home/stefanwebb/data/wikimedia/wikipedia/20231101.en'\n",
    "dataset_name = ''\n",
    "dataset_split = 'train'\n",
    "\n",
    "batch_size = 1024\n",
    "num_data_workers = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Chunking Strategy\n",
    "* Chunks are per-paragraph, split with overlap if > 510 tokens.\n",
    "* Chunks corresponding to headers and short sentences are filtered out.\n",
    "* Is this what you call Semantic chunking?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "document_encoder = SentenceTransformer(embedding_model).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: May need extra space for special tokens, hence 510 not 512\n",
    "def chunk_iterator(seq, max_seq_length=510, overlap=128):\n",
    "    \"\"\"\n",
    "    Given a list of tokens, if it is greater than maximum, break into overlapping sections\n",
    "    \"\"\"\n",
    "    \n",
    "    if len(seq) <= max_seq_length:\n",
    "        yield seq\n",
    "    else:\n",
    "        count_chunks = ((len(seq) - max_seq_length) + overlap - 1) // overlap + 1\n",
    "        for idx in range(count_chunks):\n",
    "            yield seq[(idx*overlap):min(len(seq), (idx*overlap) + max_seq_length)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chunk_by_paragraph(articles):\n",
    "    \"\"\"\n",
    "    Forms chunks by breaking document into paragraphs, and then into sub-paragraphs if necessary.\n",
    "    \"\"\"\n",
    "    chunks = []\n",
    "    for article in articles['text']:\n",
    "        # print('article', article)\n",
    "        paragraphs = article.split('\\n\\n')\n",
    "        paragraphs = [p for p in paragraphs if p != '' and p[0].isalpha()]\n",
    "        # TODO: Can I avoid tokenizing sentences twice: once here and once in model.encode?\n",
    "        # model.encode doesn't seem to be able to take tokens, only strings\n",
    "        # I could dig into source code in library to find a solution...\n",
    "        # Alternatively, just do fixed length chunking with overlap\n",
    "        \n",
    "        # Filter by string\n",
    "        if len(paragraphs) > 0:\n",
    "            tokens = document_encoder.tokenizer(paragraphs)['input_ids']\n",
    "            tokens = [x for p in tokens for x in list(chunk_iterator(p)) if len(x) > 7]\n",
    "            paragraphs = document_encoder.tokenizer.batch_decode(tokens)\n",
    "\n",
    "        chunks += paragraphs\n",
    "\n",
    "    return {'chunks': chunks}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load dataset\n",
    "* An issue I noticed is that formulae, tables, etc. are missing from data which means some paragraphs end abruptly and this info is lost\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(dataset_path, dataset_name, data_files=[\"train-00000-of-00041.parquet\"], split=dataset_split, streaming=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "chunked_ds = dataset.map(chunk_by_paragraph, batched=True, batch_size=4048, remove_columns=dataset.column_names)\n",
    "chunked_pt = DataLoader(chunked_ds, batch_size=1024, num_workers=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time Embedding of Chunks\n",
    "* Starting off with one file for testing purposes\n",
    "* Initial pause of ~40 secs is dataloader filling its buffer\n",
    "* Ignore warning \"Token indices sequence length is longer than the specified maximum sequence length\". We tokenize longer chunks before breaking them up with overlap so the ones input to model are < 510 length\n",
    "* Throughput can be further optimized so GPU is constantly occupied and not blocked by synchronous operations or data starved.\n",
    "* About 32 mins for 1/41 data files => 22 hours upper bound\n",
    "* About 4gb per Parquet file\n",
    "* I'll speed up by increasing num of data workers, and save embeddings at float16\n",
    "* TODO: Experiment with 8-bit and 1-bit embeddings from SentenceTransformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: I find this an easier way to construct schema\n",
    "table = pa.table([\n",
    "            pa.array([\n",
    "               f\"The quick brown fox jumped over the log.\",\n",
    "               f\"Colorless dreams sleep furiously.\"       \n",
    "         ]),\n",
    "            pa.array([\n",
    "               torch.randn(384).numpy(),\n",
    "               torch.randn(384).numpy()\n",
    "            ])\n",
    "         ], names=[\"chunks\", \"embeddings\"])\n",
    "\n",
    "with pq.ParquetWriter('testing.tmp.parquet', table.schema) as writer:\n",
    "    for idx, batch in enumerate(chunked_pt):\n",
    "        # Embed chunks\n",
    "        # TODO: Cast to lower float precision for saving? np.float16?\n",
    "        chunks = batch['chunks']\n",
    "        embeddings = document_encoder.encode(chunks, batch_size=1024, show_progress_bar=True)\n",
    "\n",
    "\n",
    "        # Stream output\n",
    "        # TODO: Collate batches before saving to reduce I/O ops?\n",
    "        # TODO: Is write_table synchronous or asynchronous? I.e. is this holding up GPU?\n",
    "        table = pa.table([\n",
    "            pa.array(chunks),\n",
    "            pa.array(list(embeddings))\n",
    "         ], names=[\"chunks\", \"embeddings\"])\n",
    "        writer.write_table(table)\n",
    "        \n",
    "        # DEBUG\n",
    "        # if idx == 10:\n",
    "        #     break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Embed Entire Wikipedia and Save\n",
    "* Changed my mind on quantization. I'll save as `np.float32` and do quantization in FAISS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "chunks: string\n",
      "embeddings: list<item: float>\n",
      "  child 0, item: float\n"
     ]
    }
   ],
   "source": [
    "table = pa.table([\n",
    "                pa.array([\n",
    "                f\"The quick brown fox jumped over the log.\",\n",
    "                f\"Colorless dreams sleep furiously.\"       \n",
    "            ]),\n",
    "                pa.array([\n",
    "                np.zeros((384), dtype=np.float32),\n",
    "                np.zeros((384), dtype=np.float32)\n",
    "                ])\n",
    "            ], names=[\"chunks\", \"embeddings\"])\n",
    "\n",
    "print(table.schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing train-00001-of-00041.parquet\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Token indices sequence length is longer than the specified maximum sequence length for this model (1310 > 512). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\n",
      "Processing train-00002-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\n",
      "Processing train-00003-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\n",
      "Processing train-00004-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\n",
      "Processing train-00005-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00006-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00007-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00008-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00009-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00010-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00011-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00012-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00013-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00014-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00015-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00016-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\t batch 1500 of 1530\n",
      "\t batch 1800 of 1530\n",
      "\n",
      "Processing train-00017-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00018-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00019-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\n",
      "Processing train-00020-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00021-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00022-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\n",
      "Processing train-00023-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00024-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00025-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00026-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00027-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00028-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\n",
      "Processing train-00029-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00030-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00031-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00032-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00033-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00034-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00035-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\n",
      "Processing train-00036-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\t batch 1500 of 1530\n",
      "\t batch 1800 of 1530\n",
      "\t batch 2100 of 1530\n",
      "\n",
      "Processing train-00037-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\t batch 1500 of 1530\n",
      "\t batch 1800 of 1530\n",
      "\t batch 2100 of 1530\n",
      "\t batch 2400 of 1530\n",
      "\n",
      "Processing train-00038-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\t batch 1500 of 1530\n",
      "\t batch 1800 of 1530\n",
      "\n",
      "Processing train-00039-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\t batch 1500 of 1530\n",
      "\n",
      "Processing train-00040-of-00041.parquet\n",
      "\t batch 0 of 1530\n",
      "\t batch 300 of 1530\n",
      "\t batch 600 of 1530\n",
      "\t batch 900 of 1530\n",
      "\t batch 1200 of 1530\n",
      "\t batch 1500 of 1530\n",
      "\n"
     ]
    }
   ],
   "source": [
    "files = [f\"train-{idx:05d}-of-00041.parquet\" for idx in range(41)]\n",
    "\n",
    "for file in files[1:]:\n",
    "    # Count approximate batches for progress (10 chunks per article on average)\n",
    "    metadata = pq.ParquetFile(os.path.join(dataset_path, file)).metadata\n",
    "    count_batches = math.ceil(metadata.num_rows / 1024) * 10\n",
    "    del metadata\n",
    "\n",
    "    # Load input stream\n",
    "    dataset = load_dataset(dataset_path, dataset_name, data_files=[file], split=dataset_split, streaming=True)\n",
    "    chunked_ds = dataset.map(chunk_by_paragraph, batched=True, batch_size=batch_size*4, remove_columns=dataset.column_names)\n",
    "    chunked_pt = DataLoader(chunked_ds, batch_size=batch_size)\n",
    "    \n",
    "    # Note: I find this an easier way to construct schema\n",
    "    table = pa.table([\n",
    "                pa.array([\n",
    "                f\"The quick brown fox jumped over the log.\",\n",
    "                f\"Colorless dreams sleep furiously.\"       \n",
    "            ]),\n",
    "                pa.array([\n",
    "                torch.randn(384).numpy(),\n",
    "                torch.randn(384).numpy()\n",
    "                ])\n",
    "            ], names=[\"chunks\", \"embeddings\"])\n",
    "\n",
    "    with pq.ParquetWriter(os.path.join(embeddings_path, file), table.schema) as writer:\n",
    "        print(f'Processing {file}')\n",
    "        for idx, batch in enumerate(chunked_pt):\n",
    "            # Embed chunks\n",
    "            chunks = batch['chunks']\n",
    "            embeddings = document_encoder.encode(chunks, batch_size=batch_size, show_progress_bar=False).astype(np.float32)\n",
    "\n",
    "            # Stream output\n",
    "            table = pa.table([\n",
    "                pa.array(chunks),\n",
    "                pa.array(list(embeddings))\n",
    "            ], names=[\"chunks\", \"embeddings\"])\n",
    "            writer.write_table(table)\n",
    "\n",
    "            if idx % 300 == 0:\n",
    "                print(f'\\t batch {idx} of {count_batches}')\n",
    "\n",
    "    # Not sure if following is necessary, just in case.\n",
    "    del chunked_ds\n",
    "    del chunked_pt\n",
    "\n",
    "    print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scratchpad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "chunks: string\n",
      "embeddings: list<item: halffloat>\n",
      "  child 0, item: halffloat\n"
     ]
    }
   ],
   "source": [
    "table = pa.table([\n",
    "                pa.array([\n",
    "                f\"The quick brown fox jumped over the log.\",\n",
    "                f\"Colorless dreams sleep furiously.\"       \n",
    "            ]),\n",
    "                pa.array([\n",
    "                np.zeros((384), dtype=np.float32),\n",
    "                np.zeros((384), dtype=np.float32)\n",
    "                ])\n",
    "            ], names=[\"chunks\", \"embeddings\"])\n",
    "\n",
    "print(table.schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}