PolyScribe 0.1.1__tar.gz → 0.2.0__tar.gz

{polyscribe-0.1.1 → polyscribe-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: PolyScribe
-Version: 0.1.1
+Version: 0.2.0
 Summary: Custom C++20 tokenizer.
 Author: DracoRooks
 License: MIT

polyscribe-0.2.0/include/byte_pair_encoder.hpp

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <unordered_map>
+#include <string>
+#include <vector>
+#include <cstdint>
+
+#include <./typedefs.hpp>
+
+namespace Scribe {
+    class BytePairEncoder {
+    private:
+        std::unordered_map<Token, std::vector<uint8_t>> vocab;
+        std::unordered_map<uint8_t, Token> utf8Lookup;
+        std::vector<std::pair<Pair, Token>> mergeForest;
+
+        std::vector<Token> getCodePoints(const std::string& str);
+        Pair getMostFrequentPair(const WordCounts& wordCounts);
+        Pair getBestRankedPair(const WordCounts& wordCounts);
+        void doMerge(WordCounts& wordCounts, const Pair& pairToMerge, const int newToken);
+
+    public:
+        BytePairEncoder();
+        ~BytePairEncoder() = default;
+
+        void train(const std::string& filename, int cycles, bool normalizedRanking, bool verbose);
+        std::vector<Token> encode(const std::string& data);
+        std::vector<uint8_t> decode(const std::vector<Token>& tokens);
+    };
+}

polyscribe-0.2.0/include/typedefs.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <unordered_map>
+#include <vector>
+
+namespace Scribe {
+    struct PairHash {
+        inline size_t operator()(const std::pair<int, int>& v) const {
+            return v.first * 997 + v.second;
+        }
+    };
+
+    using Count = int;
+    using Token = int;
+    using Pair = std::pair<Token, Token>;
+    using PairCounts = std::unordered_map<Pair, Count, PairHash>;
+    using WordCounts = std::vector<std::pair<std::vector<Token>, Count>>;
+}

polyscribe-0.2.0/include/word_extracter.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include <string>
+#include <unordered_map>
+
+namespace Scribe {
+    class WordExtracter
+    {
+    private:
+        int getUTFByteLength(const char ch);
+
+    public:
+        WordExtracter() = default;
+        ~WordExtracter() = default;
+
+        std::unordered_map<std::string, int> wordify(const std::string& filename);
+    };
+}

{polyscribe-0.1.1 → polyscribe-0.2.0}/pyproject.toml

@@ -21,9 +21,17 @@ cmake.version = ">=3.28"
 cmake.args = ["-D_BUILD_PY_MOD=ON"]
 ninja.make-fallback = true
 wheel.packages = []
-metadata.version.provider = "scikit_build_core.metadata.setuptools_scm"
+sdist.include = [
+    "CMakeLists.txt",
+    "src/*",
+    "include/*",
+    "python/*"
+]
+
+[[tool.dynamic-metadata]]
+provider = "scikit_build_core.metadata.setuptools_scm"
 
-[tools.setuptools_scm]
+[tool.setuptools_scm]
 
 [tool.scikit-build.cmake.define]
 _BUILD_PY_MOD = "ON"

polyscribe-0.2.0/python/bindings.cpp

@@ -0,0 +1,40 @@
+#include <nanobind/nanobind.h>
+#include "../include/byte_pair_encoder.hpp"
+
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/unordered_map.h>
+#include <nanobind/stl/vector.h>
+#include <nanobind/ndarray.h>
+
+namespace nb = nanobind;
+using namespace nb::literals;
+
+NB_MODULE(Scribe, m) {
+    m.doc() = "A fast and simple module to tokenize text data.";
+
+    nb::class_<Scribe::BytePairEncoder>(m, "BytePairEncoder", "A module that applies BPE algorithm on the given text to generate word/sub-word tokens."
+    ).def(nb::init<>()
+    ).def("train", &Scribe::BytePairEncoder::train, nb::arg("filename"), nb::arg("cycles"), nb::arg("normalizedRanking"), nb::arg("verbose"),
+        "Trains the tokenizer on text data contained in `filename` file.\n\n"
+        "Args:\n"
+        "   filename (str): Path to the training data (taken relative from the directory the function will be called in).\n"
+        "   cycles (int): Number of cycles to train the tokenizer for, total vocab size will be 256 + `cycles`.\n"
+        "   normalizedRanking (bool): If false, follows default BPE Scoring for token pairs, i.e. P(A|B)."
+        "       If true, follows WordPiece Scoring for token pairs, i.e. P(A|B)/P(A)*P(B)."
+        "   verbose (bool): If true, prints log of every merge cycle to standard output.\n"
+        "Returns:\n"
+        "   None"
+    ).def("encode", &Scribe::BytePairEncoder::encode, nb::arg("data"),
+        "Encoder function to convert text data into tokens.\n\n"
+        "Args:\n"
+        "   data (str): Text data to encode.\n"
+        "Returns:\n"
+        "   list[int]"
+    ).def("decode", &Scribe::BytePairEncoder::decode, nb::arg("tokens"),
+        "Decoder function to convert tokens into string data.\n\n"
+        "Args:\n"
+        "   tokens (list[int]): List of tokens to decode.\n"
+        "Returns:\n"
+        "   list[int]"
+    );
+}

polyscribe-0.2.0/src/byte_pair_encoder.cpp

@@ -0,0 +1,193 @@
+#include "../include/byte_pair_encoder.hpp"
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+#include <iostream>
+#include <cstdint>
+#include <fmt/format.h>
+
+#include <../include/typedefs.hpp>
+#include "../include/word_extracter.hpp"
+
+std::vector<Scribe::Token> Scribe::BytePairEncoder::getCodePoints(const std::string& str) {
+    int len = str.length();
+    int i;
+
+    std::vector<Token> codepoints(len);
+
+    for (i = 0; i < len; i++) {
+        codepoints[i] = utf8Lookup[str[i]];
+    }
+
+    return codepoints;
+}
+
+Scribe::Pair Scribe::BytePairEncoder::getMostFrequentPair(const WordCounts& wordCounts) {
+    PairCounts pairCounts;
+    Count maxCount = 0;
+    Pair mostFrequentPair = { -1, -1 };
+
+    for (const auto& [word, count] : wordCounts) {
+        int len = word.size();
+        if (len < 2) continue;
+
+        int i = 0, j = 1;
+        while (j < len) {
+            Pair pair = { word[i], word[j] };
+            pairCounts[pair] += count;
+
+            if (pairCounts[pair] > maxCount) {
+                maxCount = pairCounts[pair];
+                mostFrequentPair = pair;
+            }
+
+            i++; j++;
+        }
+    }
+
+    return mostFrequentPair;
+}
+
+Scribe::Pair Scribe::BytePairEncoder::getBestRankedPair(const WordCounts& wordCounts) {
+    PairCounts pairCounts;
+    std::unordered_map<Token, Count> tokenCounts;
+    double maxScore = 0;
+    Pair bestPair = { -1, -1 };
+
+    for (const auto& [word, count] : wordCounts) {
+        int len = word.size();
+        if (len < 2) continue;
+
+        int i = 0, j = 1;
+        tokenCounts[word[0]] += count;
+        while (j < len) {
+            Pair pair = { word[i], word[j] };
+            pairCounts[pair] += count;
+            tokenCounts[word[j]] += count;
+
+            i++; j++;
+        }
+    }
+
+    for (const auto& [word, count] : wordCounts) {
+        int len = word.size();
+        if (len < 2) continue;
+
+        int i = 0, j = 1;
+        while (j < len) {
+            Pair pair = { word[i], word[j] };
+
+            double score = (double)pairCounts[pair] / tokenCounts[word[i]];
+            score /= tokenCounts[word[j]];
+            if (score > maxScore) {
+                maxScore = score;
+                bestPair = pair;
+            }
+
+            i++; j++;
+        }
+    }
+
+    return bestPair;
+}
+
+void Scribe::BytePairEncoder::doMerge(WordCounts& wordCounts, const Pair& pairToMerge, const int newToken) {
+    for (auto& [word, _] : wordCounts) {
+        int len = word.size();
+        int read_idx = 0;
+        int write_idx = 0;
+
+        while (read_idx < len) {
+            if (read_idx + 1 < len && word[read_idx] == pairToMerge.first && word[read_idx + 1] == pairToMerge.second) {
+                word[read_idx] = newToken;
+                word[write_idx++] = std::move(word[read_idx]);
+
+                read_idx += 2;
+            } else {
+                if (write_idx != read_idx) {
+                    word[write_idx] = std::move(word[read_idx]);
+                }
+                write_idx++;
+                read_idx++;
+            }
+        }
+        word.resize(write_idx);
+    }
+}
+
+Scribe::BytePairEncoder::BytePairEncoder() {
+    for (int i = 0; i < 256; i++) {
+        const uint8_t ch = static_cast<uint8_t>(i);
+
+        vocab[i] = { ch };
+        utf8Lookup[ch] = i;
+    }
+}
+
+void Scribe::BytePairEncoder::train(const std::string& filename, int cycles, bool normalizedRanking, bool verbose) {
+    WordExtracter extracter;
+
+    std::unordered_map<std::string, Count> strWordCounts = extracter.wordify(filename);
+    // TODO: make it so that `strWordCounts` is destroyed after its use
+
+    WordCounts wordCounts;
+    wordCounts.reserve(strWordCounts.size());
+
+    for (const auto& [word, count] : strWordCounts) {
+        if (word.empty()) continue;
+
+        const std::vector<Token> initialTokens = getCodePoints(word);
+        wordCounts.emplace_back(std::move(initialTokens), count);
+    }
+
+    Pair pairToMerge;
+    for (int i = 0; i < cycles; i++) {
+        if (!normalizedRanking) pairToMerge = getMostFrequentPair(wordCounts);
+        else pairToMerge = getBestRankedPair(wordCounts);
+
+        if (pairToMerge == Pair(-1, -1)) {
+            std::cerr << "\n\n[INFO]::BYTE_PAIR_ENCODER::CYCLES: Max possible vocab size reached, ";
+            std::cerr << "{Cycles: " << cycles << "} is too high for the provided dataset.";
+            std::cerr << "\nFinishing gracefully at cycle: " << i << "." << std::endl;
+            break;
+        }
+
+        int newToken = i + 256;
+        doMerge(wordCounts, pairToMerge, newToken);
+
+        mergeForest.emplace_back(pairToMerge, newToken);
+
+        std::vector<uint8_t> newBytes = vocab[pairToMerge.first];
+        newBytes.insert(newBytes.end(), vocab[pairToMerge.second].begin(), vocab[pairToMerge.second].end());
+        vocab[newToken] = newBytes;
+
+        if (!verbose) continue;
+        std::string w1 = "", w2 = "";
+        for (auto ch : vocab[pairToMerge.first]) w1 += static_cast<char>(ch);
+        for (auto ch : vocab[pairToMerge.second]) w2 += static_cast<char>(ch);
+        std::string w3 = w1 + w2;
+        std::clog << fmt::format("[INFO]::NEW_TOKEN_{:<6}: Merged Pair {{{:^6}, {:^6}}} : {{{:^15}, {:^15}}} -> {:<40}", newToken, pairToMerge.first, pairToMerge.second, w1, w2, w3) << std::endl;
+    }
+}
+
+std::vector<Scribe::Token> Scribe::BytePairEncoder::encode(const std::string& data) {
+    std::vector<Token> codepoints = getCodePoints(data);
+    WordCounts encodedData = { std::pair<std::vector<Token>, Count>(codepoints, 1) };
+
+    for (const auto& [pair, token] : mergeForest) {
+        doMerge(encodedData, pair, token);
+    }
+
+    return encodedData[0].first;
+}
+
+std::vector<uint8_t> Scribe::BytePairEncoder::decode(const std::vector<Token>& tokens) {
+    std::vector<uint8_t> decodedData;
+
+    for (const Token token : tokens) {
+        decodedData.insert(decodedData.end(), vocab[token].begin(), vocab[token].end());
+    }
+
+    return decodedData;
+}

polyscribe-0.2.0/src/word_extracter.cpp

@@ -0,0 +1,76 @@
+#include "../include/word_extracter.hpp"
+
+#include <set>
+#include <fstream>
+#include <sstream>
+#include <string>
+
+int Scribe::WordExtracter::getUTFByteLength(const char ch) {
+    const unsigned char uc = static_cast<unsigned char>(ch);
+    if (uc < 0x80) { return 1; } // 0xxxxxxx, hence a normal ASCII char
+    else if ((uc & 0xE0) == 0xC0) { return 2; } // 110xxxxx: 2 bytes
+    else if ((uc & 0xF0) == 0xE0) { return 3; } // 1110xxxx: 3 bytes
+    else if ((uc & 0xF8) == 0xF0) { return 4; } // 11110xxx: 4 bytes
+    return 1; // for invalid utf-8 byte
+}
+
+std::unordered_map<std::string, int> Scribe::WordExtracter::wordify(const std::string& filename) {
+    std::ifstream file(filename, std::ios::binary);
+    if (!file.is_open()) {
+        throw std::runtime_error("[ERROR]::COULD_NOT_OPEN_FILE: " + filename);
+    }
+
+    std::stringstream buffer;
+    buffer << file.rdbuf();
+    std::string filedata(buffer.str());
+
+    file.close();
+
+    const std::set<char> delimiters{'!', '"', '#', '%', '&', '\'', '(', ')', '*', '+', ',', '-', '_', '.', '/', ':', ';', '<', '=', '>', '?', '[', ']', '~', '`', '\\', '^', '{', '}', '\n'};
+        
+    std::unordered_map<std::string, int> words;
+    unsigned long wordStartIndex = 0, wordLen = 0;
+
+    while (wordStartIndex < filedata.length()) {
+        const char ch = filedata[wordStartIndex + wordLen];
+        const int utfByteLen = getUTFByteLength(ch);
+
+        if (utfByteLen > 1) {
+            // Inserting the word building upto the multibyte character
+            words[std::string(&filedata[wordStartIndex], wordLen)]++;
+            wordStartIndex += wordLen;
+
+            // Inserting the multibyte character
+            words[std::string(&filedata[wordStartIndex], utfByteLen)]++;
+            wordStartIndex += utfByteLen;
+            wordLen = 0;
+
+        } else if (ch == ' ' || ch == '\0') {
+            // Inserting the word building upto the whitespace (or EOF)
+            words[std::string(&filedata[wordStartIndex], wordLen)]++;
+
+            wordStartIndex += wordLen;
+            wordLen = 1;
+
+            // TODO: Decouple ' ' and '\0' and make individual cases for them; add an EOF token for file-termination.
+        } else if (delimiters.count(ch)) {
+            // Inserting the word building upto the delimiter
+            words[std::string(&filedata[wordStartIndex], wordLen)]++;
+            wordStartIndex += wordLen;
+
+            while (delimiters.count(filedata[wordStartIndex])) {
+                // Inserting all consecutive (if) occuring delimiters
+                words[std::string(1, filedata[wordStartIndex])]++;
+                wordStartIndex++;
+            }
+
+            wordLen = 1;
+        } else {
+            wordLen++;
+        }
+
+        // TODO: Make the function work with multiple files at once.
+    }
+
+    return words;
+}