outlines-core 0.2.14

//! Post-processing operations for the tokens before they being inserted into
//! `Vocabulary`, strategies depend on the tokenizer's level.

use std::collections::HashMap;

use once_cell::sync::Lazy;
use serde::Deserialize;
use tokenizers::normalizers::Replace;
use tokenizers::{DecoderWrapper, Tokenizer};

use crate::{Error, Result};

/// Map to reverse certain UTF-8 characters back to bytes.
///
/// GPT2-like tokenizers have multibyte tokens that can have a mix of full and incomplete
/// UTF-8 characters, for example, byte ` \xf0` can be one token. These tokenizers map each
/// byte to a valid UTF-8 character, `TokenProcessor` of `ByteFallback` level will be used
/// to map back these type of characters into bytes, based on `CHAR_MAP`.
///
/// For example these should be interpreted as:
///
/// ``` ignore
/// "ĠO" = [U+0120, U+004F] => [0x20, 0x4F] = " O"
/// "Ġal" = [U+0120, U+0061, U+006C] => [0x20, 0x61, 0x6C] = " al"
/// ```
/// We'll use the following the mapping for this translation:
///
/// ``` ignore
/// 'Ā' == '\u{0100}' -> 0x00 == 0
/// 'ā' == '\u{0101}' -> 0x01 == 1
/// 'Ă' == '\u{0102}' -> 0x02 == 2
/// ...
/// 'Ğ' == '\u{011E}' -> 0x1E == 30
/// 'ğ' == '\u{011F}' -> 0x1F == 31
/// 'Ġ' == '\u{0120}' -> 0x20 == 32
/// ---
/// '!' == '\u{0021}' -> 0x21 == 33
/// '"' == '\u{0022}' -> 0x22 == 34
/// '#' == '\u{0023}' -> 0x23 == 35
/// ...
/// '|' == '\u{007C}' -> 0x7C == 124
/// '}' == '\u{007D}' -> 0x7D == 125
/// '~' == '\u{007E}' -> 0x7E == 126
/// ---
/// 'ġ' == '\u{0121}' -> 0x7F == 127
/// 'Ģ' == '\u{0122}' -> 0x80 == 128
/// 'ģ' == '\u{0123}' -> 0x81 == 129
/// ...
/// 'ŀ' == '\u{0140}' -> 0x9E == 158
/// 'Ł' == '\u{0141}' -> 0x9F == 159
/// 'ł' == '\u{0142}' -> 0xA0 == 160
/// ---
/// '¡' == '\u{00A1}' -> 0xA1 == 161
/// '¢' == '\u{00A2}' -> 0xA2 == 162
/// '£' == '\u{00A3}' -> 0xA3 == 163
/// ...
/// 'ª' == '\u{00AA}' -> 0xAA == 170
/// '«' == '\u{00AB}' -> 0xAB == 171
/// '¬' == '\u{00AC}' -> 0xAC == 172
/// ---
/// 'Ń' == '\u{0143}' -> 0xAD == 173
/// ---
/// '®' == '\u{00AE}' -> 0xAE == 174
/// '¯' == '\u{00AF}' -> 0xAF == 175
/// '°' == '\u{00B0}' -> 0xB0 == 176
/// ...
/// 'ý' == '\u{00FD}' -> 0xFD == 253
/// 'þ' == '\u{00FE}' -> 0xFE == 254
/// 'ÿ' == '\u{00FF}' -> 0xFF == 255
/// ```
static CHAR_MAP: Lazy<HashMap<char, u8>> = Lazy::new(|| {
    let mut char_map = HashMap::with_capacity(256);
    let mut key = 0x100u32;
    for byte in 0..=255u8 {
        let char = byte as char;
        if matches!(
            char, '!'..='~' | '\u{00A1}'..='\u{00AC}' | '\u{00AE}'..='\u{00FF}',
        ) {
            char_map.insert(char, byte);
        } else if let Some(ch) = char::from_u32(key) {
            char_map.insert(ch, byte);
            key += 1;
        }
    }
    char_map
});

/// Recognizes different tokenizer's levels.
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum TokenProcessorLevel {
    /// Matches byte level tokenizer (e.g., gpt2).
    Byte,
    /// Matches byte fallback tokenizer (e.g., llama), which have `<0x__>` tokens for
    /// all `__` >= `0x80` to represent incomplete UTF-8 sequences.
    ByteFallback(Mods),
}

/// Modifications to be applied by `TokenProcessor`of `ByteFallback` level.
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct Mods {
    spacechar: String,
}

impl Mods {
    /// Default string modification to be applied by `TokenProcessor` of `ByteFallback` level.
    fn apply_default(&self, token: &str) -> String {
        token.replace(&self.spacechar, " ")
    }
}

