gigatoken 0.1.0

data/src/bindings/padding.rs

@@ -0,0 +1,177 @@
+//! Padded/truncated batch assembly, serving the drop-in compatibility APIs
+//! (gigatoken.HFCompat's padding/truncation support). Kept out of lib.rs so
+//! the main encode path stays easy to read: the `encode_batch_padded`
+//! bindings there are one-call forwards into this module, and the options
+//! travel as the Rust-defined `PadTruncate` class, so per call they cost a
+//! typed downcast rather than dict-key lookups. The friendly keyword
+//! signature lives on `gigatoken.Tokenizer.encode_batch_padded`.
+
+use crate::input::file_source::DocFormat;
+use numpy::{IntoPyArray, PyArray1, PyArray2, PyArrayMethods};
+use pyo3::prelude::*;
+
+/// How `encode_batch_padded` assembles rows into a rectangular matrix.
+/// A frozen pyclass: fields are validated once at construction, and bad
+/// names or types fail there with a TypeError instead of at encode time.
+#[pyclass(frozen, get_all, from_py_object)]
+#[derive(Clone)]
+pub struct PadTruncate {
+    pub pad_id: u32,
+    /// Truncation limit and/or fixed pad width, depending on the two flags.
+    pub max_length: Option<usize>,
+    /// Pad every row to exactly `max_length` instead of the longest row.
+    pub pad_to_max_length: bool,
+    /// Keep at most `max_length` ids per row, counting prefix and suffix.
+    pub truncate: bool,
+    /// Put the padding before the tokens instead of after them.
+    pub pad_left: bool,
+    /// Drop tokens from the start of a row instead of the end.
+    pub truncate_left: bool,
+    /// Special-token ids written before / after every row's tokens; they
+    /// count toward the truncation budget, like HF post-processors.
+    pub prefix: Vec<u32>,
+    pub suffix: Vec<u32>,
+}
+
+#[pymethods]
+impl PadTruncate {
+    #[new]
+    #[allow(clippy::too_many_arguments)]
+    #[pyo3(signature = (pad_id, max_length=None, pad_to_max_length=false, truncate=false, pad_left=false, truncate_left=false, prefix=Vec::new(), suffix=Vec::new()))]
+    fn new(
+        pad_id: u32,
+        max_length: Option<usize>,
+        pad_to_max_length: bool,
+        truncate: bool,
+        pad_left: bool,
+        truncate_left: bool,
+        prefix: Vec<u32>,
+        suffix: Vec<u32>,
+    ) -> Self {
+        Self {
+            pad_id,
+            max_length,
+            pad_to_max_length,
+            truncate,
+            pad_left,
+            truncate_left,
+            prefix,
+            suffix,
+        }
+    }
+
+    fn __repr__(&self) -> String {
+        let py_bool = |b: bool| if b { "True" } else { "False" };
+        let max_length = self
+            .max_length
+            .map_or("None".to_string(), |m| m.to_string());
+        format!(
+            "PadTruncate(pad_id={}, max_length={max_length}, pad_to_max_length={}, truncate={}, \
+             pad_left={}, truncate_left={}, prefix={:?}, suffix={:?})",
+            self.pad_id,
+            py_bool(self.pad_to_max_length),
+            py_bool(self.truncate),
+            py_bool(self.pad_left),
+            py_bool(self.truncate_left),
+            self.prefix,
+            self.suffix,
+        )
+    }
+}
+
+/// The (rows x width) id matrix plus each row's real (unpadded) length —
+/// everything an attention mask needs.
+pub type PaddedMatrix<'py> = (Bound<'py, PyArray2<u32>>, Bound<'py, PyArray1<i64>>);
+
+/// Shared back-end of the `encode_batch_padded` bindings: encode like
+/// encode_batch, then pad/truncate into one (rows x width) uint32 matrix.
+pub fn encode_batch_matrix<'py>(
+    py: Python<'py>,
+    inputs: &Bound<'py, PyAny>,
+    opts: PadTruncate,
+    parallel: bool,
+    encode: impl Fn(&[&[u8]], &DocFormat) -> PyResult<(Vec<u32>, Vec<i64>)> + Send + Sync,
+) -> PyResult<PaddedMatrix<'py>> {
+    let (flat, counts) = super::bridge::encode_batch_flat(py, inputs, encode)?;
+    let (data, lengths, width) = py.detach(|| pad_truncate_matrix(&flat, &counts, &opts, parallel))?;
+    let rows = lengths.len();
+    let matrix = data.into_pyarray(py).reshape([rows, width])?;
+    Ok((matrix, lengths.into_pyarray(py)))
+}
+
+/// Assemble a ragged token batch into one row-major (rows x width) matrix in
+/// a single parallel copy pass. Each row is `prefix` + the (optionally
+/// truncated) tokens + `suffix`, padded with `pad_id` to the longest row (or
+/// to exactly `max_length` when `pad_to_max_length`). Returns the matrix
+/// data, each row's real (unpadded) length, and the width.
+fn pad_truncate_matrix(
+    flat: &[u32],
+    counts: &[i64],
+    opts: &PadTruncate,
+    parallel: bool,
+) -> PyResult<(Vec<u32>, Vec<i64>, usize)> {
+    use rayon::prelude::*;
+    let value_err = |msg: String| PyErr::new::<pyo3::exceptions::PyValueError, _>(msg);
+    let extra = opts.prefix.len() + opts.suffix.len();
+    let cap = if opts.truncate {
+        let max = opts
+            .max_length
+            .ok_or_else(|| value_err("truncate requires max_length".to_string()))?;
+        max.checked_sub(extra).ok_or_else(|| {
+            value_err(format!(
+                "max_length={max} leaves no room for the {extra} special tokens added per sequence"
+            ))
+        })?
+    } else {
+        usize::MAX
+    };
+    let mut offsets = Vec::with_capacity(counts.len());
+    let mut lengths = Vec::with_capacity(counts.len());
+    let mut longest = 0usize;
+    let mut pos = 0usize;
+    for &count in counts {
+        let kept = (count as usize).min(cap);
+        offsets.push(pos);
+        pos += count as usize;
+        lengths.push((kept + extra) as i64);
+        longest = longest.max(kept + extra);
+    }
+    let width = if opts.pad_to_max_length {
+        let max = opts
+            .max_length
+            .ok_or_else(|| value_err("pad_to_max_length requires max_length".to_string()))?;
+        if longest > max {
+            return Err(value_err(format!(
+                "a sequence is {longest} ids long but padding to max_length={max} was requested \
+                 without truncation; enable truncation or raise max_length"
+            )));
+        }
+        max
+    } else {
+        longest
+    };
+    let mut out = vec![opts.pad_id; counts.len() * width];
+    if width > 0 {
+        let fill = |(i, row): (usize, &mut [u32])| {
+            let count = counts[i] as usize;
+            let kept = count.min(cap);
+            let tokens = &flat[offsets[i]..offsets[i] + count];
+            let tokens = if opts.truncate_left {
+                &tokens[count - kept..]
+            } else {
+                &tokens[..kept]
+            };
+            let len = kept + extra;
+            let start = if opts.pad_left { width - len } else { 0 };
+            row[start..start + opts.prefix.len()].copy_from_slice(&opts.prefix);
+            row[start + opts.prefix.len()..start + opts.prefix.len() + kept].copy_from_slice(tokens);
+            row[start + len - opts.suffix.len()..start + len].copy_from_slice(&opts.suffix);
+        };
+        if parallel {
+            out.par_chunks_mut(width).enumerate().for_each(fill);
+        } else {
+            out.chunks_mut(width).enumerate().for_each(fill);
+        }
+    }
+    Ok((out, lengths, width))
+}

