gigatoken 0.1.0

data/ext/gigatoken/src/sentencepiece.rs

@@ -0,0 +1,205 @@
+//! `Gigatoken::Native::SentencePieceTokenizer`: a `gigatoken::SentencePieceBPE`
+//! plus its persistent pretoken-cache `EncodeState`, and no `WorkerPool` —
+//! the `sp_*` cores parallelize internally instead of through a pool.
+//! Mirrors the pyo3 `SentencePieceTokenizer` in the core crate's `src/lib.rs`
+//! (the `python` feature), minus the numpy/awkward-array machinery that has
+//! no Ruby analog; marshaling helpers are shared with
+//! `Gigatoken::Native::BPETokenizer` (see `crate::tokenizer`).
+//!
+//! Unlike the pyo3 binding, which trusts Python's `str` guarantee, every
+//! document/file byte region handed in from Ruby is validated as UTF-8
+//! before it reaches the `&str`-typed SentencePiece cores — a Ruby String
+//! carries no such guarantee even when tagged UTF-8 — raising
+//! `Gigatoken::Error` instead of ever calling `str::from_utf8_unchecked` on
+//! Ruby-supplied bytes.
+
+use std::cell::RefCell;
+
+use gigatoken_rs::input::file_source::DocFormat;
+use gigatoken_rs::{EncodeState, SentencePieceBPE, sp_encode_docs_ragged, sp_encode_files_docs, sp_encode_files_docs_serial};
+use magnus::{
+    Error, RArray, RClass, RHash, RModule, RString, Ruby, Value, method, prelude::*,
+    scan_args::{get_kwargs, scan_args},
+};
+
+use crate::error::raise;
+use crate::gvl::without_gvl;
+use crate::sources;
+use crate::tokenizer::{binary_string, packed_result, ragged_result};
+
+/// Validate that `bytes` is UTF-8, raising `Gigatoken::Error` otherwise.
+fn require_utf8<'a>(ruby: &Ruby, bytes: &'a [u8]) -> Result<&'a str, Error> {
+    std::str::from_utf8(bytes).map_err(|e| raise(ruby, format!("invalid UTF-8: {e}")))
+}
+
+#[magnus::wrap(class = "Gigatoken::Native::SentencePieceTokenizer", free_immediately, size)]
+pub struct SentencePieceTokenizer {
+    // `SentencePieceBPE`'s encode methods take `&self` (only `EncodeState`
+    // is mutated), so this `RefCell` is never `borrow_mut`'d — see the
+    // builder report's DISAGREEMENTS for why it's here anyway.
+    tokenizer: RefCell<SentencePieceBPE>,
+    state: RefCell<EncodeState>,
+}
+
+impl SentencePieceTokenizer {
+    pub(crate) fn from_tokenizer(tokenizer: SentencePieceBPE) -> Self {
+        Self {
+            tokenizer: RefCell::new(tokenizer),
+            state: RefCell::new(EncodeState::new()),
+        }
+    }
+
+    fn encode(ruby: &Ruby, rb_self: &Self, input: RString) -> Result<Vec<u32>, Error> {
+        // SAFETY: read-only, for the duration of this synchronous call.
+        let bytes = unsafe { input.as_slice() };
+        let text = require_utf8(ruby, bytes)?;
+        let mut ids: Vec<u32> = Vec::new();
+        let mut state = rb_self.state.borrow_mut();
+        rb_self.tokenizer.borrow().encode_raw_cb(&mut state, text, &mut |tokens| {
+            ids.extend(tokens.iter().map(|&t| u32::from(t)))
+        });
+        Ok(ids)
+    }
+
+    /// Encode a batch on the core's internal (rayon) parallelism, with the
+    /// GVL released — like `BPETokenizer#encode_batch`, always parallel
+    /// (only `encode_files` exposes a `parallel:` toggle). Every input
+    /// string is validated as UTF-8 and copied into an owned `String`
+    /// before release: nothing Ruby-managed may be touched once the GVL is
+    /// gone.
+    fn encode_batch_ragged(ruby: &Ruby, rb_self: &Self, inputs: RArray) -> Result<(Vec<u32>, Vec<i64>), Error> {
+        // SAFETY: values are read (checked-converted to `RString`, then
+        // validated and copied into owned buffers below) before anything
+        // else runs.
+        let docs: Vec<String> = unsafe { inputs.as_slice() }
+            .iter()
+            .map(|&v| {
+                let s = RString::try_convert(v)?;
+                let bytes = unsafe { s.as_slice() };
+                require_utf8(ruby, bytes).map(str::to_owned)
+            })
+            .collect::<Result<_, _>>()?;
+        let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
+        let tokenizer = rb_self.tokenizer.borrow();
+        let tokenizer: &SentencePieceBPE = &tokenizer;
+        Ok(without_gvl(|| sp_encode_docs_ragged(tokenizer, &doc_refs)))
+    }
+
+    fn encode_batch(ruby: &Ruby, rb_self: &Self, inputs: RArray) -> Result<RArray, Error> {
+        let (flat, lens) = Self::encode_batch_ragged(ruby, rb_self, inputs)?;
+        ragged_result(ruby, flat, lens)
+    }
+
+    /// The packed analog of `encode_batch`: same encode core, but the flat
+    /// token ids are handed to Ruby as one `IO::Buffer` instead of being
+    /// re-chunked into per-document Arrays (see `tokenizer::packed_result`).
+    fn encode_batch_packed(ruby: &Ruby, rb_self: &Self, inputs: RArray) -> Result<RArray, Error> {
+        let (flat, lens) = Self::encode_batch_ragged(ruby, rb_self, inputs)?;
+        packed_result(ruby, flat, lens)
+    }
+
+    /// Encode every document named by `source` in parallel, with the GVL
+    /// released for the whole run — loading the files and the encode
+    /// itself. A `separator:` on a `TextFileSource` must be valid UTF-8
+    /// (checked up front, like the pyo3 binding's own separator check):
+    /// an arbitrary byte separator could cut a document mid-character and
+    /// break the SentencePiece cores' `&str` contract. `parallel: false`
+    /// loads and encodes everything on the calling thread instead, with
+    /// identical output, never touching rayon.
+    fn encode_files_ragged(ruby: &Ruby, rb_self: &Self, args: &[Value]) -> Result<(Vec<u32>, Vec<i64>), Error> {
+        let args = scan_args::<(Value,), (), (), (), RHash, ()>(args)?;
+        let (source,) = args.required;
+        let kw = get_kwargs::<_, (), (Option<Value>,), ()>(args.keywords, &[], &["parallel"])?;
+        let (parallel,) = kw.optional;
+        let parallel = match parallel {
+            Some(v) if !v.is_nil() => bool::try_convert(v)?,
+            _ => true,
+        };
+
+        let source = sources::resolve(ruby, source)?;
+        if let DocFormat::Text { separator: Some(sep) } = &source.format {
+            if std::str::from_utf8(sep).is_err() {
+                return Err(raise(
+                    ruby,
+                    "the SentencePiece backend requires a separator that is valid UTF-8",
+                ));
+            }
+        }
+
+        let tokenizer = rb_self.tokenizer.borrow();
+        let tokenizer: &SentencePieceBPE = &tokenizer;
+        let encoded: std::io::Result<(Vec<u32>, Vec<i64>)> = without_gvl(|| {
+            sources::encode_files_ragged(&source, parallel, |files, format| {
+                for &region in files {
+                    std::str::from_utf8(region).map_err(|e| {
+                        std::io::Error::new(std::io::ErrorKind::InvalidData, format!("invalid UTF-8 in file contents: {e}"))
+                    })?;
+                }
+                Ok(if parallel {
+                    sp_encode_files_docs(tokenizer, files, format)
+                } else {
+                    sp_encode_files_docs_serial(tokenizer, files, format)
+                })
+            })
+        });
+        encoded.map_err(|e| raise(ruby, e.to_string()))
+    }
+
+    fn encode_files(ruby: &Ruby, rb_self: &Self, args: &[Value]) -> Result<RArray, Error> {
+        let (flat, lens) = Self::encode_files_ragged(ruby, rb_self, args)?;
+        ragged_result(ruby, flat, lens)
+    }
+
+    /// The packed analog of `encode_files`: same encode core (including the
+    /// `parallel:` nil-equals-omitted contract), but the flat token ids are
+    /// handed to Ruby as one `IO::Buffer` instead of being re-chunked into
+    /// per-document Arrays (see `tokenizer::packed_result`).
+    fn encode_files_packed(ruby: &Ruby, rb_self: &Self, args: &[Value]) -> Result<RArray, Error> {
+        let (flat, lens) = Self::encode_files_ragged(ruby, rb_self, args)?;
+        packed_result(ruby, flat, lens)
+    }
+
+    fn decode(ruby: &Ruby, rb_self: &Self, tokens: RArray) -> Result<RString, Error> {
+        let ids: Vec<u32> = tokens.to_vec()?;
+        let ids: Vec<_> = ids.into_iter().map(Into::into).collect();
+        let bytes = rb_self.tokenizer.borrow().decode(&ids);
+        Ok(binary_string(ruby, &bytes))
+    }
+
+    fn vocab_size(&self) -> usize {
+        self.tokenizer.borrow().vocab_size()
+    }
+
+    fn vocab(ruby: &Ruby, rb_self: &Self) -> Result<RHash, Error> {
+        let tokenizer = rb_self.tokenizer.borrow();
+        let hash = ruby.hash_new();
+        for (id, bytes) in tokenizer.vocab_entries() {
+            hash.aset(id, binary_string(ruby, bytes))?;
+        }
+        Ok(hash)
+    }
+
+    fn merges(ruby: &Ruby, rb_self: &Self) -> Result<RArray, Error> {
+        let tokenizer = rb_self.tokenizer.borrow();
+        let entries = tokenizer.merge_entries();
+        let result = ruby.ary_new_capa(entries.len());
+        for (a, b) in entries {
+            result.push((binary_string(ruby, a), binary_string(ruby, b)))?;
+        }
+        Ok(result)
+    }
+}
+
+pub fn init(ruby: &Ruby, native: RModule) -> Result<(), Error> {
+    let class: RClass = native.define_class("SentencePieceTokenizer", ruby.class_object())?;
+    class.define_method("encode", method!(SentencePieceTokenizer::encode, 1))?;
+    class.define_method("encode_batch", method!(SentencePieceTokenizer::encode_batch, 1))?;
+    class.define_method("encode_batch_packed", method!(SentencePieceTokenizer::encode_batch_packed, 1))?;
+    class.define_method("encode_files", method!(SentencePieceTokenizer::encode_files, -1))?;
+    class.define_method("encode_files_packed", method!(SentencePieceTokenizer::encode_files_packed, -1))?;
+    class.define_method("decode", method!(SentencePieceTokenizer::decode, 1))?;
+    class.define_method("vocab_size", method!(SentencePieceTokenizer::vocab_size, 0))?;
+    class.define_method("vocab", method!(SentencePieceTokenizer::vocab, 0))?;
+    class.define_method("merges", method!(SentencePieceTokenizer::merges, 0))?;
+    Ok(())
+}

