harmony_protocol/tiktoken_ext/

public_encodings.rs

use std::{
    collections::HashMap,
    fs::File,
    io::{BufReader, BufWriter, Read as _, Write as _},
    path::{Path, PathBuf},
    sync::OnceLock,
};

use base64::{prelude::BASE64_STANDARD, Engine as _};

use crate::tiktoken::{CoreBPE, Rank};
use sha1::Sha1;
use sha2::{Digest as _, Sha256};

#[derive(Debug, thiserror::Error)]
pub enum LoadError {
    #[error("the env var TIKTOKEN_ENCODINGS_BASE is not set, or invalid")]
    InvalidEncodingBaseDirEnvVar,

    #[error("unknown encoding name: {0}")]
    UnknownEncodingName(String),

    #[error("invalid tiktoken vocab file: {0}")]
    InvalidTiktokenVocabFile(#[source] std::io::Error),

    #[error("failed to create CoreBPE: {0}")]
    CoreBPECreationFailed(#[source] Box<dyn std::error::Error + Send + Sync>),

    #[error("error downloading or loading vocab file: {0}")]
    DownloadOrLoadVocabFile(
        #[source]
        #[from]
        RemoteVocabFileError,
    ),
}

#[derive(Debug, thiserror::Error)]
pub enum RemoteVocabFileError {
    #[error("failed to download or load vocab file")]
    FailedToDownloadOrLoadVocabFile(#[source] Box<dyn std::error::Error + Send + Sync>),

    #[error("an underlying IO error occurred while {0}: {1}")]
    IOError(String, #[source] std::io::Error),

    #[error("hash mismatch for remote file {file_url}")]
    HashMismatch {
        file_url: String,
        expected_hash: String,
        computed_hash: String,
    },
}

const TIKTOKEN_ENCODINGS_BASE_VAR: &str = "TIKTOKEN_ENCODINGS_BASE";
const DEFAULT_TIKTOKEN_BASE_URL: &str = "https://openaipublic.blob.core.windows.net/encodings/";

static TIKTOKEN_BASE_URL_OVERRIDE: OnceLock<String> = OnceLock::new();

pub fn set_tiktoken_base_url(base_url: impl Into<String>) {
    let mut base = base_url.into();
    if !base.ends_with('/') {
        base.push('/');
    }
    let _ = TIKTOKEN_BASE_URL_OVERRIDE.set(base);
}

fn tiktoken_base_url() -> &'static str {
    TIKTOKEN_BASE_URL_OVERRIDE
        .get()
        .map(|s| s.as_str())
        .unwrap_or(DEFAULT_TIKTOKEN_BASE_URL)
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Encoding {
    O200kBase,
    O200kHarmony,
    Cl100kBase,
}

impl Encoding {
    pub fn all() -> &'static [Self] {
        &[Self::O200kBase, Self::O200kHarmony, Self::Cl100kBase]
    }

    pub fn from_name(name: impl AsRef<str>) -> Option<Self> {
        let name_str = name.as_ref();
        for encoding in Self::all() {
            if encoding.name() == name_str {
                return Some(*encoding);
            }
        }
        None
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn load_from_name(name: impl AsRef<str>) -> Result<CoreBPE, LoadError> {
        let name = name.as_ref();
        Self::from_name(name)
            .ok_or_else(|| LoadError::UnknownEncodingName(name.to_string()))?
            .load()
    }

    #[cfg(target_arch = "wasm32")]
    pub async fn load_from_name(name: impl AsRef<str>) -> Result<CoreBPE, LoadError> {
        let name = name.as_ref();
        Self::from_name(name)
            .ok_or_else(|| LoadError::UnknownEncodingName(name.to_string()))?
            .load()
            .await
    }

    pub fn name(&self) -> &'static str {
        match self {
            Self::O200kBase => "o200k_base",
            Self::O200kHarmony => "o200k_harmony",
            Self::Cl100kBase => "cl100k_base",
        }
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn load(&self) -> Result<CoreBPE, LoadError> {
        let (vocab_file_path, check_hash) =
            if let Ok(base_dir) = std::env::var(TIKTOKEN_ENCODINGS_BASE_VAR) {
                (PathBuf::from(base_dir).join(self.vocab_file_name()), true)
            } else {
                let url = self.public_vocab_file_url();
                (
                    download_or_find_cached_file(&url, Some(self.expected_hash()))
                        .map_err(LoadError::DownloadOrLoadVocabFile)?,
                    false,
                )
            };

        match self {
            Self::O200kHarmony => {
                let mut specials: Vec<(String, Rank)> = self
                    .special_tokens()
                    .iter()
                    .map(|(s, r)| ((*s).to_string(), *r))
                    .collect();
                specials.extend((200014..=201088).map(|id| (format!("<|reserved_{id}|>"), id)));
                load_encoding_from_file(
                    vocab_file_path,
                    check_hash.then(|| self.expected_hash()),
                    specials,
                    &self.pattern(),
                )
            }
            Self::O200kBase => {
                let mut specials: Vec<(String, Rank)> = self
                    .special_tokens()
                    .iter()
                    .map(|(s, r)| ((*s).to_string(), *r))
                    .collect();
                specials.extend((199998..=201088).map(|id| (format!("<|reserved_{id}|>"), id)));
                load_encoding_from_file(
                    vocab_file_path,
                    check_hash.then(|| self.expected_hash()),
                    specials,
                    &self.pattern(),
                )
            }
            _ => load_encoding_from_file(
                vocab_file_path,
                check_hash.then(|| self.expected_hash()),
                self.special_tokens().iter().cloned(),
                &self.pattern(),
            ),
        }
    }

    #[cfg(target_arch = "wasm32")]
    pub async fn load(&self) -> Result<CoreBPE, LoadError> {
        let url = self.public_vocab_file_url();
        let vocab_bytes = download_or_find_cached_file_bytes(&url, Some(self.expected_hash()))
            .await
            .map_err(LoadError::DownloadOrLoadVocabFile)?;

        match self {
            Self::O200kHarmony => {
                let mut specials: Vec<(String, Rank)> = self
                    .special_tokens()
                    .iter()
                    .map(|(s, r)| ((*s).to_string(), *r))
                    .collect();
                specials.extend((200014..=201088).map(|id| (format!("<|reserved_{id}|>"), id)));
                load_encoding_from_bytes(&vocab_bytes, None, specials, &self.pattern())
            }
            Self::O200kBase => {
                let mut specials: Vec<(String, Rank)> = self
                    .special_tokens()
                    .iter()
                    .map(|(s, r)| ((*s).to_string(), *r))
                    .collect();
                specials.extend((199998..=201088).map(|id| (format!("<|reserved_{id}|>"), id)));
                load_encoding_from_bytes(&vocab_bytes, None, specials, &self.pattern())
            }
            _ => load_encoding_from_bytes(
                &vocab_bytes,
                None,
                self.special_tokens().iter().cloned(),
                &self.pattern(),
            ),
        }
    }

    fn public_vocab_file_url(&self) -> String {
        let base = tiktoken_base_url();
        match self {
            Self::O200kBase => format!("{base}o200k_base.tiktoken"),
            Self::O200kHarmony => format!("{base}o200k_base.tiktoken"),
            Self::Cl100kBase => format!("{base}cl100k_base.tiktoken"),
        }
    }

    fn vocab_file_name(&self) -> &'static str {
        match self {
            Self::O200kBase => "o200k_base.tiktoken",
            Self::O200kHarmony => "o200k_base.tiktoken",
            Self::Cl100kBase => "cl100k_base.tiktoken",
        }
    }

    fn expected_hash(&self) -> &'static str {
        match self {
            Self::O200kBase => "446a9538cb6c348e3516120d7c08b09f57c36495e2acfffe59a5bf8b0cfb1a2d",
            Self::O200kHarmony => {
                "446a9538cb6c348e3516120d7c08b09f57c36495e2acfffe59a5bf8b0cfb1a2d"
            }
            Self::Cl100kBase => "223921b76ee99bde995b7ff738513eef100fb51d18c93597a113bcffe865b2a7",
        }
    }

    fn special_tokens(&self) -> &'static [(&'static str, Rank)] {
        match self {
            Self::O200kBase => &[],
            Self::O200kHarmony => &[
                ("<|startoftext|>", 199998),
                ("<|endoftext|>", 199999),
                ("<|reserved_200000|>", 200000),
                ("<|reserved_200001|>", 200001),
                ("<|return|>", 200002),
                ("<|constrain|>", 200003),
                ("<|reserved_200004|>", 200004),
                ("<|channel|>", 200005),
                ("<|start|>", 200006),
                ("<|end|>", 200007),
                ("<|message|>", 200008),
                ("<|reserved_200009|>", 200009),
                ("<|reserved_200010|>", 200010),
                ("<|reserved_200011|>", 200011),
                ("<|call|>", 200012),
                ("<|reserved_200013|>", 200013),
            ],
            Self::Cl100kBase => &[
                ("<|endoftext|>", 100257),
                ("<|fim_prefix|>", 100258),
                ("<|fim_middle|>", 100259),
                ("<|fim_suffix|>", 100260),
                ("<|endofprompt|>", 100276),
            ],
        }
    }

