{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def read_corpus(corpus_path:str):\n",
    "    with open(corpus_path, 'r', encoding='utf-8') as f:\n",
    "        text = f.read()\n",
    "    return text\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "class BPEGujaratiTokenizer:\n",
    "    def __init__(self, corpus_path:str, max_vocab_size:int=5001, sample_size:int=50000):\n",
    "        self.corpus = read_corpus(corpus_path)\n",
    "        self.max_vocab_size = max_vocab_size\n",
    "        self.corpus_vocab = sorted(list(set(self.corpus)))\n",
    "        self.corpus_vocab_size = len(self.corpus_vocab)\n",
    "        self.stoi = { ch:i for i,ch in enumerate(self.corpus_vocab) }\n",
    "        self.itos = { i:ch for i,ch in enumerate(self.corpus_vocab) }\n",
    "        self.sample_size = sample_size\n",
    "\n",
    "        self.vocab, self.merges = self.train_bpe(self.corpus, self.max_vocab_size, self.sample_size)\n",
    "\n",
    "\n",
    "    def get_stats(self, ids):\n",
    "        counts = {}\n",
    "        for pair in zip(ids, ids[1:]):\n",
    "            counts[pair] = counts.get(pair, 0) + 1\n",
    "        return counts\n",
    "\n",
    "\n",
    "    def merge(self,ids, pair, idx):\n",
    "        newids = []\n",
    "        i = 0\n",
    "        while i < len(ids):\n",
    "            if i < len(ids) - 1 and ids[i] == pair[0] and ids[i+1] == pair[1]:\n",
    "                newids.append(idx)\n",
    "                i += 2\n",
    "            else:\n",
    "                newids.append(ids[i])\n",
    "                i += 1\n",
    "        return newids\n",
    "\n",
    "\n",
    "\n",
    "    def train_bpe(self, corpus, max_vocab_size, sample_size=None):\n",
    "        self.vocab = {idx: bytes([idx]) for idx in range(256)}\n",
    "        print(f\"Before Training Vocab length {len(self.vocab)}\")\n",
    "        if sample_size :\n",
    "            corpus = corpus[:sample_size]\n",
    "        num_merges = max_vocab_size - len(self.vocab)\n",
    "        print(f\"num_merges required {num_merges}\")\n",
    "        tokens = corpus.encode('utf-8')\n",
    "        tokens= list(map(int, tokens))\n",
    "        ids = list(tokens)\n",
    "        self.merges = {} # (int, int) -> int\n",
    "        print(f\"Before training: ids length: {len(ids)}\")\n",
    "        print(f\"Before training: tokens length: {len(tokens)}\")\n",
    "        print(\"Before training: merges length: \", len(self.merges))\n",
    "\n",
    "        for i in range(num_merges):\n",
    "            stats = self.get_stats(ids)\n",
    "            pair = max(stats, key=stats.get)\n",
    "            idx = len(self.vocab)+i\n",
    "            ids = self.merge(ids, pair, idx)\n",
    "            self.merges[pair] = idx\n",
    "        # merge the vocab\n",
    "        \n",
    "        for (p0, p1), idx in self.merges.items():\n",
    "            self.vocab[idx] = self.vocab[p0] + self.vocab[p1]\n",
    "        print(f\"After training: ids length: {len(ids)}\")\n",
    "        print(f\"After training: tokens length: {len(tokens)}\")\n",
    "        print(\"After training: merges length: \", len(self.merges))\n",
    "        print(f\"After Training Vocab length {len(self.vocab)}\")\n",
    "        print(f\"compression ratio: {len(tokens) / len(ids):.2f}X\")\n",
    "        return self.vocab, self.merges\n",
    "\n",
    "    def encode(self, text):\n",
    "        tokens = list(text.encode(\"utf-8\"))\n",
    "        while len(tokens) >= 2:\n",
    "            stats = self.get_stats(tokens)\n",
    "            pair = min(stats, key=lambda p: self.merges.get(p, float(\"inf\")))\n",
    "            if pair not in self.merges:\n",
    "                break # nothing else can be merged\n",
    "            idx = self.merges[pair]\n",
    "            tokens = self.merge(tokens, pair, idx)\n",
    "        return tokens\n",
    "\n",
    "    \n",
    "    def decode(self, tokens):\n",
    "        tokens = b\"\".join(self.vocab[idx] for idx in tokens)\n",
    "        text = tokens.decode(\"utf-8\", errors=\"replace\")\n",
    "        return text\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Before Training Vocab length 256\n",
      "num_merges required 4744\n",
      "Before training: ids length: 755940\n",
      "Before training: tokens length: 755940\n",
      "Before training: merges length:  0\n",
      "After training: ids length: 76306\n",
      "After training: tokens length: 755940\n",
      "After training: merges length:  4744\n",
      "After Training Vocab length 5000\n",
      "compression ratio: 9.91X\n",
      "Time taken to train: 199.02717900276184 seconds\n",
      "--------------------------------\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "start_time = time.time()\n",
    "tokenizer = BPEGujaratiTokenizer(corpus_path=\"gu_corpus.txt\", max_vocab_size=5000, sample_size=300000)\n",
    "end_time = time.time()\n",
    "print(f\"Time taken to train: {end_time - start_time} seconds\")\n",
    "print(\"--------------------------------\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[294, 307, 164, 292, 431, 325, 317, 3229, 444]\n",
      "Time taken to encode: 0.0007619857788085938 seconds\n",
      "--------------------------------\n",
      "હું તને પ્રેમ કરું છું\n",
      "Time taken to decode: 0.0004019737243652344 seconds\n",
      "--------------------------------\n",
      "original:  હું આજે ખૂબ ખુશ છું.\n",
      "encoded:  [294, 307, 1414, 853, 928, 1793, 482, 444, 46]\n",
      "decoded:  હું આજે ખૂબ ખુશ છું.\n",
      "True\n",
      "original:  તું શું કરે છે? \n",
      "encoded:  [3519, 182, 307, 391, 4339, 63, 32]\n",
      "decoded:  તું શું કરે છે? \n",
      "True\n",
      "original:  મને ચા પીવી છે. \n",
      "encoded:  [274, 292, 154, 758, 519, 269, 296, 46, 32]\n",
      "decoded:  મને ચા પીવી છે. \n",
      "True\n",
      "original:  એ બધું સરસ છે. \n",
      "encoded:  [512, 4222, 3997, 2296, 3648, 46, 32]\n",
      "decoded:  એ બધું સરસ છે. \n",
      "True\n",
      "original:  આ પુસ્તક ખૂબ રસપ્રદ છે. \n",
      "encoded:  [256, 4844, 2469, 290, 3227, 311, 4738, 345, 3648, 46, 32]\n",
      "decoded:  આ પુસ્તક ખૂબ રસપ્રદ છે. \n",
      "True\n",
      "original:  તારે ક્યારે આવવું છે? \n",
      "encoded:  [2460, 335, 484, 340, 793, 296, 63, 32]\n",
      "decoded:  તારે ક્યારે આવવું છે? \n",
      "True\n",
      "original:  આ મારો મિત્ર છે. \n",
      "encoded:  [256, 134, 309, 763, 4071, 3648, 46, 32]\n",
      "decoded:  આ મારો મિત્ર છે. \n",
      "True\n",
      "original:  હું શાકભાજી લઈ આવ્યો છું. \n",
      "encoded:  [294, 307, 182, 533, 455, 397, 666, 451, 655, 2301, 444, 46, 32]\n",
      "decoded:  હું શાકભાજી લઈ આવ્યો છું. \n",
      "True\n",
      "original:  આકાશ માં વાદળ છે. \n",
      "encoded:  [256, 134, 290, 676, 1546, 181, 390, 343, 3648, 46, 32]\n",
      "decoded:  આકાશ માં વાદળ છે. \n",
      "True\n",
      "original:  શાળા ક્યારે શરૂ થશે? \n",
      "encoded:  [332, 547, 581, 484, 3680, 165, 1168, 63, 32]\n",
      "decoded:  શાળા ક્યારે શરૂ થશે? \n",
      "True\n",
      "original:  આ પુસ્તક ખૂબ રસપ્રદ છે.\n",
      "encoded:  [256, 4844, 2469, 290, 3227, 311, 4738, 345, 3648, 46]\n",
      "decoded:  આ પુસ્તક ખૂબ રસપ્રદ છે.\n",
      "True\n",
      "Time taken to decode: 0.009686946868896484 seconds\n",
      "--------------------------------\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "print(tokenizer.encode(\"હું તને પ્રેમ કરું છું\"))\n",
    "end_time = time.time()\n",
    "print(f\"Time taken to encode: {end_time - start_time} seconds\")\n",
    "print(\"--------------------------------\")\n",
    "start_time = time.time()\n",
    "print(tokenizer.decode(tokenizer.encode(\"હું તને પ્રેમ કરું છું\")))\n",
    "end_time = time.time()\n",
    "print(f\"Time taken to decode: {end_time - start_time} seconds\")\n",
    "print(\"--------------------------------\")\n",
    "start_time = time.time()\n",
    "sentences = [\"હું આજે ખૂબ ખુશ છું.\",\"તું શું કરે છે? \",\"મને ચા પીવી છે. \",\"એ બધું સરસ છે. \",\"આ પુસ્તક ખૂબ રસપ્રદ છે. \",\"તારે ક્યારે આવવું છે? \",\"આ મારો મિત્ર છે. \",\"હું શાકભાજી લઈ આવ્યો છું. \",\"આકાશ માં વાદળ છે. \",\"શાળા ક્યારે શરૂ થશે? \",'આ પુસ્તક ખૂબ રસપ્રદ છે.']\n",
    "for sentence in sentences:\n",
    "    print(\"original: \", sentence)\n",
    "    print(\"encoded: \", tokenizer.encode(sentence))\n",
    "    print(\"decoded: \", tokenizer.decode(tokenizer.encode(sentence)))\n",
    "    print(tokenizer.decode(tokenizer.encode(sentence)) == sentence)\n",
    "end_time = time.time()\n",
    "print(f\"Time taken to decode: {end_time - start_time} seconds\")\n",
    "print(\"--------------------------------\")   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kaggle": {
   "accelerator": "none",
   "dataSources": [
    {
     "datasetId": 6426227,
     "sourceId": 10374225,
     "sourceType": "datasetVersion"
    }
   ],
   "dockerImageVersionId": 30822,
   "isGpuEnabled": false,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "venv",
   "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.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}