/// Local structure to be deserialized into from HF's `ReplaceDecoder` in order to get a replace pattern.
#[derive(Debug, Deserialize)]
struct ReplaceDecoder {
    content: String,
    pattern: ReplacePattern,
}

impl ReplaceDecoder {
    fn space_replacement(&self) -> Option<String> {
        if self.content != " " {
            return None;
        }
        match &self.pattern {
            ReplacePattern::String(pattern) => {
                let mut chars = pattern.chars();
                let char = chars.next();
                if let Some(replacement) = char {
                    if chars.next().is_none() {
                        return Some(replacement.to_string());
                    }
                }
                None
            }
        }
    }
}

#[derive(Debug, Deserialize)]
enum ReplacePattern {
    String(String),
}

/// Token processor to adjust tokens according to the tokenizer's level.
#[derive(Debug)]
pub(crate) struct TokenProcessor {
    level: TokenProcessorLevel,
}

impl TokenProcessor {
    /// Create new `TokenProcessor` with the level defined based on tokenizer's decoders.
    pub(crate) fn new(tokenizer: &Tokenizer) -> Result<Self> {
        match tokenizer.get_decoder() {
            None => Err(Error::UnsupportedByTokenProcessor),
            Some(decoder) => match decoder {
                DecoderWrapper::ByteLevel(_) => Ok(Self {
                    level: TokenProcessorLevel::Byte,
                }),
                DecoderWrapper::Sequence(decoding_sequence) => {
                    let mut is_byte_fallback = false;
                    let mut spacechar = ' '.to_string();

                    for decoder in decoding_sequence.get_decoders() {
                        match decoder {
                            DecoderWrapper::ByteFallback(_) => {
                                is_byte_fallback = true;
                            }
                            DecoderWrapper::Replace(replace) => {
                                // `Replace` decoder would replace a pattern in the output with something else,
                                // which we need to know.
                                let decoder = Self::unpack_decoder(replace)?;
                                if let Some(replacement) = decoder.space_replacement() {
                                    spacechar = replacement;
                                }
                            }
                            _ => {}
                        }
                    }

                    if is_byte_fallback {
                        Ok(Self {
                            level: TokenProcessorLevel::ByteFallback(Mods { spacechar }),
                        })
                    } else {
                        Err(Error::UnsupportedByTokenProcessor)
                    }
                }
                _ => Err(Error::UnsupportedByTokenProcessor),
            },
        }
    }

    /// Operates on each token based on the level of `TokenProcessor`.
    pub(crate) fn process(&self, token: &str) -> Result<Vec<u8>> {
        match &self.level {
            TokenProcessorLevel::Byte => token
                .chars()
                .map(|char| {
                    CHAR_MAP
                        .get(&char)
                        .copied()
                        .ok_or(Error::ByteProcessorFailed)
                })
                .collect(),
            TokenProcessorLevel::ByteFallback(mods) => {
                // If the token is of form `<0x__>`:
                if token.len() == 6 && token.starts_with("<0x") && token.ends_with('>') {
                    // Get to a single byte specified in the __ part and parse it in base 16 to a byte.
                    match u8::from_str_radix(&token[3..5], 16) {
                        Ok(byte) => Ok([byte].to_vec()),
                        Err(_) => Err(Error::ByteFallbackProcessorFailed),
                    }
                } else {
                    Ok(mods.apply_default(token).as_bytes().to_vec())
                }
            }
        }
    }