    fn pattern(&self) -> String {
        match self {
            Self::O200kBase => {
                [
                    "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
                    "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
                    "\\p{N}{1,3}",
                    " ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*",
                    "\\s*[\\r\\n]+",
                    "\\s+(?!\\S)",
                    "\\s+",
                ].join("|")
            }
            Self::O200kHarmony => {
                [
                    "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
                    "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
                    "\\p{N}{1,3}",
                    " ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*",
                    "\\s*[\\r\\n]+",
                    "\\s+(?!\\S)",
                    "\\s+",
                ].join("|")
            }
            Self::Cl100kBase => {
                "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+".to_string()
            }
        }
    }
}

fn load_tiktoken_vocab<R>(
    mut reader: R,
    expected_hash: Option<&str>,
) -> std::result::Result<HashMap<Vec<u8>, Rank>, std::io::Error>
where
    R: std::io::BufRead,
{
    let mut hasher = expected_hash.map(|_| Sha256::new());
    let mut bpe_ranks = HashMap::default();
    let mut lin_no = 0;
    let mut line_buffer = String::new();
    while reader.read_line(&mut line_buffer)? > 0 {
        lin_no += 1;
        if let Some(hasher) = hasher.as_mut() {
            hasher.update(line_buffer.as_bytes());
        }
        let line = line_buffer.trim_end();
        let (token, rank) = line.split_once(' ').ok_or_else(|| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("expected token and rank, could not split on ' ' at line {lin_no}"),
            )
        })?;
        let bytes = BASE64_STANDARD.decode(token).map_err(|e| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("failed to decode base64 token at line {lin_no}: {e}",),
            )
        })?;
        let rank = rank.parse().map_err(|e| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("failed to parse rank at line {lin_no}: {e}"),
            )
        })?;
        bpe_ranks.insert(bytes, rank);
        line_buffer.clear();
    }
    if let Some(hasher) = hasher {
        let expected_hash = expected_hash.unwrap();
        let computed_hash = format!("{:x}", hasher.finalize());
        if computed_hash != expected_hash {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("hash mismatch: computed={computed_hash}, expected={expected_hash}"),
            ));
        }
    }
    Ok(bpe_ranks)
}

