GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

openai/tiktoken

Public

mirrored from https://github.com/openai/tiktokenAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
0.13.0

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

tiktoken/src

src/lib.rs

702lines · modecode

RawDownload
Latest commit unavailable.
1use std::collections::HashSet;
2use std::num::NonZeroU64;
3use std::thread;
4
5use fancy_regex::Regex;
6#[cfg(feature = "python")]
7use pyo3::prelude::*;
8use rustc_hash::FxHashMap as HashMap;
9
10#[cfg(feature = "python")]
11mod py;
12
13pub type Rank = u32;
14
15use std::collections::BinaryHeap;
16
17#[derive(Eq, PartialEq, Clone, Copy)]
18struct Merge {
19 start: usize,
20 rank: Rank,
21}
22
23impl Ord for Merge {
24 #[inline]
25 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
26 other
27 .rank
28 .cmp(&self.rank)
29 .then_with(|| other.start.cmp(&self.start))
30 }
31}
32
33impl PartialOrd for Merge {
34 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
35 Some(self.cmp(other))
36 }
37}
38
39struct State {
40 prev: usize,
41 end: usize,
42 next_end: usize,
43 next_rank: Rank,
44 cur_rank: Rank,
45}
46
47fn _byte_pair_merge_large(ranks: &HashMap<Vec<u8>, Rank>, piece: &[u8]) -> Vec<Rank> {
48 let mut state = Vec::with_capacity(piece.len());
49 state.push(State {
50 prev: usize::MAX,
51 end: 1,
52 next_end: 2,
53 next_rank: Rank::MAX,
54 cur_rank: Rank::MAX,
55 });
56
57 let mut heap = BinaryHeap::with_capacity(piece.len());
58 for i in 0..piece.len() - 1 {
59 if let Some(&rank) = ranks.get(&piece[i..i + 2]) {
60 heap.push(Merge { start: i, rank });
61 state[i].next_rank = rank;
62 }
63 // note this is happening offset by 1
64 state.push(State {
65 prev: i,
66 end: i + 2,
67 next_end: i + 3,
68 next_rank: Rank::MAX,
69 cur_rank: Rank::MAX,
70 });
71 }
72
73 // Repeatedly find the valid merge with smallest rank. We merge the (left) token that
74 // starts at `start` and ends at `state[start].end` with the (right) token that starts at
75 // `state[start].end` and ends at `state[start].next_end`. We invalidate the old merges
76 // (the ones that started at `state[start].end` and ended at `state[start]`) and add the two
77 // new potential merges to the heap.
78
79 let potential_merge = {
80 #[inline(always)]
81 |state: &mut Vec<State>,
82 heap: &mut BinaryHeap<Merge>,
83 start: usize,
84 next_end_item: usize| {
85 state[start].next_end = next_end_item;
86 state[start].next_rank = Rank::MAX; // Always invalidate the old merge
87 if next_end_item <= piece.len()
88 && let Some(&rank) = ranks.get(&piece[start..next_end_item])
89 {
90 // We have a valid potential merge!
91 heap.push(Merge { start, rank });
92 state[start].next_rank = rank;
93 }
94 }
95 };
96
97 while let Some(left) = heap.pop() {
98 if left.rank == Rank::MAX {
99 break;
100 }
101 if left.rank != state[left.start].next_rank {
102 continue; // This merge was invalidated, ignore it
103 }
104
105 let left_start = left.start;
106 let right_start = state[left_start].end;
107 let right_end = state[left_start].next_end;
108 debug_assert!(right_end == state[right_start].end);
109 let right_next_end = state[right_start].next_end;
110
111 // Merge left and right into a single token
112 state[left_start].cur_rank = state[left_start].next_rank;
113 state[left_start].end = right_end;
114 potential_merge(&mut state, &mut heap, left_start, right_next_end);
115 if right_end < state.len() {
116 state[right_end].prev = left_start;
117 }
118 // Update the merge that ends at left_start
119 if left_start > 0 {
120 let prev_start = state[left_start].prev;
121 potential_merge(&mut state, &mut heap, prev_start, right_end);
122 }
123 // Invalidate the merge starting at right_start, so we ignore it when it comes off the heap
124 state[right_start].next_rank = Rank::MAX;
125 }
126
127 let mut result = Vec::new();
128 let mut i = 0;
129 while i < state.len() {
130 if state[i].cur_rank != Rank::MAX {
131 result.push(state[i].cur_rank);
132 } else {
133 result.push(ranks[&piece[i..state[i].end]]);
134 }
135 i = state[i].end;
136 }
137 result
138}
139
140fn _byte_pair_merge(ranks: &HashMap<Vec<u8>, Rank>, piece: &[u8]) -> Vec<(usize, Rank)> {
141 // This is a vector of (start, rank).
142 // The rank is of the pair starting at position start.
143 let mut parts = Vec::with_capacity(piece.len() + 1);
144
145 // Note that we hash bytes when indexing into `ranks`, not token pairs. As long as we train BPE
146 // the way we currently do, this is equivalent. An easy way to break this would be to decouple
147 // merge priority from token index or to prevent specific token merges.
148 let mut min_rank: (Rank, usize) = (Rank::MAX, usize::MAX);
149 for i in 0..piece.len() - 1 {
150 let rank = *ranks.get(&piece[i..i + 2]).unwrap_or(&Rank::MAX);
151 if rank < min_rank.0 {
152 min_rank = (rank, i);
153 }
154 parts.push((i, rank));
155 }
156 parts.push((piece.len() - 1, Rank::MAX));
157 parts.push((piece.len(), Rank::MAX));
158
159 let get_rank = {
160 #[inline(always)]
161 |parts: &Vec<(usize, Rank)>, i: usize| {
162 if (i + 3) < parts.len() {
163 // Similar to `piece[i..i + 2]` above. The +3 is because we haven't yet deleted
164 // parts[i + 1], see comment in the main loop.
165 *ranks
166 .get(&piece[parts[i].0..parts[i + 3].0])
167 .unwrap_or(&Rank::MAX)
168 } else {
169 Rank::MAX
170 }
171 }
172 };
173
174 // If you have n parts and m merges, this does O(mn) work.
175 // We could do something with a heap and do O(m log n) work.
176 // n is often very small so considerations like cache-locality outweigh the algorithmic
177 // complexity downsides of the `parts` vector.
178 while min_rank.0 != Rank::MAX {
179 let i = min_rank.1;
180 // Update parts[i] and parts[i - 1] before removing parts[i + 1], since
181 // `parts.remove(i + 1)` will thrash the cache.
182 if i > 0 {
183 parts[i - 1].1 = get_rank(&parts, i - 1);
184 }
185 parts[i].1 = get_rank(&parts, i);
186 parts.remove(i + 1);
187
188 min_rank = (Rank::MAX, usize::MAX);
189 for (i, &(_, rank)) in parts[..parts.len() - 1].iter().enumerate() {
190 if rank < min_rank.0 {
191 min_rank = (rank, i);
192 }
193 }
194 }
195 parts
196}
197
198pub fn byte_pair_encode(piece: &[u8], ranks: &HashMap<Vec<u8>, Rank>) -> Vec<Rank> {
199 let piece_len = piece.len();
200
201 if piece_len == 1 {
202 return vec![ranks[piece]];
203 }
204 if piece_len < 100 {
205 return _byte_pair_merge(ranks, piece)
206 .windows(2)
207 .map(|part| ranks[&piece[part[0].0..part[1].0]])
208 .collect();
209 }
210 _byte_pair_merge_large(ranks, piece)
211}
212
213pub fn byte_pair_split<'a>(piece: &'a [u8], ranks: &HashMap<Vec<u8>, Rank>) -> Vec<&'a [u8]> {
214 assert!(piece.len() > 1);
215 _byte_pair_merge(ranks, piece)
216 .windows(2)
217 .map(|part| &piece[part[0].0..part[1].0])
218 .collect()
219}
220
221// Various performance notes:
222//
223// Regex
224// =====
225// Most of the time is spent in regex. The easiest way to speed this up is by using less fancy
226// regex features. For instance, using a regex parse-able by `regex` crate is 3x faster than
227// the usual regex we use.
228//
229// However, given that we're using a regex parse-able by `regex`, there isn't much difference
230// between using the `regex` crate and using the `fancy_regex` crate.
231//
232// There is an important interaction between threading, `regex` and `fancy_regex`.
233// When using `fancy_regex`, we hit `regex.find_at`. It turns out that this causes contention on
234// some mutable scratch space inside of `regex`. This absolutely kills performance. When using plain
235// old `regex`, we don't hit this, because `find_iter` has a different code path.
236// Related: https://github.com/rust-lang/regex/blob/master/PERFORMANCE.md
237// Anyway, the way we get around this is with having a (mostly) thread local clone of the regex for
238// each thread.
239//
240// Threading
241// =========
242// I tried using `rayon`. It wasn't really faster than using Python threads and releasing the GIL.
243// So goodbye `rayon`! Let thread count etc be in control of our Python users.
244//
245// Caching
246// =======
247// The reference tokeniser has an lru cache over the equivalent of `byte_pair_encode`.
248// Originally, we had one too! Without it, we were only vaguely faster than Python.
249// I used an RWLock to protect the cache. This didn't seem to hurt single threaded performance
250// noticeably, but it did affect multi-threaded performance. Weirdly, it seemed to affect
251// multi-threaded performance even when I only had readers (maybed I messed something up?).
252// Anyway, I realised that we could get rid of the cache, if we treat the set of tokens as a cache!
253// These are exactly the set or merges that are likely to be hot. And now we don't have to think
254// about interior mutability, memory use, or cloning.
255//
256// Hashing
257// =======
258// We use FxHashMap instead of the standard HashMap. This is maybe like a 5-10% win?
259// The current implementation ends up doing a lot of hashing of bytes. In theory, this could be made
260// to be hashing of two-tuples of ints, which looks like it may also be a couple percent faster.
261
262struct FakeThreadId(NonZeroU64);
263
264fn hash_current_thread() -> usize {
265 // It's easier to use unsafe than to use nightly. Rust has this nice u64 thread id counter
266 // that works great for our use case of avoiding collisions in our array. Unfortunately,
267 // it's private. However, there are only so many ways you can layout a u64, so just transmute
268 // https://github.com/rust-lang/rust/issues/67939
269 const _: [u8; 8] = [0; std::mem::size_of::<std::thread::ThreadId>()];
270 const _: [u8; 8] = [0; std::mem::size_of::<FakeThreadId>()];
271 let x = unsafe {
272 std::mem::transmute::<std::thread::ThreadId, FakeThreadId>(thread::current().id()).0
273 };
274 u64::from(x) as usize
275}
276
277#[derive(Debug, Clone)]
278pub struct DecodeKeyError {
279 pub token: Rank,
280}
281
282impl std::fmt::Display for DecodeKeyError {
283 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
284 write!(f, "Invalid token for decoding: {}", self.token)
285 }
286}
287
288impl std::error::Error for DecodeKeyError {}
289
290#[derive(Debug, Clone)]
291pub struct DecodeError {
292 pub message: String,
293}
294
295impl std::fmt::Display for DecodeError {
296 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
297 write!(f, "Could not decode tokens: {}", self.message)
298 }
299}
300
301impl std::error::Error for DecodeError {}
302
303#[derive(Debug, Clone)]
304pub struct EncodeError {
305 pub message: String,
306}
307
308impl std::fmt::Display for EncodeError {
309 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
310 write!(f, "Could not encode string: {}", self.message)
311 }
312}
313
314impl std::error::Error for EncodeError {}
315
316const MAX_NUM_THREADS: usize = 128;
317
318#[cfg_attr(feature = "python", pyclass(frozen))]
319#[derive(Clone)]
320pub struct CoreBPE {
321 encoder: HashMap<Vec<u8>, Rank>,
322 special_tokens_encoder: HashMap<String, Rank>,
323 decoder: HashMap<Rank, Vec<u8>>,
324 special_tokens_decoder: HashMap<Rank, Vec<u8>>,
325 regex_tls: Vec<Regex>,
326 special_regex_tls: Vec<Regex>,
327 sorted_token_bytes: Vec<Vec<u8>>,
328}
329
330impl CoreBPE {
331 fn _get_tl_regex(&self) -> &Regex {
332 // See performance notes above for what this is about
333 // It's also a little janky, please make a better version of it!
334 // However, it's nice that this doesn't leak memory to short-lived threads
335 &self.regex_tls[hash_current_thread() % MAX_NUM_THREADS]
336 }
337
338 fn _get_tl_special_regex(&self) -> &Regex {
339 &self.special_regex_tls[hash_current_thread() % MAX_NUM_THREADS]
340 }
341
342 /// Decodes tokens into a list of bytes.
343 ///
344 /// The bytes are not gauranteed to be a valid utf-8 string.
345 fn decode_bytes(&self, tokens: &[Rank]) -> Result<Vec<u8>, DecodeKeyError> {
346 let mut ret = Vec::with_capacity(tokens.len() * 2);
347 for &token in tokens {
348 let token_bytes = match self.decoder.get(&token) {
349 Some(bytes) => bytes,
350 None => self
351 .special_tokens_decoder
352 .get(&token)
353 .ok_or(DecodeKeyError { token })?,
354 };
355 ret.extend(token_bytes);
356 }
357 Ok(ret)
358 }
359
360 pub fn encode_ordinary(&self, text: &str) -> Vec<Rank> {
361 // This is the core of the encoding logic; the other functions in here
362 // just make things complicated :-)
363 let regex = self._get_tl_regex();
364 let mut ret = vec![];
365 for mat in regex.find_iter(text) {
366 let piece = mat.unwrap().as_str().as_bytes();
367 match self.encoder.get(piece) {
368 Some(token) => ret.push(*token),
369 None => ret.extend(&byte_pair_encode(piece, &self.encoder)),
370 }
371 }
372 ret
373 }
374
375 pub fn encode(
376 &self,
377 text: &str,
378 allowed_special: &HashSet<&str>,
379 ) -> Result<(Vec<Rank>, usize), EncodeError> {
380 let special_regex = self._get_tl_special_regex();
381 let regex = self._get_tl_regex();
382 let mut ret = vec![];
383
384 let mut start = 0;
385 let mut last_piece_token_len = 0;
386 loop {
387 let mut next_special;
388 let mut start_find = start;
389 loop {
390 // Find the next allowed special token, if any
391 next_special = special_regex.find_from_pos(text, start_find).unwrap();
392 match next_special {
393 Some(m) => {
394 if allowed_special.contains(&text[m.start()..m.end()]) {
395 break;
396 }
397 start_find = m.start() + 1;
398 }
399 None => break,
400 }
401 }
402 let end = next_special.map_or(text.len(), |m| m.start());
403
404 // Okay, here we go, compare this logic to encode_ordinary
405 for mat_res in regex.find_iter(&text[start..end]) {
406 let mat = match mat_res {
407 Ok(m) => m,
408 Err(e) => {
409 return Err(EncodeError {
410 message: format!("Regex error while tokenizing: {e}"),
411 });
412 }
413 };
414
415 let piece = mat.as_str().as_bytes();
416 if let Some(token) = self.encoder.get(piece) {
417 last_piece_token_len = 1;
418 ret.push(*token);
419 continue;
420 }
421 let tokens = byte_pair_encode(piece, &self.encoder);
422 last_piece_token_len = tokens.len();
423 ret.extend(&tokens);
424 }
425
426 match next_special {
427 // And here we push the special token
428 Some(m) => {
429 let piece = m.as_str();
430 let token = self.special_tokens_encoder[piece];
431 ret.push(token);
432 start = m.end();
433 last_piece_token_len = 0;
434 }
435 None => break,
436 }
437 }
438
439 // last_piece_token_len is how many tokens came from the last regex split. This is used
440 // for determining unstable tokens, since you can't merge across (stable) regex splits
441 Ok((ret, last_piece_token_len))
442 }
443
444 fn _increase_last_piece_token_len(
445 &self,
446 tokens: Vec<Rank>,
447 mut last_piece_token_len: usize,
448 ) -> (Vec<Rank>, usize) {
449 // Unfortunately, the locations where our regex splits can be unstable.
450 // For the purposes of determining unstable tokens, unstable regex splitting
451 // is only a problem if a split that was present disappears, since this can
452 // lead to merging of tokens otherwise thought to be stable.
453 // cl100k_base makes our life hard by including the \s*[\r\n]+
454 // pattern. This can e.g. cause "\n" + " " to become "\n \n".
455 // Here is a quick and dirty fix:
456 {
457 let token_is_all_space = |token| {
458 self.decoder
459 .get(token)
460 .map(|token_bytes| {
461 token_bytes
462 .iter()
463 .rev()
464 .all(|&b| [b' ', b'\n', b'\t'].contains(&b))
465 })
466 .unwrap_or(false)
467 };
468 if last_piece_token_len > 0
469 && token_is_all_space(&tokens[tokens.len() - last_piece_token_len])
470 {
471 while (last_piece_token_len < tokens.len())
472 && token_is_all_space(&tokens[tokens.len() - last_piece_token_len - 1])
473 {
474 last_piece_token_len += 1;
475 }
476 }
477 }
478 debug_assert!(last_piece_token_len <= tokens.len());
479
480 (tokens, last_piece_token_len)
481 }
482
483 pub fn _encode_unstable_native(
484 &self,
485 text: &str,
486 allowed_special: &HashSet<&str>,
487 ) -> (Vec<Rank>, HashSet<Vec<Rank>>) {
488 let (tokens, last_piece_token_len) = self.encode(text, allowed_special).unwrap();
489 if last_piece_token_len == 0 {
490 // If last_piece_token_len is zero, the last token was a special token and we have
491 // no unstable bytes
492 return (tokens, HashSet::new());
493 }
494 let (mut tokens, last_piece_token_len) =
495 self._increase_last_piece_token_len(tokens, last_piece_token_len);
496
497 let unstable_bytes = self
498 .decode_bytes(&tokens[tokens.len() - last_piece_token_len..])
499 .unwrap();
500 tokens.truncate(tokens.len() - last_piece_token_len);
501
502 // TODO: we should try harder to find additional stable tokens
503 // This would reduce the amount of retokenising when determining completions
504 // Refer to the logic in an older version of this file
505
506 let mut completions = HashSet::new();
507 if unstable_bytes.is_empty() {
508 return (tokens, completions);
509 }
510
511 // This is the easy bit. Just find all single tokens that start with unstable_bytes
512 // (including tokens that exactly match unstable_bytes)
513 // Separating this from the loop below helps with performance in a common case.
514 let mut point = self
515 .sorted_token_bytes
516 .partition_point(|x| x.as_slice() < unstable_bytes.as_slice());
517 while point < self.sorted_token_bytes.len()
518 && self.sorted_token_bytes[point].starts_with(&unstable_bytes)
519 {
520 completions.insert(vec![
521 self.encoder[self.sorted_token_bytes[point].as_slice()],
522 ]);
523 point += 1;
524 }
525
526 // Now apply even more brute force. At every (other) possible position for the straddling
527 // token, concatenate additional bytes from that token (if any) to unstable_bytes,
528 // and retokenise the whole thing and see what we get.
529 for i in 1..unstable_bytes.len() {
530 let prefix = &unstable_bytes[..i];
531 let suffix = &unstable_bytes[i..];
532 let mut point = self
533 .sorted_token_bytes
534 .partition_point(|x| x.as_slice() < suffix);
535 // TODO: Perf optimisation if suffix starts with " "?
536 while point < self.sorted_token_bytes.len()
537 && self.sorted_token_bytes[point].starts_with(suffix)
538 {
539 let possibility = [prefix, self.sorted_token_bytes[point].as_slice()].concat();
540 let encoded = match std::str::from_utf8(&possibility) {
541 // Morally, this is byte_pair_encode(&possibility, &self.encoder)
542 // But we might have introduced a regex split which would prevent merges.
543 // (particularly possible in the presence of unstable regex splits)
544 // So convert to UTF-8 and do regex splitting.
545 // E.g. with cl100k_base " !" gets split to " " + " !",
546 // but byte_pair_encode(" !") != byte_pair_encode(" ")
547 Ok(s) => self.encode_ordinary(s),
548
549 // Technically, whether or not this arm is correct depends on whether there
550 // would be a regex split before the UTF-8 truncation point.
551 // Probably niche enough that no one will ever notice (after all, people didn't
552 // notice all the big holes in the previous unstable token implementation)
553 Err(_) => byte_pair_encode(&possibility, &self.encoder),
554 // Something like the following is intriguing but incorrect:
555 // Err(e) => self.encode_ordinary(unsafe {
556 // std::str::from_utf8_unchecked(&possibility[..e.valid_up_to()])
557 // }),
558 };
559 let mut seq = Vec::new();
560 let mut seq_len = 0;
561 for token in encoded {
562 seq.push(token);
563 seq_len += self.decoder[&token].len();
564 if seq_len >= unstable_bytes.len() {
565 break;
566 }
567 }
568 completions.insert(seq);
569 point += 1;
570 }
571 }
572
573 // This is also not straightforward. While we generally assume that regex splits are stable,
574 // unfortunately, they are not. That is, if adding bytes were to make a split appear in
575 // unstable_bytes, this could make tokens possible which our logic would otherwise think
576 // would be merged.
577 // For example, with gpt2, the use of \s+(?!\S) means that "\n\n" could
578 // develop a split, e.g. "\n\n0" splits into "\n"+"\n"+"0", making "\n" a possible token.
579 // Here is a quick and dirty fix:
580 // This isn't right if we ever remove \s+(?!\S)
581 if unstable_bytes.len() > 1 {
582 let last_decoded = bstr::decode_last_utf8(unstable_bytes.as_slice());
583 if unstable_bytes.len() - last_decoded.1 > 0
584 && last_decoded.0.is_some_and(|c| c.is_whitespace())
585 {
586 let mut reencoded = byte_pair_encode(
587 &unstable_bytes[..unstable_bytes.len() - last_decoded.1],
588 &self.encoder,
589 );
590 reencoded.extend(byte_pair_encode(
591 &unstable_bytes[unstable_bytes.len() - last_decoded.1..],
592 &self.encoder,
593 ));
594 completions.insert(reencoded);
595 }
596 }
597
598 (tokens, completions)
599 }
600
601 pub fn new<E, SE, NSE>(
602 encoder: E,
603 special_tokens_encoder: SE,
604 pattern: &str,
605 ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>>
606 where
607 E: IntoIterator<Item = (Vec<u8>, Rank)>,
608 SE: IntoIterator<Item = (String, Rank)>,
609 NSE: IntoIterator<Item = (String, (Rank, Rank))>,
610 {
611 Self::new_internal(
612 HashMap::from_iter(encoder),
613 HashMap::from_iter(special_tokens_encoder),
614 pattern,
615 )
616 }
617
618 fn new_internal(
619 encoder: HashMap<Vec<u8>, Rank>,
620 special_tokens_encoder: HashMap<String, Rank>,
621 pattern: &str,
622 ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
623 let regex = Regex::new(pattern)?;
624
625 let special_regex = {
626 let parts = special_tokens_encoder
627 .keys()
628 .map(|s| fancy_regex::escape(s))
629 .collect::<Vec<_>>();
630 Regex::new(&parts.join("|"))?
631 };
632
633 let decoder: HashMap<Rank, Vec<u8>> =
634 encoder.iter().map(|(k, v)| (*v, k.clone())).collect();
635
636 assert!(
637 encoder.len() == decoder.len(),
638 "Encoder and decoder must be of equal length. Encoder length: {}, decoder length: {}.\nMaybe you had duplicate token indices in your encoder?",
639 encoder.len(),
640 decoder.len()
641 );
642
643 let special_tokens_decoder: HashMap<Rank, Vec<u8>> = special_tokens_encoder
644 .iter()
645 .map(|(k, v)| (*v, k.as_bytes().to_vec()))
646 .collect();
647
648 // Clone because I don't know how to tell Rust I'm not going to change the map
649 let mut sorted_token_bytes: Vec<Vec<u8>> = encoder.keys().cloned().collect();
650 sorted_token_bytes.sort();
651
652 Ok(Self {
653 encoder,
654 special_tokens_encoder,
655 decoder,
656 special_tokens_decoder,
657 regex_tls: (0..MAX_NUM_THREADS).map(|_| regex.clone()).collect(),
658 special_regex_tls: (0..MAX_NUM_THREADS)
659 .map(|_| special_regex.clone())
660 .collect(),
661 sorted_token_bytes,
662 })
663 }
664
665 pub fn special_tokens(&self) -> HashSet<&str> {
666 self.special_tokens_encoder
667 .keys()
668 .map(|s| s.as_str())
669 .collect()
670 }
671
672 pub fn encode_with_special_tokens(&self, text: &str) -> Vec<Rank> {
673 let allowed_special = self.special_tokens();
674 self.encode(text, &allowed_special).unwrap().0
675 }
676}
677
678#[cfg(test)]
679mod tests {
680 use fancy_regex::Regex;
681 use rustc_hash::FxHashMap as HashMap;
682
683 use crate::{Rank, byte_pair_split};
684
685 fn setup_ranks() -> HashMap<Vec<u8>, Rank> {
686 HashMap::from_iter([(b"ab".to_vec(), 0), (b"cd".to_vec(), 1)])
687 }
688
689 #[test]
690 fn test_simple_characters() {
691 let ranks = setup_ranks();
692 let res = byte_pair_split(b"abcd", &ranks);
693 assert_eq!(res, vec![b"ab", b"cd"]);
694 }
695
696 #[test]
697 fn test_repeated_characters() {
698 let ranks = setup_ranks();
699 let res = byte_pair_split(b"abab", &ranks);
700 assert_eq!(res, vec![b"ab", b"ab"]);
701 }
702}
703