{"generator":"Jekyll","link":[{"@attributes":{"href":"https:\/\/guillaume-be.github.io\/feed.xml","rel":"self","type":"application\/atom+xml"}},{"@attributes":{"href":"https:\/\/guillaume-be.github.io\/","rel":"alternate","type":"text\/html"}}],"updated":"2021-09-25T08:08:25+00:00","id":"https:\/\/guillaume-be.github.io\/feed.xml","title":"Rust NLP tales","subtitle":"Natural Language Processing and Rust","author":{"name":"Guillaume Becquin","email":"guillaume.becquin@gmail.com"},"entry":[{"title":"Byte Pair Encoding and Data Structures","link":{"@attributes":{"href":"https:\/\/guillaume-be.github.io\/2021-09-16\/byte_pair_encoding","rel":"alternate","type":"text\/html","title":"Byte Pair Encoding and Data Structures"}},"published":"2021-09-16T00:00:00+00:00","updated":"2021-09-16T00:00:00+00:00","id":"https:\/\/guillaume-be.github.io\/2021-09-16\/byte_pair_encoding","content":"<p>Tokenization of input strings into sequences of words or sub-tokens is a central concept for modern Natural Language Processing techniques (NLP). This article focuses on a classic tokenization algorithm: Byte Pair Encoding (BPE) <a href=\"#bpe\">[1]<\/a>. While resources describing the working principle of the algorithm are widely available, this article focuses on its implementation, illustrating how the choice of data structures impact the performance of a real-world NLP component. This article is complemented by a working Rust-based implementation, <a href=\"https:\/\/github.com\/guillaume-be\/bpe-example\" target=\"_blank\">available on GitHub<\/a>.<\/p>\n\n<h1 id=\"introduction\">Introduction<\/h1>\n\n<p>Tokenization consists in splitting an input sequence (e.g., a sentence) into a sequence of tokens with a finite cardinality (for example words or individual characters). These tokens can then be encoded and processed further by machine learning components such as neural networks.<\/p>\n\n<p>Halfway between word-based tokenization and character input, sub-word tokenization aims at combining the benefits of both approaches:<\/p>\n<ul>\n  <li>The tokens generated carry a semantic content that is close to words (and richer than individual characters)<\/li>\n  <li>They allow for smaller vocabulary sizes, driving memory efficiency and increased representation sharing<\/li>\n<\/ul>\n\n<p>Sub-word tokenization works by splitting rare words into sub-word components instead of ignoring them or having a seldom-used dedicated representation for them. Ideally, these sub-word representations are semantically or syntactically meaningful, for example by splitting pre- or suffixes from words (<code class=\"language-plaintext highlighter-rouge\">unexpected<\/code> may become <code class=\"language-plaintext highlighter-rouge\">[un, ##expected]<\/code>). The resulting vocabulary size can be an order of magnitude smaller than for word-based segmentation while keeping common words as single entries - allowing to generate semantically-rich word embeddings for the downstream model.<\/p>\n\n<p>These models usually have two algorithmic components:<\/p>\n<ul>\n  <li>a training phase, where a vocabulary and rules to merge\/split words are learned with an unsupervised algorithm.<\/li>\n  <li>a prediction phase, where these rules are used to split an input sentence or documents into a sequence of words\/sub-tokens.<\/li>\n<\/ul>\n\n<p>While proper algorithmic design of the training phase is important to allow scaling to larger corpora, this is a one-off cost. This article focuses on the design and implementation of the prediction phase which is critical to limit the latency and operation cost of models when they are deployed and generate production. Following a previous article illustrating the <a href=\"..\/2020-05-30\/sentence_piece\" target=\"_blank\">SentencePiece unigram tokenization model<\/a>, this article will focus on another ubiquitous algorithm: Byte Pair Encoding (BPE).<\/p>\n\n<h1 id=\"1-the-byte-pair-encoding-bpe-tokenizer\">1. The Byte Pair Encoding (BPE) tokenizer<\/h1>\n\n<p>BPE is a morphological tokenizer that merges adjacent byte pairs based on their frequency in a training corpus. Based on a compression algorithm with the same name, BPE has been adapted to sub-word tokenization and can be thought of as a clustering algorithm <a href=\"#bpe-lecture\">[2]<\/a>. A starting sequence of individual characters will be aggregated bottom-up based on frequencies learned during a training phase. When no further aggregation is possible, end the process and return the sub-words generated.<\/p>\n\n<p>No unknown tokens can occur as an output of this process as long as all individual characters are present in the vocabulary (in the worst case, the sequence of individual characters will be returned). With a sufficiently large vocabulary, the most common words will be fully aggregated and the entire sequence of input characters will be returned as a single word token. Rare words will be returned as a list of sub-tokens that can no longer be aggregated.<\/p>\n\n<h2 id=\"prediction-phase-tokenization\">Prediction phase (tokenization)<\/h2>\n\n<p>Even though the BPE model needs to be trained before it can be used, we will first describe how it is used to tokenize an input sequence at prediction time in this article. This will help build an intuition on how the algorithm decomposes words into sub-tokens.<\/p>\n\n<p>The tokenization algorithm is as follows:<\/p>\n\n<p><a name=\"bpe-tokenization-naive\"><\/a><\/p>\n<pre id=\"read-0\" style=\"display:none;\">\n    \n\\begin{algorithm}\n\\caption{BPE Tokenize}\n\\begin{algorithmic}\n\\PROCEDURE{FindBestPair}{$symbols$, $merges$}\n    \\STATE ($best\\:pair$, $best\\:score$) = ($null$, 0)\n    \\FOR{$i = 0$ \\TO $len(symbols) - 1$}\n        \\STATE $score$ = \\CALL{Get}{$merges$, $(symbols[i], symbols[i+1])$}\n        \\IF{$score$ &gt; $best\\:score$}\n            \\STATE ($best\\:pair$, $best\\:score$) = ($(symbols[i], symbols[i+1])$, score)\n        \\ENDIF\n    \\ENDFOR    \n    \\RETURN ($best\\:pair$, $best\\:score$)\n\\ENDPROCEDURE\n\\STATE\n\\PROCEDURE{MergeBestPair}{$symbols$, $best\\:pair$}\n    \\FOR{$i = 0$ \\TO $len(symbols) - 1$}\n        \\IF{$(symbols[i], symbols[i+1])$ = $best\\:pair$}\n            \\STATE \\CALL{Merge}{$symbols[i], symbols[i+1]$}\n        \\ENDIF\n    \\ENDFOR    \n    \\RETURN $symbols$\n\\ENDPROCEDURE\n\\STATE\n\\PROCEDURE{Tokenize}{$text$, $merges$}\n    \\STATE $symbols = $ \\CALL{InitializeSymbols}{$text$} \\COMMENT{Pre-populate symbols with individual characters and &lt;\/w&gt; token}\n    \\WHILE {$true$}\n        \\STATE $(best\\:pair, best\\:score) = $ \\CALL{FindBestPair}{$symbols$, $merges$}\n        \\IF{$best\\:score$ is $null$}\n            \\BREAK \\COMMENT{Symbols can no longer be merged}\n        \\ENDIF\n        \\STATE $symbols$ = \\CALL{MergeBestPair}{$symbols$, $best\\:pair$}\n    \\ENDWHILE\n    \\RETURN $symbols$\n\\ENDPROCEDURE\n\\end{algorithmic}\n\\end{algorithm}\n\n<\/pre>\n<div id=\"goal-0\"><\/div>\n<script type=\"text\/javascript\">\n    var code = document.getElementById(\"read-0\").textContent;\n    var parentEl = document.getElementById(\"goal-0\");\n    var options = {\n        lineNumber: true\n    };\n    pseudocode.render(code, parentEl, options);\n<\/script>\n\n<p><code class=\"language-plaintext highlighter-rouge\">merges<\/code> is a learned mapping of symbol pairs to score that is learned during a training phase (see next section). A <code class=\"language-plaintext highlighter-rouge\">Symbol<\/code> is a sub-token of a given input, which may be a character, multiple consecutive characters that have been merged or the entire word. The following example illustrates how a merges vocabulary (mapping of symbols that may be aggregated with their frequency) is used to tokenize 2 words. The symbol pairs with higher frequencies are merged first:<\/p>\n\n<div class=\"language-json highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"p\">{<\/span><span class=\"w\">\n  <\/span><span class=\"nl\">\"e r\"<\/span><span class=\"p\">:<\/span><span class=\"w\"> <\/span><span class=\"mi\">25<\/span><span class=\"p\">,<\/span><span class=\"w\">\n  <\/span><span class=\"nl\">\"h e\"<\/span><span class=\"p\">:<\/span><span class=\"w\"> <\/span><span class=\"mi\">12<\/span><span class=\"p\">,<\/span><span class=\"w\">\n  <\/span><span class=\"nl\">\"l l\"<\/span><span class=\"p\">:<\/span><span class=\"w\"> <\/span><span class=\"mi\">11<\/span><span class=\"p\">,<\/span><span class=\"w\">\n  <\/span><span class=\"nl\">\"l o\"<\/span><span class=\"p\">:<\/span><span class=\"w\"> <\/span><span class=\"mi\">10<\/span><span class=\"p\">,<\/span><span class=\"w\">\n  <\/span><span class=\"nl\">\"he ll\"<\/span><span class=\"p\">:<\/span><span class=\"w\"> <\/span><span class=\"mi\">4<\/span><span class=\"p\">,<\/span><span class=\"w\">\n  <\/span><span class=\"nl\">\"lo w\"<\/span><span class=\"p\">:<\/span><span class=\"w\"> <\/span><span class=\"mi\">3<\/span><span class=\"p\">,<\/span><span class=\"w\">\n  <\/span><span class=\"nl\">\"hell o\"<\/span><span class=\"p\">:<\/span><span class=\"w\"> <\/span><span class=\"mi\">2<\/span><span class=\"w\"> \n<\/span><span class=\"p\">}<\/span><span class=\"w\">\n<\/span><\/code><\/pre><\/div><\/div>\n<ul>\n  <li>\n    <p>\u201clower\u201d would follow the following aggregation steps: <br \/>\n<code class=\"language-plaintext highlighter-rouge\">[l, o, w, e, r] -&gt; [l, o, w, er] -&gt; [lo, w, er] -&gt; [low, er]<\/code> (<code class=\"language-plaintext highlighter-rouge\">\"low er\"<\/code> is not in the merges vocabulary, causing an early stopping of the tokenization procedure at line <em>26<\/em>)<\/p>\n  <\/li>\n  <li>\n    <p>\u201chello\u201d would be tokenized as: <br \/>\n<code class=\"language-plaintext highlighter-rouge\">[h, e, l, l, o] -&gt; [he, l, l, o] -&gt; [he, ll, o] -&gt; [hell, o] -&gt; [hello]<\/code><\/p>\n  <\/li>\n<\/ul>\n\n<h2 id=\"training-phase\">Training phase<\/h2>\n\n<p>For computational efficiency, BPE training relies on a pre-tokenizer (e.g. whitespace tokenizer) in order to generate a dictionary with word frequencies (the original BPE algorithm does not allow cross-word merges). This word counter is used to initialize a counter of symbol pairs of adjacent characters that may be merged. After the count completes, the most frequent \u201csymbol pair\u201d is merged to create a new symbol in the dictionary, the symbol pairs counts  in the corpus are updated and the process repeats until the target vocabulary size is reached. As such, BPE is an algorithm that grows its vocabulary at each iteration (in contrast with SentencePiece unigram\u2019s model that prunes a large vocabulary at each iteration).<\/p>\n\n<p>The entire training algorithm is rather straight-forward and given below (reproduced from <a href=\"#bpe\">[1]<\/a>).<\/p>\n\n<pre id=\"read-1\" style=\"display:none;\">\n    \n\\begin{algorithm}\n\\caption{BPE Train}\n\\begin{algorithmic}\n\\PROCEDURE{UpdateMerges}{$vocab$}\n    \\STATE $merges = $\\CALL{InitializeMerges}{}\n    \\FOR{$(word, freq)$ IN $Vocab$}\n        \\STATE $symbols = $ \\CALL{SplitWhitespace}{$word$}\n        \\FOR{$i = 0$ \\TO $len(symbols)$}\n            \\STATE $merges[symbols[i], symbols[i+1]] += freq$\n        \\ENDFOR\n    \\ENDFOR\n    \\RETURN $merges$\n\\ENDPROCEDURE\n\\STATE\n\\PROCEDURE{MergeVocab}{$vocab_{in}, pair$}\n    \\STATE $vocab_{out} = $\\CALL{InitializeVocab}{}\n    \\STATE $pattern = $ \\CALL{Join}{pair}\n    \\FOR{$(w_{in}, _)$ IN $vocab_{in}$}\n        \\STATE $w_{out} = $ \\CALL{Replace}{ $w_{in}$, $pair_{0} \\sqcup pair_{1}$, $pair_{0} pair_{1}$}\n        \\STATE $vocab_{out}[w_{out}] = vocab_{in}[w_{in}]$\n    \\ENDFOR\n    \\RETURN $vocab_{out}$\n\\ENDPROCEDURE\n\\STATE\n\\PROCEDURE{Train}{$corpus$}\n    \\STATE vocab = \\CALL{CountWords}{$corpus$}\n    \\FOR{$i = 0$ \\TO $N_{merges}$}\n        \\STATE $merges = $ \\CALL{UpdateMerges}{$vocab$}\n        \\STATE $most\\:common\\:pair = $ \\CALL{MostCommon}{$merges$}\n        \\STATE $vocab = $ \\CALL{MergeVocab}{$vocab$, $most\\:common\\:pair$}\n    \\ENDFOR\n    \\RETURN $merges$\n\\ENDPROCEDURE\n\\end{algorithmic}\n\\end{algorithm}\n\n<\/pre>\n<div id=\"goal-1\"><\/div>\n<script type=\"text\/javascript\">\n    var code = document.getElementById(\"read-1\").textContent;\n    var parentEl = document.getElementById(\"goal-1\");\n    var options = {\n        lineNumber: true\n    };\n    pseudocode.render(code, parentEl, options);\n<\/script>\n\n<p>The algorithm above can be improved as described in <a href=\"#bpe\">[1]<\/a> by updating data structures instead of re-computing the \u201cmerges\u201d from scratch at each iteration. This article will however focus on the tokenization procedure called at prediction time.<\/p>\n\n<h1 id=\"2-rust-implementations-of-the-bpe-algorithm\">2. Rust implementation(s) of the BPE algorithm<\/h1>\n\n<p>The rest of this article consists of a walk-through of a number of working implementations of the BPE tokenization algorithm in Rust. Starting from a \u201cnaive\u201d implementation approach, improvements will be made to highlight pros and cons of some common data structures within the scope of BPE tokenization.<\/p>\n\n<p>4 algorithms for the BPE tokenization will be presented:<\/p>\n<ul>\n  <li><a href=\"#naive\">Naive implementation<\/a><\/li>\n  <li><a href=\"#naive-pre-split\">Naive implementation with pre-splitting<\/a><\/li>\n  <li><a href=\"#pq-bst\">Priority Queue + Binary Search Tree implementation<\/a><\/li>\n  <li><a href=\"#pq-ll\">Priority Queue + Linked List implementation<\/a><\/li>\n<\/ul>\n\n<h2 id=\"a-naive-implementation\">a. <a name=\"naive\"><\/a>Naive implementation<\/h2>\n\n<p>Let\u2019s begin with a direct implementation of <a href=\"#bpe-tokenization-naive\">[algorithm 1]<\/a>, which consists in 2 main procedures within a loop:<\/p>\n<ol>\n  <li>Find the best merge<\/li>\n  <li>Apply the best merge<\/li>\n<\/ol>\n\n<p>Let\u2019s examine the algorithm complexity, where <em>N<\/em> represents the number of characters of an input text. Assuming a lookup in the merges mapping can be done in constant time (a fair assumption for a standard hashmap implementation), the <code class=\"language-plaintext highlighter-rouge\">FindBestPair<\/code> procedure has a $O(N)$ complexity (<em>N-1<\/em> pair candidates need to be checked for the first iteration). The <code class=\"language-plaintext highlighter-rouge\">MergeBestPair<\/code> has a $O(N)$ complexity if the <code class=\"language-plaintext highlighter-rouge\">Merge<\/code> operation can be done in constant time if we allow for multiple merges per iteration, otherwise $O(1)$. The main <code class=\"language-plaintext highlighter-rouge\">Tokenize<\/code> procedure will iterate until no further merge can be found, which will result in <em>N-1<\/em> merges in the \u201cworst\u201d case (if the entire word exists in the vocabulary). This results in a combined complexity of $O(N^{2})$. This motivates the pre-tokenization (e.g. whitespace splitting) step that usually precedes BPE tokenization, so that <em>N<\/em> is the typical length of a word rather than a full sentence\/document.<\/p>\n\n<p>Let\u2019s look at a Rust implementation of this algorithm. The following code has been edited to focus on the critical parts of the algorithm. The full working version can be found at <a href=\"#bpe-code\">[3]<\/a>. Let\u2019s start by defining <code class=\"language-plaintext highlighter-rouge\">Symbols<\/code>, the data structure representing a sub-token:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"nd\">#[derive(Debug,<\/span> <span class=\"nd\">Copy,<\/span> <span class=\"nd\">Clone,<\/span> <span class=\"nd\">Eq,<\/span> <span class=\"nd\">PartialEq)]<\/span>\n<span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">Symbol<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">start_byte<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">end_byte<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>A symbol contains information related to the start and end byte positions of the sub-token. This is lighter than working with string slices (and significantly faster than string clones). It is also convenient for the operations we expect on <code class=\"language-plaintext highlighter-rouge\">Symbols<\/code> (merges and lookups). We will not manipulate <code class=\"language-plaintext highlighter-rouge\">Symbols<\/code> on their own, but rather work on a collection of <code class=\"language-plaintext highlighter-rouge\">Symbols<\/code>. The naive implementation, looping over the list of symbols indicates a Rust <code class=\"language-plaintext highlighter-rouge\">Vec<\/code> may be appropriate:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">SymbolArray<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">symbols<\/span><span class=\"p\">:<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">Symbol<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">,<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>The symbols are initialized from the characters of the input text (see algorithm 1, line <em>21<\/em>). We will create a method <code class=\"language-plaintext highlighter-rouge\">from_text<\/code> that populates the initial <code class=\"language-plaintext highlighter-rouge\">SymbolArray<\/code> from a string slice. Note that we look up character byte indices so that we can handle characters that span over multiple UTF-8 bytes correctly.<\/p>\n\n<p>We also implement a method for finding the best pair to merge (given a tokenizer\/merge dictionary) that will return an optional position of the best pair in the <code class=\"language-plaintext highlighter-rouge\">SymbolArray<\/code>. When this method return <code class=\"language-plaintext highlighter-rouge\">None<\/code>, we know that the array can no longer be merged. Finally, we implement directly a method to merge a pair of symbols mutating the <code class=\"language-plaintext highlighter-rouge\">SymbolArray<\/code> inplace. This method will take the best pair position as an input to directly insert the merged pair and remove the two parents.<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">SymbolArray<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"nf\">from_text<\/span><span class=\"p\">(<\/span><span class=\"n\">input_text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"n\">Self<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">symbols<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">Vec<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">for<\/span> <span class=\"p\">(<\/span><span class=\"n\">character_start<\/span><span class=\"p\">,<\/span> <span class=\"n\">character<\/span><span class=\"p\">)<\/span> <span class=\"n\">in<\/span> <span class=\"n\">input_text<\/span><span class=\"nf\">.char_indices<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">symbols<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"n\">Symbol<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">start_byte<\/span><span class=\"p\">:<\/span> <span class=\"n\">character_start<\/span><span class=\"p\">,<\/span>\n                <span class=\"n\">end_byte<\/span><span class=\"p\">:<\/span> <span class=\"n\">character_start<\/span> <span class=\"o\">+<\/span> <span class=\"n\">character<\/span><span class=\"nf\">.len_utf8<\/span><span class=\"p\">(),<\/span>\n            <span class=\"p\">});<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">Self<\/span> <span class=\"p\">{<\/span> <span class=\"n\">symbols<\/span> <span class=\"p\">}<\/span>\n    <span class=\"p\">}<\/span>\n\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"n\">find_best_merge<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">T<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">input_text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nb\">str<\/span><span class=\"p\">,<\/span> <span class=\"n\">tokenizer<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">T<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Option<\/span><span class=\"o\">&lt;<\/span><span class=\"nb\">usize<\/span><span class=\"o\">&gt;<\/span>\n    <span class=\"k\">where<\/span>\n        <span class=\"n\">T<\/span><span class=\"p\">:<\/span> <span class=\"n\">BpeTokenizer<\/span><span class=\"p\">,<\/span>\n    <span class=\"p\">{<\/span>\n        <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span>\n            <span class=\"nf\">.iter<\/span><span class=\"p\">()<\/span>\n            <span class=\"py\">.tuple_windows<\/span><span class=\"p\">::<\/span><span class=\"o\">&lt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">Symbol<\/span><span class=\"p\">,<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">Symbol<\/span><span class=\"p\">)<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">()<\/span>\n            <span class=\"nf\">.enumerate<\/span><span class=\"p\">()<\/span>\n            <span class=\"nf\">.filter_map<\/span><span class=\"p\">(|(<\/span><span class=\"n\">pos<\/span><span class=\"p\">,<\/span> <span class=\"p\">(<\/span><span class=\"n\">first<\/span><span class=\"p\">,<\/span> <span class=\"n\">second<\/span><span class=\"p\">))|<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">tokenizer<\/span>\n                    <span class=\"nf\">.get_merge_score<\/span><span class=\"p\">(<\/span><span class=\"n\">first<\/span><span class=\"p\">,<\/span> <span class=\"n\">second<\/span><span class=\"p\">,<\/span> <span class=\"n\">input_text<\/span><span class=\"p\">)<\/span>\n                    <span class=\"nf\">.map<\/span><span class=\"p\">(|<\/span><span class=\"n\">rank<\/span><span class=\"p\">|<\/span> <span class=\"p\">(<\/span><span class=\"n\">pos<\/span><span class=\"p\">,<\/span> <span class=\"n\">rank<\/span><span class=\"p\">))<\/span>\n            <span class=\"p\">})<\/span>\n            <span class=\"nf\">.min_by_key<\/span><span class=\"p\">(|(<\/span><span class=\"mi\">_<\/span><span class=\"p\">,<\/span> <span class=\"n\">rank<\/span><span class=\"p\">)|<\/span> <span class=\"o\">*<\/span><span class=\"n\">rank<\/span><span class=\"p\">)<\/span>\n            <span class=\"nf\">.map<\/span><span class=\"p\">(|(<\/span><span class=\"n\">pos<\/span><span class=\"p\">,<\/span> <span class=\"mi\">_<\/span><span class=\"p\">)|<\/span> <span class=\"n\">pos<\/span><span class=\"p\">)<\/span>\n    <span class=\"p\">}<\/span>\n\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"nf\">merge_symbols<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">best_pair_index<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"n\">Symbol<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">new_symbol<\/span> <span class=\"o\">=<\/span> <span class=\"n\">Symbol<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">start_byte<\/span><span class=\"p\">:<\/span> <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"p\">[<\/span><span class=\"n\">best_pair_index<\/span><span class=\"p\">]<\/span><span class=\"py\">.start_byte<\/span><span class=\"p\">,<\/span>\n            <span class=\"n\">end_byte<\/span><span class=\"p\">:<\/span> <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"p\">[<\/span><span class=\"n\">best_pair_index<\/span> <span class=\"o\">+<\/span> <span class=\"mi\">1<\/span><span class=\"p\">]<\/span><span class=\"py\">.end_byte<\/span><span class=\"p\">,<\/span>\n        <span class=\"p\">};<\/span>\n        <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"nf\">.remove<\/span><span class=\"p\">(<\/span><span class=\"n\">best_pair_index<\/span> <span class=\"o\">+<\/span> <span class=\"mi\">1<\/span><span class=\"p\">);<\/span>\n        <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"nf\">.remove<\/span><span class=\"p\">(<\/span><span class=\"n\">best_pair_index<\/span><span class=\"p\">);<\/span>\n        <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"nf\">.insert<\/span><span class=\"p\">(<\/span><span class=\"n\">best_pair_index<\/span><span class=\"p\">,<\/span> <span class=\"n\">new_symbol<\/span><span class=\"p\">);<\/span>\n        <span class=\"n\">new_symbol<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>We still need to implement the actual tokenizer and <code class=\"language-plaintext highlighter-rouge\">tokenize<\/code> procedure. This will pre-process the input text (replaces whitespaces by the corresponding <code class=\"language-plaintext highlighter-rouge\">\u2581<\/code> encoding symbol in the pre-trained vocabulary), pre-populate a <code class=\"language-plaintext highlighter-rouge\">SymbolArray<\/code> and identify\/merge best symbol pairs until no further merge is possible. It then returns string slices with the sub-tokens. At no point the tokenizer creates a copy of the string input: it solely relies on byte positions and string slices.<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">NaiveBpeTokenizer<\/span> <span class=\"p\">{<\/span>\n    <span class=\"n\">merges_vocab<\/span><span class=\"p\">:<\/span> <span class=\"n\">MergesVocab<\/span><span class=\"p\">,<\/span>\n<span class=\"p\">}<\/span>\n\n<span class=\"k\">impl<\/span> <span class=\"n\">BpeTokenizer<\/span> <span class=\"k\">for<\/span> <span class=\"n\">NaiveBpeTokenizer<\/span> <span class=\"p\">{<\/span>\n\n    <span class=\"k\">fn<\/span> <span class=\"n\">tokenize<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">input_text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"p\">,<\/span> <span class=\"n\">byte_mapping<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"nf\">.pre_process_text<\/span><span class=\"p\">(<\/span><span class=\"n\">input_text<\/span><span class=\"p\">,<\/span> <span class=\"sc\">'\u2581'<\/span><span class=\"p\">);<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">symbols<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">SymbolArray<\/span><span class=\"p\">::<\/span><span class=\"nf\">from_text<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">());<\/span>\n        <span class=\"k\">while<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">best_pair_index<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbols<\/span><span class=\"nf\">.find_best_merge<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">(),<\/span> <span class=\"k\">self<\/span><span class=\"p\">)<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">symbols<\/span><span class=\"nf\">.merge_symbols<\/span><span class=\"p\">(<\/span><span class=\"n\">best_pair_index<\/span><span class=\"p\">);<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">output<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">Vec<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">for<\/span> <span class=\"n\">symbol<\/span> <span class=\"n\">in<\/span> <span class=\"n\">symbols<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">output<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span>\n                <span class=\"o\">&amp;<\/span><span class=\"n\">input_text<\/span><span class=\"p\">[<\/span><span class=\"n\">byte_mapping<\/span><span class=\"p\">[<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">symbol<\/span><span class=\"py\">.start_byte<\/span><span class=\"p\">]<\/span><span class=\"o\">..<\/span><span class=\"n\">byte_mapping<\/span><span class=\"p\">[<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">symbol<\/span><span class=\"py\">.end_byte<\/span><span class=\"p\">]],<\/span>\n            <span class=\"p\">);<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">output<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>As mentioned previously, this algorithm has a $O(N^{2})$ complexity where N represents the number of characters in the input. This means that this algorithm will be unable to process long input text of the size of a sentence, paragraph or document. A common work-around consists in limiting the size of the input passed to the BPE tokenizer, for example by applying a whitespace splitting pre-tokenizer.<\/p>\n\n<h2 id=\"b-pre-splitting-naive-implementation\">b. <a name=\"naive-pre-split\"><\/a>Pre-splitting naive implementation<\/h2>\n\n<p>The previous implementation can easily be extended to include a pre-tokenization step. This is actually the standard BPE implementation in several widely used packages, such as subword-nmt (Python) <a href=\"#subword-nmt\">[4]<\/a> or fastBPE (C++) <a href=\"#fastbpe\">[5]<\/a>. By pre-tokenizing the sequence (for example whitespace splitting), one can effectively limit the average size of the inputs passed for BPE tokenization. A whitespace splitting will pass single words for processing. For most languages, the expected size of a word <em>M<\/em> is much smaller than the number of characters <em>N<\/em> in an average sentence or document (although after living in Germany for 10 years, the author realizes this hypothesis may be optimistic at times). The complexity of the tokenization for a word is $O(M^{2})$. Splitting the sequence into words using a simple whitespace\/punctuation rule has a $O(N)$ complexity, resulting in a number of words that is at most <em>N<\/em>, meaning the algorithm complexity is $O(N) + O(N*M^{2}) = O(N)$ if $M \\ll N$.<\/p>\n\n<p>To implement this tokenizer, the <code class=\"language-plaintext highlighter-rouge\">Tokenizer<\/code> has an additional method for whitespace\/punctuation tokenization. Punctuation marks are returned as a single character word, and other words contain the leading whitespace character if applicable. The method returns a vector of string slices:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">NaivePreSplitBpeTokenizer<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">fn<\/span> <span class=\"n\">split_whitespace_punctuation<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span>\n        <span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"p\">,<\/span>\n        <span class=\"n\">input_string<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"p\">,<\/span>\n        <span class=\"n\">whitespace_token<\/span><span class=\"p\">:<\/span> <span class=\"nb\">char<\/span><span class=\"p\">,<\/span>\n    <span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">output<\/span><span class=\"p\">:<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;&amp;<\/span><span class=\"nb\">str<\/span><span class=\"o\">&gt;<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">Vec<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">start<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span> <span class=\"o\">=<\/span> <span class=\"mi\">0<\/span><span class=\"p\">;<\/span>\n\n        <span class=\"k\">for<\/span> <span class=\"p\">(<\/span><span class=\"n\">c_pos<\/span><span class=\"p\">,<\/span> <span class=\"n\">c<\/span><span class=\"p\">)<\/span> <span class=\"n\">in<\/span> <span class=\"n\">input_string<\/span><span class=\"nf\">.char_indices<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">if<\/span> <span class=\"n\">c<\/span> <span class=\"o\">==<\/span> <span class=\"n\">whitespace_token<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">if<\/span> <span class=\"n\">start<\/span> <span class=\"o\">&lt;<\/span> <span class=\"n\">c_pos<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"n\">output<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">input_string<\/span><span class=\"p\">[<\/span><span class=\"n\">start<\/span><span class=\"o\">..<\/span><span class=\"n\">c_pos<\/span><span class=\"p\">]);<\/span>\n                <span class=\"p\">}<\/span>\n                <span class=\"n\">start<\/span> <span class=\"o\">=<\/span> <span class=\"n\">c_pos<\/span><span class=\"p\">;<\/span>\n            <span class=\"p\">}<\/span> <span class=\"k\">else<\/span> <span class=\"k\">if<\/span> <span class=\"n\">c<\/span><span class=\"nf\">.is_ascii_punctuation<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">if<\/span> <span class=\"n\">start<\/span> <span class=\"o\">&lt;<\/span> <span class=\"n\">c_pos<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"n\">output<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">input_string<\/span><span class=\"p\">[<\/span><span class=\"n\">start<\/span><span class=\"o\">..<\/span><span class=\"n\">c_pos<\/span><span class=\"p\">]);<\/span>\n                <span class=\"p\">}<\/span>\n                <span class=\"n\">output<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">input_string<\/span><span class=\"p\">[<\/span><span class=\"n\">c_pos<\/span><span class=\"o\">..<\/span><span class=\"n\">c_pos<\/span> <span class=\"o\">+<\/span> <span class=\"n\">c<\/span><span class=\"nf\">.len_utf8<\/span><span class=\"p\">()]);<\/span>\n                <span class=\"n\">start<\/span> <span class=\"o\">=<\/span> <span class=\"n\">c_pos<\/span> <span class=\"o\">+<\/span> <span class=\"n\">c<\/span><span class=\"nf\">.len_utf8<\/span><span class=\"p\">();<\/span>\n            <span class=\"p\">}<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"k\">if<\/span> <span class=\"n\">start<\/span> <span class=\"o\">&lt;<\/span> <span class=\"n\">input_string<\/span><span class=\"nf\">.len<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">output<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">input_string<\/span><span class=\"p\">[<\/span><span class=\"n\">start<\/span><span class=\"o\">..<\/span><span class=\"p\">]);<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">output<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>The previous <code class=\"language-plaintext highlighter-rouge\">tokenize<\/code> method can be modified to include the pre-tokenization step. This is essentially identical to the previous method, with the addition of the whitespace splitting and additional logic to keep track of the correct character offsets to return the sub-tokens:<\/p>\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">BpeTokenizer<\/span> <span class=\"k\">for<\/span> <span class=\"n\">NaivePreSplitBpeTokenizer<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">fn<\/span> <span class=\"nf\">get_merges_vocab<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">MergesVocab<\/span> <span class=\"p\">{<\/span>\n        <span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"py\">.merges_vocab<\/span>\n    <span class=\"p\">}<\/span>\n\n    <span class=\"k\">fn<\/span> <span class=\"n\">tokenize<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">input_text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">whitespace_token<\/span> <span class=\"o\">=<\/span> <span class=\"sc\">'\u2581'<\/span><span class=\"p\">;<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"p\">,<\/span> <span class=\"n\">byte_mapping<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"nf\">.pre_process_text<\/span><span class=\"p\">(<\/span><span class=\"n\">input_text<\/span><span class=\"p\">,<\/span> <span class=\"n\">whitespace_token<\/span><span class=\"p\">);<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">split_texts<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"nf\">.split_whitespace_punctuation<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">(),<\/span> <span class=\"n\">whitespace_token<\/span><span class=\"p\">);<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">output<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">Vec<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">offset<\/span> <span class=\"o\">=<\/span> <span class=\"mi\">0<\/span><span class=\"p\">;<\/span>\n        <span class=\"k\">for<\/span> <span class=\"n\">split_text<\/span> <span class=\"n\">in<\/span> <span class=\"n\">split_texts<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">symbols<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">SymbolArray<\/span><span class=\"p\">::<\/span><span class=\"nf\">from_text<\/span><span class=\"p\">(<\/span><span class=\"n\">split_text<\/span><span class=\"p\">);<\/span>\n            <span class=\"k\">while<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">best_pair_index<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbols<\/span><span class=\"nf\">.find_best_merge<\/span><span class=\"p\">(<\/span><span class=\"n\">split_text<\/span><span class=\"p\">,<\/span> <span class=\"k\">self<\/span><span class=\"p\">)<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">symbols<\/span><span class=\"nf\">.merge_symbols<\/span><span class=\"p\">(<\/span><span class=\"n\">best_pair_index<\/span><span class=\"p\">);<\/span>\n            <span class=\"p\">}<\/span>\n            <span class=\"k\">for<\/span> <span class=\"n\">symbol<\/span> <span class=\"n\">in<\/span> <span class=\"n\">symbols<\/span><span class=\"py\">.symbols<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">output<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span>\n                    <span class=\"o\">&amp;<\/span><span class=\"n\">input_text<\/span><span class=\"p\">[<\/span><span class=\"n\">byte_mapping<\/span><span class=\"p\">[<\/span><span class=\"o\">&amp;<\/span><span class=\"p\">(<\/span><span class=\"n\">offset<\/span> <span class=\"o\">+<\/span> <span class=\"n\">symbol<\/span><span class=\"py\">.start_byte<\/span><span class=\"p\">)]<\/span>\n                        <span class=\"o\">..<\/span><span class=\"n\">byte_mapping<\/span><span class=\"p\">[<\/span><span class=\"o\">&amp;<\/span><span class=\"p\">(<\/span><span class=\"n\">offset<\/span> <span class=\"o\">+<\/span> <span class=\"n\">symbol<\/span><span class=\"py\">.end_byte<\/span><span class=\"p\">)]],<\/span>\n                <span class=\"p\">);<\/span>\n            <span class=\"p\">}<\/span>\n            <span class=\"n\">offset<\/span> <span class=\"o\">+=<\/span> <span class=\"n\">split_text<\/span><span class=\"nf\">.len<\/span><span class=\"p\">();<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">output<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>This implementation brings the expected complexity from $O(N^{2})$ to $O(N)$ which is the best that can be done since the input needs to be scanned at least once to perform tokenization. This implementation, however, has the following drawbacks:<\/p>\n<ul>\n  <li>A whitespace\/punctuation tokenization step is required. This works well for most latin language, but may be problematic for example for some Asian languages not relying on whitespaces. This is a significant limitation and requires handling languages differently based on their morphology and use of word separators.<\/li>\n  <li>While faster, this algorithm is an approximation and prevents cross-word merges. It would for example be unable to merge words across punctuation symbols (such as \u201cmother-in-law\u201d)<\/li>\n<\/ul>\n\n<p>The following sections will present two additional implementation that will aim at reducing the computational cost of the naive algorithm without relying on a pre-tokenization step.<\/p>\n\n<h2 id=\"c-priority-queue--binary-search-tree-implementation\">c. <a name=\"pq-bst\"><\/a>Priority Queue + Binary Search Tree implementation<\/h2>\n\n<p>The nave algorithm suffers from a worst case $O(N^{2})$ complexity. This is due to up to <em>N<\/em> iterations of <code class=\"language-plaintext highlighter-rouge\">FindBestPair<\/code> and <code class=\"language-plaintext highlighter-rouge\">MergeBestPair<\/code>. Generally, an input may be made from up to <em>N-1<\/em> valid pairs that will be merged, and the outer loop complexity can therefore not be optimized.<\/p>\n\n<p>The <code class=\"language-plaintext highlighter-rouge\">MergeBestPair<\/code> complexity can be reduced to the time required to select the symbols to merge based on the best pair information if a single merge is performed at every iteration. The symbols can be compared and ordered based on their position - and a Binary Search Tree implementation of the <em>Symbols<\/em> sequence allows finding elements in $O(\\log(N))$ time instead of linear time for an array.<\/p>\n\n<p>The <code class=\"language-plaintext highlighter-rouge\">FindBestPair<\/code> can be optimized by noting that the entire symbol pair information does not need to be re-computed entirely at each iteration, but rather updated after every merge. Merging two symbols will have a local impact on the merge information: only the symbols preceding and succeeding the pair to be merged are impacted: we need to check if the preceding symbol and newly created merged symbol form a valid pair. Similarly, we need to check if the newly created merged symbol and next symbol form a valid pair. We still need a data structure to store the symbol pair information:<\/p>\n<ul>\n  <li><em>SymbolPairs<\/em> can be compared and ordered based on their \u201ccost\u201d: we need an ordered collection.<\/li>\n  <li>At each iteration, the <em>SymbolPair<\/em> with the minimum cost will be processed: we need a collection with an effective <code class=\"language-plaintext highlighter-rouge\">extract_min<\/code> operation.<\/li>\n  <li>At each iteration, up to 2 new <em>SymbolPairs<\/em> will be inserted: we need a collection with an effective <code class=\"language-plaintext highlighter-rouge\">insert<\/code> operation.<\/li>\n<\/ul>\n\n<p>These requirements effectively describe a priority queue which will be used to build and maintain the set of <em>SymbolPairs<\/em>. Using a Min Heap allows performing both <code class=\"language-plaintext highlighter-rouge\">extract_min<\/code> and <code class=\"language-plaintext highlighter-rouge\">insert<\/code> operations in $O(\\log(N))$ time.<\/p>\n\n<p>By implementing the following data structures:<\/p>\n<ul>\n  <li>Binary Search Tree for the <em>Symbols<\/em><\/li>\n  <li>Min Heap for the <em>SymbolPairs<\/em>,<\/li>\n<\/ul>\n\n<p>A worst-case complexity of $O(N(\\log(N) + \\log(N)))$ = $O(N\\log(N))$ can be achieved, significantly better than the initial $O(N^{2})$. The maintenance of a priority queue for the <em>SymbolPairs<\/em> to process is mentioned in the SentencePiece article <a href=\"#sentencepiece-paper\">[6]<\/a> and is used in the optimized BPE implementation of the C++ SentencePiece library <a href=\"#sentencepiece-bpe\">[7]<\/a>. <em>Algorithm 1<\/em> (BPE tokenization) becomes:<\/p>\n\n<pre id=\"read-2\" style=\"display:none;\">\n    \n\\begin{algorithm}\n\\caption{BPE Tokenize (Priority Queue)}\n\\begin{algorithmic}\n\\PROCEDURE{MaybeAddSymbolPair}{$symbol\\:pairs$, $merges$, $symbol_{left}$, $symbol_{right}$}\n    \\STATE $score = $\\CALL{Get}{$merges$, $(symbol_{left}, symbol_{right})$}\n    \\IF{$score \\neq null$}\n        \\STATE \\CALL{Insert}{$symbol\\:pairs$, $(symbol_{left}, symbol_{right}, score)$}\n    \\ENDIF\n\\ENDPROCEDURE\n\\STATE\n\\PROCEDURE{MergeBestPair}{$symbols$, $best\\:pair$}\n    \\STATE $new\\:symbol = $ \\CALL{Combine}{$best\\:pair_{left}$, $best\\:pair_{right}$}\n    \\STATE \\CALL{Remove}{$symbols$, $best\\:pair_{left}$}\n    \\STATE \\CALL{Remove}{$symbols$, $best\\:pair_{right}$}\n    \\STATE \\CALL{Insert}{$symbols$, $new\\:symbol$}\n    \\RETURN $new\\:symbol$\n\\ENDPROCEDURE\n\\STATE\n\\PROCEDURE{Tokenize}{$text$, $merges$}\n    \\STATE $symbols = $ \\CALL{InitializeSymbols}{$text$} \\COMMENT{Pre-populate symbols with individual characters and &lt;\/w&gt; token}\n    \\STATE $symbol\\:pairs = $ \\CALL{InitializeSymbolsPairs}{} \n    \\FOR{$i = 0$ \\TO $len(symbols) - 1$} \\COMMENT{Pre-populate symbol pairs}\n        \\STATE \\CALL{MaybeAddSymbolPair}{$symbol\\:pairs$, $merges$, $(symbols[i], symbols[i+1])$}\n    \\ENDFOR   \n    \\WHILE {$true$}\n        \\STATE $(best\\:pair, best\\:score) = $ \\CALL{Pop}{$symbol\\:pairs$}\n        \\IF{$best\\:score$ is $null$}\n            \\BREAK \\COMMENT{Symbols can no longer be merged}\n        \\ENDIF\n        \\IF{$best\\:pair_{left} \\neq null$ and $best\\:pair_{right} \\neq null$}\n            \\STATE $new\\:symbol$ = \\CALL{MergeBestPair}{$symbols$, $best\\:pair$}\n            \\STATE $prev = $ \\CALL{Left}{$symbols$, $best\\:pair_{left}$}\n            \\STATE $next = $ \\CALL{Right}{$symbols$, $best\\:pair_{right}$}\n            \\STATE \\CALL{MaybeAddSymbolPair}{$symbol\\:pairs$, $merges$, $(prev, best\\:pair_{left})$}\n            \\STATE \\CALL{MaybeAddSymbolPair}{$symbol\\:pairs$, $merges$, $(best\\:pair_{right})$, $next$}\n        \\ENDIF\n    \\ENDWHILE\n    \\RETURN $symbols$\n\\ENDPROCEDURE\n\\end{algorithmic}\n\\end{algorithm}\n\n<\/pre>\n<div id=\"goal-2\"><\/div>\n<script type=\"text\/javascript\">\n    var code = document.getElementById(\"read-2\").textContent;\n    var parentEl = document.getElementById(\"goal-2\");\n    var options = {\n        lineNumber: true\n    };\n    pseudocode.render(code, parentEl, options);\n<\/script>\n\n<p>Lines <em>17<\/em> to <em>21<\/em> initialize both the <em>Symbols<\/em> and <em>SymbolPairs<\/em> data structures. At each iteration, the best <em>SymbolPair<\/em> is popped from the priority queue (line <em>23<\/em>). The check on line <em>27<\/em> is required as instead of manually removing invalid merges following a merge (if the first 2 elements of a triplets get merged, the pair information for the last 2 elements is no longer valid), we pop pairs from the <em>SymbolPairs<\/em> and then check their validity. If the pair is still valid, we proceed to merge and check if new pairs should be added (lines <em>31<\/em> and <em>32<\/em>).<\/p>\n\n<p>Similarly to the <em>SymbolsArray<\/em>, the Rust implementation for the <em>Symbols<\/em> Binary Search Tree contains a method to mutate itself and merge two symbols:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">SymbolBTree<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">symbols<\/span><span class=\"p\">:<\/span> <span class=\"n\">BTreeSet<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">Symbol<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">,<\/span>\n<span class=\"p\">}<\/span>\n\n<span class=\"k\">impl<\/span> <span class=\"n\">SymbolBTree<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"nf\">merge_symbols<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">symbol_1<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">Symbol<\/span><span class=\"p\">,<\/span> <span class=\"n\">symbol_2<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">Symbol<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"n\">Symbol<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"nf\">.remove<\/span><span class=\"p\">(<\/span><span class=\"n\">symbol_1<\/span><span class=\"p\">);<\/span>\n        <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"nf\">.remove<\/span><span class=\"p\">(<\/span><span class=\"n\">symbol_2<\/span><span class=\"p\">);<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">new_symbol<\/span> <span class=\"o\">=<\/span> <span class=\"n\">Symbol<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">start_byte<\/span><span class=\"p\">:<\/span> <span class=\"n\">symbol_1<\/span><span class=\"py\">.start_byte<\/span><span class=\"p\">,<\/span>\n            <span class=\"n\">end_byte<\/span><span class=\"p\">:<\/span> <span class=\"n\">symbol_2<\/span><span class=\"py\">.end_byte<\/span><span class=\"p\">,<\/span>\n        <span class=\"p\">};<\/span>\n        <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"nf\">.insert<\/span><span class=\"p\">(<\/span><span class=\"n\">new_symbol<\/span><span class=\"p\">);<\/span>\n        <span class=\"n\">new_symbol<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>The Tokenizer contains methods to build and maintain a priority queue of <em>SymbolPairs<\/em> called \u201cagenda\u201d:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">PriorityQueueBpeTokenizer<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">fn<\/span> <span class=\"nf\">maybe_add_pair<\/span><span class=\"p\">(<\/span>\n        <span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"p\">,<\/span>\n        <span class=\"n\">left_symbol<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">Symbol<\/span><span class=\"p\">,<\/span>\n        <span class=\"n\">right_symbol<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">Symbol<\/span><span class=\"p\">,<\/span>\n        <span class=\"n\">input_text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nb\">str<\/span><span class=\"p\">,<\/span>\n        <span class=\"n\">agenda<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"n\">BinaryHeap<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">SymbolPair<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">,<\/span>\n    <span class=\"p\">)<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">merged_text<\/span> <span class=\"o\">=<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">input_text<\/span><span class=\"p\">[<\/span><span class=\"n\">left_symbol<\/span><span class=\"py\">.start_byte<\/span><span class=\"o\">..<\/span><span class=\"n\">right_symbol<\/span><span class=\"py\">.end_byte<\/span><span class=\"p\">];<\/span>\n        <span class=\"k\">if<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">score<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"py\">.merges_vocab<\/span><span class=\"nf\">.get<\/span><span class=\"p\">(<\/span><span class=\"n\">merged_text<\/span><span class=\"p\">)<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">agenda<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"n\">SymbolPair<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">left<\/span><span class=\"p\">:<\/span> <span class=\"o\">*<\/span><span class=\"n\">left_symbol<\/span><span class=\"p\">,<\/span>\n                <span class=\"n\">right<\/span><span class=\"p\">:<\/span> <span class=\"o\">*<\/span><span class=\"n\">right_symbol<\/span><span class=\"p\">,<\/span>\n                <span class=\"n\">score<\/span><span class=\"p\">,<\/span>\n            <span class=\"p\">})<\/span>\n        <span class=\"p\">}<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n\n<span class=\"k\">impl<\/span> <span class=\"n\">BpeTokenizer<\/span> <span class=\"k\">for<\/span> <span class=\"n\">PriorityQueueBpeTokenizer<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">fn<\/span> <span class=\"n\">tokenize<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">input_text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"p\">,<\/span> <span class=\"n\">byte_mapping<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"nf\">.pre_process_text<\/span><span class=\"p\">(<\/span><span class=\"n\">input_text<\/span><span class=\"p\">,<\/span> <span class=\"sc\">'\u2581'<\/span><span class=\"p\">);<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">symbols<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">SymbolBTree<\/span><span class=\"p\">::<\/span><span class=\"nf\">from_text<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">());<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">agenda<\/span><span class=\"p\">:<\/span> <span class=\"n\">BinaryHeap<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">SymbolPair<\/span><span class=\"o\">&gt;<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">BinaryHeap<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n\n        <span class=\"k\">for<\/span> <span class=\"p\">(<\/span><span class=\"n\">left_symbol<\/span><span class=\"p\">,<\/span> <span class=\"n\">right_symbol<\/span><span class=\"p\">)<\/span> <span class=\"n\">in<\/span> <span class=\"n\">symbols<\/span><span class=\"nf\">.iter<\/span><span class=\"p\">()<\/span><span class=\"py\">.tuple_windows<\/span><span class=\"p\">::<\/span><span class=\"o\">&lt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">Symbol<\/span><span class=\"p\">,<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">Symbol<\/span><span class=\"p\">)<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">self<\/span><span class=\"nf\">.maybe_add_pair<\/span><span class=\"p\">(<\/span><span class=\"n\">left_symbol<\/span><span class=\"p\">,<\/span> <span class=\"n\">right_symbol<\/span><span class=\"p\">,<\/span> <span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">(),<\/span> <span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"n\">agenda<\/span><span class=\"p\">);<\/span>\n        <span class=\"p\">}<\/span>\n        \n        <span class=\"k\">while<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">symbol_pair<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"n\">agenda<\/span><span class=\"nf\">.pop<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">let<\/span> <span class=\"n\">left_symbol<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbols<\/span><span class=\"nf\">.get<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">symbol_pair<\/span><span class=\"py\">.left<\/span><span class=\"p\">)<\/span><span class=\"nf\">.cloned<\/span><span class=\"p\">();<\/span>\n            <span class=\"k\">let<\/span> <span class=\"n\">right_symbol<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbols<\/span><span class=\"nf\">.get<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">symbol_pair<\/span><span class=\"py\">.right<\/span><span class=\"p\">)<\/span><span class=\"nf\">.cloned<\/span><span class=\"p\">();<\/span>\n\n            <span class=\"k\">if<\/span> <span class=\"k\">let<\/span> <span class=\"p\">(<\/span><span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">left_symbol<\/span><span class=\"p\">),<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">right_symbol<\/span><span class=\"p\">))<\/span> <span class=\"o\">=<\/span> <span class=\"p\">(<\/span><span class=\"n\">left_symbol<\/span><span class=\"p\">,<\/span> <span class=\"n\">right_symbol<\/span><span class=\"p\">)<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">let<\/span> <span class=\"n\">new_symbol<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbols<\/span><span class=\"nf\">.merge_symbols<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">left_symbol<\/span><span class=\"p\">,<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">right_symbol<\/span><span class=\"p\">);<\/span>\n                <span class=\"k\">if<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">next<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbols<\/span><span class=\"py\">.symbols<\/span><span class=\"nf\">.range<\/span><span class=\"p\">(<\/span><span class=\"n\">new_symbol<\/span><span class=\"o\">..<\/span><span class=\"p\">)<\/span><span class=\"nf\">.nth<\/span><span class=\"p\">(<\/span><span class=\"mi\">1<\/span><span class=\"p\">)<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"k\">self<\/span><span class=\"nf\">.maybe_add_pair<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">new_symbol<\/span><span class=\"p\">,<\/span> <span class=\"n\">next<\/span><span class=\"p\">,<\/span> <span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">(),<\/span> <span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"n\">agenda<\/span><span class=\"p\">);<\/span>\n                <span class=\"p\">}<\/span>\n                <span class=\"k\">if<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">prev<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbols<\/span><span class=\"py\">.symbols<\/span><span class=\"nf\">.range<\/span><span class=\"p\">(<\/span><span class=\"o\">..<\/span><span class=\"n\">new_symbol<\/span><span class=\"p\">)<\/span><span class=\"nf\">.next_back<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"k\">self<\/span><span class=\"nf\">.maybe_add_pair<\/span><span class=\"p\">(<\/span><span class=\"n\">prev<\/span><span class=\"p\">,<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">new_symbol<\/span><span class=\"p\">,<\/span> <span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">(),<\/span> <span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"n\">agenda<\/span><span class=\"p\">);<\/span>\n                <span class=\"p\">}<\/span>\n            <span class=\"p\">}<\/span>\n        <span class=\"p\">}<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">output<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">Vec<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">for<\/span> <span class=\"n\">symbol<\/span> <span class=\"n\">in<\/span> <span class=\"n\">symbols<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">output<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span>\n                <span class=\"o\">&amp;<\/span><span class=\"n\">input_text<\/span><span class=\"p\">[<\/span><span class=\"n\">byte_mapping<\/span><span class=\"p\">[<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">symbol<\/span><span class=\"py\">.start_byte<\/span><span class=\"p\">]<\/span><span class=\"o\">..<\/span><span class=\"n\">byte_mapping<\/span><span class=\"p\">[<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">symbol<\/span><span class=\"py\">.end_byte<\/span><span class=\"p\">]],<\/span>\n            <span class=\"p\">);<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">output<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>While this is a significant improvement over the naive implementation, accessing the left and right symbols of the best pair for merge (and their predecessor\/successor) could be optimized further: a total of 2 <code class=\"language-plaintext highlighter-rouge\">get<\/code> operations on the binary search tree and the construction of 2 \u201cthrowaway\u201d iterators to find the predecessor and successor likely impact the constant factor of the algorithm. The following section proposes a further optimization keeping the same asymptotic complexity but reducing the number of operations performed at each iteration.<\/p>\n\n<h2 id=\"d-priority-queue--linked-list-implementation\">d. <a name=\"pq-ll\"><\/a>Priority Queue + Linked List implementation<\/h2>\n\n<p>The previous implementation solves the problem to identifying the best symbols pair to merge without recalculating the entire list at each iteration, but suffers from inefficiencies to select the symbols to execute the merge. This involves the following operations:<\/p>\n<ol>\n  <li>Select the left symbol<\/li>\n  <li>Select the right symbol<\/li>\n  <li>Create a new symbol, combining left and right<\/li>\n  <li>Remove the left symbol<\/li>\n  <li>Remove the right symbol<\/li>\n  <li>Insert the new symbol<\/li>\n  <li>Select the left symbol\u2019s predecessor, check if it forms a valid pair with the new token<\/li>\n  <li>Select the right symbol\u2019s successor, check if it forms a valid pair with the new token<\/li>\n<\/ol>\n\n<p>These operations seem like a natural fit for a linked list: one can easily access predecessors and successors and replace a sequence of arbitrary nodes by a new element. If the <code class=\"language-plaintext highlighter-rouge\">SymbolPair<\/code> element contains pointer to its left and right node, one can even skip scanning the linked list to find the nodes, providing $O(1)$ complexity for all of the operations above.<\/p>\n\n<p>Linked lists are however deceptively simple, and are notoriously difficult to implement while both satisfying Rust\u2019s borrow checker and offering a high level of performance. An excellent article <a href=\"#too-many-inked-lists\">[8]<\/a> highlights the challenges of implementing linked lists in Rust. Instead, Rust\u2019s <code class=\"language-plaintext highlighter-rouge\">Vec<\/code> growable array data structure has been thoroughly optimized and is recommended over linked lists for better use of CPU cache in the official Rust documentation <a href=\"#rust-linked-list\">[9]<\/a>.<\/p>\n\n<p>The following implementation will implement a <code class=\"language-plaintext highlighter-rouge\">LinkedList<\/code> behaviour, storing the <em>Symbol<\/em> nodes in a Rust <code class=\"language-plaintext highlighter-rouge\">Vec<\/code>. This is feasible because while new nodes will be inserted following a merge, they will replace 2 local nodes and the linked list will only shrink following its construction. Our <em>Symbol<\/em>  data structure is modified to contain pointers to the previous and next nodes (which are effectively just a position in the \u201clinked list\/vec\u201d). This position pointer is a <code class=\"language-plaintext highlighter-rouge\">isize<\/code>, where a <code class=\"language-plaintext highlighter-rouge\">-1<\/code> value indicates that a given <em>Symbol<\/em> has no predecessor\/successor (first or last node in the list).<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"nd\">#[derive(Debug,<\/span> <span class=\"nd\">Copy,<\/span> <span class=\"nd\">Clone,<\/span> <span class=\"nd\">Eq,<\/span> <span class=\"nd\">PartialEq)]<\/span>\n<span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">Symbol<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">start_byte<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">end_byte<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">prev<\/span><span class=\"p\">:<\/span> <span class=\"nb\">isize<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">next<\/span><span class=\"p\">:<\/span> <span class=\"nb\">isize<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">size<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n<span class=\"p\">}<\/span>\n\n<span class=\"nd\">#[derive(Debug,<\/span> <span class=\"nd\">Copy,<\/span> <span class=\"nd\">Clone,<\/span> <span class=\"nd\">Eq,<\/span> <span class=\"nd\">PartialEq)]<\/span>\n<span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">SymbolPair<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">left<\/span><span class=\"p\">:<\/span> <span class=\"nb\">isize<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">right<\/span><span class=\"p\">:<\/span> <span class=\"nb\">isize<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">score<\/span><span class=\"p\">:<\/span> <span class=\"nb\">i64<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">pair_size<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>We have also added a <code class=\"language-plaintext highlighter-rouge\">size<\/code> field to both the <em>Symbol<\/em> and the <em>SymbolPair<\/em>. The reason is that following a merge, we will delete the right <em>Symbol<\/em>, and update the left <em>Symbol<\/em> in-place to represent the combined <em>Symbol<\/em> (we don\u2019t want to grow the linked list). The pointers of the predecessor and successor <em>Symbols<\/em> are updated to reflect these changes. However, there is no way to find all <em>SymbolPairs<\/em> that contain a given <em>Symbol<\/em> (the lookup only works in the other direction). This means that following a merge, some <em>SymbolPair<\/em> (that were pointing to the left element of the previous <em>SymbolPair<\/em>) are no longer valid: the <em>Symbol<\/em> at this position has changed (its size increased as a result of combining two symbols). We use the <code class=\"language-plaintext highlighter-rouge\">size<\/code> information in the <em>Symbol<\/em> and <em>SymbolPair<\/em> as a validation step when popping a new <em>SymbolPair<\/em> from the agenda: if the <code class=\"language-plaintext highlighter-rouge\">size<\/code> of the <em>SymbolPair<\/em> is smaller than the sum of the <code class=\"language-plaintext highlighter-rouge\">size<\/code> of its left and right <em>Symbols<\/em>, this <em>SymbolPair<\/em> is no longer valid and we ignore it (pop the next <em>SymbolPair<\/em>). The size for all <em>Symbols<\/em> is initialized as <code class=\"language-plaintext highlighter-rouge\">1<\/code> when the linked list is constructed from the character list of the text to tokenize.<\/p>\n\n<p><img src=\"..\/assets\/bpe\/bpe_schematics.svg\" alt=\"Linked list symbol merge\" title=\"Size validation following Symbols merge\" \/><\/p>\n\n<p>The data structure holding the <em>Symbols<\/em> is now a list, implemented using a Rust <code class=\"language-plaintext highlighter-rouge\">Vec<\/code>. It implements the method for merging two symbols (mutates itself) from two symbol positions and a size validation (the merge is executed only if the size validation described above succeeds).<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">SymbolList<\/span> <span class=\"p\">{<\/span>\n    <span class=\"n\">symbols<\/span><span class=\"p\">:<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;<\/span><span class=\"nb\">Option<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">Symbol<\/span><span class=\"o\">&gt;&gt;<\/span><span class=\"p\">,<\/span>\n<span class=\"p\">}<\/span>\n\n<span class=\"k\">impl<\/span> <span class=\"n\">SymbolList<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"nf\">merge_symbols<\/span><span class=\"p\">(<\/span>\n        <span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"k\">self<\/span><span class=\"p\">,<\/span>\n        <span class=\"n\">symbol_1_index<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n        <span class=\"n\">symbol_2_index<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n        <span class=\"n\">size_validation<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n    <span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Option<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">Symbol<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">if<\/span> <span class=\"k\">let<\/span> <span class=\"p\">(<\/span><span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">left_symbol<\/span><span class=\"p\">),<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">right_symbol<\/span><span class=\"p\">))<\/span> <span class=\"o\">=<\/span>\n        <span class=\"p\">(<\/span><span class=\"k\">self<\/span><span class=\"p\">[<\/span><span class=\"n\">symbol_1_index<\/span><span class=\"p\">],<\/span> <span class=\"k\">self<\/span><span class=\"p\">[<\/span><span class=\"n\">symbol_2_index<\/span><span class=\"p\">])<\/span>\n        <span class=\"p\">{<\/span>\n            <span class=\"k\">if<\/span> <span class=\"n\">left_symbol<\/span><span class=\"py\">.size<\/span> <span class=\"o\">+<\/span> <span class=\"n\">right_symbol<\/span><span class=\"py\">.size<\/span> <span class=\"o\">!=<\/span> <span class=\"n\">size_validation<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">return<\/span> <span class=\"nb\">None<\/span><span class=\"p\">;<\/span>\n            <span class=\"p\">}<\/span>\n            <span class=\"k\">if<\/span> <span class=\"n\">right_symbol<\/span><span class=\"py\">.next<\/span> <span class=\"o\">!=<\/span> <span class=\"o\">-<\/span><span class=\"mi\">1<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">if<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">next_next<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"nf\">.get_mut<\/span><span class=\"p\">(<\/span><span class=\"n\">right_symbol<\/span><span class=\"py\">.next<\/span> <span class=\"k\">as<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">)<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"n\">next_next<\/span><span class=\"py\">.prev<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbol_1_index<\/span> <span class=\"k\">as<\/span> <span class=\"nb\">isize<\/span><span class=\"p\">;<\/span>\n                <span class=\"p\">}<\/span>\n            <span class=\"p\">}<\/span>\n            <span class=\"k\">let<\/span> <span class=\"n\">new_symbol<\/span> <span class=\"o\">=<\/span> <span class=\"n\">Symbol<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">start_byte<\/span><span class=\"p\">:<\/span> <span class=\"n\">left_symbol<\/span><span class=\"py\">.start_byte<\/span><span class=\"p\">,<\/span>\n                <span class=\"n\">end_byte<\/span><span class=\"p\">:<\/span> <span class=\"n\">right_symbol<\/span><span class=\"py\">.end_byte<\/span><span class=\"p\">,<\/span>\n                <span class=\"n\">prev<\/span><span class=\"p\">:<\/span> <span class=\"n\">left_symbol<\/span><span class=\"py\">.prev<\/span><span class=\"p\">,<\/span>\n                <span class=\"n\">next<\/span><span class=\"p\">:<\/span> <span class=\"n\">right_symbol<\/span><span class=\"py\">.next<\/span><span class=\"p\">,<\/span>\n                <span class=\"n\">size<\/span><span class=\"p\">:<\/span> <span class=\"n\">left_symbol<\/span><span class=\"py\">.size<\/span> <span class=\"o\">+<\/span> <span class=\"n\">right_symbol<\/span><span class=\"py\">.size<\/span><span class=\"p\">,<\/span>\n            <span class=\"p\">};<\/span>\n            <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"p\">[<\/span><span class=\"n\">symbol_2_index<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"nb\">None<\/span><span class=\"p\">;<\/span>\n            <span class=\"k\">self<\/span><span class=\"py\">.symbols<\/span><span class=\"p\">[<\/span><span class=\"n\">symbol_1_index<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">new_symbol<\/span><span class=\"p\">);<\/span>\n            <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">new_symbol<\/span><span class=\"p\">)<\/span>\n        <span class=\"p\">}<\/span> <span class=\"k\">else<\/span> <span class=\"p\">{<\/span>\n            <span class=\"nb\">None<\/span>\n        <span class=\"p\">}<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>The tokenizer method looks very similar to the previous implementations. The <code class=\"language-plaintext highlighter-rouge\">maybe_add_pair<\/code> method only changes to calculate the size of a <em>SymbolPair<\/em> from the sum of its <em>Symbols<\/em> sizes and is skipped below (see the full code at <a href=\"#bpe-code\">[3]<\/a>). Accessing the left and right token of a <em>SymbolPair<\/em>, along with their predecessor and successor is now done by directly looking up the fields <code class=\"language-plaintext highlighter-rouge\">left<\/code>, <code class=\"language-plaintext highlighter-rouge\">right<\/code> of the <em>SymbolPair<\/em> and <code class=\"language-plaintext highlighter-rouge\">prev<\/code>, <code class=\"language-plaintext highlighter-rouge\">next<\/code> of the <em>Symbols<\/em>.<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">BpeTokenizer<\/span> <span class=\"k\">for<\/span> <span class=\"n\">PriorityQueueBpeLLTokenizer<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">fn<\/span> <span class=\"n\">tokenize<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">input_text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"p\">,<\/span> <span class=\"n\">byte_mapping<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"nf\">.pre_process_text<\/span><span class=\"p\">(<\/span><span class=\"n\">input_text<\/span><span class=\"p\">,<\/span> <span class=\"sc\">'\u2581'<\/span><span class=\"p\">);<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">symbols<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">SymbolList<\/span><span class=\"p\">::<\/span><span class=\"nf\">from_text<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">());<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">agenda<\/span><span class=\"p\">:<\/span> <span class=\"n\">BinaryHeap<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">SymbolPair<\/span><span class=\"o\">&gt;<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">BinaryHeap<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n\n        <span class=\"k\">for<\/span> <span class=\"n\">symbol_index<\/span> <span class=\"n\">in<\/span> <span class=\"mi\">1<\/span><span class=\"o\">..<\/span><span class=\"n\">symbols<\/span><span class=\"nf\">.len<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">self<\/span><span class=\"nf\">.maybe_add_pair<\/span><span class=\"p\">(<\/span>\n                <span class=\"n\">symbol_index<\/span> <span class=\"k\">as<\/span> <span class=\"nb\">isize<\/span> <span class=\"o\">-<\/span> <span class=\"mi\">1<\/span><span class=\"p\">,<\/span>\n                <span class=\"n\">symbol_index<\/span> <span class=\"k\">as<\/span> <span class=\"nb\">isize<\/span><span class=\"p\">,<\/span>\n                <span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">(),<\/span>\n                <span class=\"o\">&amp;<\/span><span class=\"n\">symbols<\/span><span class=\"p\">,<\/span>\n                <span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"n\">agenda<\/span><span class=\"p\">,<\/span>\n            <span class=\"p\">);<\/span>\n        <span class=\"p\">}<\/span>\n\n        <span class=\"k\">while<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">symbol_pair<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"n\">agenda<\/span><span class=\"nf\">.pop<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">let<\/span> <span class=\"n\">left_symbol_index<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbol_pair<\/span><span class=\"py\">.left<\/span><span class=\"p\">;<\/span>\n            <span class=\"k\">let<\/span> <span class=\"n\">right_symbol_index<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbol_pair<\/span><span class=\"py\">.right<\/span><span class=\"p\">;<\/span>\n            <span class=\"k\">if<\/span> <span class=\"n\">left_symbol_index<\/span> <span class=\"o\">!=<\/span> <span class=\"o\">-<\/span><span class=\"mi\">1<\/span> <span class=\"o\">&amp;&amp;<\/span> <span class=\"n\">right_symbol_index<\/span> <span class=\"o\">!=<\/span> <span class=\"o\">-<\/span><span class=\"mi\">1<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">let<\/span> <span class=\"n\">new_symbol<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbols<\/span><span class=\"nf\">.merge_symbols<\/span><span class=\"p\">(<\/span>\n                    <span class=\"n\">left_symbol_index<\/span> <span class=\"k\">as<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n                    <span class=\"n\">right_symbol_index<\/span> <span class=\"k\">as<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n                    <span class=\"n\">symbol_pair<\/span><span class=\"py\">.pair_size<\/span><span class=\"p\">,<\/span>\n                <span class=\"p\">);<\/span>\n                <span class=\"k\">if<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">new_symbol<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"n\">new_symbol<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"k\">self<\/span><span class=\"nf\">.maybe_add_pair<\/span><span class=\"p\">(<\/span>\n                        <span class=\"n\">new_symbol<\/span><span class=\"py\">.prev<\/span><span class=\"p\">,<\/span>\n                        <span class=\"n\">left_symbol_index<\/span><span class=\"p\">,<\/span>\n                        <span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">(),<\/span>\n                        <span class=\"o\">&amp;<\/span><span class=\"n\">symbols<\/span><span class=\"p\">,<\/span>\n                        <span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"n\">agenda<\/span><span class=\"p\">,<\/span>\n                    <span class=\"p\">);<\/span>\n                    <span class=\"k\">self<\/span><span class=\"nf\">.maybe_add_pair<\/span><span class=\"p\">(<\/span>\n                        <span class=\"n\">left_symbol_index<\/span><span class=\"p\">,<\/span>\n                        <span class=\"n\">new_symbol<\/span><span class=\"py\">.next<\/span><span class=\"p\">,<\/span>\n                        <span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">(),<\/span>\n                        <span class=\"o\">&amp;<\/span><span class=\"n\">symbols<\/span><span class=\"p\">,<\/span>\n                        <span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"n\">agenda<\/span><span class=\"p\">,<\/span>\n                    <span class=\"p\">);<\/span>\n                <span class=\"p\">}<\/span>\n            <span class=\"p\">}<\/span>\n        <span class=\"p\">}<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">output<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">Vec<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">for<\/span> <span class=\"n\">symbol<\/span> <span class=\"n\">in<\/span> <span class=\"n\">symbols<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">if<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">symbol<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"n\">symbol<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">output<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span>\n                    <span class=\"o\">&amp;<\/span><span class=\"n\">input_text<\/span><span class=\"p\">[<\/span><span class=\"n\">byte_mapping<\/span><span class=\"p\">[<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">symbol<\/span><span class=\"py\">.start_byte<\/span><span class=\"p\">]<\/span><span class=\"o\">..<\/span><span class=\"n\">byte_mapping<\/span><span class=\"p\">[<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">symbol<\/span><span class=\"py\">.end_byte<\/span><span class=\"p\">]],<\/span>\n                <span class=\"p\">);<\/span>\n            <span class=\"p\">}<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">output<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<h1 id=\"3-benchmarks\">3. Benchmarks<\/h1>\n\n<p>So far we have compared the implementations on a theoretical complexity level. In reality, the constants hidden in the asymptotic behaviour can have a significant impact. This section reports experimental results taking samples of varying size of Shakespeare\u2019s Hamlet <a href=\"#hamlet\">[10]<\/a>. The time taken to tokenize the first 1, 10, 100 or 1000 lines of the play is recorded for the 4 implementations presented previously:<\/p>\n\n<style>\ntable th:first-of-type {\n    width: 20%;\n}\ntable th:nth-of-type(2) {\n    width: 10%;\n}\ntable th:nth-of-type(3) {\n    width: 20%;\n}\ntable th:nth-of-type(4) {\n    width: 30%;\n}\ntable th:nth-of-type(5) {\n    width: 30%;\n}\n<\/style>\n\n<table>\n  <thead>\n    <tr>\n      <th style=\"text-align: left\">Input size<\/th>\n      <th style=\"text-align: left\">Naive<\/th>\n      <th style=\"text-align: left\">Naive <br \/>(pre-split)<\/th>\n      <th style=\"text-align: left\">Priority Queue + <br \/>Binary Search Tree<\/th>\n      <th style=\"text-align: left\">Priority Queue + Linked List<\/th>\n    <\/tr>\n  <\/thead>\n  <tbody>\n    <tr>\n      <td style=\"text-align: left\">1<\/td>\n      <td style=\"text-align: left\">27.1 $\\mu$s<\/td>\n      <td style=\"text-align: left\">9.1 $\\mu$s<\/td>\n      <td style=\"text-align: left\">14.7 $\\mu$s<\/td>\n      <td style=\"text-align: left\">8.6 $\\mu$s<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\">10<\/td>\n      <td style=\"text-align: left\">162 $\\mu$s<\/td>\n      <td style=\"text-align: left\">26.5 $\\mu$s<\/td>\n      <td style=\"text-align: left\">54.4 $\\mu$s<\/td>\n      <td style=\"text-align: left\">24.7 $\\mu$s<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\">100<\/td>\n      <td style=\"text-align: left\">85 ms<\/td>\n      <td style=\"text-align: left\">0.57 ms<\/td>\n      <td style=\"text-align: left\">1.76 ms<\/td>\n      <td style=\"text-align: left\">0.68 ms<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\">1000<\/td>\n      <td style=\"text-align: left\">18.6 s<\/td>\n      <td style=\"text-align: left\">8.6 ms<\/td>\n      <td style=\"text-align: left\">33.9 ms<\/td>\n      <td style=\"text-align: left\">12.8 ms<\/td>\n    <\/tr>\n  <\/tbody>\n<\/table>\n\n<p>This confirms the expectations derived from the algorithms earlier: the naive approach becomes unpractical for inputs lengths exceeding 100 lines. Meanwhile, all other approaches stay in the order of a millisecond even for inputs that are 1000 lines long, confirming the asymptotic benefits of the data structures investigated. We also see that the two exact solutions leveraging priority queues provide execution times that are in line with the pre-split approximation. Even though they have the same asymptotic complexity, we also note that the linked-list implementation for the <code class=\"language-plaintext highlighter-rouge\">Symbols<\/code> outperforms the binary search tree version.<\/p>\n\n<p>These results, along with asymptotic trend-lines, can be seen in the figure below:<\/p>\n\n<p><img src=\"..\/assets\/bpe\/bpe_benchmark.svg\" alt=\"BPE implementations benchmark\" title=\"BPE implementations benchmark\" \/><\/p>\n\n<h1 id=\"conclusion\">Conclusion<\/h1>\n\n<p>Byte pair Encoding is a tokenization method that is in essence very simple and effective as a pre-processing step for modern machine learning pipelines. Widely used in multiple productive libraries, its actual implementation can vary significantly from one source to another. This article gives an overview of some key implementations of the algorithm that the reader may encounter and provides a high-level intuition behind their design. It illustrates the impact that the choice of a data structure can have on a real NLP application that is used everyday by thousands of data scientists and machine learning engineers. The priority-queue \/ linked-list implementation of Byte pair Encoding has been implemented in the rust-tokenizers library <a href=\"#rust-tokenizers\">[11]<\/a>, along with other modern tokenization algorithms.<\/p>\n\n<h2 id=\"references\">References<\/h2>\n<ul>\n  <li><a name=\"bpe\"><\/a>[1] <a href=\"https:\/\/arxiv.org\/abs\/1508.07909\">Neural Machine Translation of Rare Words with Subword Units<\/a>, Rico Sennrich, Barry Haddow, Alexandra Birch, 2015<\/li>\n  <li><a name=\"bpe-lecture\"><\/a>[2] <a href=\"https:\/\/web.stanford.edu\/class\/cs224n\/slides\/cs224n-2019-lecture12-subwords.pdf\">Natural Language Processing\nwith Deep Learning, CS224N\/Ling284, lecture 12<\/a>, Christopher Manning, Stanford.<\/li>\n  <li><a name=\"bpe-code\"><\/a>[3] <a href=\"https:\/\/github.com\/guillaume-be\/bpe-example\">Byte Pair Encoding Rust Implementation<\/a>, Guillaume Becquin<\/li>\n  <li><a name=\"subword-nmt\"><\/a>[4] <a href=\"https:\/\/github.com\/rsennrich\/subword-nmt\">Subword Neural Machine Translation<\/a><\/li>\n  <li><a name=\"fastbpe\"><\/a>[5] <a href=\"https:\/\/github.com\/glample\/fastBPE\">fastBPE<\/a><\/li>\n  <li><a name=\"sentencepiece-paper\"><\/a>[6] <a href=\"https:\/\/arxiv.org\/abs\/1808.06226\">SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing<\/a>, Taku Kudo, John Richardson<\/li>\n  <li><a name=\"sentencepiece-bpe\"><\/a>[7] <a href=\"https:\/\/github.com\/google\/sentencepiece\/blob\/master\/src\/bpe_model.cc\">SentencePiece BPE model<\/a>, Taku Kudo<\/li>\n  <li><a name=\"too-many-linked-lists\"><\/a>[8] <a href=\"https:\/\/rust-unofficial.github.io\/too-many-lists\/\">Learn Rust With Entirely Too Many Linked Lists<\/a><\/li>\n  <li><a name=\"rust-linked-list\"><\/a>[9] <a href=\"https:\/\/doc.rust-lang.org\/std\/collections\/struct.LinkedList.html\">Rust std LinkedList documentation<\/a><\/li>\n  <li><a name=\"hamlet\"><\/a>[10] <a href=\"http:\/\/shakespeare.mit.edu\/hamlet\/full.html\">The Tragedy of Hamlet, Prince of Denmark<\/a>, William Shakespeare<\/li>\n  <li><a name=\"rust-tokenizers\"><\/a>[11] <a href=\"https:\/\/github.com\/guillaume-be\/rust-tokenizers\">rust-tokenizers<\/a><\/li>\n<\/ul>","author":{"name":"Guillaume"},"summary":"Tokenization of input strings into sequences of words or sub-tokens is a central concept for modern Natural Language Processing techniques (NLP). This article focuses on a classic tokenization algorithm: Byte Pair Encoding (BPE) [1]. While resources describing the working principle of the algorithm are widely available, this article focuses on its implementation, illustrating how the choice of data structures impact the performance of a real-world NLP component. This article is complemented by a working Rust-based implementation, available on GitHub."},{"title":"Accelerating text generation with Rust","link":{"@attributes":{"href":"https:\/\/guillaume-be.github.io\/2020-11-21\/generation_benchmarks","rel":"alternate","type":"text\/html","title":"Accelerating text generation with Rust"}},"published":"2020-11-21T00:00:00+00:00","updated":"2020-11-21T00:00:00+00:00","id":"https:\/\/guillaume-be.github.io\/2020-11-21\/generation_benchmarks","content":"<p>Over the past few months, text generation capabilities using Transformer-based models have been democratized by open-source efforts such as Hugging Face\u2019s Transformers <a href=\"#transformers\">[1]<\/a> library. A broad range of models and applications have been made available, including:<\/p>\n<ul>\n  <li>Summarization models fine-tuned on the CNN-DailyMail <a href=\"#cnndailymail\">[2]<\/a> or XSUM <a href=\"#xsum\">[3]<\/a> datasets, including for example BART <a href=\"#bart\">[4]<\/a> or T5 <a href=\"#t5\">[5]<\/a><\/li>\n  <li>Translation, with the availability of more than a thousand models trained by the Opus-MT team from Language Technology at the University of Helsinki <a href=\"#opusmt\">[6]<\/a><\/li>\n  <li>Text generation, generating tokens from a prompt text, including OpenAI\u2019s GPT2 model <a href=\"#gpt2\">[7]<\/a><\/li>\n<\/ul>\n\n<p>These models offer state-of-the-art performance and can be set-up with just a few lines of code using ready-to-use pipelines in the Transformers library. This allowed a wide spread of the recent advances in the field of NLP to the industry where they can be effectively used to solve business-driven use cases. As these powerful models gain a broader adoption, the issue of their computational efficiency becomes ever more critical. A lot of research has been done to reduce the size of these models with minimal impact on their accuracy, including for example distillation, quantization or pruning. This article will focus on another avenue to improve the runtime of text generation models: an implementation in the Rust programming language.<\/p>\n\n<p>The structure of the article is as follows:<\/p>\n\n<ol>\n  <li>Overview of text generation tasks and Python baseline code examples from the Transformers <a href=\"#transformers\">[1]<\/a> library<\/li>\n  <li>Description of the text generation architecture used for summarization, translation and free text generation tasks<\/li>\n  <li>Introduction to the <em>rust-bert<\/em> library, a project offering Rust-native implementation of these models in Rust<\/li>\n  <li>Benchmarks between the baseline Python implementation from Transformers and the proposed Rust implementation for these text generation tasks<\/li>\n<\/ol>\n\n<p>Readers familiar with text generation tasks and their implementation in modern language models may want to skip directly to <a href=\"#rustbert-section\">section 3<\/a>. The first two parts provide a high-level overview of the technology to better understand the scope of the proposed performance improvements and give a few references in order to dive deeper into the topic.<\/p>\n\n<h1 id=\"1-overview-of-text-generation-tasks\">1. Overview of text generation tasks<a name=\"task-overview\"><\/a><\/h1>\n\n<p>To illustrate both translation and summarization capabilities, we\u2019ll use a new article titled <a href=\"https:\/\/en.wikinews.org\/wiki\/Astronomers_find_water_vapour_in_atmosphere_of_exoplanet_K2-18b\" target=\"_blank\">Astronomers find water vapour in atmosphere of exoplanet K2-18b<\/a> from  WikiNews shared under  CC BY 2.5 license <a href=\"#wikinews\">[8]<\/a>.<\/p>\n\n<h3 id=\"summarization\">Summarization<\/h3>\n<p>In Python, this article can be summarized calling the following snippet from the Transformer\u2019s Python library <a href=\"#transformers\">[1]<\/a>, defaulting to a BART model trained on the CNN-DailyMail dataset:<\/p>\n<div class=\"language-python highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"kn\">from<\/span> <span class=\"nn\">transformers<\/span> <span class=\"kn\">import<\/span> <span class=\"n\">pipeline<\/span>\n<span class=\"n\">summarization_pipeline<\/span> <span class=\"o\">=<\/span> <span class=\"n\">pipeline<\/span><span class=\"p\">(<\/span><span class=\"s\">\"summarization\"<\/span><span class=\"p\">)<\/span>\n<span class=\"n\">summarization_pipeline<\/span><span class=\"p\">(<\/span><span class=\"n\">input_article<\/span><span class=\"p\">)<\/span>\n<\/code><\/pre><\/div><\/div>\n<p>returning:<\/p>\n<blockquote>\n  <p>K2-18b is the first such discovery in a planet in its star\u2019s habitable zone. It is not too hot and not too cold for liquid water to exist on a planet that orbits a star 110 light years from Earth. Scientists from the University of Montreal and a team from UCL found water in the atmosphere of the planet. The Montreal team used data from the NASA\u2019s Hubble telescope to assess changes in the light coming from the star as the planet passed between it and Earth.<\/p>\n<\/blockquote>\n\n<h3 id=\"translation\">Translation<\/h3>\n<p>Similarly, a translation model can be easily created to translate a sentence taken from this document to Spanish:<\/p>\n<div class=\"language-python highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"kn\">from<\/span> <span class=\"nn\">transformers<\/span> <span class=\"kn\">import<\/span> <span class=\"n\">pipeline<\/span>\n<span class=\"n\">translation_pipeline<\/span> <span class=\"o\">=<\/span> <span class=\"n\">pipeline<\/span><span class=\"p\">(<\/span><span class=\"s\">\"translation_en_to_es\"<\/span><span class=\"p\">)<\/span>\n<span class=\"n\">translation_pipeline<\/span><span class=\"p\">(<\/span><span class=\"s\">\"They found that certain wavelengths of light, which are usually <\/span><span class=\"se\">\\\n<\/span><span class=\"s\"> absorbed by water, weakened when the planet was in the way, indicating not only <\/span><span class=\"se\">\\\n<\/span><span class=\"s\"> does K2-18b have an atmosphere, but the atmosphere contains water in vapour form.\"<\/span><span class=\"p\">)<\/span>\n<\/code><\/pre><\/div><\/div>\n<p>returning:<\/p>\n<blockquote>\n  <p>Encontraron que ciertas longitudes de onda de la luz, que generalmente son absorbidas por el agua, se debilitaban cuando el planeta estaba en el camino, lo que indica que no s\u00f3lo K2-18b tiene una atm\u00f3sfera, sino que la atm\u00f3sfera contiene agua en forma de vapor.<\/p>\n<\/blockquote>\n\n<h3 id=\"text-generation\">Text generation<\/h3>\n<p>Text can be generated using a text generation pipeline:<\/p>\n<div class=\"language-python highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"kn\">from<\/span> <span class=\"nn\">transformers<\/span> <span class=\"kn\">import<\/span> <span class=\"n\">pipeline<\/span>\n<span class=\"n\">text_generator<\/span> <span class=\"o\">=<\/span> <span class=\"n\">pipeline<\/span><span class=\"p\">(<\/span><span class=\"s\">\"text-generation\"<\/span><span class=\"p\">)<\/span>\n<span class=\"k\">print<\/span><span class=\"p\">(<\/span><span class=\"n\">text_generator<\/span><span class=\"p\">(<\/span><span class=\"s\">\"The majority of crustaceans are aquatic,\"<\/span><span class=\"p\">,<\/span> \n    <span class=\"n\">max_length<\/span><span class=\"o\">=<\/span><span class=\"mi\">64<\/span><span class=\"p\">,<\/span> \n    <span class=\"n\">do_sample<\/span><span class=\"o\">=<\/span><span class=\"bp\">False<\/span><span class=\"p\">))<\/span>\n<\/code><\/pre><\/div><\/div>\n<p>returning:<\/p>\n<blockquote>\n  <p>The majority of crustaceans are aquatic, meaning they live on land, rivers, and lakes. Carnivorous crustacean species, such as those found in the Pacific Northwest, are found in all parts of the world, including the United States, Canada, Australia, New Zealand, and Japan.<\/p>\n<\/blockquote>\n\n<p>The next section will briefly describe the architecture powering these models before diving into a comparison of the baseline Python <em>Transformers<\/em> library with a proposed Rust-based implementation: <em>rust-bert<\/em> <a href=\"#rustbert\">[9]<\/a><\/p>\n\n<h1 id=\"2-overview-of-models-and-system-architecture\">2. Overview of models and system architecture<a name=\"models-overview\"><\/a><\/h1>\n\n<h3 id=\"summarization-and-translation\">Summarization and Translation<\/h3>\n\n<p>Translation and summarization both rely on a similar architecture, although the model weights naturally vary from application to application. They are essentially made of:<\/p>\n<ol>\n  <li>\n    <p>A pre-processing pipeline mostly comprising of a tokenizer (such as Byte Pair Encoding or SentencePiece\/Unigram-based) and an encoder (mapping individual tokens to vocabulary indices and other optional inputs (such as position indices)).<\/p>\n  <\/li>\n  <li>\n    <p>A transformer-based model, based on an encoder-decoder architecture. If you are not familiar with Transformer-based encoder-decoder architecture, I highly recommend the blog post <em>The Illustrated Transformer<\/em> <a href=\"#illustratedtransformer\">[10]<\/a>. The encoder is comprised of a stack of self-attention and fully connected layers and encodes the input sequence (i.e. text to be translated or summarized) into a latent space. The decoder is made of a similar stack of self-attention layers completed with cross-attention to the encoder hidden states, allowing to leverage the representations generated during encoding. The decoder takes as an input the output sequence generated so far and the encoder output to generate the next tokens. The decoder therefore generates output tokens one at a time.<\/p>\n  <\/li>\n  <li>\n    <p>A generation routine, which in its simplest form will keep calling the transformer-based models to generate tokens until the sequence is completed (output of an <code class=\"language-plaintext highlighter-rouge\">End Of Sequence<\/code> token). Note that the encoder only needs to be run once in this iterative process: its output is cached and re-used at each decoder generation step. In practice, more advanced algorithms are used to improve the quality of the generation, including beam search, sampling, length and repetition penalties. These methods are summarized in an excellent article from Hugging Face <a href=\"#generation\">[11]<\/a>. Careful design of the decoder allows to not only cache the encoder states, but also parts of the keys and values in the decoder to avoid unnecessary re-calculation and speeds-up the decoding process.<\/p>\n  <\/li>\n<\/ol>\n\n<p>This iterative process is illustrated at a high level in the figure below (with slight simplifications, especially for the end of generation condition):<\/p>\n\n<p><img src=\"..\/assets\/generation_benchmarks\/encoder_decoder.svg\" alt=\"Encoder decoder generation architecture\" title=\"Encoder decoder generation architecture\" \/><\/p>\n\n<p>This process (and in the special case of BART and Marian - the model architecture itself) is identical between translation and summarization. Only the tokenization process and the model parameters differ between the two applications, showing the high versatility of this system.<\/p>\n\n<h3 id=\"text-generation-1\">Text generation<\/h3>\n\n<p>The process for text generation using GPT2 is very similar. However, GPT2 is a decoder-only model, and does not contain the encoder part of the transformers architectures. The model uses the starting prompt (and sequence generated so far) as the only input. While it therefore does not need to compute encoder states and cache them, it still relies on an efficient caching mechanism to avoid unnecessary re-computation of activations already computed during the generation process.<\/p>\n\n<p><img src=\"..\/assets\/generation_benchmarks\/decoder.svg\" alt=\"Decoder generation architecture\" title=\"Decoder generation architecture\" width=\"70%\" \/><\/p>\n\n<h3 id=\"on-the-complexity-of-the-generation-routine\">On the complexity of the generation routine<\/h3>\n\n<p>For both architectures, one may note the complexity of the generation routine, involving a significant number of operations beyond the model forward pass, helping to improve the quality of the generated text. However, these improvements do not come for free and incur an additional computational cost in both the model forward pass (beam search increases the effective batch size) and in post-processing operations.<\/p>\n\n<p>The Python Transformer\u2019s library already leverages Rust-based tokenizers for all its ready-to-use pipelines, therefore accelerating the preprocessing part of the system. Some benchmarks on question answering <a href=\"#rustbert\">[9]<\/a> indicate that the preprocessing can amount to 20% to 30% of the processing time for simple pipelines. The post-processing, also involving operations that go beyond tensor operations, is however implemented in Python in the Transformer\u2019s library. The rest of this article assesses the impact of a high-performance implementation of the entire system in Rust (therefore covering the post-processing pipeline) using the rust-bert library <a href=\"#rustbert\">[9]<\/a>.<\/p>\n\n<h1 id=\"3-a-brief-introduction-to-rust-bert\">3. A brief introduction to rust-bert<a name=\"rustbert-section\"><\/a><\/h1>\n\n<p>Rust-bert is essentially a Rust-native port of Hugging Face\u2019s Transformers\u2019 library <a href=\"#transformers\">[1]<\/a>. Leveraging the <code class=\"language-plaintext highlighter-rouge\">rust_tokenizers<\/code> <a href=\"#rusttokenizers\">[12]<\/a> library for preprocessing, it proposes implementations for state-of-the-art transformers-based models and ready-to-use pipelines. De-noising auto-encoder (BERT, Electra), autoregressive (XLNet, GPT, GPT2) and encoder-decoder models (BART, T5) have been implemented with pre-trained sets of weights available on Hugging Face\u2019s model hub <a href=\"#modelhub\">[13]<\/a>. Any Pytorch model trained on the Transformers\u2019s library can be converted to a C-array format and used by the <code class=\"language-plaintext highlighter-rouge\">rust-bert<\/code> library.<\/p>\n\n<p>These models can be used in ready-to-use pipelines, including:<\/p>\n<ul>\n  <li>classification (e.g. sentiment analysis)<\/li>\n  <li>token classification (e.g. named entities recognition)<\/li>\n  <li>extractive question answering<\/li>\n  <li>zero-shot classification, using a natural language inference classifier<\/li>\n  <li>text generation<\/li>\n  <li>conversational model<\/li>\n  <li>translation<\/li>\n  <li>summarization<\/li>\n<\/ul>\n\n<p>The last three text generation pipelines allow for a side-by-side comparison of the Python implementation (with Rust tokenizers) and the end-to-end Rust version. The three pipelines mentioned above can also be instantiated in Rust in a few lines of code:<\/p>\n\n<h3 id=\"summarization-1\">Summarization<\/h3>\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">use<\/span> <span class=\"nn\">rust_bert<\/span><span class=\"p\">::<\/span><span class=\"nn\">pipelines<\/span><span class=\"p\">::<\/span><span class=\"nn\">summarization<\/span><span class=\"p\">::<\/span><span class=\"n\">SummarizationModel<\/span><span class=\"p\">;<\/span>\n<span class=\"k\">let<\/span> <span class=\"n\">summarization_model<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">SummarizationModel<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">(<\/span><span class=\"nn\">Default<\/span><span class=\"p\">::<\/span><span class=\"nf\">default<\/span><span class=\"p\">())<\/span><span class=\"o\">?<\/span><span class=\"p\">;<\/span>\n<span class=\"n\">summarization_model<\/span><span class=\"nf\">.summarize<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">input<\/span><span class=\"p\">);<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<h3 id=\"translation-1\">Translation<\/h3>\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">use<\/span> <span class=\"nn\">rust_bert<\/span><span class=\"p\">::<\/span><span class=\"nn\">pipelines<\/span><span class=\"p\">::<\/span><span class=\"nn\">translation<\/span><span class=\"p\">::{<\/span><span class=\"n\">Language<\/span><span class=\"p\">,<\/span> <span class=\"n\">TranslationConfig<\/span><span class=\"p\">,<\/span> <span class=\"n\">TranslationModel<\/span><span class=\"p\">};<\/span>\n<span class=\"k\">let<\/span> <span class=\"n\">translation_config<\/span> <span class=\"o\">=<\/span>\n    <span class=\"nn\">TranslationConfig<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">(<\/span><span class=\"nn\">Language<\/span><span class=\"p\">::<\/span><span class=\"n\">EnglishToSpanish<\/span><span class=\"p\">,<\/span> <span class=\"nn\">Device<\/span><span class=\"p\">::<\/span><span class=\"nf\">cuda_if_available<\/span><span class=\"p\">());<\/span>\n<span class=\"k\">let<\/span> <span class=\"n\">model<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">TranslationModel<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">(<\/span><span class=\"n\">translation_config<\/span><span class=\"p\">)<\/span><span class=\"o\">?<\/span><span class=\"p\">;<\/span>\n\n<span class=\"n\">model<\/span><span class=\"nf\">.translate<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"p\">[<\/span><span class=\"s\">\"They found that certain wavelengths of light, which are usually \n absorbed by water, weakened when the planet was in the way, indicating not only \n does K2-18b have an atmosphere, but the atmosphere contains water in vapour form.\"<\/span><span class=\"p\">]);<\/span>\n<\/code><\/pre><\/div><\/div>\n<h3 id=\"generation\">Generation<\/h3>\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">use<\/span> <span class=\"nn\">rust_bert<\/span><span class=\"p\">::<\/span><span class=\"nn\">pipelines<\/span><span class=\"p\">::<\/span><span class=\"nn\">text_generation<\/span><span class=\"p\">::<\/span><span class=\"n\">TextGenerationModel<\/span><span class=\"p\">;<\/span>\n<span class=\"k\">let<\/span> <span class=\"n\">model<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">TextGenerationModel<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">(<\/span><span class=\"nn\">Default<\/span><span class=\"p\">::<\/span><span class=\"nf\">default<\/span><span class=\"p\">())<\/span><span class=\"o\">?<\/span><span class=\"p\">;<\/span>\n<span class=\"k\">let<\/span> <span class=\"n\">input_context<\/span> <span class=\"o\">=<\/span> <span class=\"s\">\"The majority of crustaceans are aquatic,\"<\/span><span class=\"p\">;<\/span>\n<span class=\"n\">model<\/span><span class=\"nf\">.generate<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"p\">[<\/span><span class=\"n\">input_context<\/span><span class=\"p\">],<\/span> <span class=\"nb\">None<\/span><span class=\"p\">);<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>These pipelines bring state-of-the-art NLP capabilities to the Rust community. Please check <code class=\"language-plaintext highlighter-rouge\">rust-bert<\/code>\u2019s <a href=\"https:\/\/github.com\/guillaume-be\/rust-bert\">repository<\/a>, the associated paper <a href=\"#rustbert\">[9]<\/a>, or reach out to me if you are interested in learning more about the capabilities of the library. The rest of this article will focus on the performance comparison of the original Python-based text generation pipelines (using the Transformers library <a href=\"#transformers\">[1]<\/a>) and the proposed Rust-based implementation.<\/p>\n\n<h1 id=\"4-benchmarks\">4. Benchmarks<a name=\"benchmarks\"><\/a><\/h1>\n\n<p>The performance benchmarks proposed here will focus on the text generation task. Benchmarks have also been performed for simpler pipelines (for example classifications) and are available in <a href=\"#rustbert\">[9]<\/a>. For simple pipelines with low to no post-processing operations, there is little to gain from a Rust implementation. The forward pass through the neural network leverages the same backend (Rust bindings to the C++ Libtorch library <a href=\"#tch\">[14]<\/a>). Potential benefits could be gained from the pre-processing and tokenization step, but the Transformers\u2019 library uses Rust-based tokenizers <a href=\"#tokenizers\">[15]<\/a> for all its models starting from v4.0.0. The outcome is a virtually identical performance between the Rust and Python implementation for tasks such as classification, token classification or question answering.<\/p>\n\n<p>The text generation pipelines, however, do include a complex post-processing pipeline which is implemented natively in Python. Because of the iterative process involving a model forward pass and the post-processing steps, a migration of the post-processing operations to Rust and use of bindings to Python (as is the case for the tokenizers) is more difficult. This is an area where a fully Rust-native, end-to-end Rust implementation may offer benefits. This section describes a few experiments aiming at quantifying how significant these benefits may be.<\/p>\n\n<h3 id=\"experimental-setup\">Experimental setup<\/h3>\n\n<p>The experimental setup for all experiments is unchanged and described below:<\/p>\n\n<table>\n  <thead>\n    <tr>\n      <th style=\"text-align: left\">Hardware<\/th>\n      <th style=\"text-align: left\">\u00a0<\/th>\n      <th style=\"text-align: left\">\u00a0<\/th>\n      <th style=\"text-align: left\">\u00a0<\/th>\n      <th style=\"text-align: left\">Software<\/th>\n      <th style=\"text-align: left\">\u00a0<\/th>\n    <\/tr>\n  <\/thead>\n  <tbody>\n    <tr>\n      <td style=\"text-align: left\"><strong>CPU<\/strong>     \u00a0<\/td>\n      <td style=\"text-align: left\">AMD 2700X<\/td>\n      <td style=\"text-align: left\">\u00a0 \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">\u00a0<\/td>\n      <td style=\"text-align: left\"><strong>OS<\/strong> \u00a0<\/td>\n      <td style=\"text-align: left\">Windows 10 (Marian: Ubuntu 20.04)<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>GPU<\/strong>     \u00a0<\/td>\n      <td style=\"text-align: left\">Nvidia 2070 RTX<\/td>\n      <td style=\"text-align: left\">\u00a0 \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">\u00a0<\/td>\n      <td style=\"text-align: left\"><strong>CUDA<\/strong> \u00a0<\/td>\n      <td style=\"text-align: left\">10.2<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>RAM<\/strong>     \u00a0<\/td>\n      <td style=\"text-align: left\">32GB<\/td>\n      <td style=\"text-align: left\">\u00a0\u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">\u00a0<\/td>\n      <td style=\"text-align: left\"><strong>Python<\/strong> \u00a0<\/td>\n      <td style=\"text-align: left\">Python 3.7, Transformers v4.2.2<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>Drive<\/strong>     \u00a0<\/td>\n      <td style=\"text-align: left\">NVME 970 EVO<\/td>\n      <td style=\"text-align: left\">\u00a0 \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">\u00a0<\/td>\n      <td style=\"text-align: left\"><strong>Rust<\/strong> \u00a0<\/td>\n      <td style=\"text-align: left\">rust-bert v0.12.1<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\">\u00a0<\/td>\n      <td style=\"text-align: left\">\u00a0<\/td>\n      <td style=\"text-align: left\">\u00a0<\/td>\n      <td style=\"text-align: left\">\u00a0<\/td>\n      <td style=\"text-align: left\"><strong>C++<\/strong> \u00a0 \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">Opus-MT Marian Docker image<\/td>\n    <\/tr>\n  <\/tbody>\n<\/table>\n\n<p><br \/><\/p>\n\n<p>By default experiments are run in Windows 10, with the exception of Marian executed natively in Ubuntu 20.04 on the same hardware. All experiments are repeated at least 10 iterations and the mean is reported. In all benchmarks, a warm-up run is executed (loading model in the GPU buffer and running a forward pass) as the first GPU buffer allocation can be significantly slower. The source code for all benchmarks is available in the references <a href=\"#benchmark-files\">[18]<\/a> section.<\/p>\n\n<h3 id=\"translation-2\">Translation<\/h3>\n\n<p>Two models are used for benchmark purposes for translation: an English to Spanish model trained by the Opus-MT team <a href=\"#opusmt\">[6]<\/a>, and the T5-base model allowing for translation (in this case, English to French) as part of its text-to-text capabilities. For all translation tasks, the source sentences are extracted from the example provided at the beginning of this article <a href=\"#wikinews\">[8]<\/a> (and of course identical between Python and Rust).<\/p>\n\n<table>\n  <thead>\n    <tr>\n      <th style=\"text-align: left\">Setting<\/th>\n      <th style=\"text-align: left\">Value<\/th>\n    <\/tr>\n  <\/thead>\n  <tbody>\n    <tr>\n      <td style=\"text-align: left\"><strong># beams<\/strong>   \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">6<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>sampling<\/strong>   \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">false<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>early stopping<\/strong>  \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">true<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>output sequences<\/strong>  \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">1<\/td>\n    <\/tr>\n  <\/tbody>\n<\/table>\n\n<p><br \/>\nAll sentences are processed in a single batch. To illustrate the impact of the batch size and padding, a sample of 10 sentences with various lengths and a single sentence are passed to the models. Note that since translation is done with 6 beams, the effective batch size is 6x the number of input sequences.<\/p>\n\n<p>The figure below shows the results of the translation benchmark with the Marian English to Spanish model. For both input sizes, the Rust-based translation executes approximately 60% faster than its Python counterpart - regardless of the number of input sentences provided. Interestingly, the Rust and C++ (Marian) translation have the same performance, even though they do not share the same tensor operations backend (Marian uses its own optimized auto-differentiation engine <a href=\"#marian\">[16]<\/a> while the Rust version relies on bindings to the Libtorch library).<\/p>\n\n<p><img src=\"..\/assets\/generation_benchmarks\/translation_marian.svg\" alt=\"Marian Translation\" title=\"Marian translation\" width=\"80%\" \/><\/p>\n\n<p>The next figure illustrates the same experiment, taking the T5-base model (the only differences lies in the neural network architecture, the rest of the pipeline and settings remain identical). For a small effective batch size (3 input sentences), the benefits are in line with the Marian-based translation - approximately 50% reduced execution time for the Rust version. Interestingly, these benefits decrease significantly for larger effective batch sizes. This issue is still being investigated and may be caused by the handling of padding for sequences with varying length, or because T5 is a significantly larger model than Marian.<\/p>\n\n<p><img src=\"..\/assets\/generation_benchmarks\/translation_t5.svg\" alt=\"T5 Translation\" title=\"T5 translation\" width=\"80%\" \/><\/p>\n\n<h3 id=\"summarization-2\">Summarization<\/h3>\n\n<p>This section investigates the performance of summarization models using models based on a BART architecture. The article introduced at the beginning of this article is used for all experiments. Because of the higher computational cost of this pipeline, batch-wise summarization is not investigated. Note that since beam search is used, an effective batch size of length <em>#beams<\/em> is still used. In order to illustrate the impact of the model size, 3 different models (one baseline model and two distilled models) are used for benchmarks:<\/p>\n<ul>\n  <li>BART-large model finetuned on the CNN-DailyMail dataset (406M parameters)<\/li>\n  <li>DistilBART-12-6 model <a href=\"#distilbart\">[17]<\/a> with 12 encoder layers and 6 decoder layers (306M parameters)<\/li>\n  <li>DistilBART-6-6 model with 6 encoder layers and 6 decoder layers(230M parameters)<\/li>\n<\/ul>\n\n<p>The experimental set-up remains similar to that of translation, with 4 beams and no sampling.<\/p>\n\n<table>\n  <thead>\n    <tr>\n      <th style=\"text-align: left\">Setting<\/th>\n      <th style=\"text-align: left\">Value<\/th>\n    <\/tr>\n  <\/thead>\n  <tbody>\n    <tr>\n      <td style=\"text-align: left\"><strong># beams<\/strong>   \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">3<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>sampling<\/strong>   \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">false<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>early stopping<\/strong>  \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">true<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>output sequences<\/strong>  \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">1<\/td>\n    <\/tr>\n  <\/tbody>\n<\/table>\n\n<p><br \/><\/p>\n\n<p>The figure below shows the execution time for each model for both Python and Rust. Rust-based translation consistently runs approximately 2x faster than its Python counterpart. The Rust-based version of the BART-large 406M parameters runs faster than the smallest distilled model with 230M parameters in Python, showing that while research efforts aimed at reducing the model size are critical to ensure sustainable use of NLP technologies, engineering-driven improvements allow reaching comparable speed-ups. More interestingly, this consistent 50%+ execution time reduction shows that the two approaches are complementary. Combining distillation and implementation in Rust, the summarization of a document can be accelerated by a factor of close to 5, from 2.57s down to less than 600ms.<\/p>\n\n<p><img src=\"..\/assets\/generation_benchmarks\/summarization_benchmark.svg\" alt=\"Summarization\" title=\"Summarization\" width=\"95%\" \/><\/p>\n\n<h3 id=\"text-generation-2\">Text generation<\/h3>\n\n<p>The last experiment investigates the performance benefits of a Rust implementation for free text generation using a GPT2 language model. The architecture for this pipeline is slightly different (relies on a decoder model) and the model size is significantly smaller than for the previous pipelines. Sampling is turned on for these pipelines. In order to ensure a fair comparison (the sampling leads to non-deterministic output), the sequence output size is fixed to 64 tokens by fixing the sequence minimum and maximum lengths to this value. 5 output sequences with 5 beams are generated for each prompt, leading to an effective batch size equal to 25x the number of prompts. The text generation benchmark is illustrated for both a single prompt and 3 prompts processed in a single batch.<\/p>\n\n<table>\n  <thead>\n    <tr>\n      <th style=\"text-align: left\">Setting<\/th>\n      <th style=\"text-align: left\">Value<\/th>\n    <\/tr>\n  <\/thead>\n  <tbody>\n    <tr>\n      <td style=\"text-align: left\"><strong># beams<\/strong>   \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">5<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>sampling<\/strong>   \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">true<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>early stopping<\/strong>  \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">true<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>output sequences<\/strong>  \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">5<\/td>\n    <\/tr>\n    <tr>\n      <td style=\"text-align: left\"><strong>sequence length<\/strong>  \u00a0\u00a0<\/td>\n      <td style=\"text-align: left\">64<\/td>\n    <\/tr>\n  <\/tbody>\n<\/table>\n\n<p><br \/><\/p>\n\n<p>The benefits of the Rust implementation for text generation are significantly higher than for the previous experiments, with the Rust pipeline running roughly 4x faster than its Python equivalent. Note that these experiments do not use the ready-to-use pipeline from the Transformers library as this does not support batched input yet. Instead, the inputs for the Python experiment have been manually encoded and padded to the left with <code class=\"language-plaintext highlighter-rouge\">&lt;eos&gt;<\/code> tokens. The inputs are subsequently processed in a single batch in both the Python and Rust pipelines to allow for a fair comparison.<\/p>\n\n<p>The text generation pipeline uses sampling, while the previous two experiments covering summarization and translation did not (a deterministic behaviour is expected in this case). This additional post-processing operation is likely to be the cause of the significant increase in performance difference between Python and Rust. For validation purposes, this sampling was turned off for both frameworks and the experiment repeated. The experiments without sampling validate this assumption, and Rust is approximately 2x faster than Python for deterministic generation, in line with the past benchmark.<\/p>\n\n<p>This last experiment provides two insights:<\/p>\n<ul>\n  <li>First, sampling for text generation is expensive. It slows down the Python generation by a factor of ~4 and Rust by a factor of ~3 and should therefore be used with its additional computational cost in mind.<\/li>\n  <li>Second, this illustrates that as the post-processing pipeline increases in complexity, the Rust benefits become more significant. The benefits from faster beam search (~2x speedup) seem to be comparable to the benefit from faster sampling (additional ~2x speedup), resulting in ~4x speedup for the complete post-processing operation.<\/li>\n<\/ul>\n\n<p><img src=\"..\/assets\/generation_benchmarks\/generation.svg\" alt=\"Text generation\" title=\"Text generation\" width=\"100%\" \/><\/p>\n\n<h1 id=\"final-thoughts\">Final thoughts<\/h1>\n\n<p>These results highlight the potential of high-performance languages, including Rust, for serving text generation models under low latency. Bringing performance benefits in line with C++, its memory safety, concurrency and accessibility to Machine Learning engineers make it a powerful additional choice for the deployment of performant, machine-learning powered applications.<\/p>\n\n<p>Research efforts aimed at reducing the computational cost of deep learning models have translated into significant gains in execution speed, at only a marginal cost in the performance of these models. The proposed Rust implementation synergizes very well with this work: while techniques such as distillation or quantization are effective at reducing the cost of the forward pass through the neural network, a Rust implementation can significantly speed up the auxiliary operations (whose relative cost increases as the neural network gets optimized). Combined with Rust safe concurrency capabilities, the combination of these techniques enables a significant acceleration of text generation pipelines using state-of-the-art models.<\/p>\n\n<p><em>last revision: 2021\/01\/31, updated benchmark with Transformers v4.2.2<\/em><\/p>\n\n<h2 id=\"references\">References<\/h2>\n<ul>\n  <li><a name=\"transformers\"><\/a>[1] <a href=\"https:\/\/www.aclweb.org\/anthology\/2020.emnlp-demos.6\/\">Transformers: State-of-the-Art Natural Language Processing\n<\/a>, Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, Alexander Rush.<\/li>\n  <li><a name=\"cnndailymail\"><\/a>[2] <a href=\"http:\/\/arxiv.org\/abs\/1704.04368\">Get To The Point: Summarization with Pointer-Generator Networks<\/a>, Abigail See, Peter J. Liu, Christopher D. Manning<\/li>\n  <li><a name=\"xsum\"><\/a>[3] <a href=\"https:\/\/www.aclweb.org\/anthology\/D18-1206\/\">Don\u2019t Give Me the Details, Just the Summary! Topic-Aware Convolutional Neural Networks for Extreme Summarization<\/a>, Shashi Narayan, Shay B. Cohen, Mirella Lapata<\/li>\n  <li><a name=\"bart\"><\/a>[4] <a href=\"https:\/\/www.aclweb.org\/anthology\/2020.acl-main.703\/\">BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension<\/a>, Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, Luke Zettlemoyer<\/li>\n  <li><a name=\"t5\"><\/a>[5] <a href=\"https:\/\/arxiv.org\/abs\/1910.10683\">Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer<\/a>, Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu<\/li>\n  <li><a name=\"opusmt\"><\/a>[6] <a href=\"https:\/\/github.com\/Helsinki-NLP\/Opus-MT\">Open neural machine translation models and web services<\/a>, The Opus-MT team<\/li>\n  <li><a name=\"gpt2\"><\/a>[7] <a href=\"https:\/\/d4mucfpksywv.cloudfront.net\/better-language-models\/language-models.pdf\">Language Models are Unsupervised Multitask Learners<\/a>, Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever<\/li>\n  <li><a name=\"wikinews\"><\/a>[8] <a href=\"https:\/\/en.wikinews.org\/wiki\/Astronomers_find_water_vapour_in_atmosphere_of_exoplanet_K2-18b\">Astronomers find water vapour in atmosphere of exoplanet K2-18b<\/a>, WikiNews<\/li>\n  <li><a name=\"rustbert\"><\/a>[9] <a href=\"https:\/\/www.aclweb.org\/anthology\/2020.nlposs-1.4\/\">End-to-end NLP Pipelines in Rust<\/a>, Becquin, Guillaume<\/li>\n  <li><a name=\"illustratedtransformer\"><\/a>[10] <a href=\"http:\/\/jalammar.github.io\/illustrated-transformer\/\">The Illustrated Transformer<\/a>, Alammar, Jay<\/li>\n  <li><a name=\"generation\"><\/a>[11] <a href=\"https:\/\/huggingface.co\/blog\/how-to-generate\">How to generate text: using different decoding methods for language generation with Transformers<\/a>, von Platen, Patrick<\/li>\n  <li><a name=\"rusttokenizers\"><\/a>[12] <a href=\"https:\/\/github.com\/guillaume-be\/rust-tokenizers\">rust_tokenizers<\/a><\/li>\n  <li><a name=\"modelhub\"><\/a>[13] <a href=\"https:\/\/huggingface.co\/models?filter=rust\">Hugging Face model hub<\/a><\/li>\n  <li><a name=\"tch\"><\/a>[14] <a href=\"https:\/\/github.com\/LaurentMazare\/tch-rs\">tch-rs crate<\/a>, Mazare, Laurent<\/li>\n  <li><a name=\"tokenizers\"><\/a>[15] <a href=\"https:\/\/github.com\/huggingface\/tokenizers\">Tokenizers: Fast State-of-the-Art Tokenizers optimized for Research and Production<\/a>, The Huggingface team<\/li>\n  <li><a name=\"marian\"><\/a>[16] <a href=\"https:\/\/www.aclweb.org\/anthology\/P18-4020\/\">Marian: Fast Neural Machine Translation in C++<\/a>, Marcin Junczys-Dowmunt, Roman Grundkiewicz, Tomasz Dwojak, Hieu Hoang, Kenneth Heafield, Tom Neckermann, Frank Seide, Ulrich Germann, Alham Fikri Aji, Nikolay Bogoychev, Andr\u00e9 F. T. Martins, Alexandra Birch<\/li>\n  <li><a name=\"distilbart\"><\/a>[17] <a href=\"https:\/\/arxiv.org\/abs\/2010.13002\">Pre-trained Summarization Distillation<\/a>, Sam Shleifer, Alexander M. Rush<\/li>\n  <li><a name=\"benchmark-files\"><\/a>[18] Source code for benchmarks: <a href=\"https:\/\/gist.github.com\/guillaume-be\/3c103bce3a23336b2bdde308639a42c6\" target=\"_blank\">translation (Python)<\/a>, <a href=\"https:\/\/gist.github.com\/guillaume-be\/56a6c42b939fcd1f7770765cb530d313\" target=\"_blank\">summarization (Python)<\/a>, <a href=\"https:\/\/gist.github.com\/guillaume-be\/75b5909ee2b28fe24c5ba120abcd5101\" target=\"_blank\">text generation (Python)<\/a>, <a href=\"https:\/\/github.com\/guillaume-be\/rust-bert\/tree\/master\/benches\" target=\"_blank\">Rust<\/a><\/li>\n<\/ul>","author":{"name":"Guillaume"},"summary":"Over the past few months, text generation capabilities using Transformer-based models have been democratized by open-source efforts such as Hugging Face\u2019s Transformers [1] library. A broad range of models and applications have been made available, including: Summarization models fine-tuned on the CNN-DailyMail [2] or XSUM [3] datasets, including for example BART [4] or T5 [5] Translation, with the availability of more than a thousand models trained by the Opus-MT team from Language Technology at the University of Helsinki [6] Text generation, generating tokens from a prompt text, including OpenAI\u2019s GPT2 model [7]"},{"title":"A Rust SentencePiece implementation","link":{"@attributes":{"href":"https:\/\/guillaume-be.github.io\/2020-05-30\/sentence_piece","rel":"alternate","type":"text\/html","title":"A Rust SentencePiece implementation"}},"published":"2020-05-30T00:00:00+00:00","updated":"2020-05-30T00:00:00+00:00","id":"https:\/\/guillaume-be.github.io\/2020-05-30\/sentence_piece","content":"<h1 id=\"abstract\">Abstract<\/h1>\n\n<p>Tokenization into words or sub-word units is a key component of Natural Language Processing pipeline. Modern approaches such as Byte Pair Encoding (<a href=\"https:\/\/arxiv.org\/abs\/1508.07909\">Sennrich et al., 2015<\/a>), WordPiece or SentencePiece (<a href=\"https:\/\/arxiv.org\/abs\/1808.06226\">Kudo et al., 2018<\/a>) segment rare words into sub-tokens in order to limit the size of the resulting vocabulary, which in turn results in more compact embedding matrices, reduced memory footprint and better generalization. These would for example split the word <code class=\"language-plaintext highlighter-rouge\">unaffable<\/code> into <code class=\"language-plaintext highlighter-rouge\">[un, ##aff, ##able]<\/code> where the prefix <code class=\"language-plaintext highlighter-rouge\">un<\/code> could share its representation with other negating prefixes for other words, for example <code class=\"language-plaintext highlighter-rouge\">[un, ##expected]<\/code>.<\/p>\n\n<p>I believe that familiarity with the implementation details, benefits and limitations of these algorithms is key to understanding the same for the end-to-end models using these tokenizers. For example, while Byte Pair Encoding is a morphological tokenizer agglomerating common character pairs into subtokens, the SentencePiece unigram tokenizer is a statistical model that uses a unigram language model to return the statistically most likely segmentation of an input. While information on the popular Byte Pair Encoding (BPE) is <a href=\"https:\/\/web.stanford.edu\/class\/cs224n\/slides\/cs224n-2019-lecture12-subwords.pdf\">available<\/a> with detailed information on the encoding algorithms, I felt working examples (with acceptable real world performance) of the sentence piece unigram encoding algorithm were missing. Not meant at being a replacement for the <a href=\"https:\/\/arxiv.org\/abs\/1804.10959\">reference article<\/a>, this article aims at providing an illustrated walk-through the Sentence Piece unigram model encoding algorithm. However, the scale of the project (and potential lack of familiarity with C++) can make a deep-dive through the source code a daunting task for the interested data scientist willing to better understand the algorithm.<\/p>\n\n<p><a href=\"https:\/\/www.rust-lang.org\">Rust<\/a> has been shown to be a language that is adequate for high-performance tokenization (for example in my <a href=\"https:\/\/crates.io\/crates\/rust_tokenizers\">rust-tokenizers<\/a> project or Hugging Face\u2019s <a href=\"https:\/\/crates.io\/crates\/tokenizers\">tokenizers<\/a>). The Rust programming language offers a number of features making it a great language for natural language processing (robust and modern string handling, memory safety, effective multi-threading application design, speed and readability). This article provides two example implementations for the SentencePiece ending step in Rust. Importantly, I will show that the raw performance of a language should not replace careful design of the algorithms. This article proposes two implementations that are an order of magnitude apart in terms of execution speed, with the fastest being comparable to the reference C++ code. With both implementation fitting in around 150 lines of code, they make a good learning reference for both Rust and the SentencePiece algorithm itself.<\/p>\n\n<h1 id=\"an-overview-of-the-sentencepiece-unigram-model\">An overview of the SentencePiece unigram model<\/h1>\n\n<p>SentencePiece refers to both a probabilistic word segmentation algorithm (usually the unigram model) and a library (implemented in C++). The SentencePiece unigram model decomposes an input into a sequence of tokens that would have the highest likelihood (probability) to occur in an unigram language model, i.e. the decomposition that maximizes the product of the sub-tokens probability (or more conveniently the sum of their log probability).<\/p>\n\n<p>Mathematically, the objective of the tokenization is to identify the subtokens \\((x_1, x_2, \\dots, x_n)\\)  from the set of segmentation candidates \\(S(V)\\) for an input \\(X\\) such as\n\\(\\newcommand{\\argmax}{\\arg\\!\\max}\\)<\/p>\n\n\\[P(X) = \\prod\\limits_{i=1}^{M} P(x_i)\\]\n\n\\[(x_1, x_2, \\dots, x_n) = \\argmax_{x \\in S(V)} P(X)\\]\n\n<p>The individual token probabilities are built such as:<\/p>\n\n\\[\\sum_{Vocabulary}P(x) = 1\\]\n\n<p>(from <a href=\"https:\/\/arxiv.org\/abs\/1808.06226\">Kudo et al., 2018<\/a>). Note that log probabilities are usually used rather than the direct probabilities so that the most likely sequence can be derived from the sum of log probabilities rather than the product of probabilities.<\/p>\n\n<p>The algorithm consists of two macro steps: the training on a large corpus and the encoding of sentences at inference time. A detailed description of the training mechanism of the unigram model is beyond the scope of this article, and the focus will rather be on the encoding mechanism.<\/p>\n\n<h2 id=\"1-training\">1. Training<\/h2>\n\n<p>The SentencePiece unigram model aims at maximizing the likelihood of a unigram language model by starting f which is pruned iteratively using the Expectation Maximization algorithm.<\/p>\n\n<p>The SentencePiece tokenization aims at finding the tokenization that maximizes the likelihood of a language model trained using this tokenization. Because these depend on each other, the Expectation Maximization algorithm is used to build this tokenization model with the following steps:<\/p>\n\n<ol>\n  <li>Initiate a reasonably large seed vocabulary (generated for example from the (<a href=\"https:\/\/ieeexplore.ieee.org\/document\/4976463\">Suffix Array algorithm, Nong et al.<\/a>) from a training corpus. Define a desired vocabulary size that is smaller than the initial vocabulary (<em>note: this differs from BPE training where the vocabulary grows until it reaches the target size!<\/em>)<\/li>\n  <li>Repeat the following Expectation Maximization steps until convergence:\n    <ol>\n      <li>Estimate the probability of each vocabulary token using frequency counts in the training corpus<\/li>\n      <li>Use this probability to tokenize the text using the encoding algirthm described in the following section (<a href=\"https:\/\/en.wikipedia.org\/wiki\/Viterbi_algorithm\">Viterbi decoding algorithm<\/a>).<\/li>\n      <li>Compute the loss for every vocabulary token, loss being the impact of dropping said token from the vocabulary on the overall likelihood \\(\\mathcal{L}\\).<\/li>\n      <li>Prune the vocabulary by keeping the top \\(x\\% (e.g. 80\\%)\\)) tokens with the lowest loss. In order to eliminate the risk of an out-of-vocabulary character item, single characters are usually excluded from this pruning step.<\/li>\n    <\/ol>\n  <\/li>\n<\/ol>\n\n<h2 id=\"2-encoding\">2. Encoding<\/h2>\n\n<h3 id=\"algorithm\">Algorithm<\/h3>\n\n<p>Encoding (the focus of this article) is used during the training maximization step of EM and at inference in order to encode a new input text. The encoding is done using the <a href=\"https:\/\/en.wikipedia.org\/wiki\/Viterbi_algorithm\">Viterbi decoding algorithm<\/a> consisting of 2 macro steps: a forward step (where the possible sub-tokens are identified) and a backward step (where the most likely decoding sequence is identified). These steps are described in detail in this <a href=\"https:\/\/everdark.github.io\/k9\/notebooks\/ml\/natural_language_understanding\/subword_units\/subword_units.nb.html\">excellent article<\/a>.<\/p>\n\n<p>The forward steps identifies the sub-token with the highest probability at each character position of the word. The result at the end of the forward pass is a list of best sub-tokens ending at each character position of the input string. Note that this best sub-token at a given ending character position takes into account the raw likelihood of the sub-token <em>and<\/em> the likelihood of the best token sequence leading to it (as opposed to the probability provided in the vocabulary that does not consider the surrounding context). The Viterbi algorithm uses the fact that the likelihood of a token ending at character position <code class=\"language-plaintext highlighter-rouge\">end_idx<\/code> and starting at position <code class=\"language-plaintext highlighter-rouge\">start_idx<\/code>, is given by \\(\\mathcal{L}(Best\\:sequence\\:up\\:to\\:start\\:idx) + \\mathcal{L}(subtoken)\\). Two methods for this forward pass will be described in the following with a proposed implementation in Rust.<\/p>\n\n<p>In both case the high level algorithm forward pass accomplishes the following:<\/p>\n\n<p><a name=\"viterbi-forward\"><\/a><\/p>\n<pre id=\"read-0\" style=\"display:none;\">\n    \n\\begin{algorithm}\n\\caption{Forward pass}\n\\begin{algorithmic}\n\\PROCEDURE{Forward}{$input$, $scores$}\n    \\STATE $best\\:ending\\:subtokens$ = [nil] * \\CALL{Len}{$input$}\n    \\FOR{$i = 0$ \\TO $len(input)$}\n        \\STATE $best\\:ending\\:subtokens[i]$ = \\CALL{BestPrefix}{$input$, $i$, $scores$}\n    \\ENDFOR    \n    \\RETURN ($best\\:ending\\:subtokens$)\n\\ENDPROCEDURE\n\\end{algorithmic}\n\\end{algorithm}\n\n<\/pre>\n<div id=\"goal-0\"><\/div>\n<script type=\"text\/javascript\">\n    var code = document.getElementById(\"read-0\").textContent;\n    var parentEl = document.getElementById(\"goal-0\");\n    var options = {\n        lineNumber: true\n    };\n    pseudocode.render(code, parentEl, options);\n<\/script>\n\n<p>the <em>BestPrefix<\/em> procedure is described in the following sections of this article.<\/p>\n\n<p>The backward pass takes this list of most likely tokens ending at each position to decode the sequence starting from the end and building its way back to the beginning.<\/p>\n\n<p><a name=\"viterbi-backward\"><\/a><\/p>\n<pre id=\"read-1\" style=\"display:none;\">\n    \n\\begin{algorithm}\n\\caption{Backward pass}\n\\begin{algorithmic}\n\\PROCEDURE{Backward}{$input$, $best\\:ending\\:subtokens$}\n    \\STATE $end$ = \\CALL{Len}{$input$}\n    \\STATE $subtokens$ = []\n    \\WHILE{$end &gt; 0$}\n        \\STATE $best\\:token$ = $best\\:ending\\:subtokens[end]$\n        \\STATE \\CALL{Push}{$subtokens$, $best\\:token$}\n        \\STATE $end$ = $end$ - \\CALL{Len}{$best\\:token$}\n    \\ENDWHILE\n    \\RETURN \\CALL{Reverse}{$subtokens$}\n\\ENDPROCEDURE\n\\end{algorithmic}\n\\end{algorithm}\n\n<\/pre>\n<div id=\"goal-1\"><\/div>\n<script type=\"text\/javascript\">\n    var code = document.getElementById(\"read-1\").textContent;\n    var parentEl = document.getElementById(\"goal-1\");\n    var options = {\n        lineNumber: true\n    };\n    pseudocode.render(code, parentEl, options);\n<\/script>\n\n<h3 id=\"example\">Example<\/h3>\n<p>The illustrations are heavily inspired by another recommended read on the SentencePiece unigram algorithm by <a href=\"https:\/\/everdark.github.io\/k9\/notebooks\/ml\/natural_language_understanding\/subword_units\/subword_units.nb.html\">Kyle Chung<\/a>. Let\u2019s take the following example word to tokenize. :<\/p>\n<blockquote>\n  <p>Sesquipedalophobia (fear of long words)<\/p>\n<\/blockquote>\n\n<p>The reference SentencePiece library with a pretrained tokenizer used by <a href=\"https:\/\/arxiv.org\/abs\/1906.08237\">XLNet<\/a> tokenizes this long and rare word as :<\/p>\n<div class=\"language-python highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"kn\">import<\/span> <span class=\"nn\">sentencepiece<\/span>\n<span class=\"n\">sp_model<\/span> <span class=\"o\">=<\/span> <span class=\"n\">sentencepiece<\/span><span class=\"p\">.<\/span><span class=\"n\">SentencePieceProcessor<\/span><span class=\"p\">()<\/span>\n<span class=\"n\">sp_model<\/span><span class=\"p\">.<\/span><span class=\"n\">Load<\/span><span class=\"p\">(<\/span><span class=\"s\">'xlnet-base-cased-spiece.model'<\/span><span class=\"p\">)<\/span>\n\n<span class=\"n\">sp_model<\/span><span class=\"p\">.<\/span><span class=\"n\">EncodeAsPieces<\/span><span class=\"p\">(<\/span><span class=\"s\">\"Sesquipedalophobia\"<\/span><span class=\"p\">)<\/span>\n<\/code><\/pre><\/div><\/div>\n<p><code class=\"language-plaintext highlighter-rouge\">['\u2581Se', 's', 'qui', 'ped', 'alo', 'phobia']<\/code><\/p>\n\n<p>The forward pass identifies the most likely token ending at each character position (more on that to follow):<\/p>\n\n<p><img src=\"..\/assets\/sentencepiece\/spiece1.png\" alt=\"Forward pass example\" title=\"Forward pass example\" \/><\/p>\n\n<p>By construction, the forward pass provides the token with the highest likelihood to occur at a given position. The decoding then happens backwards: by starting at the last character, we are sure that the token ending at this position is the most likely. We then only have to walk the proposed sub-tokens from end to start:<\/p>\n\n<p><img src=\"..\/assets\/sentencepiece\/spiece2.png\" alt=\"Forward pass example\" title=\"backward pass example\" \/><\/p>\n\n<p>The following will walk through two implementations of the algorithm in Rust.<\/p>\n\n<h1 id=\"first-implementation\">First implementation<\/h1>\n\n<p>As a reminder the <a href=\"#viterbi-forward\">forward pass<\/a> of the algorithm requires identifying the most likely token at each character position of the input. One potential implementation of the <em>BestPrefix<\/em> method is to take slices of the substring ending at <code class=\"language-plaintext highlighter-rouge\">end_idx<\/code> with incremental start position and calculating the likelihood for each slice. This results in a nested <code class=\"language-plaintext highlighter-rouge\">for<\/code> loop to generate these slices. For each slice, the log probability or the token is equal to the log probability of the slice itself added to the probability of the best sequence leading to it. To implement this Viterbi algorithm the best slices and best scores are stored for each ending position and can be re-used to populate the most likely tokens ending at subsequent positions. This is the proposed implementation by <a href=\"https:\/\/everdark.github.io\/k9\/notebooks\/ml\/natural_language_understanding\/subword_units\/subword_units.nb.html\">Kyle Chung<\/a>:<\/p>\n\n<pre id=\"read-2\" style=\"display:none;\">\n    \n\\begin{algorithm}\n\\caption{Forward pass (nested loops)}\n\\begin{algorithmic}\n\\PROCEDURE{Forward}{$input$, $scores$}\n    \\STATE $best\\:ending\\:subtokens$ = [nil] * \\CALL{Len}{$input$}\n    \\STATE $best\\:ending\\:scores$ = [$-\\infty$] * \\CALL{Len}{$input$}\n    \\FOR{$end = 0$ \\TO $len(input)$}\n        \\FOR{$start = 0$ \\TO $end$}\n            \\STATE $substring$ = $input[start..end]$\n            \\STATE $score_{local}$ = $best\\:ending\\:scores[start]$ + $scores[substring]$\n            \\IF{$score_{local}$ &gt; $best\\:ending\\:scores[end]$}\n                \\STATE $best\\:ending\\:subtokens[end]$ = $substring$\n                \\STATE $best\\:ending\\:scores[end]$ = $score_{local}$\n            \\ENDIF            \n        \\ENDFOR\n    \\ENDFOR    \n    \\RETURN ($best\\:ending\\:subtokens$)\n\\ENDPROCEDURE\n\\end{algorithmic}\n\\end{algorithm}\n\n<\/pre>\n<div id=\"goal-2\"><\/div>\n<script type=\"text\/javascript\">\n    var code = document.getElementById(\"read-2\").textContent;\n    var parentEl = document.getElementById(\"goal-2\");\n    var options = {\n        lineNumber: true\n    };\n    pseudocode.render(code, parentEl, options);\n<\/script>\n\n<p>The first step is to define a SentencePiece model holding the pre-trained vocabulary and tokens log-likelihood:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">vocab<\/span><span class=\"p\">:<\/span> <span class=\"n\">HashMap<\/span><span class=\"o\">&lt;<\/span><span class=\"nb\">String<\/span><span class=\"p\">,<\/span> <span class=\"p\">(<\/span><span class=\"nb\">f32<\/span><span class=\"p\">,<\/span> <span class=\"nb\">i64<\/span><span class=\"p\">)<\/span><span class=\"o\">&gt;<\/span>\n<span class=\"p\">}<\/span>\n\n<span class=\"k\">impl<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"nf\">from_file<\/span><span class=\"p\">(<\/span><span class=\"n\">path<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">f<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">File<\/span><span class=\"p\">::<\/span><span class=\"nf\">open<\/span><span class=\"p\">(<\/span><span class=\"n\">path<\/span><span class=\"p\">)<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">contents<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">Vec<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n        <span class=\"n\">f<\/span><span class=\"nf\">.read_to_end<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"n\">contents<\/span><span class=\"p\">)<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">();<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"n\">proto<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">parse_from_bytes<\/span><span class=\"p\">::<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">ModelProto<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"n\">contents<\/span><span class=\"nf\">.as_slice<\/span><span class=\"p\">())<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">vocab<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">HashMap<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">for<\/span> <span class=\"p\">(<\/span><span class=\"n\">idx<\/span><span class=\"p\">,<\/span> <span class=\"n\">piece<\/span><span class=\"p\">)<\/span> <span class=\"n\">in<\/span> <span class=\"n\">proto<\/span><span class=\"nf\">.get_pieces<\/span><span class=\"p\">()<\/span><span class=\"nf\">.iter<\/span><span class=\"p\">()<\/span><span class=\"nf\">.enumerate<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">vocab<\/span><span class=\"nf\">.insert<\/span><span class=\"p\">(<\/span><span class=\"n\">piece<\/span><span class=\"nf\">.get_piece<\/span><span class=\"p\">()<\/span><span class=\"nf\">.to_owned<\/span><span class=\"p\">(),<\/span> <span class=\"p\">(<\/span><span class=\"n\">piece<\/span><span class=\"nf\">.get_score<\/span><span class=\"p\">(),<\/span> <span class=\"n\">idx<\/span> <span class=\"k\">as<\/span> <span class=\"nb\">i64<\/span><span class=\"p\">));<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span> <span class=\"n\">vocab<\/span> <span class=\"p\">}<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n<p>This reads a pretrained SentencePiece model using the <a href=\"https:\/\/github.com\/google\/sentencepiece\/blob\/master\/src\/sentencepiece_model.proto\">associated Protobuf<\/a>. The vocabulary is stored in a <code class=\"language-plaintext highlighter-rouge\">HashMap<\/code> with the token strings as keys and their <code class=\"language-plaintext highlighter-rouge\">(log score, index)<\/code> as keys.<\/p>\n\n<p>To hold the results of the forward pass, a utility structure is created, called <code class=\"language-plaintext highlighter-rouge\">Node<\/code> (that would be part of the forward Viterbi lattice). This hold the token score, vocabulary index, start and end character position and a <strong>reference<\/strong> to the token string (to avoid unnecessary copying):<\/p>\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"nd\">#[derive(Debug,<\/span> <span class=\"nd\">Clone,<\/span> <span class=\"nd\">Copy)]<\/span>\n<span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">Node<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">score<\/span><span class=\"p\">:<\/span> <span class=\"nb\">f32<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">index<\/span><span class=\"p\">:<\/span> <span class=\"nb\">i64<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">start<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">end<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>The forward pass outputs a list of optional Nodes.<\/p>\n<ol>\n  <li>The best node and best score are populated with <code class=\"language-plaintext highlighter-rouge\">None<\/code> and <code class=\"language-plaintext highlighter-rouge\">std::f32::NEG_INFINITY<\/code> respectively.<\/li>\n  <li>Slices are taken by looping through <code class=\"language-plaintext highlighter-rouge\">end_idx<\/code> and <code class=\"language-plaintext highlighter-rouge\">start_idx<\/code>.<\/li>\n  <li>Each slice is looked-up in the vocabulary. If it is found, its raw score is retrieved.<\/li>\n  <li>Its true local score (taken into account the probability of the best sequence leading to that token) is calculated by looking up the best score for the sequence ending at the start index for this token.<\/li>\n  <li>If the local score is greater than the store value, the best token and best score are updated with the current value<\/li>\n  <li>If after looping through all possible start values for a given end value the score is still <code class=\"language-plaintext highlighter-rouge\">std::f32::NEG_INFINITY<\/code>, this means the character does not exist in the vocabulary. In this case an unknown <code class=\"language-plaintext highlighter-rouge\">Node<\/code> is created. The score for the sequence ending at this unknown character is reset to 0 to allow for decoding.<\/li>\n<\/ol>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n    <span class=\"c\">\/\/...<\/span>\n\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"n\">decode_forward<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;<\/span><span class=\"nb\">Option<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">Node<\/span><span class=\"o\">&gt;&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">char_positions<\/span> <span class=\"o\">=<\/span> <span class=\"n\">text<\/span>\n            <span class=\"nf\">.char_indices<\/span><span class=\"p\">()<\/span>\n            <span class=\"nf\">.map<\/span><span class=\"p\">(|(<\/span><span class=\"n\">pos<\/span><span class=\"p\">,<\/span> <span class=\"mi\">_<\/span><span class=\"p\">)|<\/span> <span class=\"n\">pos<\/span><span class=\"p\">)<\/span>\n            <span class=\"nf\">.collect_vec<\/span><span class=\"p\">();<\/span>\n        <span class=\"n\">char_positions<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"nf\">.len<\/span><span class=\"p\">());<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">results<\/span> <span class=\"o\">=<\/span> <span class=\"nd\">vec!<\/span><span class=\"p\">(<\/span><span class=\"nb\">None<\/span><span class=\"p\">;<\/span> <span class=\"n\">char_positions<\/span><span class=\"nf\">.len<\/span><span class=\"p\">());<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">scores<\/span> <span class=\"o\">=<\/span> <span class=\"nd\">vec!<\/span><span class=\"p\">(<\/span><span class=\"nn\">std<\/span><span class=\"p\">::<\/span><span class=\"nn\">f32<\/span><span class=\"p\">::<\/span><span class=\"n\">NEG_INFINITY<\/span><span class=\"p\">;<\/span> <span class=\"n\">char_positions<\/span><span class=\"nf\">.len<\/span><span class=\"p\">());<\/span>\n        <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"mi\">0<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"mf\">0f32<\/span><span class=\"p\">;<\/span>\n    \n        <span class=\"k\">for<\/span> <span class=\"n\">char_end<\/span> <span class=\"n\">in<\/span> <span class=\"mi\">0<\/span><span class=\"o\">..<\/span><span class=\"n\">char_positions<\/span><span class=\"nf\">.len<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">for<\/span> <span class=\"n\">char_start<\/span> <span class=\"n\">in<\/span> <span class=\"mi\">0<\/span><span class=\"o\">..<\/span><span class=\"n\">char_end<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">let<\/span> <span class=\"n\">sub_text<\/span> <span class=\"o\">=<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">text<\/span><span class=\"p\">[<\/span><span class=\"n\">char_positions<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span><span class=\"p\">]<\/span><span class=\"o\">..<\/span><span class=\"n\">char_positions<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]];<\/span>\n                <span class=\"k\">if<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">subtoken<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"py\">.vocab<\/span><span class=\"nf\">.get<\/span><span class=\"p\">(<\/span><span class=\"n\">sub_text<\/span><span class=\"p\">)<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"k\">let<\/span> <span class=\"n\">local_score<\/span> <span class=\"o\">=<\/span> <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span><span class=\"p\">]<\/span> <span class=\"o\">+<\/span> <span class=\"n\">subtoken<\/span><span class=\"na\">.0<\/span><span class=\"p\">;<\/span>\n                    <span class=\"k\">if<\/span> <span class=\"n\">local_score<\/span> <span class=\"o\">&gt;<\/span> <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]<\/span> <span class=\"p\">{<\/span>\n                        <span class=\"n\">results<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">Node<\/span> <span class=\"p\">{<\/span>\n                            <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"n\">sub_text<\/span><span class=\"p\">,<\/span>\n                            <span class=\"n\">score<\/span><span class=\"p\">:<\/span> <span class=\"n\">local_score<\/span><span class=\"p\">,<\/span>\n                            <span class=\"n\">index<\/span><span class=\"p\">:<\/span> <span class=\"n\">subtoken<\/span><span class=\"na\">.1<\/span><span class=\"p\">,<\/span>\n                            <span class=\"n\">start<\/span><span class=\"p\">:<\/span> <span class=\"n\">char_start<\/span><span class=\"p\">,<\/span>\n                            <span class=\"n\">end<\/span><span class=\"p\">:<\/span> <span class=\"n\">char_end<\/span>\n                        <span class=\"p\">});<\/span>\n                        <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"n\">local_score<\/span><span class=\"p\">;<\/span>\n                    <span class=\"p\">}<\/span>\n                <span class=\"p\">}<\/span>\n            <span class=\"p\">}<\/span>\n            <span class=\"k\">if<\/span> <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]<\/span> <span class=\"o\">&lt;=<\/span> <span class=\"nn\">std<\/span><span class=\"p\">::<\/span><span class=\"nn\">f32<\/span><span class=\"p\">::<\/span><span class=\"n\">MIN<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">results<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">Node<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">text<\/span><span class=\"p\">[<\/span><span class=\"n\">char_positions<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span> <span class=\"o\">-<\/span> <span class=\"mi\">1<\/span><span class=\"p\">]<\/span><span class=\"o\">..<\/span><span class=\"n\">char_positions<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]],<\/span>\n                    <span class=\"n\">score<\/span><span class=\"p\">:<\/span> <span class=\"nn\">std<\/span><span class=\"p\">::<\/span><span class=\"nn\">f32<\/span><span class=\"p\">::<\/span><span class=\"n\">MIN<\/span><span class=\"p\">,<\/span>\n                    <span class=\"n\">index<\/span><span class=\"p\">:<\/span> <span class=\"mi\">0<\/span><span class=\"p\">,<\/span>\n                    <span class=\"n\">start<\/span><span class=\"p\">:<\/span> <span class=\"n\">char_end<\/span> <span class=\"o\">-<\/span> <span class=\"mi\">1<\/span><span class=\"p\">,<\/span>\n                    <span class=\"n\">end<\/span><span class=\"p\">:<\/span> <span class=\"n\">char_end<\/span>\n                <span class=\"p\">});<\/span>\n                <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"mf\">0f32<\/span><span class=\"p\">;<\/span>\n            <span class=\"p\">}<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">results<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>From this vector of <code class=\"language-plaintext highlighter-rouge\">Nodes<\/code>, the backward pass is far simpler and navigates the array starting from the last node. The next position is stored in the node itself. The backward pass returns a re-ordered list of best nodes for this tokenization. Here again references are used rather than the actual objects to avoid unnecessary copies:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n    <span class=\"c\">\/\/...<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"n\">decode_backward<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">nodes<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;<\/span><span class=\"nb\">Option<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">Node<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;&gt;&gt;<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"n\">Node<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">next_node<\/span> <span class=\"o\">=<\/span> <span class=\"n\">nodes<\/span><span class=\"nf\">.last<\/span><span class=\"p\">()<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">best_sequence<\/span> <span class=\"o\">=<\/span> <span class=\"nd\">vec!<\/span><span class=\"p\">();<\/span>\n\n        <span class=\"k\">while<\/span> <span class=\"n\">next_node<\/span><span class=\"nf\">.is_some<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">let<\/span> <span class=\"n\">node_value<\/span> <span class=\"o\">=<\/span> <span class=\"n\">next_node<\/span><span class=\"nf\">.as_ref<\/span><span class=\"p\">()<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">();<\/span>\n            <span class=\"n\">best_sequence<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"n\">node_value<\/span><span class=\"p\">);<\/span>\n            <span class=\"n\">next_node<\/span> <span class=\"o\">=<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">nodes<\/span><span class=\"p\">[<\/span><span class=\"n\">node_value<\/span><span class=\"py\">.start<\/span><span class=\"p\">];<\/span>\n        <span class=\"p\">};<\/span>\n        <span class=\"n\">best_sequence<\/span><span class=\"nf\">.reverse<\/span><span class=\"p\">();<\/span>\n        <span class=\"n\">best_sequence<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>The final tokenization combines the two steps and returns a string:<\/p>\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n    <span class=\"c\">\/\/...<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"nf\">tokenize<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;<\/span><span class=\"nb\">String<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">text<\/span> <span class=\"o\">=<\/span> <span class=\"n\">text<\/span><span class=\"nf\">.replace<\/span><span class=\"p\">(<\/span><span class=\"sc\">' '<\/span><span class=\"p\">,<\/span> <span class=\"s\">\"<\/span><span class=\"err\">\\<\/span><span class=\"s\">u{2581}\"<\/span><span class=\"p\">);<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">text<\/span> <span class=\"o\">=<\/span> <span class=\"n\">text<\/span><span class=\"nf\">.as_str<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">output<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"nf\">.decode_forward<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"p\">);<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">decoded<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"nf\">.decode_backward<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">output<\/span><span class=\"p\">);<\/span>\n        <span class=\"n\">decoded<\/span><span class=\"nf\">.into_iter<\/span><span class=\"p\">()<\/span><span class=\"nf\">.map<\/span><span class=\"p\">(|<\/span><span class=\"n\">node<\/span><span class=\"p\">|<\/span> <span class=\"n\">node<\/span><span class=\"py\">.text<\/span><span class=\"nf\">.to_string<\/span><span class=\"p\">())<\/span><span class=\"nf\">.collect<\/span><span class=\"p\">()<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>Let\u2019s try this out:<\/p>\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">fn<\/span> <span class=\"nf\">main<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">let<\/span> <span class=\"n\">text<\/span> <span class=\"o\">=<\/span> <span class=\"s\">\" Sesquipedalophobia\"<\/span><span class=\"p\">;<\/span>\n    <span class=\"k\">let<\/span> <span class=\"n\">tokenizer<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">SentencePieceModel<\/span><span class=\"p\">::<\/span><span class=\"nf\">from_file<\/span><span class=\"p\">(<\/span><span class=\"s\">\"E:\/Coding\/notebooks\/xlnet-base-cased-spiece.model\"<\/span><span class=\"p\">);<\/span>\n\n    <span class=\"k\">let<\/span> <span class=\"n\">output<\/span> <span class=\"o\">=<\/span> <span class=\"n\">tokenizer<\/span><span class=\"nf\">.tokenize<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"p\">);<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n<p><code class=\"language-plaintext highlighter-rouge\">[\"\u2581Se\", \"s\", \"qui\", \"ped\", \"alo\", \"phobia\"]<\/code><\/p>\n\n<p>All good, we get the expected output! The actual algorithm involves optional lower-casing, NFKC normalization and replaces the whitespaces by <code class=\"language-plaintext highlighter-rouge\">\\u{2581}<\/code>. These have been ignored in this simple example.<\/p>\n\n<p>It works but how fast does it run compared to the C++ implementation? Let\u2019s benchmark it from its Python bindings:<\/p>\n<div class=\"language-python highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"o\">%<\/span><span class=\"n\">timeit<\/span> <span class=\"o\">-<\/span><span class=\"n\">n<\/span> <span class=\"mi\">100<\/span> <span class=\"n\">sp_model<\/span><span class=\"p\">.<\/span><span class=\"n\">EncodeAsPieces<\/span><span class=\"p\">(<\/span><span class=\"s\">\"The quick brown fox jumps over the lazy dog\"<\/span><span class=\"p\">)<\/span>\n<span class=\"o\">%<\/span><span class=\"n\">timeit<\/span> <span class=\"o\">-<\/span><span class=\"n\">n<\/span> <span class=\"mi\">100<\/span> <span class=\"n\">sp_model<\/span><span class=\"p\">.<\/span><span class=\"n\">EncodeAsPieces<\/span><span class=\"p\">(<\/span><span class=\"s\">\"All human beings are born free and equal in dignity and rights. They are endowed with reason and conscience and should act towards one another in a spirit of brotherhood.\"<\/span><span class=\"p\">)<\/span> \n<\/code><\/pre><\/div><\/div>\n<p>(<a href=\"https:\/\/www.un.org\/en\/universal-declaration-human-rights\/\">article 1 or the Universal Declaration of Human Rights<\/a>)<\/p>\n\n<p>This C++ implementation takes <strong>9.7\u00b5s<\/strong> and <strong>29.8\u00b5s<\/strong>, respectively. The Rust version (with all compiler optimizations) takes <strong>29\u00b5s<\/strong> for the first and <strong>614.5\u00b5s<\/strong> for the second sentence.<\/p>\n\n<p><em>Oops!<\/em> That is 3 times slower than C++ for the short sentence and more than an order of magnitude slower for the longer sentence.<\/p>\n\n<p>Looking back at the implementation, one can notice the inefficiency during the forward pass: the nested loop will causes <code class=\"language-plaintext highlighter-rouge\">N\u00b2<\/code> lookups into a very large <code class=\"language-plaintext highlighter-rouge\">HashMap<\/code>. Using the <a href=\"https:\/\/github.com\/rust-lang\/hashbrown\">hashbrown<\/a> port of  Google\u2019s  <a href=\"https:\/\/abseil.io\/blog\/20180927-swisstables\">SwissTable<\/a> allows to take these times down to 14.2\u00b5s and 238.9\u00b5s. This is better but we definitely still have a problem for longer inputs because of this <code class=\"language-plaintext highlighter-rouge\">N\u00b2<\/code> complexity!<\/p>\n\n<h1 id=\"second-implementation\">Second implementation<\/h1>\n\n<p>The nested <code class=\"language-plaintext highlighter-rouge\">for<\/code> loop over the entire input length is too costly, especially for longer inputs. Intuitively, as the <code class=\"language-plaintext highlighter-rouge\">end_idx<\/code> cursor moves towards the end of the input, one can guess that low values for <code class=\"language-plaintext highlighter-rouge\">start_idx<\/code> of the substring are unlikely to be in the vocabulary. Therefore, one should only explore part of the space of substrings ending (or starting) at a given input character position. The <a href=\"https:\/\/github.com\/google\/sentencepiece\/blob\/master\/src\/unigram_model.cc\">reference C++ implementation<\/a> relies on <a href=\"https:\/\/en.wikipedia.org\/wiki\/Directed_acyclic_graph\">Directed Acrylic Graphs<\/a>.<\/p>\n\n<p>The main idea is to build a character-based directed acrylic graph (DAG) from the pre-trained tokens and their scores. Every time we insert a token, we also add the character-based path to the token if it does not yet exist. A difference is made between the DAG nodes that do exist in the SentencePiece model (marked as leaves) and the nodes that do not have a matching. Every token is pushed by being added to the children of the token containing all characters except the last one:<\/p>\n<ol>\n  <li>Get the characters of the token to insert<\/li>\n  <li>Start at the DAG root node. If the first character is in the children of the root, move to that node. Otherwise create it as a non-leaf node.<\/li>\n  <li>Select the character node as the current node.<\/li>\n  <li>Continue iterating through the characters, checking if the next character in the token to insert is in the children and adding if required.<\/li>\n  <li>When the last character is reached, mark the current node as a leaf node and store the token score (log probability) on the node.<\/li>\n<\/ol>\n\n<p>The following figure illustrates the population of the DAG assuming the first 3 words are <code class=\"language-plaintext highlighter-rouge\">cat<\/code>, <code class=\"language-plaintext highlighter-rouge\">carp<\/code> and <code class=\"language-plaintext highlighter-rouge\">car<\/code>.<\/p>\n\n<p><img src=\"..\/assets\/sentencepiece\/dag.png\" alt=\"DAG population\" title=\"DAG population\" \/><\/p>\n\n<p>When \u201ccat\u201d is fist inserted, the DAG is empty and therefore the intermediate nodes \u201cc\u201d, \u201cca\u201d and \u201ccat\u201d need to be created. At the last character position (node \u201ccat\u201d), the node is marked as a leaf. When \u201ccarp\u201d is inserted, the nodes \u201cc\u201d and \u201cca\u201d already exist and the graph is navigated. The child node \u201ccar\u201d cannot be found in \u201cca\u201d children, and is therefore created. At this point it is marked as a non-leaf node, and carp\u201d is added to its children as a leaf node. When the token \u201ccar\u201d is inserted to the DAG, all nodes already exist. We navigate the graph to the existing node \u201ccar\u201d, mark it as a leaf and store its log probability.<\/p>\n\n<p>This DAG allows to quickly get all of the possible token that share the same prefix for a given substring. However, one can note that while this was done by iteratively checking <em>every<\/em> substring of a given input in the previous method (therefore requiring <code class=\"language-plaintext highlighter-rouge\">N<\/code> operations for a substring of length <code class=\"language-plaintext highlighter-rouge\">N<\/code>), we can now navigate the DAG by feeding the character sequence from the substring query. We keep track of all of the leaf node encountered along the way, and stop early when we cannot find the next character in the current node children. This means that for any substring, we do at most <code class=\"language-plaintext highlighter-rouge\">k<\/code> operations (where k is the depth of the DAG) and in most cases far less.<\/p>\n\n<p>By asking for all DAG nodes sharing a prefix with \u201ccarpooling\u201d, we\u2019d obtain the result <code class=\"language-plaintext highlighter-rouge\">[\"car\", \"carp\"]<\/code> in 5 operations as opposed to 10 (length of carpooling), as illustrated below:<\/p>\n\n<p><img src=\"..\/assets\/sentencepiece\/dag2.png\" alt=\"DAG query\" title=\"DAG query\" \/><\/p>\n\n<p>The Viterbi forward pass can be re-written to take advantage of this efficient way of recovering tokens sharing a common prefix:<\/p>\n\n<pre id=\"read-3\" style=\"display:none;\">\n    \n\\begin{algorithm}\n\\caption{Forward pass (common prefix)}\n\\begin{algorithmic}\n\\PROCEDURE{Forward}{$input$, $scores$}\n    \\STATE $best\\:ending\\:subtokens$ = [nil] * \\CALL{Len}{$input$}\n    \\STATE $best\\:ending\\:scores$ = [$-\\infty$] * \\CALL{Len}{$input$}\n    \\FOR{$end = 0$ \\TO $len(input)$}\n        \\STATE $substring$ = $input[start..]$\n        \\STATE $matches$ = \\CALL{FindCommonPrefix}{$substring$, $scores$}\n        \\FOR{$match$ in $matches$}\n            \\STATE $score_{local}$ = $best\\:ending\\:scores[start]$ + $match.score$\n            \\IF{$score_{local}$ &gt; $best\\:ending\\:scores[end]$}\n                \\STATE $best\\:ending\\:subtokens[end]$ = $substring$\n                \\STATE $best\\:ending\\:scores[end]$ = $score_{local}$\n            \\ENDIF\n        \\ENDFOR\n    \\ENDFOR    \n    \\RETURN ($best\\:ending\\:subtokens$)\n\\ENDPROCEDURE\n\\end{algorithmic}\n\\end{algorithm}\n\n<\/pre>\n<div id=\"goal-3\"><\/div>\n<script type=\"text\/javascript\">\n    var code = document.getElementById(\"read-3\").textContent;\n    var parentEl = document.getElementById(\"goal-3\");\n    var options = {\n        lineNumber: true\n    };\n    pseudocode.render(code, parentEl, options);\n<\/script>\n\n<p>The nested <code class=\"language-plaintext highlighter-rouge\">for<\/code> loop over the length of the input has been replaced by the more effective <code class=\"language-plaintext highlighter-rouge\">find_common_prefix<\/code> illustrated above that depends on the DAG depth. The decoding algorithm remains identical.<\/p>\n\n<p>How would a Rust implementation look like? Let\u2019s start by defining a few structures to store the nodes from the DAG. A node contains a String, a length (UTF-8 code points), a score, vocabulary index, a leaf flag and a children Hashmap initialized to be empty.<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">DagNode<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"nb\">String<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">len<\/span><span class=\"p\">:<\/span> <span class=\"nb\">usize<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">score<\/span><span class=\"p\">:<\/span> <span class=\"nb\">f32<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">index<\/span><span class=\"p\">:<\/span> <span class=\"nb\">i32<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">leaf<\/span><span class=\"p\">:<\/span> <span class=\"nb\">bool<\/span><span class=\"p\">,<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">children<\/span><span class=\"p\">:<\/span> <span class=\"n\">BrownHashMap<\/span><span class=\"o\">&lt;<\/span><span class=\"nb\">char<\/span><span class=\"p\">,<\/span> <span class=\"n\">DagNode<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">,<\/span>\n<span class=\"p\">}<\/span>\n\n<span class=\"k\">impl<\/span> <span class=\"n\">DagNode<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"nf\">new<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"nb\">String<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"n\">DagNode<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">len<\/span> <span class=\"o\">=<\/span> <span class=\"n\">text<\/span><span class=\"nf\">.chars<\/span><span class=\"p\">()<\/span><span class=\"nf\">.count<\/span><span class=\"p\">();<\/span>\n        <span class=\"n\">DagNode<\/span> <span class=\"p\">{<\/span> <span class=\"n\">text<\/span><span class=\"p\">,<\/span> <span class=\"n\">len<\/span><span class=\"p\">,<\/span> <span class=\"n\">score<\/span><span class=\"p\">:<\/span> <span class=\"mf\">0.0<\/span><span class=\"p\">,<\/span> <span class=\"n\">index<\/span><span class=\"p\">:<\/span> <span class=\"mi\">0<\/span><span class=\"p\">,<\/span> <span class=\"n\">leaf<\/span><span class=\"p\">:<\/span> <span class=\"k\">false<\/span><span class=\"p\">,<\/span> <span class=\"n\">children<\/span><span class=\"p\">:<\/span> <span class=\"nn\">BrownHashMap<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">()<\/span> <span class=\"p\">}<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>Our SentencePiece model now uses a DAG (referenced by the root <code class=\"language-plaintext highlighter-rouge\">\"\"<\/code> node) instead of a vocabulary Hashmap. This is populated following the method described previously using an <code class=\"language-plaintext highlighter-rouge\">insert<\/code> function:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">pub<\/span> <span class=\"k\">struct<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"n\">root<\/span><span class=\"p\">:<\/span> <span class=\"n\">DagNode<\/span>\n<span class=\"p\">}<\/span>\n\n<span class=\"k\">impl<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"nf\">from_file<\/span><span class=\"p\">(<\/span><span class=\"n\">path<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">f<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">File<\/span><span class=\"p\">::<\/span><span class=\"nf\">open<\/span><span class=\"p\">(<\/span><span class=\"n\">path<\/span><span class=\"p\">)<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">contents<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">Vec<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">();<\/span>\n        <span class=\"n\">f<\/span><span class=\"nf\">.read_to_end<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"n\">contents<\/span><span class=\"p\">)<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">();<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"n\">proto<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">parse_from_bytes<\/span><span class=\"p\">::<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">ModelProto<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"n\">contents<\/span><span class=\"nf\">.as_slice<\/span><span class=\"p\">())<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">root<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">DagNode<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">(<\/span><span class=\"s\">\"\"<\/span><span class=\"nf\">.to_string<\/span><span class=\"p\">());<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">model<\/span> <span class=\"o\">=<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span> <span class=\"n\">root<\/span> <span class=\"p\">};<\/span>\n        <span class=\"k\">for<\/span> <span class=\"p\">(<\/span><span class=\"n\">idx<\/span><span class=\"p\">,<\/span> <span class=\"n\">piece<\/span><span class=\"p\">)<\/span> <span class=\"n\">in<\/span> <span class=\"n\">proto<\/span><span class=\"nf\">.get_pieces<\/span><span class=\"p\">()<\/span><span class=\"nf\">.iter<\/span><span class=\"p\">()<\/span><span class=\"nf\">.enumerate<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">model<\/span><span class=\"nf\">.insert<\/span><span class=\"p\">(<\/span><span class=\"n\">piece<\/span><span class=\"nf\">.get_piece<\/span><span class=\"p\">(),<\/span> <span class=\"n\">piece<\/span><span class=\"nf\">.get_score<\/span><span class=\"p\">(),<\/span> <span class=\"n\">idx<\/span> <span class=\"k\">as<\/span> <span class=\"nb\">i32<\/span><span class=\"p\">);<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">model<\/span>\n\n    <span class=\"p\">}<\/span>\n\n    <span class=\"k\">fn<\/span> <span class=\"nf\">insert<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">word<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nb\">str<\/span><span class=\"p\">,<\/span> <span class=\"n\">score<\/span><span class=\"p\">:<\/span> <span class=\"nb\">f32<\/span><span class=\"p\">,<\/span> <span class=\"n\">index<\/span><span class=\"p\">:<\/span> <span class=\"nb\">i32<\/span><span class=\"p\">)<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"n\">char_count<\/span> <span class=\"o\">=<\/span> <span class=\"n\">word<\/span><span class=\"nf\">.chars<\/span><span class=\"p\">()<\/span><span class=\"nf\">.count<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">node<\/span> <span class=\"o\">=<\/span> <span class=\"o\">&amp;<\/span><span class=\"k\">mut<\/span> <span class=\"k\">self<\/span><span class=\"py\">.root<\/span><span class=\"p\">;<\/span>\n\n        <span class=\"k\">for<\/span> <span class=\"p\">(<\/span><span class=\"n\">idx<\/span><span class=\"p\">,<\/span> <span class=\"n\">character<\/span><span class=\"p\">)<\/span> <span class=\"n\">in<\/span> <span class=\"n\">word<\/span><span class=\"nf\">.chars<\/span><span class=\"p\">()<\/span><span class=\"nf\">.enumerate<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">if<\/span> <span class=\"o\">!<\/span><span class=\"n\">node<\/span><span class=\"py\">.children<\/span><span class=\"nf\">.contains_key<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">character<\/span><span class=\"p\">)<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">text<\/span> <span class=\"o\">=<\/span> <span class=\"n\">node<\/span><span class=\"py\">.text<\/span><span class=\"nf\">.clone<\/span><span class=\"p\">();<\/span>\n                <span class=\"n\">text<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"n\">character<\/span><span class=\"p\">);<\/span>\n                <span class=\"k\">let<\/span> <span class=\"n\">new_node<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">DagNode<\/span><span class=\"p\">::<\/span><span class=\"nf\">new<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"p\">);<\/span>\n                <span class=\"n\">node<\/span><span class=\"py\">.children<\/span><span class=\"nf\">.insert<\/span><span class=\"p\">(<\/span><span class=\"n\">character<\/span><span class=\"p\">,<\/span> <span class=\"n\">new_node<\/span><span class=\"p\">);<\/span>\n            <span class=\"p\">}<\/span>\n            <span class=\"n\">node<\/span> <span class=\"o\">=<\/span> <span class=\"n\">node<\/span><span class=\"py\">.children<\/span><span class=\"nf\">.get_mut<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">character<\/span><span class=\"p\">)<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">();<\/span>\n            <span class=\"k\">if<\/span> <span class=\"n\">idx<\/span> <span class=\"o\">==<\/span> <span class=\"n\">char_count<\/span> <span class=\"o\">-<\/span> <span class=\"mi\">1<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">node<\/span><span class=\"py\">.leaf<\/span> <span class=\"o\">=<\/span> <span class=\"k\">true<\/span><span class=\"p\">;<\/span>\n                <span class=\"n\">node<\/span><span class=\"py\">.score<\/span> <span class=\"o\">=<\/span> <span class=\"n\">score<\/span><span class=\"p\">;<\/span>\n                <span class=\"n\">node<\/span><span class=\"py\">.index<\/span> <span class=\"o\">=<\/span> <span class=\"n\">index<\/span><span class=\"p\">;<\/span>\n            <span class=\"p\">}<\/span>\n        <span class=\"p\">}<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>The common prefix search method returns a vector of references to the nodes of the DAG sharing the same prefix of a query string slice:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n    <span class=\"c\">\/\/...<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"n\">common_prefix_search<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;&amp;<\/span><span class=\"n\">DagNode<\/span><span class=\"o\">&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">results<\/span> <span class=\"o\">=<\/span> <span class=\"nd\">vec!<\/span><span class=\"p\">();<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">characters<\/span> <span class=\"o\">=<\/span> <span class=\"n\">text<\/span><span class=\"nf\">.chars<\/span><span class=\"p\">();<\/span>\n\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">node<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"py\">.root.children<\/span><span class=\"nf\">.get<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">characters<\/span><span class=\"nf\">.next<\/span><span class=\"p\">()<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">());<\/span>\n        <span class=\"k\">if<\/span> <span class=\"n\">node<\/span><span class=\"nf\">.is_some<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">if<\/span> <span class=\"n\">node<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">()<\/span><span class=\"py\">.leaf<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">results<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"n\">node<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">());<\/span>\n            <span class=\"p\">}<\/span>\n        <span class=\"p\">}<\/span> <span class=\"k\">else<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">return<\/span> <span class=\"nd\">vec!<\/span><span class=\"p\">();<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"k\">while<\/span> <span class=\"k\">let<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">character<\/span><span class=\"p\">)<\/span> <span class=\"o\">=<\/span> <span class=\"n\">characters<\/span><span class=\"nf\">.next<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n            <span class=\"n\">node<\/span> <span class=\"o\">=<\/span> <span class=\"n\">node<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">()<\/span><span class=\"py\">.children<\/span><span class=\"nf\">.get<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">character<\/span><span class=\"p\">);<\/span>\n            <span class=\"k\">if<\/span> <span class=\"n\">node<\/span><span class=\"nf\">.is_some<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">if<\/span> <span class=\"n\">node<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">()<\/span><span class=\"py\">.leaf<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"n\">results<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"n\">node<\/span><span class=\"nf\">.unwrap<\/span><span class=\"p\">());<\/span>\n                <span class=\"p\">}<\/span>\n            <span class=\"p\">}<\/span> <span class=\"k\">else<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">break<\/span><span class=\"p\">;<\/span>\n            <span class=\"p\">}<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">results<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>Finally the forward method looks very similar to the first implementation but is updated to leverage this decoding, and creates substring slices by moving a <code class=\"language-plaintext highlighter-rouge\">start_idx<\/code> rather than a <code class=\"language-plaintext highlighter-rouge\">end_char<\/code> cursor:<\/p>\n\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">impl<\/span> <span class=\"n\">SentencePieceModel<\/span> <span class=\"p\">{<\/span>\n    <span class=\"c\">\/\/...<\/span>\n    <span class=\"k\">pub<\/span> <span class=\"k\">fn<\/span> <span class=\"n\">decode_forward_dag<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"k\">self<\/span><span class=\"p\">,<\/span> <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"nv\">'a<\/span> <span class=\"nb\">str<\/span><span class=\"p\">)<\/span> <span class=\"k\">-&gt;<\/span> <span class=\"nb\">Vec<\/span><span class=\"o\">&lt;<\/span><span class=\"nb\">Option<\/span><span class=\"o\">&lt;<\/span><span class=\"n\">Node<\/span><span class=\"o\">&lt;<\/span><span class=\"nv\">'a<\/span><span class=\"o\">&gt;&gt;&gt;<\/span> <span class=\"p\">{<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">char_positions<\/span> <span class=\"o\">=<\/span> <span class=\"n\">text<\/span>\n            <span class=\"nf\">.char_indices<\/span><span class=\"p\">()<\/span>\n            <span class=\"nf\">.map<\/span><span class=\"p\">(|(<\/span><span class=\"n\">pos<\/span><span class=\"p\">,<\/span> <span class=\"mi\">_<\/span><span class=\"p\">)|<\/span> <span class=\"n\">pos<\/span><span class=\"p\">)<\/span>\n            <span class=\"nf\">.collect_vec<\/span><span class=\"p\">();<\/span>\n        <span class=\"n\">char_positions<\/span><span class=\"nf\">.push<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"nf\">.len<\/span><span class=\"p\">());<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">results<\/span> <span class=\"o\">=<\/span> <span class=\"nd\">vec!<\/span><span class=\"p\">(<\/span><span class=\"nb\">None<\/span><span class=\"p\">;<\/span> <span class=\"n\">char_positions<\/span><span class=\"nf\">.len<\/span><span class=\"p\">());<\/span>\n        <span class=\"k\">let<\/span> <span class=\"k\">mut<\/span> <span class=\"n\">scores<\/span> <span class=\"o\">=<\/span> <span class=\"nd\">vec!<\/span><span class=\"p\">(<\/span><span class=\"nn\">std<\/span><span class=\"p\">::<\/span><span class=\"nn\">f32<\/span><span class=\"p\">::<\/span><span class=\"n\">NEG_INFINITY<\/span><span class=\"p\">;<\/span> <span class=\"n\">char_positions<\/span><span class=\"nf\">.len<\/span><span class=\"p\">());<\/span>\n        <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"mi\">0<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"mf\">0f32<\/span><span class=\"p\">;<\/span>\n\n        <span class=\"k\">for<\/span> <span class=\"n\">char_start<\/span> <span class=\"n\">in<\/span> <span class=\"mi\">0<\/span><span class=\"o\">..<\/span><span class=\"n\">char_positions<\/span><span class=\"nf\">.len<\/span><span class=\"p\">()<\/span> <span class=\"o\">-<\/span> <span class=\"mi\">1<\/span> <span class=\"p\">{<\/span>\n            <span class=\"k\">let<\/span> <span class=\"n\">matches<\/span> <span class=\"o\">=<\/span> <span class=\"k\">self<\/span><span class=\"nf\">.common_prefix_search<\/span><span class=\"p\">(<\/span><span class=\"o\">&amp;<\/span><span class=\"n\">text<\/span><span class=\"p\">[<\/span><span class=\"n\">char_positions<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span><span class=\"p\">]<\/span><span class=\"o\">..<\/span><span class=\"p\">]);<\/span>\n            <span class=\"k\">for<\/span> <span class=\"n\">node<\/span> <span class=\"n\">in<\/span> <span class=\"n\">matches<\/span> <span class=\"p\">{<\/span>\n                <span class=\"k\">let<\/span> <span class=\"n\">local_score<\/span> <span class=\"o\">=<\/span> <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span><span class=\"p\">]<\/span> <span class=\"o\">+<\/span> <span class=\"n\">node<\/span><span class=\"py\">.score<\/span><span class=\"p\">;<\/span>\n                <span class=\"k\">let<\/span> <span class=\"n\">char_end<\/span> <span class=\"o\">=<\/span> <span class=\"n\">char_start<\/span> <span class=\"o\">+<\/span> <span class=\"n\">node<\/span><span class=\"py\">.len<\/span><span class=\"p\">;<\/span>\n                <span class=\"k\">if<\/span> <span class=\"n\">local_score<\/span> <span class=\"o\">&gt;<\/span> <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"n\">results<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">Node<\/span> <span class=\"p\">{<\/span>\n                        <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">text<\/span><span class=\"p\">[<\/span><span class=\"n\">char_positions<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span><span class=\"p\">]<\/span><span class=\"o\">..<\/span><span class=\"n\">char_positions<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]],<\/span>\n                        <span class=\"n\">score<\/span><span class=\"p\">:<\/span> <span class=\"n\">local_score<\/span><span class=\"p\">,<\/span>\n                        <span class=\"n\">index<\/span><span class=\"p\">:<\/span> <span class=\"n\">node<\/span><span class=\"py\">.index<\/span><span class=\"p\">,<\/span>\n                        <span class=\"n\">start<\/span><span class=\"p\">:<\/span> <span class=\"n\">char_start<\/span><span class=\"p\">,<\/span>\n                        <span class=\"n\">end<\/span><span class=\"p\">:<\/span> <span class=\"n\">char_end<\/span><span class=\"p\">,<\/span>\n                    <span class=\"p\">});<\/span>\n                    <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_end<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"n\">local_score<\/span><span class=\"p\">;<\/span>\n                <span class=\"p\">}<\/span>\n            <span class=\"p\">}<\/span>\n            <span class=\"k\">if<\/span> <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span> <span class=\"o\">+<\/span> <span class=\"mi\">1<\/span><span class=\"p\">]<\/span> <span class=\"o\">&lt;=<\/span> <span class=\"nn\">std<\/span><span class=\"p\">::<\/span><span class=\"nn\">f32<\/span><span class=\"p\">::<\/span><span class=\"n\">MIN<\/span> <span class=\"p\">{<\/span>\n                <span class=\"n\">results<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span> <span class=\"o\">+<\/span> <span class=\"mi\">1<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"nf\">Some<\/span><span class=\"p\">(<\/span><span class=\"n\">Node<\/span> <span class=\"p\">{<\/span>\n                    <span class=\"n\">text<\/span><span class=\"p\">:<\/span> <span class=\"o\">&amp;<\/span><span class=\"n\">text<\/span><span class=\"p\">[<\/span><span class=\"n\">char_positions<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span><span class=\"p\">]<\/span><span class=\"o\">..<\/span><span class=\"n\">char_positions<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span> <span class=\"o\">+<\/span> <span class=\"mi\">1<\/span><span class=\"p\">]],<\/span>\n                    <span class=\"n\">score<\/span><span class=\"p\">:<\/span> <span class=\"nn\">std<\/span><span class=\"p\">::<\/span><span class=\"nn\">f32<\/span><span class=\"p\">::<\/span><span class=\"n\">MIN<\/span><span class=\"p\">,<\/span>\n                    <span class=\"n\">index<\/span><span class=\"p\">:<\/span> <span class=\"mi\">0<\/span><span class=\"p\">,<\/span>\n                    <span class=\"n\">start<\/span><span class=\"p\">:<\/span> <span class=\"n\">char_start<\/span><span class=\"p\">,<\/span>\n                    <span class=\"n\">end<\/span><span class=\"p\">:<\/span> <span class=\"n\">char_start<\/span> <span class=\"o\">+<\/span> <span class=\"mi\">1<\/span><span class=\"p\">,<\/span>\n                <span class=\"p\">});<\/span>\n                <span class=\"n\">scores<\/span><span class=\"p\">[<\/span><span class=\"n\">char_start<\/span> <span class=\"o\">+<\/span> <span class=\"mi\">1<\/span><span class=\"p\">]<\/span> <span class=\"o\">=<\/span> <span class=\"mf\">0f32<\/span><span class=\"p\">;<\/span>\n            <span class=\"p\">}<\/span>\n        <span class=\"p\">}<\/span>\n        <span class=\"n\">results<\/span>\n    <span class=\"p\">}<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n\n<p>Let\u2019s try this out!<\/p>\n<div class=\"language-rust highlighter-rouge\"><div class=\"highlight\"><pre class=\"highlight\"><code><span class=\"k\">fn<\/span> <span class=\"nf\">main<\/span><span class=\"p\">()<\/span> <span class=\"p\">{<\/span>\n    <span class=\"k\">let<\/span> <span class=\"n\">text<\/span> <span class=\"o\">=<\/span> <span class=\"s\">\" Sesquipedalophobia\"<\/span><span class=\"p\">;<\/span>\n    <span class=\"k\">let<\/span> <span class=\"n\">tokenizer<\/span> <span class=\"o\">=<\/span> <span class=\"nn\">SentencePieceModel<\/span><span class=\"p\">::<\/span><span class=\"nf\">from_file<\/span><span class=\"p\">(<\/span><span class=\"s\">\"E:\/Coding\/notebooks\/xlnet-base-cased-spiece.model\"<\/span><span class=\"p\">);<\/span>\n\n    <span class=\"k\">let<\/span> <span class=\"n\">output<\/span> <span class=\"o\">=<\/span> <span class=\"n\">tokenizer<\/span><span class=\"nf\">.tokenize<\/span><span class=\"p\">(<\/span><span class=\"n\">text<\/span><span class=\"p\">);<\/span>\n<span class=\"p\">}<\/span>\n<\/code><\/pre><\/div><\/div>\n<p><code class=\"language-plaintext highlighter-rouge\">[\"\u2581Se\", \"s\", \"qui\", \"ped\", \"alo\", \"phobia\"]<\/code><\/p>\n\n<p>The output is still as expected. Do we observe performance benefits? Running the same benchmark as before, we now obtain a tokenization speed of <strong>12\u00b5s<\/strong> for the first sentence and <strong>44.8\u00b5s<\/strong> for the longer second sentence. This is still slightly slower (although in the same order of magnitude) than the reference C++ implementation, but more than 13x faster than our original implementation!<\/p>\n\n<p>There is still room for optimization, but this implementation of a SentencePiece encoding is readable, fits in about 150 lines of code (compared to a much larger codebase for the C++ counterpart) and runs at a speed that is comparable to the reference. The supporting code used for this article is available in the following <a href=\"https:\/\/github.com\/guillaume-be\/SentencePiece-Rust-example\">repository<\/a>. Please note this code is designed for educational purposes. A more comprehensive implementation of this tokenization that includes vocabulary encoding and begin\/end offsets of the tokens with respect to the original strings is available on the <a href=\"https:\/\/crates.io\/crates\/rust_tokenizers\">rust_tokenizers<\/a> crate.<\/p>\n\n<h2 id=\"references\">References<\/h2>\n<ul>\n  <li><a href=\"https:\/\/arxiv.org\/abs\/1508.07909\">Neural Machine Translation of Rare Words with Subword Units, 2015<\/a>, Rico Sennrich, Barry Haddow, Alexandra Birch<\/li>\n  <li><a href=\"https:\/\/arxiv.org\/abs\/1808.06226\">SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing, 2018<\/a>, Taku Kudo, John Richardson<\/li>\n  <li><a href=\"https:\/\/web.stanford.edu\/class\/cs224n\/slides\/cs224n-2019-lecture12-subwords.pdf\">Natural Language Processing with Deep Learning - CS224N<\/a>, Christopher Manning, Stanford University<\/li>\n  <li><a href=\"https:\/\/ieeexplore.ieee.org\/document\/4976463\">Linear Suffix Array Construction by Almost Pure Induced-Sorting, 2009<\/a>, Ge Nong, Sen Zhang, Wai Hong Chan<\/li>\n  <li><a href=\"https:\/\/everdark.github.io\/k9\/notebooks\/ml\/natural_language_understanding\/subword_units\/subword_units.nb.html\">On Subword Units: Segmentation for Natural Language Modeling, 2019<\/a>, Kyle Chung<\/li>\n  <li><a href=\"https:\/\/github.com\/google\/sentencepiece\">SentencePiece library<\/a><\/li>\n  <li><a href=\"https:\/\/github.com\/tensorflow\/tfjs-models\/tree\/master\/universal-sentence-encoder\/src\/tokenizer\">Tensorflow.js Universal Sentence Encoder repository<\/a>, Tensorflow Team<\/li>\n  <li><a href=\"https:\/\/www.rust-lang.org\">The Rust programming language<\/a><\/li>\n<\/ul>","author":{"name":"Guillaume"},"summary":"Abstract"}]}