pub fn load_tiktoken_vocab_file<P>(
    path: P,
    expected_hash: Option<&str>,
) -> std::result::Result<HashMap<Vec<u8>, Rank>, std::io::Error>
where
    P: AsRef<Path>,
{
    let file = std::fs::File::open(path)?;
    let reader = std::io::BufReader::new(file);
    load_tiktoken_vocab(reader, expected_hash)
}

pub fn load_encoding_from_file<P, S, TS>(
    file_path: P,
    expected_hash: Option<&str>,
    special_tokens: S,
    pattern: &str,
) -> Result<CoreBPE, LoadError>
where
    P: AsRef<Path>,
    S: IntoIterator<Item = (TS, Rank)>,
    TS: Into<String>,
{
    let encoder = load_tiktoken_vocab_file(file_path, expected_hash)
        .map_err(LoadError::InvalidTiktokenVocabFile)?;
    CoreBPE::new(
        encoder,
        special_tokens.into_iter().map(|(k, v)| (k.into(), v)),
        pattern,
    )
    .map_err(LoadError::CoreBPECreationFailed)
}

#[cfg(target_arch = "wasm32")]
pub fn load_encoding_from_bytes<S, TS>(
    vocab_bytes: &[u8],
    expected_hash: Option<&str>,
    special_tokens: S,
    pattern: &str,
) -> Result<CoreBPE, LoadError>
where
    S: IntoIterator<Item = (TS, Rank)>,
    TS: Into<String>,
{
    let reader = std::io::BufReader::new(vocab_bytes);
    let encoder = load_tiktoken_vocab(reader, expected_hash)
        .map_err(LoadError::InvalidTiktokenVocabFile)?;
    CoreBPE::new(
        encoder,
        special_tokens.into_iter().map(|(k, v)| (k.into(), v)),
        pattern,
    )
    .map_err(LoadError::CoreBPECreationFailed)
}

#[cfg(not(target_arch = "wasm32"))]
fn download_or_find_cached_file(
    url: &str,
    expected_hash: Option<&str>,
) -> Result<PathBuf, RemoteVocabFileError> {
    let cache_dir = resolve_cache_dir()?;
    let cache_path = resolve_cache_path(&cache_dir, url);
    if cache_path.exists() {
        if verify_file_hash(&cache_path, expected_hash)? {
            return Ok(cache_path);
        }
        let _ = std::fs::remove_file(&cache_path);
    }
    let hash = load_remote_file(url, &cache_path)?;
    if let Some(expected_hash) = expected_hash {
        if hash != expected_hash {
            let _ = std::fs::remove_file(&cache_path);
            return Err(RemoteVocabFileError::HashMismatch {
                file_url: url.to_string(),
                expected_hash: expected_hash.to_string(),
                computed_hash: hash,
            });
        }
    }
    Ok(cache_path)
}