data/src/bindings/pretokenize.rs

@@ -0,0 +1,53 @@
+//! Python bindings exposing the pretokenizer directly (an iterator over one
+//! document's pretokens and a parallel pretoken-count function) — used for
+//! inspection and tests, not by the encode path.
+
+use crate::pretokenize;
+use pyo3::prelude::*;
+use pyo3::types::PyBytes;
+
+#[pyclass]
+pub(crate) struct PretokenizerIter {
+    /// Byte offset into `bytes`; the pretokenizer is stateless beyond this, so
+    /// each `__next__` resumes a fresh `FastR50kPretokenizer` at this position.
+    pos: usize,
+    bytes: Py<PyBytes>,
+}
+
+#[pymethods]
+impl PretokenizerIter {
+    fn __iter__<'py>(slf: PyRef<'py, Self>) -> PyRef<'py, PretokenizerIter> {
+        slf
+    }
+
+    fn __next__<'py>(&'py mut self, py: Python<'py>) -> Option<&'py [u8]> {
+        let bytes: &'py [u8] = self.bytes.as_bytes(py);
+        let mut iter = pretokenize::FastR50kPretokenizer::with_pos(bytes, self.pos);
+        let result = iter.next();
+        self.pos = iter.pos();
+        Some(result?.0)
+    }
+}
+
+#[pyfunction]
+pub(crate) fn pretokenizer<'py>(text: Bound<'py, PyBytes>) -> PyResult<PretokenizerIter> {
+    Ok(PretokenizerIter {
+        pos: 0,
+        bytes: text.into(),
+    })
+}
+
+#[pyfunction]
+#[pyo3(signature = (text, separator = None))]
+pub(crate) fn pretokenized_counts<'py>(
+    text: Bound<'py, PyBytes>,
+    separator: Option<&[u8]>,
+) -> PyResult<Vec<(Bound<'py, PyBytes>, usize)>> {
+    let separator = separator.unwrap_or(pretokenize::DEFAULT_SEPARATOR);
+    let tokens_counts = pretokenize::pretokenize_par_bytes(text.as_bytes(), separator);
+    let tokens_counts = tokens_counts
+        .into_iter()
+        .map(|(k, v)| (PyBytes::new(text.py(), k.as_ref()), v))
+        .collect::<Vec<_>>();
+    Ok(tokens_counts)
+}

