wordchipper 0.9.2

HPC Rust LLM Tokenizer Library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

//! # Parallel Decoder

use crate::{
    TokenType,
    WCResult,
    alloc::sync::Arc,
    decoders::{
        BatchDecodeResult,
        DecodeResult,
        TokenDecoder,
    },
    prelude::*,
};

/// Batch-Level Parallel Decoder Wrapper.
///
/// Enables ``rayon`` decoding of batches when available.
pub struct ParallelRayonDecoder<T: TokenType> {
    /// Wrapped decoder.
    pub inner: Arc<dyn TokenDecoder<T>>,

    _marker: std::marker::PhantomData<T>,
}

impl<T> ParallelRayonDecoder<T>
where
    T: TokenType,
{
    /// Create a new parallel token decoders.
    ///
    /// ## Arguments
    /// * `inner` - The token decoder to wrap.
    ///
    /// ## Returns
    /// A new `ParallelRayonDecoder` instance.
    pub fn new(inner: Arc<dyn TokenDecoder<T>>) -> Self {
        Self {
            inner,
            _marker: std::marker::PhantomData,
        }
    }
}

impl<T> TokenDecoder<T> for ParallelRayonDecoder<T>
where
    T: TokenType,
{
    fn try_decode_to_bytes(
        &self,
        tokens: &[T],
    ) -> WCResult<DecodeResult<Vec<u8>>> {
        self.inner.try_decode_to_bytes(tokens)
    }

    fn try_decode_batch_to_bytes(
        &self,
        batch: &[&[T]],
    ) -> WCResult<BatchDecodeResult<Vec<u8>>> {
        use rayon::prelude::*;

        batch
            .par_iter()
            .map(|tokens| self.try_decode_to_bytes(tokens))
            .collect::<WCResult<Vec<_>>>()
            .map(BatchDecodeResult::from)
    }

    fn try_decode_batch_to_strings(
        &self,
        batch: &[&[T]],
    ) -> WCResult<BatchDecodeResult<String>> {
        use rayon::prelude::*;

        batch
            .par_iter()
            .map(|tokens| self.try_decode_to_string(tokens))
            .collect::<WCResult<Vec<_>>>()
            .map(BatchDecodeResult::from)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        TokenDecoderOptions,
        UnifiedTokenVocab,
        decoders::utility::testing::common_decoder_tests,
        pretrained::openai::OA_CL100K_BASE_PATTERN,
        spanners::TextSpanningConfig,
        vocab::utility::testing::{
            build_test_shift_byte_vocab,
            build_test_vocab,
        },
    };

    #[test]
    fn test_rayon_decoder() {
        type T = u16;

        let vocab: Arc<UnifiedTokenVocab<T>> = build_test_vocab(
            build_test_shift_byte_vocab(10),
            TextSpanningConfig::from_pattern(OA_CL100K_BASE_PATTERN),
        )
        .into();

        let inner = TokenDecoderOptions::default()
            .with_parallel(false)
            .build(vocab.clone());
        let decoder = ParallelRayonDecoder::new(inner);

        common_decoder_tests(vocab, Arc::new(decoder));
    }
}