#[cfg(target_arch = "wasm32")]
async fn download_or_find_cached_file_bytes(
    url: &str,
    expected_hash: Option<&str>,
) -> Result<Vec<u8>, RemoteVocabFileError> {
    let bytes = load_remote_file_bytes(url).await?;
    if let Some(expected_hash) = expected_hash {
        let computed_hash = format!("{:x}", Sha256::digest(&bytes));
        if computed_hash != expected_hash {
            return Err(RemoteVocabFileError::HashMismatch {
                file_url: url.to_string(),
                expected_hash: expected_hash.to_string(),
                computed_hash,
            });
        }
    }
    Ok(bytes)
}

fn resolve_cache_dir() -> Result<PathBuf, RemoteVocabFileError> {
    let cache_dir_override = std::env::var("TIKTOKEN_RS_CACHE_DIR").ok();
    if let Some(cache_dir_override) = cache_dir_override {
        Ok(PathBuf::from(cache_dir_override))
    } else {
        let cache_dir = std::env::temp_dir().join("tiktoken-rs-cache");
        std::fs::create_dir_all(&cache_dir).map_err(|e| {
            RemoteVocabFileError::IOError(format!("creating cache dir {cache_dir:?}"), e)
        })?;
        Ok(cache_dir)
    }
}

fn resolve_cache_path(cache_dir: &Path, url: &str) -> PathBuf {
    let mut hasher = Sha1::new();
    hasher.update(url.as_bytes());
    let cache_key = format!("{:x}", hasher.finalize());
    cache_dir.join(cache_key)
}

fn verify_file_hash(
    file_path: &Path,
    expected_hash: Option<&str>,
) -> Result<bool, RemoteVocabFileError> {
    let Some(expected_hash) = expected_hash else {
        return Ok(true);
    };
    let file = File::open(file_path)
        .map_err(|e| RemoteVocabFileError::IOError(format!("opening file {file_path:?}"), e))?;
    let mut reader = BufReader::new(file);
    let mut hasher = Sha256::new();
    std::io::copy(&mut reader, &mut hasher).map_err(|e| {
        RemoteVocabFileError::IOError(format!("copying file {file_path:?} contents to hasher"), e)
    })?;
    let computed_hash = format!("{:x}", hasher.finalize());
    Ok(computed_hash == expected_hash)
}

#[cfg(not(target_arch = "wasm32"))]
fn load_remote_file(url: &str, destination: &Path) -> Result<String, RemoteVocabFileError> {
    let client = reqwest::blocking::Client::new();
    let mut response = client
        .get(url)
        .send()
        .and_then(|r| r.error_for_status())
        .map_err(|e| RemoteVocabFileError::FailedToDownloadOrLoadVocabFile(Box::new(e)))?;

    let file = File::create(destination)
        .map_err(|e| RemoteVocabFileError::IOError(format!("creating file {destination:?}"), e))?;
    let mut dest = BufWriter::new(file);
    let mut hasher = Sha256::new();
    let mut buffer = [0u8; 8192];
    loop {
        let bytes_read = response.read(&mut buffer).map_err(|e| {
            RemoteVocabFileError::IOError(format!("reading from response {url}"), e)
        })?;
        if bytes_read == 0 {
            break;
        }
        dest.write_all(&buffer[..bytes_read]).map_err(|e| {
            RemoteVocabFileError::IOError(format!("writing to file {destination:?}"), e)
        })?;
        hasher.update(&buffer[..bytes_read]);
    }
    Ok(format!("{:x}", hasher.finalize()))
}

#[cfg(target_arch = "wasm32")]
async fn load_remote_file_bytes(url: &str) -> Result<Vec<u8>, RemoteVocabFileError> {
    use reqwest::Client;

    let client = Client::new();
    let response = client
        .get(url)
        .send()
        .await
        .and_then(|r| r.error_for_status())
        .map_err(|e| RemoteVocabFileError::FailedToDownloadOrLoadVocabFile(Box::new(e)))?;
    let bytes = response
        .bytes()
        .await
        .map_err(|e| RemoteVocabFileError::FailedToDownloadOrLoadVocabFile(Box::new(e)))?;
    Ok(bytes.to_vec())
}