split_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from transformers import (\n",
    "    DebertaV2Tokenizer,\n",
    "    DebertaV2ForSequenceClassification,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    "    get_polynomial_decay_schedule_with_warmup,\n",
    ")\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CFG:\n",
    "    split_text = True\n",
    "    split_text_improved = True\n",
    "    max_length = 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = DebertaV2Tokenizer.from_pretrained(\n",
    "    \"microsoft/deberta-v3-base\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.read_parquet(\"./large_dataset/data.parquet\")\n",
    "print(data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.drop_duplicates(subset=[\"text\"], inplace=True)\n",
    "print(data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = data[[\"text\", \"source\", \"prompt_id\"]]\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if CFG.split_text_improved:\n",
    "\n",
    "    def split_text_improved(row, max_length=CFG.max_length - 10):\n",
    "        tokens = tokenizer.encode(row[\"text\"], add_special_tokens=False)\n",
    "        if len(tokens) <= max_length:\n",
    "            return [row.to_dict()]\n",
    "\n",
    "        new_rows = []\n",
    "        current_chunk = []\n",
    "        current_length = 0\n",
    "\n",
    "        for token in tokens:\n",
    "            current_chunk.append(token)\n",
    "            current_length += 1\n",
    "\n",
    "            if tokenizer.decode([token]) in [\".\"]:\n",
    "                if current_length >= max_length:\n",
    "                    text_chunk = tokenizer.decode(\n",
    "                        current_chunk,\n",
    "                        clean_up_tokenization_spaces=True\n",
    "                    )\n",
    "                    new_row = row.to_dict()\n",
    "                    new_row[\"text\"] = text_chunk\n",
    "                    new_rows.append(new_row)\n",
    "\n",
    "                    current_chunk = []\n",
    "                    current_length = 0\n",
    "\n",
    "        if current_chunk:\n",
    "            text_chunk = tokenizer.decode(\n",
    "                current_chunk,\n",
    "                clean_up_tokenization_spaces=True\n",
    "            )\n",
    "            new_row = row.to_dict()\n",
    "            new_row[\"text\"] = text_chunk\n",
    "            new_rows.append(new_row)\n",
    "\n",
    "        return new_rows\n",
    "\n",
    "    new_rows = data.apply(split_text_improved, axis=1).tolist()\n",
    "    flattened_rows = [item for sublist in new_rows for item in sublist]\n",
    "    new_df = pd.DataFrame(flattened_rows)\n",
    "    new_df = new_df.reset_index(drop=True)\n",
    "\n",
    "    print(new_df.shape)\n",
    "    print(new_df.head())\n",
    "    new_df.to_csv(\"./large_dataset/split_text_improved.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if CFG.split_text:\n",
    "    def split_text(row, max_length=CFG.max_length + 10):\n",
    "        tokens = tokenizer.encode(row[\"text\"], add_special_tokens=False)\n",
    "        if len(tokens) <= max_length:\n",
    "            return [row.to_dict()]\n",
    "\n",
    "        chunks = [tokens[i:i + max_length]\n",
    "                  for i in range(0, len(tokens), max_length)]\n",
    "        new_rows = []\n",
    "        for chunk in chunks:\n",
    "            text_chunk = tokenizer.decode(\n",
    "                chunk,\n",
    "                clean_up_tokenization_spaces=True\n",
    "            )\n",
    "            new_row = row.to_dict()\n",
    "            new_row[\"text\"] = text_chunk\n",
    "            new_rows.append(new_row)\n",
    "        return new_rows\n",
    "\n",
    "    new_rows = data.apply(split_text, axis=1).tolist()\n",
    "    flattened_rows = [item for sublist in new_rows for item in sublist]\n",
    "    new_df = pd.DataFrame(flattened_rows)\n",
    "\n",
    "    new_df = new_df.reset_index(drop=True)\n",
    "\n",
    "    print(new_df.shape)\n",
    "    print(new_df.head())\n",
    "    new_df.to_csv(\"./large_dataset/split_text_simple.csv\", index=False)"
   ]
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}