implement review comments

Author: marinegor

crates/bpe/src/byte_pair_encoding.rs

@@ -555,11 +555,33 @@ impl BytePairEncoding {
     /// This function computes the shortest possible encoding sequence which will usually differ from the
     /// tokenization produced by the original BPE algorithm.
     #[cfg(feature = "rand")]
-    pub fn encode_minimal_dropout(&self, text: &[u8], dropout: f32) -> Vec<u32> {
+    pub fn encode_minimal_dropout(&self, text: &[u8], dropout: f32, seed: Option<u64>) -> Vec<u32> {
+        use rand::rngs::StdRng;
         use rand::Rng;
+        use rand::SeedableRng;
+
+        fn get_rng(seed: Option<u64>) -> StdRng {
+            match seed {
+                Some(num) => {
+                    // Expand the u64 seed to 32 bytes
+                    let mut seed_bytes = [0u8; 32];
+                    seed_bytes[..8].copy_from_slice(&num.to_le_bytes());
+                    StdRng::from_seed(seed_bytes)
+                }
+                None => {
+                    // Seed StdRng with a random 32-byte array from ThreadRng
+                    let mut thread_rng = rand::rng();
+                    let mut seed_bytes = [0u8; 32];
+                    thread_rng.fill(&mut seed_bytes);
+                    StdRng::from_seed(seed_bytes)
+                }
+            }
+        }
+
+        let mut rng = get_rng(seed);
+
         assert!(0.0 <= dropout);
         assert!(dropout <= 1.0);
-        let mut rng = rand::rng();
 
         let mut last_token: Vec<(u32, u32)> = Vec::with_capacity(text.len());
         let mut state = self.overlapping_searcher.start_state();
@@ -572,8 +594,9 @@ impl BytePairEncoding {
                     best = (m.value(), 1);
                     break;
                 } else if last_token[m.start() - 1].1 + 1 < best.1 {
-                    best = (m.value(), last_token[m.start() - 1].1 + 1);
                     if rng.random_range(0.0..=1.0) < dropout {
+                        best = (m.value(), last_token[m.start() - 1].1 + 1);
+                    } else {
                         best = (m.value(), 1);
                     }
                 }

crates/bpe/tests/Cargo.toml

@@ -8,3 +8,6 @@ bpe-openai = { path = "../../bpe-openai" }
 itertools = "0.14"
 rand = "0.9"
 tiktoken-rs = "0.9"
+
+[dev-dependencies]
+rand_chacha = { version = "0.9" }

crates/bpe/tests/src/lib.rs

@@ -149,9 +149,9 @@ mod tests {
             for _ in 0..8 {
                 let input = create_test_bytes(bpe, bytes);
                 let encoded = bpe.encode_minimal(&input);
-                let encoded_d_min = bpe.encode_minimal_dropout(&input, 0.2);
-                let encoded_d_max = bpe.encode_minimal_dropout(&input, 0.9);
-                let encoded_d_max_again = bpe.encode_minimal_dropout(&input, 0.9);
+                let encoded_d_min = bpe.encode_minimal_dropout(&input, 0.2, Some(0));
+                let encoded_d_max = bpe.encode_minimal_dropout(&input, 0.9, Some(1));
+                let encoded_d_max_again = bpe.encode_minimal_dropout(&input, 0.9, Some(2));
                 let decoded = bpe.decode_tokens(&encoded);
                 let decoded_min = bpe.decode_tokens(&encoded_d_min);
                 let decoded_max = bpe.decode_tokens(&encoded_d_max);