Blame for notebooks/load_data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e2289881",
   "metadata": {},
   "source": [
    "Downloading datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef193d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "_ = load_dataset(\"ahazeemi/iwslt14-en-fr\", cache_dir=\"./data/raw\")\n",
    "_ = load_dataset(\"Helsinki-NLP/europarl\", \"en-fr\", cache_dir=\"./data/raw\")\n",
    "_ = load_dataset(\"wmt/wmt14\", \"fr-en\", cache_dir=\"./data/raw\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb48dbf1",
   "metadata": {},
   "source": [
    "Loading datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e0838f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "small_dataset = load_dataset(\"./data/raw/ahazeemi___iwslt14-en-fr\")\n",
    "medium_dataset = load_dataset(\"./data/raw/Helsinki-NLP___europarl\")\n",
    "large_dataset = load_dataset(\"./data/raw/wmt___wmt14\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb452f91",
   "metadata": {},
   "source": [
    "Preprocessing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca4d1867",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "import yaml\n",
    "from datasets import DatasetDict\n",
    "\n",
    "from mini_transformer.utils import set_global_seed\n",
    "\n",
    "with open(\"configs/runtime/local.yaml\", encoding=\"utf-8\") as cfg_file:\n",
    "    runtime_cfg = yaml.safe_load(cfg_file) or {}\n",
    "seed = int(runtime_cfg.get(\"seed\", 0))\n",
    "set_global_seed(seed)\n",
    "\n",
    "\n",
    "def preprocess_data(example, min_len=2, max_len=256, length_ratio=2.0):\n",
    "    \"\"\"\n",
    "    Clean and filter a parallel translation pair (English-French).\n",
    "\n",
    "    Args:\n",
    "        example (dict): {\"translation\": {\"en\": \"...\", \"fr\": \"...\"}}\n",
    "        min_len (int): Minimum tokenized length allowed.\n",
    "        max_len (int): Maximum tokenized length allowed.\n",
    "        length_ratio (float): Max src/tgt length ratio to keep.\n",
    "\n",
    "    Returns:\n",
    "        dict | None: Cleaned translation dict, or None if filtered out.\n",
    "    \"\"\"\n",
    "\n",
    "    src = example[\"translation\"][\"en\"]\n",
    "    tgt = example[\"translation\"][\"fr\"]\n",
    "\n",
    "    # ----------------\n",
    "    # Basic cleanup\n",
    "    # ----------------\n",
    "    def clean_text(text: str) -> str:\n",
    "        # Remove HTML tags\n",
    "        text = re.sub(r\"<.*?>\", \" \", text)\n",
    "\n",
    "        # Remove URLs and emails\n",
    "        text = re.sub(r\"http\\S+|www\\.\\S+\", \" \", text)\n",
    "        text = re.sub(r\"\\S+@\\S+\", \" \", text)\n",
    "\n",
    "        # Normalize quotes and punctuation\n",
    "        text = re.sub(r\"[“”]\", '\"', text)\n",
    "        text = re.sub(r\"[‘’]\", \"'\", text)\n",
    "        text = re.sub(r\"([!?.,])\\1+\", r\"\\1\", text)  # collapse repeated punct\n",
    "\n",
    "        # Remove non-alphabetic junk (keep accents, digits, basic punct)\n",
    "        text = re.sub(r\"[^a-zA-ZÀ-ÖØ-öø-ÿ0-9\\s.,!?']\", \" \", text)\n",
    "\n",
    "        # Collapse whitespace\n",
    "        text = re.sub(r\"\\s+\", \" \", text).strip()\n",
    "        return text.lower()\n",
    "\n",
    "    src = clean_text(src)\n",
    "    tgt = clean_text(tgt)\n",
    "\n",
    "    # ----------------\n",
    "    # Length filtering\n",
    "    # ----------------\n",
    "    src_len, tgt_len = len(src.split()), len(tgt.split())\n",
    "    if src_len < min_len or tgt_len < min_len:\n",
    "        return None\n",
    "    if src_len > max_len or tgt_len > max_len:\n",
    "        return None\n",
    "    if not (1 / length_ratio <= src_len / tgt_len <= length_ratio):\n",
    "        return None\n",
    "\n",
    "    # ----------------\n",
    "    # Optional sanity check: drop if empty after cleaning\n",
    "    # ----------------\n",
    "    if not src or not tgt:\n",
    "        return None\n",
    "\n",
    "    return {\"en\": src, \"fr\": tgt}\n",
    "\n",
    "\n",
    "processed_small_dataset = small_dataset.map(preprocess_data).filter(lambda x: x is not None)\n",
    "processed_medium_dataset = medium_dataset.map(preprocess_data).filter(lambda x: x is not None)\n",
    "processed_large_dataset = large_dataset.map(preprocess_data).filter(lambda x: x is not None)\n",
    "\n",
    "splits = processed_medium_dataset[\"train\"].train_test_split(test_size=0.1, seed=seed)\n",
    "train_split = splits[\"train\"]\n",
    "temp_split = splits[\"test\"].train_test_split(test_size=0.5, seed=seed)\n",
    "processed_medium_dataset = DatasetDict(\n",
    "    {\n",
    "        \"train\": train_split,\n",
    "        \"validation\": temp_split[\"train\"],\n",
    "        \"test\": temp_split[\"test\"],\n",
    "    }\n",
    ")\n",
    "\n",
    "processed_small_dataset.save_to_disk(\"./data/processed/processed_small_dataset\")\n",
    "processed_medium_dataset.save_to_disk(\"./data/processed/processed_medium_dataset\")\n",
    "# processed_large_dataset.save_to_disk(\"./data/processed/processed_large_dataset\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d782d9db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from hydra import compose, initialize\n",
    "from hydra.core.global_hydra import GlobalHydra\n",
    "from omegaconf import OmegaConf\n",
    "from transformers import PreTrainedTokenizerFast\n",
    "\n",
    "from mini_transformer.configs import TokenizerCfg\n",
    "\n",
    "\n",
    "def load_tokenizer(name: str) -> PreTrainedTokenizerFast:\n",
    "    \"\"\"Load a tokenizer defined in configs/tokenizer/<name>.yaml.\"\"\"\n",
    "    GlobalHydra.instance().clear()\n",
    "    initialize(config_path=\"./configs/tokenizer\", version_base=None)\n",
    "    cfg = compose(config_name=\"bpe_4k\")\n",
    "    scfg_temp = OmegaConf.merge(OmegaConf.structured(TokenizerCfg), cfg)\n",
    "    tokenizer_cfg: TokenizerCfg = OmegaConf.to_object(scfg_temp)\n",
    "    tokenizer_path = tokenizer_cfg.path\n",
    "    tokenizer = PreTrainedTokenizerFast(\n",
    "        tokenizer_file=str(tokenizer_path),\n",
    "        bos_token=cfg.get(\"bos_token\"),\n",
    "        eos_token=cfg.get(\"eos_token\"),\n",
    "        unk_token=cfg.get(\"unk_token\"),\n",
    "        pad_token=cfg.get(\"pad_token\"),\n",
    "    )\n",
    "    tokenizer.model_max_length = cfg.get(\"max_seq_len\", tokenizer.model_max_length)\n",
    "    return tokenizer\n",
    "\n",
    "\n",
    "def _resolve_tokenizer(tokenizers, field: str) -> PreTrainedTokenizerFast:\n",
    "    if isinstance(tokenizers, dict):\n",
    "        tok = tokenizers.get(field)\n",
    "        if tok is None:\n",
    "            raise KeyError(f\"No tokenizer provided for field '{field}'\")\n",
    "        return tok\n",
    "    if tokenizers is None:\n",
    "        raise ValueError(\"Tokenizer or field->tokenizer mapping required\")\n",
    "    return tokenizers\n",
    "\n",
    "\n",
    "def _tokenize_text(\n",
    "    tokenizer: PreTrainedTokenizerFast, text: str, include_special_tokens: bool\n",
    ") -> list[str]:\n",
    "    encoded = tokenizer(\n",
    "        text,\n",
    "        add_special_tokens=include_special_tokens,\n",
    "        return_attention_mask=False,\n",
    "        return_token_type_ids=False,\n",
    "    )\n",
    "    input_ids = encoded[\"input_ids\"]\n",
    "    if input_ids and isinstance(input_ids[0], list):\n",
    "        input_ids = input_ids[0]\n",
    "    return tokenizer.convert_ids_to_tokens(input_ids)\n",
    "\n",
    "\n",
    "def _subset_dataset(ds, sample_size: int | None):\n",
    "    if sample_size is None or sample_size >= len(ds):\n",
    "        return ds\n",
    "    return ds.select(range(sample_size))\n",
    "\n",
    "\n",
    "def show_top_k_tokens(\n",
    "    dataset_dict,\n",
    "    tokenizers,\n",
    "    fields=(\"en\", \"fr\"),\n",
    "    k: int = 20,\n",
    "    sample_size: int | None = 2000,\n",
    "    include_special_tokens: bool = False,\n",
    "):\n",
    "    \"\"\"Print the top-k most frequent tokens using configured tokenizer(s).\"\"\"\n",
    "    for split, ds in dataset_dict.items():\n",
    "        ds_view = _subset_dataset(ds, sample_size)\n",
    "        counters = {field: Counter() for field in fields}\n",
    "        for example in ds_view:\n",
    "            for field in fields:\n",
    "                tokenizer = _resolve_tokenizer(tokenizers, field)\n",
    "                tokens = _tokenize_text(tokenizer, example[field], include_special_tokens)\n",
    "                counters[field].update(tokens)\n",
    "        print(f\"\\nSplit: {split}\")\n",
    "        for field in fields:\n",
    "            print(f\"  Top {k} tokens for '{field}':\")\n",
    "            for token, freq in counters[field].most_common(k):\n",
    "                print(f\"    {token!r}: {freq}\")\n",
    "\n",
    "\n",
    "def plot_token_length_barchart(\n",
    "    dataset_dict,\n",
    "    tokenizers,\n",
    "    fields=(\"en\", \"fr\"),\n",
    "    sample_size: int | None = 2000,\n",
    "    include_special_tokens: bool = False,\n",
    "):\n",
    "    \"\"\"Display bar charts of token-length distributions using configured tokenizer(s).\"\"\"\n",
    "    for split, ds in dataset_dict.items():\n",
    "        ds_view = _subset_dataset(ds, sample_size)\n",
    "        length_counters = {field: Counter() for field in fields}\n",
    "        for example in ds_view:\n",
    "            for field in fields:\n",
    "                tokenizer = _resolve_tokenizer(tokenizers, field)\n",
    "                count = len(_tokenize_text(tokenizer, example[field], include_special_tokens))\n",
    "                length_counters[field][count] += 1\n",
    "        support = sorted({length for counter in length_counters.values() for length in counter})\n",
    "        if not support:\n",
    "            continue\n",
    "        bin_size = 50\n",
    "        bin_starts = list(range(0, support[-1] + bin_size, bin_size))\n",
    "        bin_labels = [f\"{start}-{start + bin_size - 1}\" for start in bin_starts]\n",
    "        bin_counts = {field: [0] * len(bin_starts) for field in fields}\n",
    "        for field in fields:\n",
    "            for length, freq in length_counters[field].items():\n",
    "                idx = min(len(bin_starts) - 1, length // bin_size)\n",
    "                bin_counts[field][idx] += freq\n",
    "        x = np.arange(len(bin_starts))\n",
    "        width = 0.8 / max(1, len(fields))\n",
    "        fig, ax = plt.subplots(figsize=(8, 4))\n",
    "        for idx_field, field in enumerate(fields):\n",
    "            counts = bin_counts[field]\n",
    "            ax.bar(x + idx_field * width, counts, width=width, label=field)\n",
    "        ax.set_xticks(x + width * (len(fields) - 1) / 2)\n",
    "        ax.set_xticklabels(bin_labels)\n",
    "        ax.set_xlabel(\"Token count per example\")\n",
    "        ax.set_ylabel(\"Number of examples\")\n",
    "        ax.set_title(f\"Example token lengths in '{split}' split\")\n",
    "        ax.legend()\n",
    "        fig.tight_layout()\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8ed32b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizers = {\n",
    "    \"en\": load_tokenizer(\"bpe_8k\"),\n",
    "    \"fr\": load_tokenizer(\"bpe_8k\"),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04f27397",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Small dataset statistics\n",
    "show_top_k_tokens(processed_small_dataset, tokenizers, k=20, sample_size=2000)\n",
    "plot_token_length_barchart(processed_small_dataset, tokenizers, sample_size=2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7819046",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Medium dataset statistics\n",
    "show_top_k_tokens(processed_medium_dataset, tokenizers, k=20, sample_size=2000)\n",
    "plot_token_length_barchart(processed_medium_dataset, tokenizers, sample_size=2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16332b57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Large dataset statistics\n",
    "show_top_k_tokens(processed_large_dataset, tokenizers, k=20, sample_size=2000)\n",
    "plot_token_length_barchart(processed_large_dataset, tokenizers, sample_size=2000)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mini-transformer",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}