7.2 Data_Preparation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f5f15f90-68bc-4d41-add1-aeaee8a9d21a",
   "metadata": {},
   "source": [
    "# Data Preparation\n",
    "\n",
    "Here we'll carry out some of the data cleaning steps required to prepare data for pretraining. In the video, Sung mentioned an Upstage tool called **Dataverse** which can help us with data cleaning. We can checkout the features of Dataverse at [this link](https://github.com/UpstageAI/dataverse)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f72caad8-34d6-4deb-a7f4-49d8d67925cd",
   "metadata": {
    "height": 47
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4772e5df",
   "metadata": {},
   "source": [
    "## 1. Sourcing datasets for pretraining\n",
    "\n",
    "1. Download an existing dataset from Hugging Face\n",
    "2. Create a dataset of python scripts sourced from Github\n",
    "\n",
    "In both cases the result will be a Hugging Face `Dataset` object, part of the `Datasets` library. We can read more about the properties of Datasets and how to work with them on the [Hugging Face website](https://huggingface.co/docs/datasets/en/index).\n",
    "\n",
    "### Download data from Hugging face\n",
    "\n",
    "The dataset you download here is a subset of a much larger dataset called **Red Pajama**. The full, 1 trillion token dataset is available on Hugging Face at [this link](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f7bc396b-faaa-452a-be0e-7b17a79b8e0e",
   "metadata": {
    "height": 98
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef87430a27104ea2967b97fa19a49632",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/150M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4e3a7c7fe8f24f0ba30f7e36afefc720",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import datasets\n",
    "pretraining_dataset = datasets.load_dataset(\n",
    "    \"upstage/Pretraining_Dataset\",\n",
    "    split=\"train\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ae796e0e-c719-4b39-be29-286aa5a436fb",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['text', 'meta'],\n",
      "    num_rows: 60000\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "print(pretraining_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36d4266d",
   "metadata": {},
   "source": [
    "Only work with the `text` column:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3765e9c6-5f18-4325-8157-206e3bb199fa",
   "metadata": {
    "height": 64
   },
   "outputs": [],
   "source": [
    "pretraining_dataset = pretraining_dataset.select_columns(\n",
    "    [\"text\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d592d6c5",
   "metadata": {},
   "source": [
    "Print a sample:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "84aaf8af-c20b-49ea-ba75-65ea7f87c146",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In 1793 Zaman Shah, a grandson of Ahmad Shah Durrani, won a brief war of succession to become ruler of Afghanistan. The support of Painda Khan, chief of the Baraksai branch of the Durrani tribe, was decisive in his victory. In the next fifty year., the brothers of Zaman shah and the sons of Painda Khan were to dominate the affairs of Afghanistan. The Durrani tribe was very large with several branches and numerous clans. 1 Abmad Shah and his successors belonged to the Sadozai clan, but other clan\n"
     ]
    }
   ],
   "source": [
    "print(pretraining_dataset[0][\"text\"][:500])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56ca9d93",
   "metadata": {},
   "source": [
    "### Compare pretraining and fine-tuning datasets\n",
    "We'll download a fine-tuning dataset to contrast with the pretraining dataset you loaded above. You can read more about the Alpaca model and instruction tuning dataset [here](https://crfm.stanford.edu/2023/03/13/alpaca.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5fb48e62-11fa-47a3-8e2c-8b9ae4011c7f",
   "metadata": {
    "height": 98
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "96d0b59dd01740a0a524830a1868fb6e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/1.39k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b1975b34ad6f46b0bb974ddf9733c0d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/43.4M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51514cd9732b4d3da63d14a667d053ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/52002 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['instruction', 'input', 'output'],\n",
      "    num_rows: 52002\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "instruction_dataset = datasets.load_dataset(\n",
    "    \"c-s-ale/alpaca-gpt4-data\",\n",
    "    split='train'\n",
    ")\n",
    "print(instruction_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "768b3d25-2bca-4a93-b3b0-bd05c55bf70e",
   "metadata": {
    "height": 81
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Instruction: Give three tips for staying healthy.\n",
      "Input: \n",
      "Output: 1. Eat a balanced and nutritious diet: Make sure your meals are inclusive of a variety of fruits and vegetables, lean protein, whole grains, and healthy fats. This helps to provide your body with the essential nutrients to function at its best and can help prevent chronic diseases.\n",
      "\n",
      "2. Engage in regular physical activity: Exercise is crucial for maintaining strong bones, muscles, and cardiovascular health. Aim for at least 150 minutes of moderate aerobic exercise or 75 minutes of vigorous exercise each week.\n",
      "\n",
      "3. Get enough sleep: Getting enough quality sleep is crucial for physical and mental well-being. It helps to regulate mood, improve cognitive function, and supports healthy growth and immune function. Aim for 7-9 hours of sleep each night.\n"
     ]
    }
   ],
   "source": [
    "i=0\n",
    "print(\"Instruction: \" + instruction_dataset[i][\"instruction\"] \n",
    "      + \"\\nInput: \" + instruction_dataset[i][\"input\"] \n",
    "      + \"\\nOutput: \" + instruction_dataset[i][\"output\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42ef4d82",
   "metadata": {},
   "source": [
    "Notice how in contrast to the pretraining data, which is just raw text, fine-tuning datasets are structured into question-answer pairs or instruction-response sets that can include additional input context if required. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6288d0f",
   "metadata": {},
   "source": [
    "### Scrape python code from Github\n",
    "Here, we'll download a selection of python scripts from Github and then prepare them as a Hugging Face `Dataset` object to use in training. \n",
    "\n",
    "The same pattern here will work for preparing any text scraped from the web."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c14d65d1-1a52-4777-95f8-4a2f84a3c9ac",
   "metadata": {
    "height": 79
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "\n",
    "code_dir = \"./code\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7e4a0b0b-2e76-4567-9ef5-f3b6d8020851",
   "metadata": {
    "height": 200
   },
   "outputs": [],
   "source": [
    "urls = [\n",
    "    \"https://raw.githubusercontent.com/TheAlgorithms/Python/master/searches/double_linear_search_recursion.py\",\n",
    "    \"https://raw.githubusercontent.com/KosingZhu/tensorflow/master/tensorflow/python/tools/module_util.py\",\n",
    "    \"https://raw.githubusercontent.com/EricRemmerswaal/tensorflow/master/tensorflow/python/distribute/distribute_coordinator_context.py\",\n",
    "    \"https://raw.githubusercontent.com/computationalartist/tensorflow/master/tensorflow/python/ops/numpy_ops/integration_test/benchmarks/numpy_mlp.py\",\n",
    "    \"https://raw.githubusercontent.com/Van-an/tensorflow/master/tensorflow/python/distribute/coordinator/values.py\",\n",
    "    \"https://raw.githubusercontent.com/nkgwer/tensorflow/master/tensorflow/lite/tools/visualize.py\",\n",
    "    \"https://raw.githubusercontent.com/gitblazer/youtube-dl/master/youtube_dl/version.py\",\n",
    "    \"https://raw.githubusercontent.com/Joshua-Barawa/My-Photos/master/venv/lib/python3.8/site-packages/django/contrib/messages/__init__.py\",\n",
    "    \"https://raw.githubusercontent.com/PaliC/pytorch/master/test/fx/test_subgraph_rewriter.py\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1818ac26",
   "metadata": {},
   "source": [
    "Retrieve the python scripts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4bc904ed-da23-40c0-b600-bcb1feebbd31",
   "metadata": {
    "height": 149
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working on url: https://raw.githubusercontent.com/TheAlgorithms/Python/master/searches/double_linear_search_recursion.py\n",
      "Working on url: https://raw.githubusercontent.com/KosingZhu/tensorflow/master/tensorflow/python/tools/module_util.py\n",
      "Working on url: https://raw.githubusercontent.com/EricRemmerswaal/tensorflow/master/tensorflow/python/distribute/distribute_coordinator_context.py\n",
      "Working on url: https://raw.githubusercontent.com/computationalartist/tensorflow/master/tensorflow/python/ops/numpy_ops/integration_test/benchmarks/numpy_mlp.py\n",
      "Working on url: https://raw.githubusercontent.com/Van-an/tensorflow/master/tensorflow/python/distribute/coordinator/values.py\n",
      "Working on url: https://raw.githubusercontent.com/nkgwer/tensorflow/master/tensorflow/lite/tools/visualize.py\n",
      "Working on url: https://raw.githubusercontent.com/gitblazer/youtube-dl/master/youtube_dl/version.py\n",
      "Working on url: https://raw.githubusercontent.com/Joshua-Barawa/My-Photos/master/venv/lib/python3.8/site-packages/django/contrib/messages/__init__.py\n",
      "Working on url: https://raw.githubusercontent.com/PaliC/pytorch/master/test/fx/test_subgraph_rewriter.py\n"
     ]
    }
   ],
   "source": [
    "for url in urls:\n",
    "    print(f\"Working on url: {url}\")\n",
    "    response = requests.get(url)\n",
    "    file_name = os.path.basename(url)\n",
    "    file_path = os.path.join(code_dir, file_name)\n",
    "    \n",
    "    with open(file_path, \"wb\") as file:\n",
    "        file.write(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "807f9164-ed85-4df8-bb01-e5ed5a41cba5",
   "metadata": {
    "height": 64
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "double_linear_search_recursion.py\n",
      "module_util.py\n",
      "distribute_coordinator_context.py\n",
      "numpy_mlp.py\n",
      "values.py\n",
      "visualize.py\n",
      "version.py\n",
      "__init__.py\n",
      "test_subgraph_rewriter.py\n"
     ]
    }
   ],
   "source": [
    "files = os.listdir(code_dir)\n",
    "for file in files:\n",
    "    print(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7561587",
   "metadata": {},
   "source": [
    "Concatenate scripts into a list:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0c9b603c-9fd4-410b-b013-802a6d055c49",
   "metadata": {
    "height": 98
   },
   "outputs": [],
   "source": [
    "code_dataset = []\n",
    "for file in os.listdir(code_dir):\n",
    "    code_dataset.append(\n",
    "        {'text': open(os.path.join(code_dir, file), 'r').read()}\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01df9274",
   "metadata": {},
   "source": [
    "Convert list to Hugging Face `Dataset` object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "cfa9b259-5434-4d2c-96a1-d0c68149a3cd",
   "metadata": {
    "height": 47
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['text'],\n",
      "    num_rows: 9\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "code_dataset = datasets.Dataset.from_list(code_dataset)\n",
    "print(code_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d33dce90",
   "metadata": {},
   "source": [
    "Combine the python code dataset with the pretraining dataset you downloaded above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3d2b129e-4307-4b8c-a94c-b348fd66a7fe",
   "metadata": {
    "height": 81
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['text'],\n",
      "    num_rows: 60009\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "dataset = datasets.concatenate_datasets(\n",
    "    [pretraining_dataset, code_dataset]\n",
    ")\n",
    "print(dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73cf07e4-5be3-4e1f-99e0-c14013d10fd3",
   "metadata": {},
   "source": [
    "## 2. Data cleaning\n",
    "\n",
    "**Bad Data →** duplicated, typos & spelling mistakes, inconsistent, hallucinations, toxic, etc.<br>\n",
    "**Data Cleaning →** Deduplication, quality filtering (remove text of other languages except eng), content filters (remove toxic words like stupid, idiot), privacy reduction (remove names, phone no and replace with [Name], [address]).\n",
    "\n",
    "1. Filter out samples that are too short\n",
    "2. Remove repetitions within a single text example\n",
    "3. Remove duplicated documents\n",
    "4. Quality filter to remove non-English texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "185ba49c-c777-4429-b01c-f0b03f33a4df",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "60009"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.num_rows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c657a5ba",
   "metadata": {},
   "source": [
    "### Remove examples that are too short"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b6f293b4-486d-4361-ba9f-6038864e67de",
   "metadata": {
    "height": 200
   },
   "outputs": [],
   "source": [
    "import heapq\n",
    "\n",
    "def paragraph_length_filter(x):\n",
    "    \"\"\"Returns False iff a page has too few lines or lines are too short.\"\"\"\n",
    "    lines = x['text'].split('\\n')\n",
    "    if (\n",
    "        len(lines) < 3\n",
    "        or min(heapq.nlargest(3, [len(line) for line in lines])) < 3\n",
    "    ):\n",
    "        return False\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2755ea1a-a9d3-4771-950e-73304d21252b",
   "metadata": {
    "height": 81
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64a8a978f69f4bf09d923a884659c0a0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/60009 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset = dataset.filter(\n",
    "    paragraph_length_filter,\n",
    "    load_from_cache_file=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a87c47c4-9331-4b44-a5e2-6b81dfbd7bf1",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "52357"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.num_rows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eced32d0",
   "metadata": {},
   "source": [
    "### Remove repeated text within training examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6b0a79c4-e714-42d4-bdb7-d3930d74c5de",
   "metadata": {
    "height": 268
   },
   "outputs": [],
   "source": [
    "def find_duplicates(paragraphs):\n",
    "    \"\"\"\n",
    "    Use this function to find the number of repetitions \n",
    "    in the paragraphs.\n",
    "    \"\"\"\n",
    "    unique_x = set()\n",
    "    duplicate_chars = 0\n",
    "    duplicate_elements = 0\n",
    "    for element in paragraphs:\n",
    "        if element in unique_x:\n",
    "            duplicate_chars += len(element)\n",
    "            duplicate_elements += 1\n",
    "        else:\n",
    "            unique_x.add(element)\n",
    "    return duplicate_elements, duplicate_chars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "f5dc415c-35e1-4d63-95bc-74775bfdaba4",
   "metadata": {
    "height": 251
   },
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "def paragraph_repetition_filter(x):\n",
    "    \"\"\"\n",
    "    Returns False iff a page has too many repetitions.\n",
    "    \"\"\"\n",
    "    text = x['text']\n",
    "    paragraphs = re.compile(r\"\\n{2,}\").split(text.strip())                # Split by paragraphs (2 or more newlines)\n",
    "    paragraphs_duplicates, char_duplicates = find_duplicates(paragraphs)  # Find number of duplicates in paragraphs\n",
    "    if paragraphs_duplicates / len(paragraphs) > 0.3:\n",
    "        return False\n",
    "    if char_duplicates / len(text) > 0.2:\n",
    "        return False\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "522708f3-58fa-4f7e-bf34-a3f282b309e4",
   "metadata": {
    "height": 81
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "415dc4b71f7d4282853d86ecf7525162",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/52357 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset = dataset.filter(\n",
    "    paragraph_repetition_filter,\n",
    "    load_from_cache_file=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e81f0202-b3ef-4953-8203-5062ad20402d",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "52327"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.num_rows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "515b4f79",
   "metadata": {},
   "source": [
    "### Deduplication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b8dcd316-20f1-4def-959e-ce3dfe9426c4",
   "metadata": {
    "height": 268
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c93298fa581b4a1abb4617a0c00be195",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/52327 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def deduplication(ds):\n",
    "    def dedup_func(x):\n",
    "        \"\"\"Use this function to remove duplicate entries\"\"\"\n",
    "        if x['text'] in unique_text:\n",
    "            return False\n",
    "        else:\n",
    "            unique_text.add(x['text'])\n",
    "            return True\n",
    "\n",
    "    unique_text = set()\n",
    "\n",
    "    ds = ds.filter(dedup_func, load_from_cache_file=False, num_proc=1)\n",
    "    return ds\n",
    "\n",
    "dataset = deduplication(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "754444d7-def6-499c-9da8-d18136a4ee6b",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "43598"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.num_rows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72eb3a2c",
   "metadata": {},
   "source": [
    "### Quality filter - Language\n",
    "\n",
    "Remove any text examples that are in a language other than English. The code here uses a language detection model called fastText. Read about fastText [here](https://fasttext.cc/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "64ae5c3e-ae16-495c-a651-0b0200169569",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "# !pip install fasttext"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f11bc056-faeb-470d-acc2-9c1bae47c3c1",
   "metadata": {
    "height": 319
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Parameter 'function'=<function english_language_filter.<locals>.is_english at 0x7f1ac010ad40> of the transform datasets.arrow_dataset.Dataset.filter@2.0.1 couldn't be hashed properly, a random hash was used instead. Make sure your transforms and parameters are serializable with pickle or dill for the dataset fingerprinting and caching to work. If you reuse this transform, the caching mechanism will consider it to be different from the previous calls and recompute everything. This warning is only showed once. Subsequent hashing failures won't be showed.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b9a07bbffe647199dac05f97eeeb5b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/43598 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import urllib\n",
    "from fasttext.FastText import _FastText\n",
    "\n",
    "def english_language_filter(ds):\n",
    "    # load language detection model\n",
    "    model = _FastText('./models/L2_language_model.bin')\n",
    "    \n",
    "    def is_english(x):\n",
    "        # Predict language of the text and probability\n",
    "        language, score = model.predict(x['text'].replace(\"\\n\", \"\"))\n",
    "\n",
    "        language = language[0].split(\"__\")[2]\n",
    "        return score > 0.4 and language == \"en\" # change code here if building a model in another language\n",
    "\n",
    "    ds = ds.filter(is_english, load_from_cache_file=False, num_proc=1)\n",
    "    return ds\n",
    "\n",
    "dataset = english_language_filter(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "f71f18f5-53b6-42fa-bb3a-d321904e3dc1",
   "metadata": {
    "height": 30
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "40474"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.num_rows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "428dcbc5",
   "metadata": {},
   "source": [
    "## 3. Save the dataset to disk\n",
    "\n",
    "Read more about the parquet data format [here](https://parquet.apache.org/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "21e68cc7-058a-4783-9dd5-96c2ef9a752e",
   "metadata": {
    "height": 47
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e8442618517748d29441e2b198a5ec58",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/41 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "197101804"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file_path = \"./data/preprocessed_dataset.parquet\"\n",
    "dataset.to_parquet(file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fd2cb60-e263-4f0a-8c41-eea29a210c1e",
   "metadata": {
    "height": 30
   },
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}