data/src/bindings/sources.rs

@@ -0,0 +1,273 @@
+//! The FileSource/BytesSource Python classes and the helpers that turn an
+//! encode_files or train_bpe argument into loaded, format-tagged file
+//! contents.
+
+use crate::input::Resource;
+use crate::input::file_source::{DocFormat, LoadedFile, detect_default_format, load_file};
+use pyo3::prelude::*;
+use pyo3::pybacked::PyBackedBytes;
+use std::path::PathBuf;
+
+/// Base class for file sources. Not directly constructible from Python —
+/// use `TextFileSource` or `JsonlFileSource`, which pin down the document
+/// format and its parameters. Compression (.gz/.zst) is always detected
+/// from the file extension, independent of the source type.
+#[pyclass(subclass, from_py_object)]
+#[derive(Clone)]
+pub(crate) struct FileSource {
+    pub(crate) paths: Vec<PathBuf>,
+    pub(crate) format: DocFormat,
+}
+
+#[pymethods]
+impl FileSource {
+    fn __repr__(&self) -> String {
+        let n = self.paths.len();
+        match &self.format {
+            DocFormat::Jsonl { field } => {
+                format!("JsonlFileSource(paths=[{n} files], field={field:?})")
+            }
+            DocFormat::Text {
+                separator: Some(sep),
+            } => format!(
+                "TextFileSource(paths=[{n} files], separator={:?})",
+                String::from_utf8_lossy(sep)
+            ),
+            DocFormat::Text { separator: None } => {
+                format!("TextFileSource(paths=[{n} files])")
+            }
+            DocFormat::Parquet { column } => {
+                format!("ParquetFileSource(paths=[{n} files], column={column:?})")
+            }
+        }
+    }
+}
+
+/// A `separator` argument: bytes used as-is, or str encoded to its UTF-8
+/// bytes.
+#[derive(FromPyObject)]
+pub(crate) enum Separator {
+    Bytes(Vec<u8>),
+    Str(String),
+}
+
+impl Separator {
+    fn into_bytes(self) -> Vec<u8> {
+        match self {
+            Separator::Bytes(b) => b,
+            Separator::Str(s) => s.into_bytes(),
+        }
+    }
+}
+
+/// Plain-text files. With `separator` (str or bytes), documents are the
+/// pieces between separator occurrences (the separator itself belongs to no
+/// document); without one, each file is a single document.
+#[pyclass(extends = FileSource)]
+pub(crate) struct TextFileSource;
+
+#[pymethods]
+impl TextFileSource {
+    #[new]
+    #[pyo3(signature = (paths, separator = None))]
+    fn new(paths: Vec<PathBuf>, separator: Option<Separator>) -> PyClassInitializer<Self> {
+        PyClassInitializer::from(FileSource {
+            paths,
+            format: DocFormat::Text {
+                separator: separator.map(Separator::into_bytes),
+            },
+        })
+        .add_subclass(Self)
+    }
+}
+
+/// JSON Lines files: one document per line, text taken from `field`.
+#[pyclass(extends = FileSource)]
+pub(crate) struct JsonlFileSource;
+
+#[pymethods]
+impl JsonlFileSource {
+    #[new]
+    #[pyo3(signature = (paths, field = "text"))]
+    fn new(paths: Vec<PathBuf>, field: &str) -> PyClassInitializer<Self> {
+        PyClassInitializer::from(FileSource {
+            paths,
+            format: DocFormat::Jsonl {
+                field: field.to_string(),
+            },
+        })
+        .add_subclass(Self)
+    }
+}
+
+/// Parquet files: one document per row, text taken from `column` (a string
+/// or binary column; null rows become empty documents, so results stay
+/// row-aligned with the table). Parquet handles compression internally, so
+/// no .gz/.zst detection applies.
+#[pyclass(extends = FileSource)]
+pub(crate) struct ParquetFileSource;
+
+#[pymethods]
+impl ParquetFileSource {
+    #[new]
+    #[pyo3(signature = (paths, column = "text"))]
+    fn new(paths: Vec<PathBuf>, column: &str) -> PyClassInitializer<Self> {
+        PyClassInitializer::from(FileSource {
+            paths,
+            format: DocFormat::Parquet {
+                column: column.to_string(),
+            },
+        })
+        .add_subclass(Self)
+    }
+}
+
+/// In-memory bytes for encode_batch, the buffer analog of TextFileSource:
+/// with `separator` (str or bytes), each buffer's documents are the pieces
+/// between separator occurrences (the separator itself belongs to no
+/// document, and empty pieces are skipped); without one, each buffer is a
+/// single document. The buffers are borrowed, not copied, and split during
+/// the (parallel) encode itself — handing a whole corpus over as a few
+/// buffers plus a separator is much faster than pre-splitting it into
+/// per-document Python objects.
+#[pyclass(frozen)]
+pub(crate) struct BytesSource {
+    pub(crate) buffers: Vec<PyBackedBytes>,
+    pub(crate) format: DocFormat,
+}
+
+#[pymethods]
+impl BytesSource {
+    #[new]
+    #[pyo3(signature = (data, separator = None))]
+    fn new(data: &Bound<'_, PyAny>, separator: Option<Separator>) -> PyResult<Self> {
+        let buffers = if let Ok(buffer) = data.extract::<PyBackedBytes>() {
+            vec![buffer]
+        } else {
+            data.extract::<Vec<PyBackedBytes>>().map_err(|_| {
+                PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!(
+                    "expected bytes or a list of bytes, got {}",
+                    data.get_type()
+                ))
+            })?
+        };
+        Ok(Self {
+            buffers,
+            format: DocFormat::Text {
+                separator: separator.map(Separator::into_bytes),
+            },
+        })
+    }
+
+    fn __repr__(&self) -> String {
+        let n = self.buffers.len();
+        let total: usize = self.buffers.iter().map(|b| b.len()).sum();
+        match &self.format {
+            DocFormat::Text {
+                separator: Some(sep),
+            } => format!(
+                "BytesSource(data=[{n} buffers, {total} bytes], separator={:?})",
+                String::from_utf8_lossy(sep)
+            ),
+            _ => format!("BytesSource(data=[{n} buffers, {total} bytes])"),
+        }
+    }
+}
+
+/// Resolve an encode_files argument: a FileSource (TextFileSource /
+/// JsonlFileSource), a single path, or a list of paths. Bare paths get a
+/// default format from the first path's extension — all inputs in a batch
+/// are assumed to be of the same type.
+pub(crate) fn resolve_files_source(obj: &Bound<'_, PyAny>) -> PyResult<(Vec<PathBuf>, DocFormat)> {
+    if let Ok(fs) = obj.extract::<FileSource>() {
+        return Ok((fs.paths, fs.format));
+    }
+    if let Ok(path) = obj.extract::<PathBuf>() {
+        let format = detect_default_format(&path);
+        return Ok((vec![path], format));
+    }
+    if let Ok(paths) = obj.extract::<Vec<PathBuf>>() {
+        let format = paths
+            .first()
+            .map(|p| detect_default_format(p))
+            .unwrap_or(DocFormat::Text { separator: None });
+        return Ok((paths, format));
+    }
+    Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!(
+        "expected a TextFileSource/JsonlFileSource/ParquetFileSource, a path, or a list of paths, got {}",
+        obj.get_type()
+    )))
+}
+
+/// Shared scaffold of the encode_files pymethods: resolve the source
+/// argument, load the files with the GIL released, hand their contents and
+/// document format to `encode` (still detached), and return the ragged
+/// result as an awkward Array. The per-backend encoding lives in
+/// `batch::encode_files_docs` / `batch::sp_encode_files_docs`.
+pub(crate) fn encode_files_ragged<'py>(
+    py: Python<'py>,
+    source: &Bound<'py, PyAny>,
+    parallel: bool,
+    encode: impl FnOnce(&[&[u8]], &DocFormat) -> (Vec<u32>, Vec<i64>) + Send,
+) -> PyResult<Bound<'py, PyAny>> {
+    let (paths, format) = resolve_files_source(source)?;
+    let (flat, counts) = py.detach(|| -> PyResult<_> {
+        // Parquet rows can't be split out of the raw file bytes, so they are
+        // materialized as owned documents here and encoded through the
+        // whole-document path (each buffer one document).
+        if let DocFormat::Parquet { column } = &format {
+            let docs = load_parquet_docs(&paths, column, parallel)?;
+            let bytes: Vec<&[u8]> = docs.iter().map(|d| d.as_slice()).collect();
+            return Ok(encode(&bytes, &DocFormat::Text { separator: None }));
+        }
+        let files = load_files(&paths, parallel)?;
+        let bytes: Vec<&[u8]> = files.iter().map(|f| f.as_bytes()).collect();
+        Ok(encode(&bytes, &format))
+    })?;
+    super::bridge::ragged_to_python(py, flat, counts)
+}
+
+/// 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,
+) -> PyResult<Vec<Vec<u8>>> {
+    use crate::input::parquet::read_docs;
+    use rayon::prelude::*;
+    // input::parquet errors already carry the offending path.
+    let to_pyerr = |e: std::io::Error| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string());
+    let per_file: Vec<Vec<Vec<u8>>> = if parallel {
+        paths
+            .par_iter()
+            .map(|p| read_docs(p, column, true).map_err(to_pyerr))
+            .collect::<PyResult<_>>()?
+    } else {
+        paths
+            .iter()
+            .map(|p| read_docs(p, column, false).map_err(to_pyerr))
+            .collect::<PyResult<_>>()?
+    };
+    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).
+pub(crate) fn load_files(paths: &[PathBuf], parallel: bool) -> PyResult<Vec<LoadedFile>> {
+    use rayon::prelude::*;
+    let load = |p: &PathBuf| {
+        load_file(p).map_err(|e| {
+            PyErr::new::<pyo3::exceptions::PyIOError, _>(format!("{}: {e}", p.display()))
+        })
+    };
+    if parallel {
+        paths.par_iter().map(load).collect()
+    } else {
+        paths.iter().map(load).collect()
+    }
+}

