toklen 0.2.0

use fancy_regex::Regex;
use serde::Deserialize;

use crate::pre_tokenized::{PreTokenizedString, PtSplit};

use super::Error;

use crate::byte_level::BYTE_TO_UTF8;
use crate::byte_level::BYTE_TO_UTF8_LEN;

/// Encode an entire byte slice into GPT-2 byte-level characters.
///
/// # Safety
/// `out` must have at least `src.len() * 2` bytes of remaining capacity.
unsafe fn encode_bytes_bulk(src: &[u8], out: &mut Vec<u8>) {
    let mut pos = out.len();
    let base = out.as_mut_ptr();
    for &b in src {
        let utf8 = BYTE_TO_UTF8[b as usize];
        let len = BYTE_TO_UTF8_LEN[b as usize] as usize;
        unsafe {
            std::ptr::copy_nonoverlapping(utf8.as_ptr(), base.add(pos), len);
        }
        pos += len;
    }
    unsafe { out.set_len(pos) };
}

/// Append the byte-level encoding of `s` to `out` without checking capacity.
///
/// # Safety
/// `out` must have at least `s.len() * 2` bytes of spare capacity.
unsafe fn encode_bytes_into_unchecked(s: &str, out: &mut String) {
    unsafe { encode_bytes_bulk(s.as_bytes(), out.as_mut_vec()) };
}

/// Append the byte-level encoding of `s` to `out`.
fn encode_bytes_into(s: &str, out: &mut String) {
    unsafe {
        let v = out.as_mut_vec();
        v.reserve(s.len() << 1);
        encode_bytes_bulk(s.as_bytes(), v);
    }
}

/// GPT-2 pretokenization regex.
const GPT2_PATTERN: &str = concat!(
    r"'(?i:[sdmt])",
    r"|'(?i:ll|ve|re)",
    r"| ?\p{L}+",
    r"| ?\p{N}+",
    r"| ?[^\s\p{L}\p{N}]+",
    r"|\s+(?!\S)",
    r"|\s+",
);

#[rustfmt::skip]
#[inline(always)]
const fn default_true() -> bool { true }

/// Raw deserialization helper for [`ByteLevel`].
#[derive(Deserialize)]
struct ByteLevelRaw {
    #[serde(default = "default_true")]
    add_prefix_space: bool,
    #[serde(default = "default_true")]
    trim_offsets: bool,
    #[serde(default = "default_true")]
    use_regex: bool,
}

/// A compiled ByteLevel pre-tokenizer.
#[derive(Clone, Debug, Deserialize)]
#[serde(try_from = "ByteLevelRaw")]
pub struct ByteLevel {
    regex: Option<Regex>,
    add_prefix_space: bool,
    #[allow(dead_code)]
    trim_offsets: bool,
}

impl TryFrom<ByteLevelRaw> for ByteLevel {
    type Error = Error;

    fn try_from(raw: ByteLevelRaw) -> Result<Self, Error> {
        let regex = if raw.use_regex {
            Some(Regex::new(GPT2_PATTERN)?)
        } else {
            None
        };
        Ok(Self {
            regex,
            add_prefix_space: raw.add_prefix_space,
            trim_offsets: raw.trim_offsets,
        })
    }
}

impl ByteLevel {
    /// Pre-tokenize in place using byte-level encoding.
    pub fn pre_tokenize(&self, pts: &mut PreTokenizedString) -> Result<(), Error> {
        if self.regex.is_none() && !self.add_prefix_space {
            return self.pre_tokenize_bulk(pts);
        }

        let old_buf = pts.buffer.as_str();
        let old_splits = &pts.splits;
        let mut new_buf = String::with_capacity(old_buf.len().saturating_mul(2));
        let mut new_splits = Vec::with_capacity(old_splits.len().saturating_mul(4));

        for split in &pts.splits {
            let text = pts.split_text(split);

            if split.token_id.is_some() {
                let start = new_buf.len();
                encode_bytes_into(text, &mut new_buf);
                let end = new_buf.len();
                new_splits.push(PtSplit {
                    range: start..end,
                    token_id: split.token_id,
                });
                continue;
            }

            if text.is_empty() {
                continue;
            }

            let prefixed;
            let text = if self.add_prefix_space && !text.starts_with(' ') {
                prefixed = format!(" {text}");
                prefixed.as_str()
            } else {
                text
            };

            match &self.regex {
                Some(re) => {
                    for m in re.find_iter(text) {
                        let m = m?;
                        if m.start() < m.end() {
                            let start = new_buf.len();
                            encode_bytes_into(&text[m.start()..m.end()], &mut new_buf);
                            let end = new_buf.len();
                            new_splits.push(PtSplit {
                                range: start..end,
                                token_id: None,
                            });
                        }
                    }
                }
                None => {
                    let start = new_buf.len();
                    encode_bytes_into(text, &mut new_buf);
                    let end = new_buf.len();
                    if start < end {
                        new_splits.push(PtSplit {
                            range: start..end,
                            token_id: None,
                        });
                    }
                }
            }
        }

        pts.buffer = new_buf;
        pts.splits = new_splits;
        Ok(())
    }

    /// Encode the buffer and remap split ranges without per-split overhead.
    fn pre_tokenize_bulk(&self, pts: &mut PreTokenizedString) -> Result<(), Error> {
        let old_buf = pts.buffer.as_str();
        let old_splits = &pts.splits;
        let mut new_buf = String::with_capacity(old_buf.len() << 1);
        let mut new_splits = Vec::with_capacity(old_splits.len());

        for split in old_splits {
            let text = pts.split_text(split);
            if text.is_empty() && split.token_id.is_none() {
                continue;
            }
            let start = new_buf.len();
            unsafe { encode_bytes_into_unchecked(text, &mut new_buf) };
            let end = new_buf.len();
            new_splits.push(PtSplit {
                range: start..end,
                token_id: split.token_id,
            });
        }

        pts.buffer = new_buf;
        pts.splits = new_splits;
        Ok(())
    }
}