{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "VRdIio8gyD5F" }, "source": [ "# **Assignment 0: Tokenization!**\n", "\n", "# **Code Submission Instructions**\n", "1. Set `SUBMISSION_READY = True`\n", "2. In Google Colab, please click File > Download > Download .py\n", "3. Upload the .py file to Gradescope\n", "# **Introduction**\n", "\n", "This notebook implements a **byte-level tokenizer** using a **trie-based vocabulary** and an iterative **byte-pair merging algorithm** (similar to BPE).\n", "\n", "The goal is to learn a **compact and efficient vocabulary** from raw text data, which can then be used to tokenize and encode text into integer IDs for NLP applications.\n", "\n", "Key features:\n", "- Works at the **byte level**, so it can handle any Unicode text without pre-tokenization.\n", "- Uses a **maximum-length greedy tokenization** heuristic to match the longest token in the trie.\n", "- Supports **configurable vocabulary size**, maximum token length, and control over merging across spaces.\n", "- Can **save/load** the vocabulary for reuse in other tasks or models.\n", "\n", "This notebook contains:\n", "1. Class definitions for `TokenizerLearner` and `Tokenizer`.\n", "2. Dataset loading and preprocessing.\n", "3. Vocabulary learning loop with adjacency counting.\n", "4. Vocabulary saving and verification.\n", "5. Example encoding and decoding to test the tokenizer.\n", "\n", "There is no need for a GPU for this assignment." ] }, { "cell_type": "markdown", "metadata": { "id": "MyQeJ7iYyHIp" }, "source": [ "# **Imports and Dependencies**" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tlF0lpw7uRmW" }, "outputs": [], "source": [ "# If you are not using Google Colab, ensure you have Python 3.8+\n", "!pip install pygtrie" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "id": "dK7ek49OuJOv" }, "outputs": [], "source": [ "import pygtrie\n", "from collections import Counter\n", "import datasets\n", "import json\n", "import itertools\n", "\n", "SUBMISSION_READY = False" ] }, { "cell_type": "markdown", "metadata": { "id": "z-DJLa5uuk49" }, "source": [ "# **TokenizerLearner**\n", "The `TokenizerLearner` class builds a vocabulary using byte-pair merges." ] }, { "cell_type": "code", "execution_count": 12, "metadata": { "id": "mJt8X2GLuxoc" }, "outputs": [], "source": [ "class TokenizerLearner:\n", " def __init__(self, data_iterator, vocab_size=65536, docs_per_iter=10000, max_token_length=30, no_subwords_across_space=True):\n", " print(f\"\\nInitializing TokenizerLearner:\")\n", " print(f\" - vocab_size: {vocab_size}\")\n", " print(f\" - docs_per_iter: {docs_per_iter}\")\n", " print(f\" - max_token_length: {max_token_length}\")\n", " self.vocab_size = vocab_size\n", " self.data = data_iterator\n", " self.no_subwords_across_space = no_subwords_across_space\n", " self.data_iterator = iter(data_iterator)\n", " self.max_token_length = max_token_length\n", " self.docs_per_iter = docs_per_iter\n", " self.vocab = None\n", "\n", " # Byte-level vs unicode literals to control token boundaries\n", " self.space_char = list(b' ')[0]\n", " self.newline_char = list(b'\\n')[0]\n", "\n", " def maybe_add(self, most_common_adjacencies):\n", " for most_common_adjacency in most_common_adjacencies:\n", " if self.space_char in most_common_adjacency[0][1] and self.no_subwords_across_space:\n", " continue\n", " if len(most_common_adjacency[0][0]) + len(most_common_adjacency[0][1]) > self.max_token_length:\n", " continue\n", " if (most_common_adjacency[0][0] == self.newline_char) + (most_common_adjacency[0][1] == self.newline_char) == 1:\n", " continue\n", " new_token = most_common_adjacency[0][0] + most_common_adjacency[0][1]\n", " new_string = bytes(new_token).decode('utf-8', errors='ignore')\n", " print(f\" Most common adjacency: '{new_token}' (count: {most_common_adjacency[1]})\")\n", " print(f\" Most common adjacency: '{new_string}' (count: {most_common_adjacency[1]})\")\n", " return new_token\n", " return None\n", "\n", " def learn(self):\n", " print(\"\\nStarting vocabulary learning...\")\n", " iteration = 0\n", "\n", " # Make a tokenizer\n", " print(\" Creating tokenizer...\")\n", " self.tokenizer = Tokenizer(vocab=self.vocab, max_token_length=self.max_token_length)\n", "\n", " # Initialize with all one-length byte strings\n", " self.tokenizer.update_trie([(x,) for x in range(256)])\n", "\n", " while len(self.tokenizer.trie) < self.vocab_size:\n", " iteration += 1\n", " print(f\"\\nIteration {iteration}:\")\n", " print(f\"Current vocab size: {len(self.tokenizer.trie)}\")\n", "\n", " text_docs = []\n", " for i in range(self.docs_per_iter):\n", " try:\n", " doc = next(self.data_iterator)\n", " except StopIteration:\n", " self.data_iterator = iter(self.data)\n", " doc = next(self.data_iterator)\n", " text = doc['text'].encode('utf-8')\n", " text_docs.append(text)\n", "\n", " # In this section,\n", " # (1) iterate through a batch of text documents, tokenizing\n", " # each one and counting token pair adjacences.\n", " # use self.tokenizer._tokenize(doc) to tokenize (so you'll need to\n", " # implement that first.)\n", " # (2) next, go through the sorted token adjacenies pair and\n", " # use the self.maybe_add function to get which token should be added\n", " # (3) update the tokenizer's trie with the new token.\n", " # --------------------------------- BEGIN STUDENT TODO\n", "\n", "\n", "\n", " # --------------------------------- END STUDENT TODO\n", "\n", " def save(self, path):\n", " print(f\"\\nSaving vocabulary to {path}\")\n", " with open(path, 'w') as f:\n", " for token in sorted(self.tokenizer.trie):\n", " f.write(json.dumps([token])+'\\n')\n", " print(f\"Saved {len(self.tokenizer.trie)} tokens\")\n" ] }, { "cell_type": "markdown", "metadata": { "id": "c-8yYry-u3kf" }, "source": [ "# **Tokenizer**\n", "The `Tokenizer` class performs encoding/decoding with the learned vocabulary.\n" ] }, { "cell_type": "code", "execution_count": 14, "metadata": { "id": "Ork3fez7uO7v" }, "outputs": [], "source": [ "class Tokenizer:\n", " def __init__(self, vocab_path=None, vocab=None, max_token_length=30, partial_trie=None):\n", " print(f\"\\nInitializing Tokenizer:\")\n", " print(f\" - vocab_path: {vocab_path}\")\n", " print(f\" - vocab size: {len(vocab) if vocab else 'None'}\")\n", " print(f\" - max_token_length: {max_token_length}\")\n", " self.vocab_path = vocab_path\n", " self.id_to_tok = []\n", " self.trie = pygtrie.Trie()\n", " self.trie = self.trie if partial_trie is None else partial_trie\n", " self.max_token_length = max_token_length\n", "\n", " if vocab_path or vocab:\n", " self.update_trie(vocab)\n", "\n", " def update_trie(self, new_vocab=None):\n", " print(\"\\nUpdating trie...\")\n", " if new_vocab is None and self.vocab_path:\n", " print(f\"Loading from vocab file: {self.vocab_path}\")\n", " with open(self.vocab_path, 'r') as f:\n", " for i, line in enumerate(f):\n", " token = tuple(json.loads(line)[0])\n", " self.id_to_tok.append(token)\n", " self.trie[token] = i\n", " elif new_vocab:\n", " for token in new_vocab:\n", " print(token)\n", " self.id_to_tok.append(token)\n", " self.trie[token] = len(self.trie)\n", "\n", " def encode(self, text):\n", " return self._tokenize(text, return_ids=True)\n", "\n", " def decode(self, tokens):\n", " tokens = [self.id_to_tok[x] for x in tokens]\n", " return bytes(itertools.chain.from_iterable(tokens)).decode('utf-8', errors='ignore')\n", "\n", " def _tokenize(self, text, return_ids=False):\n", " # In this section,\n", " # (1) encode the text to receive a bytestring using\n", " # text.encode('utf-8', errors='ignore')\n", " # (2) tokenize the string using the trie we're developing\n", " # As a hint, consider how to use the self.max_token_length to\n", " # efficiently query the trie, and note that we use the maximum-length\n", " # greedy tokenization heuristic.\n", " # (3) if return_ids=True, then return a list of integer ids. Otherwise,\n", " # return a list of byte lists.\n", " # --------------------------------- BEGIN STUDENT TODO\n", " pass\n", "\n", "\n", " # --------------------------------- END STUDENT TODO" ] }, { "cell_type": "markdown", "metadata": { "id": "QjuMD9lVycvt" }, "source": [ "# **Dataset Loading, Tokenizer Training, and Saving Vocabulary**\n", "The dataset lives [here](https://huggingface.co/datasets/coms4705-hewitt/fineweb-linuxlike/tree/main). It should download automatically.\n", "\n", "Runtime roughly scales with vocab size. Feel free to play around with it. What happens when it is less than 256?" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Gfk8KrWsvNGc" }, "outputs": [], "source": [ "print(\"Starting tokenizer test...\")\n", "print(\"Loading dataset...\")\n", "dataset = datasets.load_dataset('coms4705-hewitt/fineweb-linuxlike', 'default', streaming=True)['train']\n", "print(\"Dataset loaded\")\n", "\n", "print(\"\\nCreating TokenizerLearner...\")\n", "# learner = TokenizerLearner(dataset, vocab_size=65536, docs_per_iter=20, no_subwords_across_space=False)\n", "learner = TokenizerLearner(dataset, vocab_size=600, docs_per_iter=20, no_subwords_across_space=True)\n", "print(\"Starting learning process...\")\n", "if not SUBMISSION_READY:\n", " learner.learn()\n", "print(\"\\nLearning complete!\")\n", "learner.save('vocab-65k-fw-byte-sas.txt')" ] }, { "cell_type": "markdown", "metadata": { "id": "BCqhVlJQzFPZ" }, "source": [ "# **Testing the Tokenizer**\n", "\n", "Once the vocabulary is learned, we can test the tokenizer by encoding some example strings and decoding them back to verify correctness." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "VYnAIRaYzJQy" }, "outputs": [], "source": [ "examples = [\n", " \"Hello, world!\",\n", " \"The quick brown fox jumps over the lazy dog.\",\n", " \"def tokenize(text): return text.split()\",\n", " \"🌟 Unicode characters work too! 🚀\",\n", " \"What happens when you decode a \\u2603?\"\n", "]\n", "\n", "print(\"Encoding and decoding examples:\")\n", "for text in examples:\n", " print(\"\\nOriginal text: \", text)\n", " encoded = learner.tokenizer.encode(text)\n", " print(\"Encoded bytes: \", encoded)\n", " token_strings = [learner.tokenizer.decode([token]) for token in encoded]\n", " print(\"Individual 'decoded' token strings: \", token_strings)\n", " print(f\"Broke {len(text)} characters into {len(encoded)} tokens!\")\n", " decoded = learner.tokenizer.decode(encoded)\n", " print(\"Decoded text:\", decoded)\n", " print('Is decoded same as original text?: ', decoded==text)" ] } ], "metadata": { "colab": { "provenance": [] }, "kernelspec": { "display_name": "base", "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.9.7" } }, "nbformat": 4, "nbformat_minor": 0 }