data/ext/gigatoken/src/sources.rs

@@ -0,0 +1,207 @@
+//! `Gigatoken::Native::{Text,Jsonl,Parquet}FileSource`: paths plus how to
+//! split their bytes into documents, for `BPETokenizer#encode_files`.
+//! Mirrors the pyo3 `FileSource` family (`src/bindings/sources.rs`) —
+//! same constructor options, same `DocFormat` payload — but each format
+//! gets its own Ruby class instead of a `FileSource` base class, since
+//! magnus ties one wrapped Rust type to exactly one Ruby class.
+//!
+//! `encode_files_ragged` below mirrors the pyo3 path's loading scaffold
+//! (`src/bindings/sources.rs::encode_files_ragged`), now that the crate root
+//! re-exports `batch::encode_files_docs`/`_serial` (`src/lib.rs`): parquet
+//! rows can't be split out of raw file bytes, so they're read directly as
+//! owned documents via `input::parquet::read_docs` and encoded through the
+//! whole-document path; text/JSONL files are loaded whole (mmapped, or
+//! decompressed into memory for .gz/.zst — `load_file` handles both) and
+//! handed straight to the fused chunk+extract+encode core as byte regions,
+//! with no separate single-threaded document-splitting pass.
+
+use std::path::PathBuf;
+
+use gigatoken_rs::input::file_source::{DocFormat, LoadedFile, load_file};
+use gigatoken_rs::input::parquet;
+use magnus::{
+    prelude::*,
+    scan_args::{get_kwargs, scan_args},
+    Error, RClass, RHash, RModule, RString, Ruby, TryConvert, Value,
+};
+
+use crate::error::raise;
+
+/// Paths plus how to split their bytes into documents — the payload shared
+/// by the three FileSource classes below.
+#[derive(Clone)]
+pub(crate) struct FileSource {
+    pub(crate) paths: Vec<PathBuf>,
+    pub(crate) format: DocFormat,
+}
+
+fn rstring_bytes(s: RString) -> Vec<u8> {
+    // SAFETY: read-only, for the duration of this synchronous call.
+    unsafe { s.as_slice() }.to_vec()
+}
+
+/// Plain-text files. With `separator`, documents are the pieces between
+/// separator occurrences; without one, each file is a single document.
+#[magnus::wrap(class = "Gigatoken::Native::TextFileSource", free_immediately, size)]
+pub struct TextFileSource(pub(crate) FileSource);
+
+impl TextFileSource {
+    fn new(args: &[Value]) -> Result<Self, Error> {
+        let args = scan_args::<(Vec<PathBuf>,), (), (), (), RHash, ()>(args)?;
+        let (paths,) = args.required;
+        let kw = get_kwargs::<_, (), (Option<Option<RString>>,), ()>(args.keywords, &[], &["separator"])?;
+        let (separator,) = kw.optional;
+        let separator = separator.flatten();
+        Ok(Self(FileSource {
+            paths,
+            format: DocFormat::Text {
+                separator: separator.map(rstring_bytes),
+            },
+        }))
+    }
+}
+
+/// JSON Lines files: one document per line, text taken from `field`
+/// (default `"text"`).
+#[magnus::wrap(class = "Gigatoken::Native::JsonlFileSource", free_immediately, size)]
+pub struct JsonlFileSource(pub(crate) FileSource);
+
+impl JsonlFileSource {
+    fn new(args: &[Value]) -> Result<Self, Error> {
+        let args = scan_args::<(Vec<PathBuf>,), (), (), (), RHash, ()>(args)?;
+        let (paths,) = args.required;
+        let kw = get_kwargs::<_, (), (Option<Option<String>>,), ()>(args.keywords, &[], &["field"])?;
+        let (field,) = kw.optional;
+        let field = field.flatten();
+        Ok(Self(FileSource {
+            paths,
+            format: DocFormat::Jsonl {
+                field: field.unwrap_or_else(|| "text".to_string()),
+            },
+        }))
+    }
+}
+
+/// Parquet files: one document per row, text taken from `column` (default
+/// `"text"`); null rows become empty documents.
+#[magnus::wrap(class = "Gigatoken::Native::ParquetFileSource", free_immediately, size)]
+pub struct ParquetFileSource(pub(crate) FileSource);
+
+impl ParquetFileSource {
+    fn new(args: &[Value]) -> Result<Self, Error> {
+        let args = scan_args::<(Vec<PathBuf>,), (), (), (), RHash, ()>(args)?;
+        let (paths,) = args.required;
+        let kw = get_kwargs::<_, (), (Option<Option<String>>,), ()>(args.keywords, &[], &["column"])?;
+        let (column,) = kw.optional;
+        let column = column.flatten();
+        Ok(Self(FileSource {
+            paths,
+            format: DocFormat::Parquet {
+                column: column.unwrap_or_else(|| "text".to_string()),
+            },
+        }))
+    }
+}
+
+/// Resolve an `encode_files` argument to its `FileSource`: a
+/// `TextFileSource`, `JsonlFileSource`, or `ParquetFileSource`. Bare paths
+/// are wrapped into a `TextFileSource` on the Ruby side
+/// (`Gigatoken::Tokenizer#encode_files`), so this only needs to handle the
+/// three native classes.
+pub(crate) fn resolve(ruby: &Ruby, source: Value) -> Result<FileSource, Error> {
+    if let Ok(s) = <&TextFileSource>::try_convert(source) {
+        return Ok(s.0.clone());
+    }
+    if let Ok(s) = <&JsonlFileSource>::try_convert(source) {
+        return Ok(s.0.clone());
+    }
+    if let Ok(s) = <&ParquetFileSource>::try_convert(source) {
+        return Ok(s.0.clone());
+    }
+    Err(raise(
+        ruby,
+        format!(
+            "expected a TextFileSource, JsonlFileSource, or ParquetFileSource, got {}",
+            source.class()
+        ),
+    ))
+}
+
+/// Load every document named by `source`, in file/row order, and hand its
+/// contents and document format to `encode` — the fused chunk+extract+encode
+/// core (`batch::encode_files_docs`/`_serial`, called by `tokenizer.rs`).
+/// Parquet rows can't be split out of raw file bytes, so they are
+/// materialized as owned documents here and encoded through the
+/// whole-document path (each buffer one document); other formats are loaded
+/// whole and left for `encode` to split. `encode` itself returns a
+/// `std::io::Result` (not a `magnus::Error`) since this whole call runs with
+/// the GVL released — the SentencePiece backend's UTF-8 validation surfaces
+/// its failure this way too, converted to `Gigatoken::Error` only after
+/// `without_gvl` returns (see `sentencepiece.rs`).
+pub(crate) fn encode_files_ragged(
+    source: &FileSource,
+    parallel: bool,
+    encode: impl FnOnce(&[&[u8]], &DocFormat) -> std::io::Result<(Vec<u32>, Vec<i64>)>,
+) -> std::io::Result<(Vec<u32>, Vec<i64>)> {
+    if let DocFormat::Parquet { column } = &source.format {
+        let docs = load_parquet_docs(&source.paths, column, parallel)?;
+        let bytes: Vec<&[u8]> = docs.iter().map(Vec::as_slice).collect();
+        return encode(&bytes, &DocFormat::Text { separator: None });
+    }
+    let files = load_files(&source.paths, parallel)?;
+    let bytes: Vec<&[u8]> = files.iter().map(LoadedFile::as_bytes).collect();
+    encode(&bytes, &source.format)
+}
+
+/// Load `column` of every parquet file as one owned document per row, files
+/// in argument order, rows in row order. Parallel across files and row
+/// groups with rayon, or fully on the calling thread when `parallel` is
+/// false (the sequential encode paths must never touch the rayon pool).
+fn load_parquet_docs(
+    paths: &[PathBuf],
+    column: &str,
+    parallel: bool,
+) -> std::io::Result<Vec<Vec<u8>>> {
+    use rayon::prelude::*;
+    let per_file: Vec<Vec<Vec<u8>>> = if parallel {
+        paths
+            .par_iter()
+            .map(|p| parquet::read_docs(p, column, true))
+            .collect::<std::io::Result<_>>()?
+    } else {
+        paths
+            .iter()
+            .map(|p| parquet::read_docs(p, column, false))
+            .collect::<std::io::Result<_>>()?
+    };
+    Ok(per_file.into_iter().flatten().collect())
+}
+
+/// Load all files: mmap when stored uncompressed, decompress .gz/.zst into
+/// memory otherwise (parallel chunking needs random access). In parallel
+/// with rayon, or serially on the calling thread when `parallel` is false
+/// (the sequential encode paths must never touch the rayon pool).
+fn load_files(paths: &[PathBuf], parallel: bool) -> std::io::Result<Vec<LoadedFile>> {
+    use rayon::prelude::*;
+    let load = |p: &PathBuf| {
+        load_file(p).map_err(|e| std::io::Error::new(e.kind(), format!("{}: {e}", p.display())))
+    };
+    if parallel {
+        paths.par_iter().map(load).collect()
+    } else {
+        paths.iter().map(load).collect()
+    }
+}
+
+pub fn init(ruby: &Ruby, native: RModule) -> Result<(), Error> {
+    let text: RClass = native.define_class("TextFileSource", ruby.class_object())?;
+    text.define_singleton_method("new", magnus::function!(TextFileSource::new, -1))?;
+
+    let jsonl: RClass = native.define_class("JsonlFileSource", ruby.class_object())?;
+    jsonl.define_singleton_method("new", magnus::function!(JsonlFileSource::new, -1))?;
+
+    let parquet: RClass = native.define_class("ParquetFileSource", ruby.class_object())?;
+    parquet.define_singleton_method("new", magnus::function!(ParquetFileSource::new, -1))?;
+
+    Ok(())
+}