{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.dataframe as dd\n",
    "from dask.diagnostics import ProgressBar\n",
    "import os\n",
    "\n",
    "directory_path = '/Users/fionachow/Documents/NYU/CDS/Fall 2023/CSCI - GA 2271 - Computer Vision/Project/'\n",
    "\n",
    "file_prefix = 'part'\n",
    "\n",
    "def list_files_with_prefix(directory, prefix):\n",
    "    file_paths = []\n",
    "\n",
    "    for root, _, files in os.walk(directory):\n",
    "        for file in files:\n",
    "            if file.startswith(prefix):\n",
    "                absolute_path = os.path.join(root, file)\n",
    "                file_paths.append(absolute_path)\n",
    "\n",
    "    return file_paths\n",
    "\n",
    "laion_file_paths = list_files_with_prefix(directory_path, file_prefix)\n",
    "\n",
    "dataframes = [dd.read_parquet(file) for file in laion_file_paths]\n",
    "combined_dataframe = dd.multi.concat(dataframes)\n",
    "\n",
    "with ProgressBar():\n",
    "    row_count = combined_dataframe.shape[0].compute()\n",
    "    print(row_count)\n",
    "\n",
    "filtered_df = combined_dataframe[combined_dataframe['NSFW'] == \"UNLIKELY\"]\n",
    "\n",
    "num_samples = 225_000\n",
    "selected_rows = filtered_df.sample(frac=num_samples / filtered_df.shape[0].compute())\n",
    "\n",
    "with ProgressBar():\n",
    "    sampled_rows = selected_rows.compute()\n",
    "\n",
    "print(len(sampled_rows))\n",
    "\n",
    "with ProgressBar():\n",
    "    selected_rows.to_parquet('laion_sampled', write_index=False)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bloom",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}