notebooks/save_hf_datasets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cd95b974",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "80"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tokenizers import Tokenizer\n",
    "from transformers import PreTrainedTokenizerFast\n",
    "from tokenizers.models import WordLevel, BPE\n",
    "from tokenizers.pre_tokenizers import Whitespace,Split,ByteLevel, WhitespaceSplit\n",
    "from tokenizers.normalizers import Lowercase, NFKC\n",
    "import os\n",
    "import polars as pl\n",
    "from joblib import Parallel, delayed\n",
    "import multiprocessing\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "import json\n",
    "from tokenizers.trainers import BpeTrainer, UnigramTrainer, WordLevelTrainer, WordPieceTrainer\n",
    "import gc\n",
    "from transformers import AutoConfig, AutoTokenizer, AutoModel, DataCollatorWithPadding\n",
    "import datasets\n",
    "from datasets import Dataset\n",
    "\n",
    "multiprocessing.cpu_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "effa5d3d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.876520295627415 GB\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (98_415_610, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>smiles</th><th>BRD4</th><th>HSA</th><th>sEH</th></tr><tr><td>str</td><td>bool</td><td>bool</td><td>bool</td></tr></thead><tbody><tr><td>&quot;C#CCOc1ccc(CNc…</td><td>false</td><td>false</td><td>false</td></tr><tr><td>&quot;C#CCOc1ccc(CNc…</td><td>false</td><td>false</td><td>false</td></tr><tr><td>&quot;C#CCOc1ccc(CNc…</td><td>false</td><td>false</td><td>false</td></tr><tr><td>&quot;C#CCOc1ccc(CNc…</td><td>false</td><td>false</td><td>false</td></tr><tr><td>&quot;C#CCOc1ccc(CNc…</td><td>false</td><td>false</td><td>false</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>&quot;[N-]=[N+]=NCCC…</td><td>false</td><td>false</td><td>false</td></tr><tr><td>&quot;[N-]=[N+]=NCCC…</td><td>false</td><td>false</td><td>false</td></tr><tr><td>&quot;[N-]=[N+]=NCCC…</td><td>false</td><td>false</td><td>false</td></tr><tr><td>&quot;[N-]=[N+]=NCCC…</td><td>false</td><td>false</td><td>false</td></tr><tr><td>&quot;[N-]=[N+]=NCCC…</td><td>false</td><td>false</td><td>false</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (98_415_610, 4)\n",
       "┌───────────────────────────────────┬───────┬───────┬───────┐\n",
       "│ smiles                            ┆ BRD4  ┆ HSA   ┆ sEH   │\n",
       "│ ---                               ┆ ---   ┆ ---   ┆ ---   │\n",
       "│ str                               ┆ bool  ┆ bool  ┆ bool  │\n",
       "╞═══════════════════════════════════╪═══════╪═══════╪═══════╡\n",
       "│ C#CCOc1ccc(CNc2nc(NCC3CCCN3c3ccc… ┆ false ┆ false ┆ false │\n",
       "│ C#CCOc1ccc(CNc2nc(NCc3cccc(Br)n3… ┆ false ┆ false ┆ false │\n",
       "│ C#CCOc1ccc(CNc2nc(NCc3ccc(OCC#C)… ┆ false ┆ false ┆ false │\n",
       "│ C#CCOc1ccc(CNc2nc(NCCNC(=O)C(=C)… ┆ false ┆ false ┆ false │\n",
       "│ C#CCOc1ccc(CNc2nc(NCC(=O)NCC=C)n… ┆ false ┆ false ┆ false │\n",
       "│ …                                 ┆ …     ┆ …     ┆ …     │\n",
       "│ [N-]=[N+]=NCCC[C@H](Nc1nc(Nc2noc… ┆ false ┆ false ┆ false │\n",
       "│ [N-]=[N+]=NCCC[C@H](Nc1nc(Nc2nnc… ┆ false ┆ false ┆ false │\n",
       "│ [N-]=[N+]=NCCC[C@H](Nc1nc(Nc2nnc… ┆ false ┆ false ┆ false │\n",
       "│ [N-]=[N+]=NCCC[C@H](Nc1nc(Nc2nnn… ┆ false ┆ false ┆ false │\n",
       "│ [N-]=[N+]=NCCC[C@H](Nc1nc(Nc2noc… ┆ false ┆ false ┆ false │\n",
       "└───────────────────────────────────┴───────┴───────┴───────┘"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df = pl.scan_csv('/home/dangnh36/datasets/competitions/leash_belka/processed/train_v2.csv').select(\n",
    "        pl.col('molecule').alias('smiles'),\n",
    "#         pl.col('bb1', 'bb2', 'bb3').cast(pl.UInt16),\n",
    "        pl.col('BRD4', 'HSA', 'sEH').cast(pl.Boolean),\n",
    "    ).collect()\n",
    "print(train_df.estimated_size('gb'), 'GB')\n",
    "train_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a334264",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e406e5c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.06128192972391844 GB\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (878_022, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>smiles</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;C#CCCC[C@H](Nc…</td></tr><tr><td>&quot;C#CCCC[C@H](Nc…</td></tr><tr><td>&quot;C#CCCC[C@H](Nc…</td></tr><tr><td>&quot;C#CCCC[C@H](Nc…</td></tr><tr><td>&quot;C#CCCC[C@H](Nc…</td></tr><tr><td>&hellip;</td></tr><tr><td>&quot;Cn1ncc2cc(Nc3n…</td></tr><tr><td>&quot;[N-]=[N+]=NCCC…</td></tr><tr><td>&quot;COC(=O)c1ccnc(…</td></tr><tr><td>&quot;COC1CCC(CCNc2n…</td></tr><tr><td>&quot;[N-]=[N+]=NCCC…</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (878_022, 1)\n",
       "┌───────────────────────────────────┐\n",
       "│ smiles                            │\n",
       "│ ---                               │\n",
       "│ str                               │\n",
       "╞═══════════════════════════════════╡\n",
       "│ C#CCCC[C@H](Nc1nc(Nc2ccc(C=C)cc2… │\n",
       "│ C#CCCC[C@H](Nc1nc(Nc2ccc(C=C)cc2… │\n",
       "│ C#CCCC[C@H](Nc1nc(NCC2(O)CCCC2(C… │\n",
       "│ C#CCCC[C@H](Nc1nc(Nc2ccc(C=C)cc2… │\n",
       "│ C#CCCC[C@H](Nc1nc(NCC2CCC(SC)CC2… │\n",
       "│ …                                 │\n",
       "│ Cn1ncc2cc(Nc3nc(Nc4nncs4)nc(N[C@… │\n",
       "│ [N-]=[N+]=NCCC[C@H](Nc1nc(NCC2CC… │\n",
       "│ COC(=O)c1ccnc(Nc2nc(Nc3noc4ccc(F… │\n",
       "│ COC1CCC(CCNc2nc(Nc3noc4ccc(F)cc3… │\n",
       "│ [N-]=[N+]=NCCC[C@H](Nc1nc(NCc2cc… │\n",
       "└───────────────────────────────────┘"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_df = pl.scan_csv('/home/dangnh36/datasets/competitions/leash_belka/processed/test_v4.csv').select(\n",
    "        pl.col('molecule').alias('smiles'),\n",
    "#         pl.col('bb1', 'bb2', 'bb3').cast(pl.UInt16),\n",
    "        # pl.col('BRD4', 'HSA', 'sEH').cast(pl.UInt8),\n",
    "    ).collect()\n",
    "print(test_df.estimated_size('gb'), 'GB')\n",
    "test_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d28fafc4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "41151195",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !rm -r /home/dangnh36/datasets/competitions/leash_belka/processed/hf/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "766ebd85",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "56b8cd22",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9faa632ac0f74b60a37992a948ceb1db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/17 shards):   0%|          | 0/98415610 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset = Dataset.from_polars(train_df)\n",
    "dataset.save_to_disk(\n",
    "        dataset_path = '/home/dangnh36/datasets/competitions/leash_belka/processed/hf/datasets/train',\n",
    "        max_shard_size = None,\n",
    "        num_shards = None,\n",
    "        num_proc = None,\n",
    "        storage_options = None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c50d587d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd4f9e78bca44d53a8b4bf1d21457997",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/878022 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset = Dataset.from_polars(test_df)\n",
    "dataset.save_to_disk(\n",
    "        dataset_path = '/home/dangnh36/datasets/competitions/leash_belka/processed/hf/datasets/test',\n",
    "        max_shard_size = None,\n",
    "        num_shards = None,\n",
    "        num_proc = None,\n",
    "        storage_options = None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "369c23fe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d636a8ac",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "199a493a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_from_disk, concatenate_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3d1fe892",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e92f665c40444abe8d5d65e1afbb7826",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading dataset from disk:   0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "98415610"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_ds = load_from_disk('/home/dangnh36/datasets/competitions/leash_belka/processed/hf/datasets/train',\n",
    "                          keep_in_memory=False)\n",
    "len(train_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "444bf443",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "878022"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_ds = load_from_disk('/home/dangnh36/datasets/competitions/leash_belka/processed/hf/datasets/test',\n",
    "                          keep_in_memory=False)\n",
    "len(test_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1a6f9667",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "99293632"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = concatenate_datasets([train_ds, test_ds])\n",
    "len(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b49dcc60",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21175e446feb4114a71b0b631b59a030",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/17 shards):   0%|          | 0/99293632 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ds.save_to_disk(\n",
    "        dataset_path = '/home/dangnh36/datasets/competitions/leash_belka/processed/hf/datasets/all',\n",
    "        max_shard_size = None,\n",
    "        num_shards = None,\n",
    "        num_proc = None,\n",
    "        storage_options = None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ffec0d8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2467604e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1e7b55b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eebcfbc0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}