refactor(tests): move heap bpe into test utils

Author: ChieloNewctle

src/inc_impl.rs src/inc_bpe.rssrc/inc_impl.rs renamed to src/inc_bpe.rs

@@ -213,7 +213,7 @@ mod tests {
 
     use crate::{
         Dictionary, IncBpeToken, IncBpeTokenChainIter, IncBpeTokenizer, NormalizedDict, TokenId,
-        Vocab, sp_impl::sentence_piece_impl,
+        Vocab, test_utils::bpe_with_heap,
     };
 
     fn inc_bpe_short_any_case(vocab: &[&str], rules: &[(&str, &str)], sequences: &[&str]) {
@@ -244,7 +244,7 @@ mod tests {
 
     fn validate(dict: &Dictionary, seq: &[TokenId], inc_res: &[IncBpeToken]) {
         for i in 0..seq.len() {
-            let expected = sentence_piece_impl::<false>(dict, &seq[0..i + 1]);
+            let expected = bpe_with_heap::<false>(dict, &seq[0..i + 1]);
             let output = IncBpeTokenChainIter::new(inc_res, i).token_ids();
             let output = output.chain(std::iter::repeat(TokenId::MAX));
             assert!(expected.into_iter().rev().zip(output).all(|(i, j)| i == j));

src/lib.rs

@@ -1,19 +1,20 @@
 mod aho_corasick;
 mod centroid;
 mod dict;
-mod heap;
-mod inc_impl;
+mod inc_bpe;
 mod normalize;
-mod sp_impl;
 mod successor;
 mod suf_suc;
 mod typed_vec;
 mod vocab;
 
 pub use crate::{
     dict::{DictBuildError, Dictionary, Rule, RuleId, UnknownToken},
-    inc_impl::{IncBpeToken, IncBpeTokenChainIter, IncBpeTokenization, IncBpeTokenizer},
+    inc_bpe::{IncBpeToken, IncBpeTokenChainIter, IncBpeTokenization, IncBpeTokenizer},
     normalize::{NormalizedDict, NormalizedDictBuildError},
     successor::SkipLen,
     vocab::{Token, TokenId, Vocab, VocabBuildError},
 };
+
+#[cfg(test)]
+mod test_utils;

src/normalize.rs

@@ -163,7 +163,7 @@ impl NormalizedDict {
 
 #[cfg(test)]
 mod tests {
-    use crate::{Dictionary, NormalizedDict, RuleId, Vocab, sp_impl::sentence_piece_impl};
+    use crate::{Dictionary, NormalizedDict, RuleId, Vocab, test_utils::bpe_with_heap};
 
     fn build_dict<T: AsRef<[u8]>, R: IntoIterator<Item = (T, T)>>(
         vocab: &Vocab,
@@ -178,7 +178,7 @@ mod tests {
             let token_id = rule.merged;
             assert!(!dict.is_single(token_id));
             let seq = &dict[token_id];
-            let res = sentence_piece_impl::<false>(&dict, dict.split_bytes_to_tokens(seq, 0usize));
+            let res = bpe_with_heap::<false>(&dict, dict.split_bytes_to_tokens(seq, 0usize));
             assert!(dict.is_useful(token_id) ^ (res != vec![token_id]));
         }
         dict
@@ -196,7 +196,7 @@ mod tests {
                 }
             };
             assert!(!dict.is_single(token_id));
-            let res = sentence_piece_impl::<false>(&dict, dict.split_utf8_to_tokens(seq, 0usize));
+            let res = bpe_with_heap::<false>(&dict, dict.split_utf8_to_tokens(seq, 0usize));
             assert!(dict.is_useful(token_id) ^ (res != vec![token_id]));
         }
         dict

src/sp_impl.rs src/test_utils/bpe.rssrc/sp_impl.rs renamed to src/test_utils/bpe.rs

@@ -1,15 +1,14 @@
-#![cfg(test)]
 use crate::{
     Dictionary, RuleId, TokenId,
-    heap::AdjustableHeap,
+    test_utils::heap::AdjustableHeap,
     typed_vec::{TypedVec, typed_vec_index},
 };
 
 typed_vec_index!(pub(super) InputTextPos, u32);
 
 type Heap = AdjustableHeap<InputTextPos, RuleId>;
 
-pub fn sentence_piece_impl<const ALLOW_IMPROPER_RULES: bool>(
+pub fn bpe_with_heap<const ALLOW_IMPROPER_RULES: bool>(
     dict: &Dictionary,
     seq: impl Into<Vec<TokenId>>,
 ) -> Vec<TokenId> {
@@ -98,7 +97,7 @@ pub fn sentence_piece_impl<const ALLOW_IMPROPER_RULES: bool>(
 
 #[cfg(test)]
 mod tests {
-    use crate::{Dictionary, TokenId, Vocab, sp_impl::sentence_piece_impl};
+    use crate::{Dictionary, TokenId, Vocab, test_utils::bpe::bpe_with_heap};
 
     fn build_dict<T: AsRef<[u8]>, R: IntoIterator<Item = (T, T)>>(
         vocab: &Vocab,
@@ -114,7 +113,7 @@ mod tests {
     ) {
         let tokens: Vec<_> = tokens.into_iter().map(I::into).collect();
         let inputs = dict.split_bytes_to_tokens(seq.as_ref(), 0usize);
-        assert_eq!(sentence_piece_impl::<true>(dict, inputs), tokens);
+        assert_eq!(bpe_with_heap::<true>(dict, inputs), tokens);
         assert!(dict.is_proper_in_bytes().is_ok());
         check_properly_in_bytes(dict, seq, tokens);
     }
@@ -126,7 +125,7 @@ mod tests {
     ) {
         let tokens: Vec<_> = tokens.into_iter().map(I::into).collect();
         let inputs = dict.split_utf8_to_tokens(seq.as_ref(), 0usize);
-        assert_eq!(sentence_piece_impl::<true>(dict, inputs), tokens);
+        assert_eq!(bpe_with_heap::<true>(dict, inputs), tokens);
         assert!(dict.is_proper_in_utf8().is_ok());
         check_properly_in_utf8(dict, seq, tokens);
     }
@@ -138,7 +137,7 @@ mod tests {
     ) {
         let tokens: Vec<_> = tokens.into_iter().map(I::into).collect();
         let inputs = dict.split_bytes_to_tokens(seq.as_ref(), 0usize);
-        assert_eq!(sentence_piece_impl::<false>(dict, inputs), tokens);
+        assert_eq!(bpe_with_heap::<false>(dict, inputs), tokens);
     }
 
     fn check_properly_in_utf8<S: AsRef<str>, I: Into<TokenId>, T: IntoIterator<Item = I>>(
@@ -148,11 +147,11 @@ mod tests {
     ) {
         let tokens: Vec<_> = tokens.into_iter().map(I::into).collect();
         let inputs = dict.split_utf8_to_tokens(seq.as_ref(), 0usize);
-        assert_eq!(sentence_piece_impl::<false>(dict, inputs), tokens);
+        assert_eq!(bpe_with_heap::<false>(dict, inputs), tokens);
     }
 
     #[test]
-    fn test_sp_impl() {
+    fn test_bpe_with_heap() {
         let vocab = Vocab::new([
             b"<unk>" as &[_],
             b"a",

src/heap.rs src/test_utils/heap.rssrc/heap.rs renamed to src/test_utils/heap.rs

@@ -1,12 +1,11 @@
-#![cfg(test)]
 use crate::typed_vec::{TypedVec, TypedVecIndex, typed_vec_index};
 
 typed_vec_index!(NodeId, u32);
 
 const ROOT: NodeId = NodeId::ZERO.next();
 
 #[derive(Debug)]
-pub(crate) struct AdjustableHeap<Pos, Key> {
+pub(super) struct AdjustableHeap<Pos, Key> {
     nodes: TypedVec<NodeId, (Key, Pos)>,
     pos_to_node_id: TypedVec<Pos, NodeId>,
 }
@@ -186,7 +185,7 @@ mod tests {
     use rand::{Rng, SeedableRng, rngs::StdRng};
 
     use crate::{
-        heap::AdjustableHeap,
+        test_utils::heap::AdjustableHeap,
         typed_vec::{TypedVec, typed_vec_index},
     };
 

src/test_utils/mod.rs

@@ -0,0 +1,4 @@
+mod bpe;
+mod heap;
+
+pub use bpe::bpe_with_heap;