    /// Since all fields of HF's `Replace` are private with no getters, it needs to be unpacked
    /// into local `ReplaceDecoder` structure.
    #[cfg(not(tarpaulin_include))]
    fn unpack_decoder(decoder: &Replace) -> Result<ReplaceDecoder> {
        match serde_json::to_value(decoder) {
            Err(_) => Err(Error::DecoderUnpackingFailed),
            Ok(value) => match serde_json::from_value(value) {
                Ok(d) => Ok(d),
                Err(_) => Err(Error::DecoderUnpackingFailed),
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn byte_level_processor() {
        let model = "openai-community/gpt2";
        let tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");
        let processor = TokenProcessor::new(&tokenizer).expect("Processor failed");

        assert_eq!(processor.level, TokenProcessorLevel::Byte);

        for (ch, byte) in [
            ('Ā', 0x00),
            ('ā', 0x01),
            ('Ă', 0x02),
            ('Ğ', 0x1E),
            ('ğ', 0x1F),
            ('Ġ', 0x20),
            ('!', 0x21),
            ('"', 0x22),
            ('#', 0x23),
            ('|', 0x7C),
            ('}', 0x7D),
            ('~', 0x7E),
            ('ġ', 0x7F),
            ('Ģ', 0x80),
            ('ģ', 0x81),
            ('ŀ', 0x9E),
            ('Ł', 0x9F),
            ('ł', 0xA0),
            ('¡', 0xA1),
            ('¢', 0xA2),
            ('£', 0xA3),
            ('ª', 0xAA),
            ('«', 0xAB),
            ('¬', 0xAC),
            ('Ń', 0xAD),
            ('®', 0xAE),
            ('¯', 0xAF),
            ('°', 0xB0),
            ('ý', 0xFD),
            ('þ', 0xFE),
            ('ÿ', 0xFF),
        ] {
            let processed = processor.process(&ch.to_string()).expect("Not processed");
            assert_eq!(processed, [byte]);
        }
    }

    #[test]
    fn byte_fallback_level_processor() {
        let model = "hf-internal-testing/llama-tokenizer";
        let tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");
        let processor = TokenProcessor::new(&tokenizer).expect("Processor failed");
        let spacechar = '▁'.to_string();
        let mods = Mods {
            spacechar: spacechar.clone(),
        };

        assert_eq!(processor.level, TokenProcessorLevel::ByteFallback(mods));

        for (input, expected) in [
            ("abc", vec![0x61, 0x62, 0x63]),
            ("<0x61>", vec![0x61]),
            ("<0x61>a", vec![0x3C, 0x30, 0x78, 0x36, 0x31, 0x3E, 0x61]),
            (&spacechar, vec![0x20]),
            (
                &format!("{}{}abc", spacechar, spacechar),
                vec![0x20, 0x20, 0x61, 0x62, 0x63],
            ),
            (
                &format!("{}{}{}", spacechar, spacechar, spacechar),
                vec![0x20, 0x20, 0x20],
            ),
        ] {
            let processed = processor.process(input).expect("Not processed");
            assert_eq!(processed, expected);
        }
    }

    #[test]
    fn unsupported_tokenizer_error() {
        let model = "hf-internal-testing/tiny-random-XLMRobertaXLForCausalLM";
        let tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");

        let result = TokenProcessor::new(&tokenizer);
        match result {
            Err(Error::UnsupportedByTokenProcessor) => {}
            _ => unreachable!(),
        }
    }

    #[test]
    fn byte_processor_error() {
        let model = "openai-community/gpt2";
        let tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");
        let processor = TokenProcessor::new(&tokenizer).expect("Processor failed");

        for token in ["𝒜𝒷𝒸𝒟𝓔", "🦄🌈🌍🔥🎉", "京东购物"] {
            let result = processor.process(token);
            match result {
                Err(Error::ByteProcessorFailed) => {}
                _ => unreachable!(),
            }
        }
    }

    #[test]
    fn byte_fallback_processor_error() {
        let model = "hf-internal-testing/llama-tokenizer";
        let tokenizer = Tokenizer::from_pretrained(model, None).expect("Tokenizer failed");
        let processor = TokenProcessor::new(&tokenizer).expect("Processor failed");

        let result = processor.process("<0x6y>");
        match result {
            Err(Error::ByteFallbackProcessorFailed) => {}
            _ => unreachable!(),
        }
    }

    #[test]
    fn only_get_spacechar_replacement() {
        let one_char = "_".to_string();
        let pattern = ReplacePattern::String(one_char);
        let not_spacechar = "-".to_string();
        let decoder = ReplaceDecoder {
            content: not_spacechar,
            pattern,
        };
        assert!(decoder.space_replacement().is_none());
    }

    #[test]
    fn only_one_pattern_char_for_spacechar_replacement() {
        let two_chars = "_*".to_string();
        let pattern = ReplacePattern::String(two_chars);
        let spacechar = " ".to_string();
        let decoder = ReplaceDecoder {
            content: spacechar,
            pattern,
        };
        assert!(decoder.space_replacement().is_none());
    }

    #[test]
    fn tokenizer_without_decoders_is_unsupported() {
        use tokenizers::models::bpe::BPE;

        let tokenizer = Tokenizer::new(BPE::default());
        let result = TokenProcessor::new(&tokenizer);
        match result {
            Err(Error::UnsupportedByTokenProcessor) => {}
            _ => unreachable!(),
        }
    }

    #[test]
    fn tokenizer_without_supported_decoders_in_sequence_is_unsupported() {
        use tokenizers::decoders::sequence::Sequence;
        use tokenizers::decoders::wordpiece::WordPiece;
        use tokenizers::models::bpe::BPE;

        let mut tokenizer = Tokenizer::new(BPE::default());
        let decoder = WordPiece::default();
        let sequence = Sequence::new(vec![DecoderWrapper::WordPiece(decoder)]);
        let decoder_sequence = DecoderWrapper::Sequence(sequence);
        tokenizer.with_decoder(Some(decoder_sequence));

        let result = TokenProcessor::new(&tokenizer);
        match result {
            Err(Error::UnsupportedByTokenProcessor) => {}
            _ => unreachable!(),
        }
    }
}