{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e360283a",
   "metadata": {},
   "source": [
    "imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47d34b96",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import concatenate_datasets, load_from_disk\n",
    "from hydra import compose, initialize\n",
    "from hydra.core.global_hydra import GlobalHydra\n",
    "from omegaconf import OmegaConf\n",
    "from tokenizers import Tokenizer, models, pre_tokenizers, trainers\n",
    "from tokenizers.processors import TemplateProcessing\n",
    "\n",
    "from mini_transformer.configs import TokenizerCfg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a0c380c",
   "metadata": {},
   "source": [
    "Choosing tokenizer config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f13dc2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "GlobalHydra.instance().clear()\n",
    "initialize(config_path=\"./configs/tokenizer\", version_base=None)\n",
    "cfg = compose(config_name=\"bpe_16k\")\n",
    "scfg_temp = OmegaConf.merge(OmegaConf.structured(TokenizerCfg), cfg)\n",
    "tokenizer_cfg: TokenizerCfg = OmegaConf.to_object(scfg_temp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5ee3d16",
   "metadata": {},
   "source": [
    "Loading datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d880e59e",
   "metadata": {},
   "outputs": [],
   "source": [
    "small_dataset = load_from_disk(\"./data/processed/processed_small_dataset\")\n",
    "medium_dataset = load_from_disk(\"./data/processed/processed_medium_dataset\")\n",
    "\n",
    "concatenated_dataset = concatenate_datasets(\n",
    "    [\n",
    "        small_dataset[\"train\"],\n",
    "        small_dataset[\"validation\"],\n",
    "        small_dataset[\"test\"],\n",
    "        medium_dataset[\"train\"],\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6b53e40",
   "metadata": {},
   "source": [
    "preparing corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f60ac4b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(tokenizer_cfg.corpus, \"w\", encoding=\"utf-8\") as f:\n",
    "    for example in globals()[tokenizer_cfg.dataset]:\n",
    "        en = example[\"translation\"][\"en\"]\n",
    "        fr = example[\"translation\"][\"fr\"]\n",
    "        if en and fr:\n",
    "            f.write(en + \"\\n\")\n",
    "            f.write(fr + \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7a50455",
   "metadata": {},
   "source": [
    "training the tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29fc0f7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONTINUING_SUBWORD_PREFIX = \"##\"\n",
    "\n",
    "tokenizer = Tokenizer(\n",
    "    models.BPE(\n",
    "        unk_token=tokenizer_cfg.unk_token,\n",
    "        continuing_subword_prefix=CONTINUING_SUBWORD_PREFIX,\n",
    "    )\n",
    ")\n",
    "tokenizer.pre_tokenizer = pre_tokenizers.Sequence(\n",
    "    [pre_tokenizers.Whitespace(), pre_tokenizers.Punctuation()]\n",
    ")\n",
    "trainer = trainers.BpeTrainer(\n",
    "    vocab_size=tokenizer_cfg.vocab_size,\n",
    "    special_tokens=tokenizer_cfg.special_tokens,\n",
    "    continuing_subword_prefix=CONTINUING_SUBWORD_PREFIX,\n",
    ")\n",
    "tokenizer.train([tokenizer_cfg.corpus], trainer)\n",
    "tokenizer.post_processor = TemplateProcessing(\n",
    "    single=f\"{tokenizer_cfg.bos_token} $A {tokenizer_cfg.eos_token}\",\n",
    "    special_tokens=[\n",
    "        (tokenizer_cfg.bos_token, tokenizer_cfg.bos_id),\n",
    "        (tokenizer_cfg.eos_token, tokenizer_cfg.eos_id),\n",
    "    ],\n",
    ")\n",
    "tokenizer.save(tokenizer_cfg.path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "537b9d9d",
   "metadata": {},
   "source": [
    "Loading and testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "443ed7ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tokenizers import Tokenizer\n",
    "\n",
    "tokenizer = Tokenizer.from_file(tokenizer_cfg.path)\n",
    "print(\n",
    "    tokenizer.encode(\"i hereby annouce you as the next president\", add_special_tokens=False).tokens\n",
    ")\n",
    "vocab = tokenizer.get_vocab()\n",
    "\n",
    "print(\"PAD:\", vocab[\"<pad>\"])\n",
    "print(\"BOS:\", vocab[\"<bos>\"])\n",
    "print(\"EOS:\", vocab[\"<eos>\"])\n",
    "print(\"UNK:\", vocab[\"<unk>\"])"
   ]
  }
 ],
 "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
}