generate dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import math\n",
    "import multiprocessing\n",
    "import warnings\n",
    "\n",
    "import bs4\n",
    "import numpy as np\n",
    "import torch\n",
    "from bs4 import MarkupResemblesLocatorWarning\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=MarkupResemblesLocatorWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def text_with_newdataset(elem):\n",
    "    text = \"\"\n",
    "    for e in elem.descendants:\n",
    "        if isinstance(e, str):\n",
    "            text += e\n",
    "        elif e.name == \"br\" or e.name == \"p\":\n",
    "            text += \"\\n\"\n",
    "\n",
    "    return text\n",
    "\n",
    "\n",
    "def parse_line(line):\n",
    "    data = json.loads(line)\n",
    "    posts_text = []\n",
    "\n",
    "    for post in data.get(\"posts\", []):\n",
    "        if \"com\" in post:\n",
    "            soup = bs4.BeautifulSoup(post[\"com\"], \"lxml\")\n",
    "            post_text = text_with_newdataset(soup).strip()\n",
    "        else:\n",
    "            post_text = \"\"\n",
    "\n",
    "        post_text = f\"--- {post['no']}\\n{post_text}\"\n",
    "        posts_text.append(post_text)\n",
    "\n",
    "    return \"\\n\".join(posts_text) + \"\\n-----\\n\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_str = []\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-neox-20b\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with multiprocessing.Pool() as pool:\n",
    "    with open(\"dataset.ndjson\") as file:\n",
    "        # imap is fine too, but it's slower\n",
    "        for data in pool.map(parse_line, file):\n",
    "            dataset_str.append(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since huggingface's tokenizer only support int64, we have to work arround it's BS before the tensor gets too big\n",
    "\n",
    "I only got 256 GB ram to work with :("
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chunk_size = 100000\n",
    "num_chunks = math.ceil(len(dataset_str) / chunk_size)\n",
    "\n",
    "dataset_token = []\n",
    "\n",
    "for i in range(num_chunks):\n",
    "    start_index = i * chunk_size\n",
    "    end_index = min((i + 1) * chunk_size, len(dataset_str))\n",
    "\n",
    "    chunk = dataset_str[start_index:end_index]\n",
    "    chunk = '\\n'.join(chunk) + \"\\n\"\n",
    "\n",
    "    chunk = tokenizer(chunk, return_tensors=\"pt\")[\"input_ids\"].squeeze()\n",
    "    dataset_token.append(chunk.type(torch.uint16))\n",
    "\n",
    "del dataset_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_token = torch.cat(dataset_token)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will store memmap in uint16 since our vocab is ~50k and uint16 goes up to 65535\n",
    "\n",
    "Using int64 is just wasting space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mp = np.memmap(\"dataset.dat\", dtype=np.uint16, mode=\"w+\", shape=dataset_token.numel())\n",
    "mp[:] = dataset_token[:]\n",
    "mp.flush()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mamba-4chan",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}