data/src/bindings/train.rs

@@ -0,0 +1,125 @@
+//! Python binding for BPE training: the train_bpe function and the
+//! conversion of its result into Python vocab/merges objects.
+
+use super::sources::FileSource;
+use crate::bpe_train;
+use crate::input::file_source::FileSourceSpec;
+use crate::input::{MmappedFile, Resource};
+use crate::pretokenize;
+use itertools::Itertools;
+use pyo3::prelude::*;
+use pyo3::types::{IntoPyDict, PyBytes, PyDict};
+use std::path::PathBuf;
+
+/// Vocab dict plus ordered merge pairs, as Python objects.
+type PyVocabAndMerges<'py> = (
+    Bound<'py, PyDict>,
+    Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>,
+);
+
+fn bpe_result_to_python<'py>(
+    py: Python<'py>,
+    result: bpe_train::BPEResult,
+) -> PyResult<PyVocabAndMerges<'py>> {
+    let vocab_py = result
+        .vocab
+        .into_iter()
+        .map(|(k, v)| (k, PyBytes::new(py, &v)))
+        .sorted_by(|e1, e2| Ord::cmp(&e1.0, &e2.0))
+        .into_py_dict(py);
+    let merges_py: Vec<_> = result
+        .merges
+        .into_iter()
+        .map(|(k, v)| (PyBytes::new(py, &k), PyBytes::new(py, &v)))
+        .collect();
+    Ok((vocab_py?, merges_py))
+}
+
+fn parse_tie_breaking(s: &str) -> PyResult<bpe_train::TieBreaking> {
+    match s {
+        "huggingface" => Ok(bpe_train::TieBreaking::HuggingFace),
+        "raw_token_ids" => Ok(bpe_train::TieBreaking::RawTokenIds),
+        "assembled_bytes" => Ok(bpe_train::TieBreaking::AssembledBytes),
+        other => Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
+            "tie_breaking must be 'huggingface', 'raw_token_ids', or 'assembled_bytes', got {other:?}"
+        ))),
+    }
+}
+
+#[pyfunction]
+#[allow(clippy::type_complexity)]
+#[pyo3(signature = (in_data, vocab_size, special_tokens, tie_breaking = "huggingface", separator = None))]
+pub(crate) fn train_bpe<'py>(
+    py: Python<'py>,
+    in_data: Bound<'py, PyAny>,
+    vocab_size: usize,
+    special_tokens: Vec<String>,
+    tie_breaking: &str,
+    separator: Option<&[u8]>,
+) -> PyResult<(
+    Bound<'py, PyDict>,
+    Vec<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)>,
+)> {
+    assert!(
+        vocab_size <= 2_usize.pow(32),
+        "vocab_size must be less than 2^32"
+    );
+    let tie_breaking = parse_tie_breaking(tie_breaking)?;
+    let separator = separator.unwrap_or(pretokenize::DEFAULT_SEPARATOR);
+
+    // --- FileSource: multi-file parallel processing ---
+    if let Ok(file_source) = in_data.extract::<FileSource>() {
+        let spec = FileSourceSpec {
+            paths: file_source.paths,
+            format: file_source.format,
+        };
+        let counts = spec.pretokenize().map_err(|e| {
+            PyErr::new::<pyo3::exceptions::PyIOError, _>(format!(
+                "FileSource processing failed: {}",
+                e
+            ))
+        })?;
+        let result = bpe_train::train_bpe(counts, vocab_size, special_tokens, tie_breaking);
+        return bpe_result_to_python(py, result);
+    }
+
+    // --- Single bytes or file path ---
+    let mmap_resource;
+    let bytes: &[u8] = if in_data.is_instance_of::<PyBytes>() {
+        in_data.extract::<&[u8]>()?
+    } else if let Ok(path) = in_data.extract::<PathBuf>() {
+        if let Some(ext) = path.extension()
+            && ext == "parquet"
+        {
+            // A bare path takes the default column "text", matching
+            // detect_default_format; use ParquetFileSource to choose another
+            // column.
+            let spec = FileSourceSpec {
+                paths: vec![path],
+                format: crate::input::file_source::DocFormat::Parquet {
+                    column: "text".to_string(),
+                },
+            };
+            let counts = spec
+                .pretokenize()
+                .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;
+            let result = bpe_train::train_bpe(counts, vocab_size, special_tokens, tie_breaking);
+            return bpe_result_to_python(py, result);
+        }
+        mmap_resource = MmappedFile::open(&path).map_err(|e| {
+            PyErr::new::<pyo3::exceptions::PyIOError, _>(format!(
+                "Failed to open file {:?}: {}",
+                path, e
+            ))
+        })?;
+        mmap_resource.as_bytes()
+    } else {
+        return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
+            "in_data must be bytes, a path, or a FileSource",
+        ));
+    };
+
+    let counts = pretokenize::pretokenize_par_bytes(bytes, separator);
+    let result = bpe_train::train_bpe(counts, vocab_size, special_tokens, tie_breaking);
+    bpe_result_to_python(py, result)
+}