gigatoken 0.1.0

data/src/load_tokenizer/hub.rs

@@ -0,0 +1,344 @@
+//! Minimal HuggingFace Hub file download into the standard HF cache.
+//!
+//! Rust port of the former pure-Python `gigatoken._load.hub`: same endpoint
+//! and URL layout as `huggingface_hub.hf_hub_download`, same token discovery
+//! (HF_TOKEN env var, then the token file written by `hf auth login`), and
+//! the same cache directory resolution, without requiring huggingface_hub,
+//! tokenizers, or transformers. Files already present in the standard HF
+//! cache are served with a pure-filesystem lookup (no network); on a miss the
+//! file is downloaded straight into the shared cache — under the commit hash
+//! the Hub reports via `x-repo-commit`, with the branch ref recorded — so
+//! huggingface_hub and later lookups serve it from the same place.
+
+use eyre::{Context, Result, bail};
+use std::fmt;
+use std::io;
+use std::path::PathBuf;
+
+/// Filename suffixes of local tokenizer files (tokenizer.json contents and
+/// raw sentencepiece models — the formats the hf loader reads from disk).
+/// A name ending in one of these is never treated as a Hub repo id, so a
+/// mistyped local path fails fast instead of hitting the network. Keep in
+/// sync with `gigatoken._load.hub.TOKENIZER_FILE_SUFFIXES`.
+pub const TOKENIZER_FILE_SUFFIXES: &[&str] = &[".json", ".model"];
+
+/// Whether `name` is shaped like a HuggingFace Hub repo id: `org/name`, or a
+/// bare legacy repo name like `gpt2`. At most one slash, and not something
+/// that is obviously a filesystem path to a local tokenizer file.
+pub fn looks_like_repo_id(name: &str) -> bool {
+    let word = |c: char| c.is_ascii_alphanumeric() || matches!(c, '_' | '.' | '-');
+    let part_ok = |part: &str, first_alnum: bool| {
+        let mut chars = part.chars();
+        let Some(first) = chars.next() else {
+            return false;
+        };
+        let first_ok = if first_alnum { first.is_ascii_alphanumeric() } else { word(first) };
+        first_ok && chars.all(word)
+    };
+    let mut parts = name.split('/');
+    let (org, rest) = (parts.next().unwrap_or(""), parts.next());
+    parts.next().is_none()
+        && part_ok(org, true)
+        && rest.is_none_or(|r| part_ok(r, false))
+        && !TOKENIZER_FILE_SUFFIXES.iter().any(|s| name.ends_with(s))
+}
+
+fn env_nonempty(key: &str) -> Option<String> {
+    std::env::var(key).ok().filter(|v| !v.is_empty())
+}
+
+/// $HF_HOME, defaulting to $XDG_CACHE_HOME/huggingface then
+/// ~/.cache/huggingface — the root for both the hub cache and the token file.
+fn hf_home() -> PathBuf {
+    env_nonempty("HF_HOME").map(PathBuf::from).unwrap_or_else(|| {
+        env_nonempty("XDG_CACHE_HOME")
+            .map(PathBuf::from)
+            .unwrap_or_else(|| std::env::home_dir().expect("home dir").join(".cache"))
+            .join("huggingface")
+    })
+}
+
+/// The standard HuggingFace hub cache directory, resolved like
+/// huggingface_hub does it: HF_HUB_CACHE, then $HF_HOME/hub, then
+/// $XDG_CACHE_HOME/huggingface/hub, then ~/.cache/huggingface/hub.
+pub fn hf_hub_cache_dir() -> PathBuf {
+    env_nonempty("HF_HUB_CACHE")
+        .map(PathBuf::from)
+        .unwrap_or_else(|| hf_home().join("hub"))
+}
+
+/// The HuggingFace access token, discovered like huggingface_hub does it:
+/// the HF_TOKEN (or legacy HUGGING_FACE_HUB_TOKEN) environment variable,
+/// then the token file (HF_TOKEN_PATH, default $HF_HOME/token).
+pub fn get_hf_token() -> Option<String> {
+    if let Some(token) = env_nonempty("HF_TOKEN")
+        .or_else(|| env_nonempty("HUGGING_FACE_HUB_TOKEN"))
+        .map(|t| t.trim().to_owned())
+        .filter(|t| !t.is_empty())
+    {
+        return Some(token);
+    }
+    let token_path = env_nonempty("HF_TOKEN_PATH")
+        .map(PathBuf::from)
+        .unwrap_or_else(|| hf_home().join("token"));
+    let token = std::fs::read_to_string(token_path).ok()?;
+    let token = token.trim();
+    (!token.is_empty()).then(|| token.to_owned())
+}
+
+/// A full git commit hash: cache snapshot directories are named by these.
+fn is_commit_hash(revision: &str) -> bool {
+    revision.len() == 40 && revision.bytes().all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase())
+}
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub enum RepoType {
+    #[default]
+    Model,
+    Dataset,
+    Space,
+}
+
+impl RepoType {
+    /// The cache directory prefix (`models--org--name` etc.).
+    fn cache_prefix(self) -> &'static str {
+        match self {
+            RepoType::Model => "models",
+            RepoType::Dataset => "datasets",
+            RepoType::Space => "spaces",
+        }
+    }
+
+    /// The URL path prefix before the repo id.
+    fn url_prefix(self) -> &'static str {
+        match self {
+            RepoType::Model => "",
+            RepoType::Dataset => "datasets/",
+            RepoType::Space => "spaces/",
+        }
+    }
+}
+
+impl std::str::FromStr for RepoType {
+    type Err = eyre::Report;
+
+    fn from_str(s: &str) -> Result<Self> {
+        match s {
+            "model" => Ok(RepoType::Model),
+            "dataset" => Ok(RepoType::Dataset),
+            "space" => Ok(RepoType::Space),
+            _ => Err(eyre::eyre!("unknown repo_type {s:?}: expected \"model\", \"dataset\", or \"space\"")),
+        }
+    }
+}
+
+/// The cache directory of a repo (`models--org--name` etc.).
+fn repo_cache_dir(repo_type: RepoType, repo_id: &str) -> PathBuf {
+    hf_hub_cache_dir().join(format!("{}--{}", repo_type.cache_prefix(), repo_id.replace('/', "--")))
+}
+
+/// Model-repo [`cached_hub_file_in`].
+pub fn cached_hub_file(repo_id: &str, filename: &str, revision: &str) -> Option<PathBuf> {
+    cached_hub_file_in(RepoType::Model, repo_id, filename, revision)
+}
+
+/// Path of `filename` in the local HF cache, or None when not cached.
+///
+/// A pure-filesystem lookup — no request is made. `revision` may be a commit
+/// hash (used directly as the snapshot name) or a branch/tag name (followed
+/// through the cached ref).
+pub fn cached_hub_file_in(repo_type: RepoType, repo_id: &str, filename: &str, revision: &str) -> Option<PathBuf> {
+    let repo_dir = repo_cache_dir(repo_type, repo_id);
+    let commit_owned;
+    let commit = if is_commit_hash(revision) {
+        revision
+    } else {
+        commit_owned = std::fs::read_to_string(repo_dir.join("refs").join(revision)).ok()?;
+        commit_owned.trim()
+    };
+    let path = repo_dir.join("snapshots").join(commit).join(filename);
+    path.is_file().then_some(path)
+}
+
+/// Download failure with a definite HTTP cause, kept as a typed error so the
+/// Python bindings can raise the matching exception (FileNotFoundError for
+/// 404, PermissionError for 401/403).
+#[derive(Debug)]
+pub enum FetchError {
+    /// 404: no such repo, revision, or file.
+    NotFound { url: String },
+    /// 401/403: private or gated repo.
+    Unauthorized { url: String, status: u16, had_token: bool },
+}
+
+impl fmt::Display for FetchError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            FetchError::NotFound { url } => {
+                write!(f, "{url}: HTTP 404 — no such repo with that file, and no such local file either")
+            }
+            FetchError::Unauthorized { url, status, had_token } => {
+                let token = if *had_token {
+                    "the request used the discovered token"
+                } else {
+                    "no token was found"
+                };
+                write!(
+                    f,
+                    "{url}: HTTP {status} — the repo may be private or gated ({token}; set HF_TOKEN or run \
+                     `hf auth login`, and accept the repo's terms on huggingface.co if it is gated)"
+                )
+            }
+        }
+    }
+}
+
+impl std::error::Error for FetchError {}
+
+/// Model-repo [`hub_file_in`].
+pub fn hub_file(repo_id: &str, filename: &str, revision: &str) -> Result<PathBuf> {
+    hub_file_in(RepoType::Model, repo_id, filename, revision)
+}
+
+/// Path of `filename` from Hub repo `repo_id` at `revision`, served from the
+/// standard HF cache, downloading into it first when absent.
+pub fn hub_file_in(repo_type: RepoType, repo_id: &str, filename: &str, revision: &str) -> Result<PathBuf> {
+    if let Some(path) = cached_hub_file_in(repo_type, repo_id, filename, revision) {
+        return Ok(path);
+    }
+    download_into_cache(repo_type, repo_id, filename, revision)
+        .wrap_err_with(|| format!("downloading {filename} from Hub repo {repo_id} at revision {revision}"))
+}
+
+/// GET `endpoint/repo/resolve/revision/filename` and stream the body into the
+/// cache snapshot named by the `x-repo-commit` response header, recording the
+/// branch ref so later lookups (ours and huggingface_hub's) resolve it.
+fn download_into_cache(repo_type: RepoType, repo_id: &str, filename: &str, revision: &str) -> Result<PathBuf> {
+    let endpoint = env_nonempty("HF_ENDPOINT").unwrap_or_else(|| "https://huggingface.co".to_owned());
+    let url = format!(
+        "{}/{}{repo_id}/resolve/{revision}/{filename}",
+        endpoint.trim_end_matches('/'),
+        repo_type.url_prefix()
+    );
+    let token = get_hf_token();
+
+    // Redirects are followed by hand: resolve/ URLs answer with the
+    // `x-repo-commit` header and a redirect to a CDN for LFS files, and the
+    // Authorization header must not travel to the other host.
+    let agent = ureq::Agent::config_builder()
+        .max_redirects(0)
+        .http_status_as_error(false)
+        .build()
+        .new_agent();
+    let mut request = agent.get(&url).header("User-Agent", "gigatoken");
+    if let Some(token) = &token {
+        request = request.header("Authorization", format!("Bearer {token}"));
+    }
+    let mut response = request.call().wrap_err_with(|| format!("requesting {url}"))?;
+    let header = |resp: &ureq::http::Response<ureq::Body>, name: &str| {
+        resp.headers().get(name).and_then(|v| v.to_str().ok()).map(str::to_owned)
+    };
+    let commit = header(&response, "x-repo-commit");
+
+    let mut hops = 0;
+    while response.status().is_redirection() {
+        hops += 1;
+        ensure_status(&url, response.status().as_u16(), token.is_some())?;
+        if hops > 10 {
+            bail!("{url}: too many redirects");
+        }
+        let location = header(&response, "location")
+            .ok_or_else(|| eyre::eyre!("{url}: redirect with no Location header"))?;
+        let next_url = absolutize(&location, &url);
+        // No Authorization here: the redirect target is a presigned CDN URL
+        // on another host (requests/huggingface_hub drop the header too).
+        response = agent
+            .get(&next_url)
+            .header("User-Agent", "gigatoken")
+            .call()
+            .wrap_err_with(|| format!("requesting {next_url}"))?;
+    }
+    ensure_status(&url, response.status().as_u16(), token.is_some())?;
+
+    // Snapshot directory: the commit the Hub reported, falling back to the
+    // requested revision (e.g. a plain file server behind HF_ENDPOINT).
+    let commit = commit.unwrap_or_else(|| revision.to_owned());
+    let repo_dir = repo_cache_dir(repo_type, repo_id);
+    let target = repo_dir.join("snapshots").join(&commit).join(filename);
+    let dir = target.parent().expect("snapshot file has a parent");
+    std::fs::create_dir_all(dir).wrap_err_with(|| format!("creating {}", dir.display()))?;
+
+    // Stream to a sibling temp file, then rename: concurrent downloaders
+    // race benignly and readers never observe a partial file.
+    let tmp = dir.join(format!(".{}.{}.tmp", target.file_name().unwrap().to_string_lossy(), std::process::id()));
+    let result = (|| -> Result<()> {
+        let mut file = std::fs::File::create(&tmp)?;
+        io::copy(&mut response.body_mut().as_reader(), &mut file)?;
+        file.sync_all()?;
+        std::fs::rename(&tmp, &target)?;
+        Ok(())
+    })();
+    if result.is_err() {
+        let _ = std::fs::remove_file(&tmp);
+    }
+    result.wrap_err_with(|| format!("writing {}", target.display()))?;
+
+    if !is_commit_hash(revision) && revision != commit {
+        let refs_dir = repo_dir.join("refs");
+        std::fs::create_dir_all(&refs_dir)
+            .and_then(|_| std::fs::write(refs_dir.join(revision), &commit))
+            .wrap_err_with(|| format!("recording ref {revision} -> {commit}"))?;
+    }
+    Ok(target)
+}
+
+fn ensure_status(url: &str, status: u16, had_token: bool) -> Result<()> {
+    match status {
+        200..=399 => Ok(()),
+        404 => Err(FetchError::NotFound { url: url.to_owned() }.into()),
+        401 | 403 => Err(FetchError::Unauthorized { url: url.to_owned(), status, had_token }.into()),
+        _ => Err(eyre::eyre!("{url}: HTTP {status}")),
+    }
+}
+
+/// A redirect Location resolved against the request URL: absolute URLs pass
+/// through, host-relative (`/x/y`) and path-relative ones join the base.
+fn absolutize(location: &str, base: &str) -> String {
+    if location.contains("://") {
+        return location.to_owned();
+    }
+    let origin_len = base.find("://").map(|i| i + 3).unwrap_or(0);
+    let origin_end = base[origin_len..].find('/').map_or(base.len(), |i| origin_len + i);
+    if location.starts_with('/') {
+        format!("{}{location}", &base[..origin_end])
+    } else {
+        let dir_end = base.rfind('/').unwrap_or(base.len());
+        format!("{}/{location}", &base[..dir_end.max(origin_end)])
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn repo_id_shapes() {
+        assert!(looks_like_repo_id("gpt2"));
+        assert!(looks_like_repo_id("openai-community/gpt2"));
+        assert!(looks_like_repo_id("Qwen/Qwen3.5-9B"));
+        assert!(!looks_like_repo_id("data/tokenizers/gpt2.json"));
+        assert!(!looks_like_repo_id("./gpt2"));
+        assert!(!looks_like_repo_id("/abs/path"));
+        assert!(!looks_like_repo_id("gpt2_tokenizer.json"));
+        assert!(!looks_like_repo_id("subdir/tokenizer.model"));
+        assert!(!looks_like_repo_id(""));
+        assert!(!looks_like_repo_id("org/"));
+    }
+
+    #[test]
+    fn location_resolution() {
+        assert_eq!(absolutize("https://cdn.example/x", "https://huggingface.co/a/b"), "https://cdn.example/x");
+        assert_eq!(absolutize("/api/x", "https://huggingface.co/a/b"), "https://huggingface.co/api/x");
+        assert_eq!(absolutize("y", "https://huggingface.co/a/b"), "https://huggingface.co/a/y");
+    }
+}

data/src/load_tokenizer/mod.rs

@@ -0,0 +1,3 @@
+pub mod hf;
+pub mod hub;
+pub mod tiktoken;

data/src/load_tokenizer/tiktoken.rs

@@ -0,0 +1,87 @@
+use crate::bpe::Tokenizer;
+use crate::pretokenize::PretokenizerType;
+use eyre::{Context, Result, ensure};
+use std::path::Path;
+
+/// The base64-per-line mergeable ranks of a .tiktoken/tiktoken.model file,
+/// in rank order (merges are reconstructed from this list).
+fn load_tiktoken_ranks(file_path: impl AsRef<Path>) -> Result<Vec<Vec<u8>>> {
+    use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
+    use base64::prelude::*;
+    use std::io::Read;
+    let mut buf = String::new();
+    std::fs::File::open(&file_path)
+        .with_context(|| format!("Failed to read {}", file_path.as_ref().display()))?
+        .read_to_string(&mut buf)?;
+    buf.lines()
+        .enumerate()
+        .map(|(i, line)| {
+            let (base64_token, id_str) = line
+                .split_once(' ')
+                .ok_or_else(|| eyre::eyre!("line {i} has no rank field"))?;
+            let id = id_str.trim().parse::<u32>()?;
+            ensure!(id == i as u32, "rank {id} at line {i}: ranks must be dense");
+            Ok(BASE64_STANDARD.decode(base64_token)?)
+        })
+        .collect()
+}
+
+pub fn load_tiktoken(file_path: impl AsRef<Path>) -> Result<Tokenizer> {
+    let rank_vocab = load_tiktoken_ranks(file_path)?;
+    let n_ranks = rank_vocab.len() as u32;
+    let mut tokenizer = Tokenizer::from_ranks(rank_vocab)?;
+    // Tiktoken vocab files carry no special tokens; GPT-2-family vocabs
+    // (gpt2/r50k) place <|endoftext|> at the id right after the mergeable
+    // ranks. Register it so tiktoken- and tokenizer.json-loaded tokenizers
+    // encode and decode identically.
+    tokenizer.add_special_token(b"<|endoftext|>".to_vec(), n_ranks.into());
+    Ok(tokenizer)
+}
+
+/// The fields of a HuggingFace tokenizer_config.json a tiktoken-rank repo
+/// needs: the special tokens (there is no tokenizer.json to carry them).
+#[derive(serde::Deserialize)]
+struct TokenizerConfigJson {
+    #[serde(default)]
+    added_tokens_decoder: std::collections::BTreeMap<String, AddedTokenJson>,
+}
+
+#[derive(serde::Deserialize)]
+struct AddedTokenJson {
+    content: String,
+}
+
+/// Load a tokenizer from a tiktoken rank file plus a HuggingFace
+/// `tokenizer_config.json` whose `added_tokens_decoder` carries the special
+/// tokens — the layout of repos that ship no tokenizer.json (e.g. the
+/// moonshotai Kimi/Moonlight line's `tiktoken.model`). The pretokenizer
+/// scheme is passed by the caller: such repos define their split regex only
+/// in remote code, so no shipped file can name it.
+pub fn load_tiktoken_model(
+    model_path: impl AsRef<Path>,
+    config_path: impl AsRef<Path>,
+    pretokenizer: PretokenizerType,
+) -> Result<Tokenizer> {
+    let rank_vocab = load_tiktoken_ranks(model_path)?;
+    let n_ranks = rank_vocab.len() as u32;
+    let mut tokenizer = Tokenizer::from_ranks(rank_vocab)?;
+    tokenizer.set_pretokenizer_type(pretokenizer);
+    let config_path = config_path.as_ref();
+    let config: TokenizerConfigJson = sonic_rs::from_slice(
+        &std::fs::read(config_path)
+            .with_context(|| format!("Failed to read {}", config_path.display()))?,
+    )
+    .map_err(|e| eyre::eyre!("Failed to parse {}: {e}", config_path.display()))?;
+    for (id, token) in &config.added_tokens_decoder {
+        let id = id.parse::<u32>().with_context(|| {
+            format!("added_tokens_decoder id {id:?} in {}", config_path.display())
+        })?;
+        ensure!(
+            id >= n_ranks,
+            "added token {:?} (id {id}) overlaps the {n_ranks} mergeable ranks",
+            token.content
+        );
+        tokenizer.add_special_token(token.content.clone().into_bytes(), id.into());
+    }
+    Ok(tokenizer)
+}

data/src/main.rs

@@ -0,0 +1,95 @@
+#![feature(portable_simd)]
+
+mod bpe;
+mod bpe_train;
+mod input;
+mod load_tokenizer;
+mod pretokenize;
+#[cfg(test)]
+mod test_hub;
+mod token;
+
+use input::MmappedFile;
+use input::Resource;
+use input::jsonl::JsonLinesSlice;
+use std::path::PathBuf;
+
+#[allow(deprecated)]
+pub fn main() {
+    let args: Vec<String> = std::env::args().collect();
+
+    let tokenizer_path = args.get(1).map(PathBuf::from).unwrap_or_else(|| {
+        PathBuf::from(concat!(
+            env!("CARGO_MANIFEST_DIR"),
+            "/tests/scripts/tinyllama_tokenizer.json"
+        ))
+    });
+
+    let input_path = args.get(2).map(PathBuf::from).unwrap_or_else(|| {
+        std::env::home_dir()
+            .unwrap()
+            .join("data/dclm-baseline/shard_00000000_processed.jsonl")
+    });
+
+    let field = args.get(3).map(|s| s.as_str()).unwrap_or("text");
+
+    eprintln!("Tokenizer:  {}", tokenizer_path.display());
+    eprintln!("Input:      {}", input_path.display());
+    eprintln!("Field:      {}", field);
+
+    // Load tokenizer
+    let tokenizer = load_tokenizer::hf::load_hf_sentencepiece(&tokenizer_path)
+        .expect("Failed to load tokenizer");
+    eprintln!("Loaded tokenizer: {:?}", tokenizer);
+
+    // Mmap input file
+    let mmap = MmappedFile::open(&input_path).expect("Failed to open input file");
+    let bytes = mmap.as_bytes();
+    eprintln!("Input size: {:.1} MB", bytes.len() as f64 / 1e6);
+
+    // Parse JSONL to get documents
+    let start = std::time::Instant::now();
+    let docs: Vec<_> = JsonLinesSlice::new(bytes, field)
+        .map(|doc| {
+            let b = doc.as_ref();
+            // Convert to string for encoding
+            unsafe { std::str::from_utf8_unchecked(b) }.to_string()
+        })
+        .collect();
+    let parse_time = start.elapsed();
+    let total_chars: usize = docs.iter().map(|d| d.len()).sum();
+    let total_text_mb = total_chars as f64 / 1e6;
+    eprintln!(
+        "Parsed {} docs ({:.1} MB text) in {:.2}s",
+        docs.len(),
+        total_text_mb,
+        parse_time.as_secs_f64()
+    );
+    eprintln!();
+
+    // Load r50k tokenizer for GPT-2 style encoding
+    let r50k_path = args.get(4).map(PathBuf::from).unwrap_or_else(|| {
+        std::env::home_dir()
+            .unwrap()
+            .join("data/tokenizers/r50k_base.tiktoken")
+    });
+    let mut r50k =
+        load_tokenizer::tiktoken::load_tiktoken(&r50k_path).expect("Failed to load r50k tokenizer");
+    eprintln!("Loaded r50k: {:?}", r50k);
+
+    // Encode with r50k (GPT-2 style: pretokenize + memoized BPE)
+    let start = std::time::Instant::now();
+    let mut total_tokens: usize = 0;
+    for doc in &docs {
+        r50k.memoized_encode(pretokenize::pretokenize_as_iter(doc.as_bytes()), |tokens| {
+            total_tokens += tokens.len();
+        });
+    }
+    let elapsed = start.elapsed();
+    eprintln!(
+        "r50k encode: {} tokens in {:.2}s ({:.1} MB/s)",
+        total_tokens,
+        elapsed.as_secs_f64(),
+        total_text_mb / elapsed.as_secs_f64(),
+    );
+}