{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f0871aab-843d-4a9b-9e19-98334a965bca",
   "metadata": {},
   "source": [
    "# Data Loading"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f347d1b-70b4-49c9-bdae-8eb39af7e701",
   "metadata": {},
   "source": [
    "## 1. dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "1a4927dd-af99-4e83-9e26-bb77907d5c52",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from pathlib import Path\n",
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bda9048-445f-4d3f-ad35-369f05e7e099",
   "metadata": {},
   "source": [
    "### Connect to Educational fineweb dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "1df6e0e5-9472-445f-b5a9-65885239d915",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 1 Connect\n",
    "fineweb = load_dataset(\n",
    "    \"HuggingFaceFW/fineweb-edu\",\n",
    "    name=\"sample-10BT\",\n",
    "    split=\"train\",\n",
    "    streaming=True,\n",
    ")\n",
    "\n",
    "    # add to our corpus"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "361e75d6-72c5-408d-bdf5-d7b8194dfa03",
   "metadata": {},
   "source": [
    "### Get a sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "50a81aba-88cf-4bcd-ba29-d9a7d9dac528",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 150 documents\n"
     ]
    }
   ],
   "source": [
    "fineweb_sample = []\n",
    "\n",
    "for row in fineweb:\n",
    "    \n",
    "    fineweb_sample.append(row)\n",
    "    \n",
    "    if len(fineweb_sample) >= 150:\n",
    "        break\n",
    "\n",
    "print(f\"Loaded {len(fineweb_sample):,} documents\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb48c2a4-c524-47fe-a3e4-dc6bae563f34",
   "metadata": {},
   "source": [
    "### Calculate dataset statistics "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "6319b538-50cf-4f31-8616-2f9c7fdd1b20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Documents:          150\n",
      "Characters/tokens:  533,917\n",
      "Average length:     3,559\n",
      "Minimum length:     335\n",
      "Maximum length:     18,639\n"
     ]
    }
   ],
   "source": [
    "# Cell 5 — Calculate dataset statistics\n",
    "\n",
    "num_documents = 0\n",
    "total_characters = 0\n",
    "min_length = None\n",
    "max_length = 0\n",
    "\n",
    "for row in fineweb_sample:\n",
    "    \n",
    "    text = row[\"text\"]\n",
    "\n",
    "    length = len(text)\n",
    "\n",
    "    num_documents += 1\n",
    "    total_characters += length\n",
    "\n",
    "    if min_length is None:\n",
    "        min_length = length\n",
    "    else:\n",
    "        min_length = min(min_length, length)\n",
    "\n",
    "    max_length = max(max_length, length)\n",
    "\n",
    "print(f\"Documents:          {num_documents:,}\")\n",
    "print(f\"Characters/tokens:  {total_characters:,}\")\n",
    "print(f\"Average length:     {total_characters / num_documents:,.0f}\")\n",
    "print(f\"Minimum length:     {min_length:,}\")\n",
    "print(f\"Maximum length:     {max_length:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1402db28-f863-420b-b9be-ab2ed999392b",
   "metadata": {},
   "source": [
    "### Compare corpus lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "100cdb47-f15e-4eac-9505-ecbd05cf7263",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shakespeare text_length : 1115394\n",
      "Tiny Stories text_length : 299987\n"
     ]
    }
   ],
   "source": [
    "# Shakespear\n",
    "sp_corpus = \"../data/tiny_shakespear.txt\"\n",
    "sp_data_path = Path(sp_corpus)\n",
    "sp_text = sp_data_path.read_text(encoding=\"utf-8\",)\n",
    "\n",
    "sp_text_length = len(sp_text)\n",
    "print(f\"Shakespeare text_length : {sp_text_length}\")\n",
    "\n",
    "# Tiny Stories\n",
    "ts_corpus = \"../data/tinystories_300k.txt\"\n",
    "ts_data_path = Path(ts_corpus)\n",
    "ts_text = ts_data_path.read_text(encoding=\"utf-8\",)\n",
    "\n",
    "ts_text_length = len(ts_text)\n",
    "print(f\"Tiny Stories text_length : {ts_text_length}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "3bafffc9-0044-42f4-8f17-7e8a16b45b9b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'\\nOnce upon a time' 221\n",
      "'\\nOne day' 107\n"
     ]
    }
   ],
   "source": [
    "#print(repr(sp_text[:2_000]))\n",
    "#print(repr(ts_text[:2_000]))\n",
    "for marker in [\"\\nOnce upon a time\", \"\\nOne day\"]:\n",
    "    print(repr(marker), ts_text.count(marker))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a10cc60-80ab-41f7-9270-3e6fd3a4bfe7",
   "metadata": {},
   "source": [
    "### Save to parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "3e69c0e5-9ac2-4dd4-bcef-59e410dd92a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved 150 documents\n",
      "Path: ../data/fineweb-edu-sample.parquet\n"
     ]
    }
   ],
   "source": [
    "output_path = Path(\"../data/fineweb-edu-sample.parquet\")\n",
    "output_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "table = pa.Table.from_pylist(fineweb_sample)\n",
    "\n",
    "pq.write_table(\n",
    "    table,\n",
    "    output_path,\n",
    "    )\n",
    "\n",
    "print(f\"Saved {len(fineweb_sample):,} documents\")\n",
    "print(f\"Path: {output_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4692e4c6-92ea-4eef-96ef-3ee602c52aaa",
   "metadata": {},
   "source": [
    "## Read from parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "01f76dc2-3a9f-48ab-8953-a571acd6a485",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rows: 150\n",
      "Columns: ['text', 'id', 'dump', 'url', 'file_path', 'language', 'language_score', 'token_count', 'score', 'int_score']\n",
      "text: string\n",
      "id: string\n",
      "dump: string\n",
      "url: string\n",
      "file_path: string\n",
      "language: string\n",
      "language_score: double\n",
      "token_count: int64\n",
      "score: double\n",
      "int_score: int64\n",
      "================================================================================\n",
      "Loaded 150 documents\n",
      "The Independent Jane\n",
      "For all the love, romance and scandal in Jane Austen’s books, what they are really about is freedom and independence. Independence of thought and the freedom to choose.\n",
      "Elizabeth’s refusal of Mr. Collins offer of marriage showed an independence seldom seen in heroines of the day. Her refusal of Mr. Darcy while triggered by anger showed a level of independence that left him shocked and stunned.\n",
      "The freedom she exhibited in finally accepting him in direct defiance of Lady Cath\n"
     ]
    }
   ],
   "source": [
    "import pyarrow.parquet as pq\n",
    "\n",
    "table = pq.read_table(\"../data/fineweb-edu-sample.parquet\")\n",
    "\n",
    "print(f\"Rows: {table.num_rows:,}\")\n",
    "print(f\"Columns: {table.column_names}\")\n",
    "print(table.schema)\n",
    "\n",
    "print(\"=\" * 80)\n",
    "\n",
    "texts = table[\"text\"].to_pylist()\n",
    "\n",
    "print(f\"Loaded {len(texts):,} documents\")\n",
    "print(texts[0][:500])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ded286b8-45bd-4ac1-b495-5bae27d7859d",
   "metadata": {},
   "source": [
    "## Join em"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c7d5710-6b11-485a-8b8e-87935b5f5274",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from pathlib import Path\n",
    "# import pandas as pd\n",
    "\n",
    "# DATA_DIR = Path(\"data\")\n",
    "\n",
    "# rows = []\n",
    "\n",
    "# # 1. Load Parquet files\n",
    "# for path in DATA_DIR.glob(\"*.parquet\"):\n",
    "#     df = pd.read_parquet(path)\n",
    "\n",
    "#     # Adapt this to your actual column names\n",
    "#     if \"text\" in df.columns:\n",
    "#         for text in df[\"text\"].dropna():\n",
    "#             rows.append({\n",
    "#                 \"text\": str(text),\n",
    "#                 \"source\": path.name,\n",
    "#             })\n",
    "\n",
    "# # 2. Load TXT files\n",
    "# for path in DATA_DIR.glob(\"*.txt\"):\n",
    "#     text = path.read_text(encoding=\"utf-8\", errors=\"replace\").strip()\n",
    "\n",
    "#     if text:\n",
    "#         rows.append({\n",
    "#             \"text\": text,\n",
    "#             \"source\": path.name,\n",
    "#         })\n",
    "\n",
    "# # 3. Create unified dataset\n",
    "# dataset = pd.DataFrame(rows)\n",
    "\n",
    "# # 4. Basic cleaning\n",
    "# dataset[\"text\"] = (\n",
    "#     dataset[\"text\"]\n",
    "#     .str.replace(r\"\\s+\", \" \", regex=True)\n",
    "#     .str.strip()\n",
    "# )\n",
    "\n",
    "# dataset = dataset[dataset[\"text\"].str.len() > 0]\n",
    "\n",
    "# # 5. Remove exact duplicates\n",
    "# dataset = dataset.drop_duplicates(subset=[\"text\"])\n",
    "\n",
    "# # 6. Shuffle\n",
    "# dataset = dataset.sample(frac=1, random_state=42).reset_index(drop=True)\n",
    "\n",
    "# # 7. Save\n",
    "# dataset.to_parquet(\"training_dataset.parquet\", index=False)\n",
    "\n",
    "# print(dataset.shape)\n",
    "# print(dataset.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76e50a4e-596f-44ce-b832-acef209df30d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from datasets import load_dataset\n",
    "\n",
    "# tiny_stories_ds = load_dataset(\"roneneldan/TinyStories\", split=\"train\", streaming=True,)\n",
    "\n",
    "# # for story in tiny_stories_ds.take(3):\n",
    "# #     print(\"=\" * 80)\n",
    "# #     print(story[\"text\"])\n",
    "    \n",
    "# TARGET = 300_000\n",
    "# chunks = []\n",
    "# n = 0\n",
    "\n",
    "# for row in ds:\n",
    "#     story = row[\"text\"].strip() + \"\\n\"\n",
    "    \n",
    "#     if n + len(story) > TARGET:\n",
    "#         break\n",
    "    \n",
    "#     chunks.append(story)\n",
    "#     n += len(story)\n",
    "\n",
    "# with open(\"../data/tinystories_300k.txt\", \"w\", encoding=\"utf-8\") as f:\n",
    "#     f.write(\"\".join(chunks))\n",
    "\n",
    "# print(f\"Characters: {n:,}\")\n",
    "# print(f\"Stories: {len(chunks):,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfdc3c53-8bf3-4a5b-a80b-31551d452fed",
   "metadata": {},
   "outputs": [],
   "source": [
    "## JSONL"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00a6e19f-cc40-4e56-98f1-86e44b78d15b",
   "metadata": {},
   "source": [
    "## Parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47a64994-be0d-4b77-9d7b-564d9c794d9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datatrove.pipeline.readers import ParquetReader\n",
    "\n",
    "# limit determines how many documents will be streamed (remove for all)\n",
    "data_reader = ParquetReader(\"hf://datasets/HuggingFaceFW/fineweb-edu\", glob_pattern=\"data/*/*.parquet\", limit=1000)\n",
    "# or to fetch a specific dump CC-MAIN-2024-10,  eplace \"CC-MAIN-2024-10\" with \"sample/100BT\" to use the 100BT sample\n",
    "data_reader = ParquetReader(\"hf://datasets/HuggingFaceFW/fineweb-edu/CC-MAIN-2024-10\", limit=1000) \n",
    "for document in data_reader():\n",
    "    # do something with document\n",
    "    print(document)\n",
    "\n",
    "###############################    \n",
    "# OR for a processing pipeline:\n",
    "###############################\n",
    "\n",
    "from datatrove.executor import LocalPipelineExecutor\n",
    "from datatrove.pipeline.readers import ParquetReader\n",
    "from datatrove.pipeline.filters import LambdaFilter\n",
    "from datatrove.pipeline.writers import JsonlWriter\n",
    "\n",
    "pipeline_exec = LocalPipelineExecutor(\n",
    "    pipeline=[\n",
    "        # replace \"CC-MAIN-2024-10\" with \"sample/100BT\" to use the 100BT sample\n",
    "        ParquetReader(\"hf://datasets/HuggingFaceFW/fineweb-edu/CC-MAIN-2024-10\", limit=1000),\n",
    "        LambdaFilter(lambda doc: \"hugging\" in doc.text),\n",
    "        JsonlWriter(\"some-output-path\")\n",
    "    ],\n",
    "    tasks=10\n",
    ")\n",
    "pipeline_exec.run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63abdde7-0abf-4a00-8f20-c56774e0d5bd",
   "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.14.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}