|
2 | 2 |
|
3 | 3 | > The current Index is a naive implementation. It means for a given DFA build from a regex it will 'bruteforce'
|
4 | 4 | > each state encountered during progression in the graph with all the tokens in order to build the tokens transitions table.
|
5 |
| -> This results in a complexity proportional to the size of the model vocabulary, the average size of the tokens in bytes and the complexity of the regex. (The complexity of a regex will be defined later.) |
| 5 | +> This results in a complexity proportional to the size of the model vocabulary, the average size of the tokens in bytes and the complexity of the regex. |
6 | 6 | > The following is the will of build an approach that takes the behaviors of DFA for regexes and extends them to the token scale in order to be less burdened by the complexity of regexes and the size of vocabularies.
|
7 | 7 | >
|
8 | 8 | > At the end, the V2Index has much better compile-time performance than its predecessor, much better performance in serving the list of allowed tokens for each state, and takes up less memory in most cases.
|
9 | 9 | ---
|
10 | 10 |
|
11 |
| - ## TokensDFA |
| 11 | + ## A. TokensDFA : Description |
12 | 12 |
|
13 | 13 | This new version of Index includes a TokensDFA object.
|
14 | 14 | This TokenDFA can be seen as an extension of DFA in that it leverages DFA optimizations to reduce the computational complexity of constructing the tokens transitions table.
|
@@ -58,7 +58,161 @@ We will use and abuse of these classes.
|
58 | 58 | We take the ByteClasses of the DFA and we construct the class of each token by concating the classes of each of their byte.
|
59 | 59 | In other world, if the range of bytes `[a-z]` has the class `[a]`, the token `'man'` will have the class `[a][a][a]` like all the
|
60 | 60 | tokens of 3 letters.
|
61 |
| -We reduce the size |
| 61 | +So we put all the tokens behind their classes which allows us to only consider the classes for the construction of the transition table. |
| 62 | + |
| 63 | +### 3. Prefix-Based Graph |
| 64 | + |
| 65 | +After grouping tokens by their regex byte classes, we construct directed prefix-based graphs to efficiently represent token hierarchies and optimize pattern matching traversal. |
| 66 | +``` |
| 67 | +[a] |
| 68 | + ↳ [a,b] |
| 69 | + ↳ [a,b,c] |
| 70 | +
|
| 71 | +[b] |
| 72 | + ↳ [b,c] |
| 73 | + ↳ [b,c,d] |
| 74 | + ↳ [b,c,e] |
| 75 | +``` |
| 76 | +```rust |
| 77 | +let eos_class_id = init_classes_and_graph_optimized( |
| 78 | + vocabulary.tokens(), |
| 79 | + &additionnal_tokens, |
| 80 | + &mut token_classes_graph, |
| 81 | + &mut transitions_table, |
| 82 | + byte_classes, |
| 83 | + &mut dead_byte_classes, |
| 84 | + eos_token_id); |
| 85 | +``` |
| 86 | +By traversing the DFA transition table with each prefix-based graph, this allows us to quickly discriminate entire sections of tokens as soon as one of their prefixes encounters a dead state. |
| 87 | + |
| 88 | +### 4. Good old Parallelization |
| 89 | + |
| 90 | +The previous optimisation, a bunch of graphs which have no intersection, unlock the possibilities to to go through the DFA in parallel, with a thread by graph. |
| 91 | +```rust |
| 92 | +use rayon::prelude::*; |
| 93 | +let roots = read_only_graph.get_roots(); |
| 94 | + roots.par_iter() |
| 95 | + .for_each(|root| { |
| 96 | + ... |
| 97 | + } |
| 98 | +``` |
| 99 | + |
| 100 | +### 5. Ultima Optima : Mute Literals and coalescence |
| 101 | + |
| 102 | +At this stage of optimization, the compilation times were already pretty good for sample regexes benchmark. |
| 103 | +But it was weak for JSON structure : |
| 104 | + |
| 105 | + |
| 106 | + |
| 107 | + |
| 108 | +After investigation it turns out that the problem comes from the literals ! |
| 109 | +Literals are worst nightmare for DFA (and by extension, TokensDFA). |
| 110 | +It's easy to understand why. If we reconsidered our last regex `"^a[a-z]$"`, the char 'a' is a literal. |
| 111 | +With classification, the char 'a' will not have the same class as the other letters. |
| 112 | +By extension, every token for a given size, with a letter 'a' will not have the same classe as the other tokens with exact same size. |
| 113 | +If we take two classes `'a' -> [a]` and `'b-z' -> [b]`, the words "hand", "five" and "ante" respectively have the classes |
| 114 | +'[b][a][b][b]' , '[b][b][b][b]' and '[a][b][b][b]'. It increases drastically the size of the alphabet, the number of transitions and the number of reached state. |
| 115 | +And the big issue is that there is a lot of literals in JSON structures. (Every keys of attributes at least, every symboles {, ",}, etc...) |
| 116 | +
|
| 117 | +The best example is the 'HTTPS' regex. |
| 118 | +| Regular Expression | V2Index Time | Index Time | |
| 119 | +| ------------------ | ------------ | ---------- | |
| 120 | +| `(https?:\/\/)?([\da-z\.-]+)\.([a-z\.]{2,6})([\/\w \.-]*)*\/?` | 27.683738s | 22.3142975s | |
| 121 | +
|
| 122 | +Here, 'https' is a literal but also 'http', 'h', 't' and 'p'. It a huge stab in the previous optimisation. |
| 123 | +Now, if we transform the 'https' determinist sequence by two 'ghost' symbols. (one for 'http', the other for 's' because 's' is optionnal with '?') : |
| 124 | +
|
| 125 | +| Regular Expression | V2Index Time | Index Time | |
| 126 | +| ------------------ | ------------ | ---------- | |
| 127 | +| `(∟1(∟2)?:\/\/)?([\da-z\.-]+)\.([a-z\.]{2,6})([\/\w \.-]*)*\/?` | 1.41s | 22.3142975s | |
| 128 | +
|
| 129 | +Yes, it's a huge improvment. Again, literals are the worst nightmare of Regexes. |
| 130 | + |
| 131 | +So, at the beginning, we add an other static analysis of the regex to extract every literals (or 'determinist sequence') with alphanumeric chars. |
| 132 | +```rust |
| 133 | +let (muted_regex, muted_list) = mute_literals(regex, vocabulary, &mut additionnal_tokens); |
| 134 | +``` |
| 135 | +
|
| 136 | +For each of them, we will find the best combination of tokens to express them. This is where **coalescence** takes place. |
| 137 | +If we extract the literal 'filename', we can express it with tokens 'file', 'name', 'f', 'i', 'l', 'e', 'n', 'a', 'm', 'e'. |
| 138 | +Then, we find the smallest combination, here, the tokens 'file' and 'name'. For these tokens, we create two 'ghost' symbols. |
| 139 | +'Ghost' tokens are choosen with char which have small probabilities to appear in the regex and zero probabilities to be a prefix of real tokens. |
| 140 | +
|
| 141 | +So, every 'Ghost' tokens begins by the char "\x1C" which is the File separator (Very Rare) then we concate with iteration index. |
| 142 | +In our example, 'file' will be [28, 49] (byte values for "\x1C1") and 'name' will be [28,50] (byte values for "\x1C2"). |
| 143 | +We affect to 'ghost' tokens same ids than their respective real token and we create new regex with ghost tokens combination instead of the literals. |
| 144 | +
|
| 145 | +
|
| 146 | +
|
| 147 | +### 6 Minimize Transitions Table |
| 148 | +
|
| 149 | +We use the same structure as the CompressIndex here : https://github.com/agourdel/outlines-core/tree/opt/new-index-struct |
| 150 | +to reduce the index size on average after compilation and increase the performance to serve the allowed tokens. |
| 151 | +When we reduce, we replace the ghost tokens by the real tokens. |
| 152 | +
|
| 153 | +```rust |
| 154 | +transitions_table.reduce(muted_list); |
| 155 | +``` |
| 156 | +
|
| 157 | +Bitset Masks of allowed tokens are already initiate for every state. |
| 158 | +
|
| 159 | +
|
| 160 | +## B - Compilations Benchmark (From Rust) |
| 161 | +
|
| 162 | + |
| 163 | +
|
| 164 | +## C - Memory Sizes Benchmark (From Rust) |
| 165 | +
|
| 166 | + |
| 167 | +
|
| 168 | +## D - Average Time to Inner Mask (From Python) |
| 169 | +
|
| 170 | +*Using mask reference as parameter* |
| 171 | + |
| 172 | +
|
| 173 | +## E - Ready-To-Use |
| 174 | +
|
| 175 | +With this branch, the V2Index is directly integrated into the Index python class without any breaking changes. |
| 176 | +It's ready to use. |
| 177 | +```python |
| 178 | +class Guide: |
| 179 | + [...] |
| 180 | + def get_tokens(self, mask:Optional[array.array]) -> List[int]: |
| 181 | + """Gets the list of allowed tokens for the current state.""" |
| 182 | + ... |
| 183 | + def advance(self, token_id: int, mask: Optional[array.array]) -> List[int]: |
| 184 | + """Guide moves to the next state provided by the token id and returns a list of allowed tokens.""" |
| 185 | + [...] |
| 186 | +``` |
| 187 | +The 'get_tokens()' and 'advance()' functions can be used as previous version. |
| 188 | + |
| 189 | +```python |
| 190 | +from outlines_core import Guide, Index, Vocabulary |
| 191 | + |
| 192 | +v2_index = Index(regex, vocab) |
| 193 | +v2_guide = Guide(v2_index) |
| 194 | + |
| 195 | +list_tokens = v2_guide.get_tokens() |
| 196 | +new_list_tokens = v2_guide.advance(list_tokens[0]) |
| 197 | + |
| 198 | +``` |
| 199 | + |
| 200 | +Or, they can be used with a reference to a mask. (Much faster) |
| 201 | + |
| 202 | +```python |
| 203 | +from outlines_core import Guide, Index, Vocabulary |
| 204 | + |
| 205 | +v2_index = Index(regex, vocab) |
| 206 | +v2_guide = Guide(v2_index) |
| 207 | +mask : array.array = create_mask(vocab.size()) |
| 208 | +v2_guide.get_tokens(mask) |
| 209 | +v2_guide.advance(mask) |
| 210 | + |
| 211 | +``` |
| 212 | + |
| 213 | +## TODO |
| 214 | + |
| 215 | + |
62 | 216 |
|
63 | 217 |
|
64 | 218 |
|
|
0 commit comments