PolyScribe 0.1.2__tar.gz → 0.2.0__tar.gz

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

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

{polyscribe-0.1.2 → polyscribe-0.2.0}/include/byte_pair_encoder.hpp

@@ -5,19 +5,9 @@
 #include <vector>
 #include <cstdint>
 
-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>>;
+#include <./typedefs.hpp>
 
+namespace Scribe {
     class BytePairEncoder {
     private:
         std::unordered_map<Token, std::vector<uint8_t>> vocab;
@@ -26,13 +16,14 @@ namespace Scribe {
 
         std::vector<Token> getCodePoints(const std::string& str);
         Pair getMostFrequentPair(const WordCounts& wordCounts);
-        void doMerge(WordCounts& wordCounts, const Pair& mostFrequentPair, const int newToken);
+        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 verbose);
+        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.1.2 → polyscribe-0.2.0}/python/bindings.cpp

@@ -14,11 +14,13 @@ NB_MODULE(Scribe, m) {
 
     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("verbose"),
+    ).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"

{polyscribe-0.1.2 → polyscribe-0.2.0}/src/byte_pair_encoder.cpp

@@ -7,6 +7,7 @@
 #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) {
@@ -48,14 +49,57 @@ Scribe::Pair Scribe::BytePairEncoder::getMostFrequentPair(const WordCounts& word
     return mostFrequentPair;
 }
 
-void Scribe::BytePairEncoder::doMerge(WordCounts& wordCounts, const Pair& mostFrequentPair, const int newToken) {
+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] == mostFrequentPair.first && word[read_idx + 1] == mostFrequentPair.second) {
+            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]);
 
@@ -81,7 +125,7 @@ Scribe::BytePairEncoder::BytePairEncoder() {
     }
 }
 
-void Scribe::BytePairEncoder::train(const std::string& filename, int cycles, bool verbose) {
+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);
@@ -97,9 +141,12 @@ void Scribe::BytePairEncoder::train(const std::string& filename, int cycles, boo
         wordCounts.emplace_back(std::move(initialTokens), count);
     }
 
+    Pair pairToMerge;
     for (int i = 0; i < cycles; i++) {
-        Pair mostFrequentPair = getMostFrequentPair(wordCounts);
-        if (mostFrequentPair == Pair(-1, -1)) {
+        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;
@@ -107,20 +154,20 @@ void Scribe::BytePairEncoder::train(const std::string& filename, int cycles, boo
         }
 
         int newToken = i + 256;
-        doMerge(wordCounts, mostFrequentPair, newToken);
+        doMerge(wordCounts, pairToMerge, newToken);
 
-        mergeForest.emplace_back(mostFrequentPair, newToken);
+        mergeForest.emplace_back(pairToMerge, newToken);
 
-        std::vector<uint8_t> newBytes = vocab[mostFrequentPair.first];
-        newBytes.insert(newBytes.end(), vocab[mostFrequentPair.second].begin(), vocab[mostFrequentPair.second].end());
+        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[mostFrequentPair.first]) w1 += static_cast<char>(ch);
-        for (auto ch : vocab[mostFrequentPair.second]) w2 += static_cast<char>(ch);
+        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, mostFrequentPair.first, mostFrequentPair.second, w1, w2, w3) << std::endl;
+        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;
     }
 }