{ "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/.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 }