#[pyo3::pymodule(gil_used = false)]
mod _lib {
use std::{collections::BTreeMap, path::PathBuf, sync::{Arc, mpsc::sync_channel}};
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1};
use ordermap::OrderMap;
use pyo3::{prelude::*, pymethods, types::{PyAny, PyIterator}};
use crate::{MyError, MyResult, bpe::{BpeEncoder, BpeModel, BpeTrainer, BpeTrainerConfig, CharIdx, CharSplit, Character, Idx, IdxLike, InitialAlphabet, TieBreak, Word, encoder::BpeBuilder, utils::ToWord}, counter::SourceBatchOptions, pretokenizer::{BoundaryMode, ChunkHint, ChunkOptions, UnicodeBigramMixedBoundary, UnicodeBigramSelection as CoreUnicodeBigramSelection, parse_unicode_bigrams, unicode_bigram_to_string}, spec::{Spec, gpt2::Gpt2Spec, unitoken::UnitokenSpec}, traits::{CanEncode, CanStrToWord, Encoder, Train as _}};
#[pyclass(subclass)]
pub struct BpeTrainerBase;
#[pyclass]
pub struct UnicodeBigramSelection {
#[pyo3(get)]
pub bigrams: Vec<String>,
#[pyo3(get)]
pub cutoff_freq: Option<i64>,
#[pyo3(get)]
pub max_excluded_freq: Option<i64>,
}
impl From<CoreUnicodeBigramSelection> for UnicodeBigramSelection {
fn from(selection: CoreUnicodeBigramSelection) -> Self {
let mut bigrams = selection.bigrams.into_iter().collect::<Vec<_>>();
bigrams.sort_unstable();
Self {
bigrams: bigrams.into_iter().map(unicode_bigram_to_string).collect(),
cutoff_freq: selection.cutoff_freq,
max_excluded_freq: selection.max_excluded_freq,
}
}
}
enum BpeModelInner {
Byte(BpeModel<u8, Idx>),
Unicode(BpeModel<Character, CharIdx>),
}
#[pyclass]
pub struct BpeModelBase {
inner: BpeModelInner,
}
#[allow(dead_code)]
/// this is just a reference for impl blocks, not directly used
pub trait BpeTrainerBaseImpl: Sized {
fn new_py(special_tokens: Vec<String>) -> (Self, BpeTrainerBase);
fn add_words(&mut self, py: Python, words: Vec<(String, i64)>);
fn vocab_size(&self) -> usize;
fn last_merge_freq(&self) -> Option<i64>;
fn init_training(&mut self, py: Python) -> PyResult<()>;
fn train_until(&mut self, py: Python, vocab_size: usize) -> PyResult<i64>;
fn step(&mut self, py: Python) -> PyResult<i64>;
fn get_vocab(&self) -> Vocabulary;
fn validate_model(&self, py: Python) -> PyResult<BpeModelBase>;
}
fn save_model_vocab<C, I>(model: &BpeModel<C, I>, path: &PathBuf, spec: &dyn Spec<C, I>) -> MyResult<()> {
let mut file = std::fs::File::create(path)?;
let mut writer = std::io::BufWriter::new(&mut file);
model.save_vocab_json(spec, &mut writer)?;
std::io::Write::flush(&mut writer)?;
Ok(())
}
fn save_model_merges<C, I>(model: &BpeModel<C, I>, path: &PathBuf, spec: &dyn Spec<C, I>) -> MyResult<()> {
let mut file = std::fs::File::create(path)?;
let mut writer = std::io::BufWriter::new(&mut file);
model.save_merges_txt(spec, &mut writer)?;
std::io::Write::flush(&mut writer)?;
Ok(())
}
fn map_model_error(error: MyError) -> PyErr {
let message = error.to_string();
match error {
MyError::Io(_) => pyo3::exceptions::PyIOError::new_err(message),
_ => pyo3::exceptions::PyValueError::new_err(message),
}
}
#[pymethods]
impl BpeModelBase {
#[getter]
/// Primary segmentation unit used by this model.
pub fn unit(&self) -> &'static str {
match self.inner {
BpeModelInner::Byte(_) => "byte",
BpeModelInner::Unicode(_) => "unicode",
}
}
#[getter]
/// Frequency of the final pair merge, if the model contains one.
pub fn last_merge_freq(&self) -> Option<i64> {
match &self.inner {
BpeModelInner::Byte(model) => model.last_merge_freq(),
BpeModelInner::Unicode(model) => model.last_merge_freq(),
}
}
#[getter]
/// Reserved special tokens in vocabulary order.
pub fn special_tokens(&self) -> Vec<String> {
match &self.inner {
BpeModelInner::Byte(model) => model.special_tokens().to_vec(),
BpeModelInner::Unicode(model) => model.special_tokens().to_vec(),
}
}
/// Build an encoder directly from the validated model.
#[pyo3(signature = (pat_str=None, unicode_bigrams=None, unicode_bigram_mixed_boundary="keep", split_on_vocab_bigrams=true))]
pub fn encoder(
&self,
py: Python,
pat_str: Option<String>,
unicode_bigrams: Option<Vec<String>>,
unicode_bigram_mixed_boundary: &str,
split_on_vocab_bigrams: bool,
) -> PyResult<BpeEncoderBase> {
py.detach(|| {
let inner: Arc<dyn Encoder<Idx> + Send + Sync> = match &self.inner {
BpeModelInner::Byte(model) => Arc::new(configure_encoder(
model.to_encoder_with_options(pat_str.as_deref(), split_on_vocab_bigrams)?,
unicode_bigrams.as_deref(),
unicode_bigram_mixed_boundary,
)?),
BpeModelInner::Unicode(model) => Arc::new(configure_encoder(
model.to_encoder_with_options(pat_str.as_deref(), split_on_vocab_bigrams)?,
unicode_bigrams.as_deref(),
unicode_bigram_mixed_boundary,
)?),
};
Ok(BpeEncoderBase(inner))
}).map_err(map_model_error)
}
/// Return a view of the validated vocabulary.
pub fn get_vocab(&self) -> Vocabulary {
let inner: Box<dyn VocabularyImpl + Send + Sync> = match &self.inner {
BpeModelInner::Byte(model) => Box::new(VocabularyInner::new(model.vocab())),
BpeModelInner::Unicode(model) => Box::new(VocabularyInner::new(model.vocab())),
};
Vocabulary { inner }
}
/// Save the validated vocabulary JSON.
pub fn save_vocab(&self, py: Python, path: PathBuf, format: &str) -> PyResult<()> {
py.detach(|| match (&self.inner, format) {
(BpeModelInner::Byte(model), "gpt2") => save_model_vocab(model, &path, &Gpt2Spec),
(BpeModelInner::Byte(model), "unitoken") => save_model_vocab(model, &path, &UnitokenSpec),
(BpeModelInner::Unicode(model), "unitoken") => save_model_vocab(model, &path, &UnitokenSpec),
_ => Err(MyError::SpecError(format!("format {format} is not compatible with unit {}", self.unit()))),
}).map_err(map_model_error)
}
/// Save the validated merge list.
pub fn save_merges_txt(&self, py: Python, path: PathBuf, format: &str) -> PyResult<()> {
py.detach(|| match (&self.inner, format) {
(BpeModelInner::Byte(model), "gpt2") => save_model_merges(model, &path, &Gpt2Spec),
(BpeModelInner::Byte(model), "unitoken") => save_model_merges(model, &path, &UnitokenSpec),
(BpeModelInner::Unicode(model), "unitoken") => save_model_merges(model, &path, &UnitokenSpec),
_ => Err(MyError::SpecError(format!("format {format} is not compatible with unit {}", self.unit()))),
}).map_err(map_model_error)
}
}
fn trainer_config(
initial_alphabet: Option<&str>,
tie_break: Option<&str>,
parallel_merge_min_occurs_in: Option<usize>,
hot_pair_window_size: Option<usize>,
bigram_cutoff_freq: Option<i64>,
) -> PyResult<BpeTrainerConfig> {
let initial_alphabet = match initial_alphabet.unwrap_or("raw") {
"raw" => InitialAlphabet::RawBytes,
"byte_level" => InitialAlphabet::ByteLevel,
value => return Err(pyo3::exceptions::PyValueError::new_err(format!("Unknown initial_alphabet: {value}"))),
};
let tie_break = match tie_break.unwrap_or("smallest_pair_id") {
"smallest_pair_id" => TieBreak::SmallestPairId,
"largest_content" => TieBreak::LargestContent,
value => return Err(pyo3::exceptions::PyValueError::new_err(format!("Unknown tie_break: {value}"))),
};
if hot_pair_window_size == Some(0) {
return Err(pyo3::exceptions::PyValueError::new_err(
"hot_pair_window_size must be positive",
));
}
if bigram_cutoff_freq.is_some_and(|cutoff| cutoff <= 0) {
return Err(pyo3::exceptions::PyValueError::new_err(
"bigram_cutoff_freq must be positive",
));
}
Ok(BpeTrainerConfig {
initial_alphabet,
tie_break,
parallel_merge_min_occurs_in,
hot_pair_window_size,
bigram_cutoff_freq,
})
}
fn validate_primary_vocab_ratio(primary_vocab_ratio: f64) -> PyResult<()> {
if !primary_vocab_ratio.is_finite() || !(0.0..=1.0).contains(&primary_vocab_ratio) {
return Err(pyo3::exceptions::PyValueError::new_err(
"primary_vocab_ratio must be finite and between 0 and 1",
));
}
Ok(())
}
fn trainer_memory_usage<C, I>(trainer: &BpeTrainer<C, I>) -> BTreeMap<&'static str, u64> {
let usage = trainer.memory_usage();
BTreeMap::from([
("estimated_persistent_bytes", usage.estimated_persistent_bytes as u64),
("word_entries", usage.word_entries as u64),
("word_storage_bytes", usage.word_storage_bytes as u64),
("pair_entries", usage.pair_entries as u64),
("pair_table_bytes", usage.pair_table_bytes as u64),
("occurrence_set_header_bytes", usage.occurrence_set_header_bytes as u64),
("occurrence_capacity_entries", usage.occurrence_capacity_entries as u64),
("occurrence_capacity_bytes", usage.occurrence_capacity_bytes as u64),
("merge_heap_bytes", usage.merge_heap_bytes as u64),
("merge_storage_bytes", usage.merge_storage_bytes as u64),
("vocab_entries", usage.vocab_entries as u64),
("vocab_token_bytes", usage.vocab_token_bytes as u64),
])
}
fn chunk_options(chunk_size: u64, boundary: &str) -> PyResult<ChunkOptions> {
let boundary = BoundaryMode::parse(boundary)
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
if chunk_size == 0 {
return Err(pyo3::exceptions::PyValueError::new_err("chunk_size must be at least 1"));
}
Ok(ChunkOptions {
hint: ChunkHint::Size(chunk_size),
boundary,
})
}
#[allow(non_camel_case_types)]
#[pyclass(extends = BpeTrainerBase)]
pub struct BpeTrainer_u8_Idx {
pub inner: BpeTrainer<u8, Idx>,
}
#[pymethods]
impl BpeTrainer_u8_Idx {
#[new]
/// Create a new BPE trainer (byte-level) for Python.
///
/// Returns `(trainer, base)` where `base` enables Python-side subclassing.
#[pyo3(signature = (special_tokens, initial_alphabet=None, tie_break=None, parallel_merge_min_occurs_in=None, hot_pair_window_size=None, bigram_cutoff_freq=None))]
pub fn new_py(
special_tokens: Vec<String>,
initial_alphabet: Option<&str>,
tie_break: Option<&str>,
parallel_merge_min_occurs_in: Option<usize>,
hot_pair_window_size: Option<usize>,
bigram_cutoff_freq: Option<i64>,
) -> PyResult<(Self, BpeTrainerBase)> {
let config = trainer_config(
initial_alphabet,
tie_break,
parallel_merge_min_occurs_in,
hot_pair_window_size,
bigram_cutoff_freq,
)?;
Ok((
Self {
inner: BpeTrainer::new_with_config(vec![], special_tokens, config),
},
BpeTrainerBase {},
))
}
/// Add `(word, frequency)` pairs to the trainer's inventory.
pub fn add_words(&mut self, py: Python, words: Vec<(String, i64)>) {
py.detach(||
self.inner.add_words(&mut words.iter().map(|(w, f)| (w.as_str(), *f)))
)
}
/// Replace the trainer inventory by consuming a word counter.
pub fn add_word_counter(&mut self, py: Python, mut counter: PyRefMut<'_, WordCounter>) {
let counts = counter.take_counts();
drop(counter);
py.detach(||
self.inner.add_words(&mut counts.iter().map(|(word, frequency)| (word.as_str(), *frequency)))
)
}
/// Current vocabulary size.
pub fn vocab_size(&self) -> usize {
self.inner.vocab_size()
}
#[getter]
/// Frequency of the most recently completed pair merge.
pub fn last_merge_freq(&self) -> Option<i64> {
self.inner.last_merge_freq()
}
#[getter]
/// Diagnostics for the bounded pair-posting window, if enabled.
pub fn hot_pair_window_stats(&self) -> Option<BTreeMap<&'static str, u64>> {
let stats = self.inner.hot_pair_window_stats()?;
Some(BTreeMap::from([
("hydration_scans", stats.hydration_scans),
("hydrated_word_entries", stats.hydrated_word_entries),
("batch_prunes", stats.batch_prunes),
("prune_evictions", stats.prune_evictions),
("peak_resident_pairs", stats.peak_resident_pairs as u64),
("resident_pairs", self.inner.hot_resident_pairs() as u64),
("occurrence_capacity", self.inner.hot_occurrence_capacity() as u64),
]))
}
#[getter]
/// Capacity-backed breakdown of persistent trainer storage.
///
/// This attributes trainer-owned allocations and is not process RSS; use it
/// to identify why bounded occurrence postings do or do not dominate.
pub fn memory_usage(&self) -> BTreeMap<&'static str, u64> {
trainer_memory_usage(&self.inner)
}
/// Initialize internal training state.
pub fn init_training(&mut self, py: Python) -> PyResult<()> {
py.detach(|| self.inner.init_training());
Ok(())
}
/// Train until the vocabulary reaches `vocab_size` or the pair cutoff.
///
/// Returns the updated vocabulary size.
pub fn train_until(&mut self, py: Python, vocab_size: usize) -> PyResult<i64> {
py.detach(|| self.inner.train_until(vocab_size)).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
Ok(self.inner.vocab_size() as i64)
}
/// Perform one training step.
///
/// Returns the updated vocabulary size.
pub fn step(&mut self, py: Python) -> PyResult<i64> {
py.detach(|| self.inner.step()).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
Ok(self.inner.vocab_size() as i64)
}
/// Return a view of the current vocabulary.
pub fn get_vocab(&self) -> Vocabulary {
Vocabulary {
inner: Box::new(VocabularyInner::new(&self.inner.vocab)),
}
}
/// Validate the current state and return an immutable model snapshot.
pub fn validate_model(&self, py: Python) -> PyResult<BpeModelBase> {
py.detach(|| self.inner.validate_model())
.map(|model| BpeModelBase { inner: BpeModelInner::Byte(model) })
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}
}
#[allow(non_camel_case_types)]
#[pyclass(extends = BpeTrainerBase)]
pub struct BpeTrainer_Character_CharIdx {
pub inner: BpeTrainer<Character, CharIdx>,
}
#[pymethods]
impl BpeTrainer_Character_CharIdx {
#[new]
/// Create a new BPE trainer (character-level) for Python.
///
/// Returns `(trainer, base)` where `base` enables Python-side subclassing.
#[pyo3(signature = (special_tokens, initial_alphabet=None, tie_break=None, parallel_merge_min_occurs_in=None, hot_pair_window_size=None, bigram_cutoff_freq=None))]
pub fn new_py(
special_tokens: Vec<String>,
initial_alphabet: Option<&str>,
tie_break: Option<&str>,
parallel_merge_min_occurs_in: Option<usize>,
hot_pair_window_size: Option<usize>,
bigram_cutoff_freq: Option<i64>,
) -> PyResult<(Self, BpeTrainerBase)> {
let config = trainer_config(
initial_alphabet,
tie_break,
parallel_merge_min_occurs_in,
hot_pair_window_size,
bigram_cutoff_freq,
)?;
Ok((
Self {
inner: BpeTrainer::new_with_config(vec![], special_tokens, config),
},
BpeTrainerBase {},
))
}
/// Add `(word, frequency)` pairs to the trainer's inventory.
pub fn add_words(&mut self, py: Python, words: Vec<(String, i64)>) {
py.detach(||
self.inner.add_words(&mut words.iter().map(|(w, f)| (w.as_str(), *f)))
)
}
/// Replace the trainer inventory by consuming a word counter.
pub fn add_word_counter(&mut self, py: Python, mut counter: PyRefMut<'_, WordCounter>) {
let counts = counter.take_counts();
drop(counter);
py.detach(||
self.inner.add_words(&mut counts.iter().map(|(word, frequency)| (word.as_str(), *frequency)))
)
}
/// Current vocabulary size.
pub fn vocab_size(&self) -> usize {
self.inner.vocab_size()
}
#[getter]
/// Frequency of the most recently completed pair merge.
pub fn last_merge_freq(&self) -> Option<i64> {
self.inner.last_merge_freq()
}
#[getter]
/// Diagnostics for the bounded pair-posting window, if enabled.
pub fn hot_pair_window_stats(&self) -> Option<BTreeMap<&'static str, u64>> {
let stats = self.inner.hot_pair_window_stats()?;
Some(BTreeMap::from([
("hydration_scans", stats.hydration_scans),
("hydrated_word_entries", stats.hydrated_word_entries),
("batch_prunes", stats.batch_prunes),
("prune_evictions", stats.prune_evictions),
("peak_resident_pairs", stats.peak_resident_pairs as u64),
("resident_pairs", self.inner.hot_resident_pairs() as u64),
("occurrence_capacity", self.inner.hot_occurrence_capacity() as u64),
]))
}
#[getter]
/// Capacity-backed breakdown of persistent trainer storage.
///
/// This attributes trainer-owned allocations and is not process RSS; use it
/// to identify why bounded occurrence postings do or do not dominate.
pub fn memory_usage(&self) -> BTreeMap<&'static str, u64> {
trainer_memory_usage(&self.inner)
}
/// Initialize internal training state.
pub fn init_training(&mut self, py: Python) -> PyResult<()> {
py.detach(|| self.inner.init_training());
Ok(())
}
/// Train until the vocabulary reaches `vocab_size` or the pair cutoff.
///
/// Returns the updated vocabulary size.
pub fn train_until(&mut self, py: Python, vocab_size: usize) -> PyResult<i64> {
py.detach(|| self.inner.train_until(vocab_size)).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
Ok(self.inner.vocab_size() as i64)
}
/// Train to `vocab_size` with a terminal byte-BPE fallback phase.
///
/// Returns the updated vocabulary size.
#[pyo3(signature = (vocab_size, *, primary_vocab_ratio=0.9))]
pub fn train_until_with_bbpe_fallback(
&mut self,
py: Python,
vocab_size: usize,
primary_vocab_ratio: f64,
) -> PyResult<i64> {
validate_primary_vocab_ratio(primary_vocab_ratio)?;
py.detach(|| self.inner.train_until_with_bbpe_fallback(vocab_size, primary_vocab_ratio))
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
Ok(self.inner.vocab_size() as i64)
}
/// Perform one training step.
///
/// Returns the updated vocabulary size.
pub fn step(&mut self, py: Python) -> PyResult<i64> {
py.detach(|| self.inner.step()).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
Ok(self.inner.vocab_size() as i64)
}
/// Return a view of the current vocabulary.
pub fn get_vocab(&self) -> Vocabulary {
Vocabulary {
inner: Box::new(VocabularyInner::new(&self.inner.vocab)),
}
}
/// Validate the current state and return an immutable model snapshot.
pub fn validate_model(&self, py: Python) -> PyResult<BpeModelBase> {
py.detach(|| self.inner.validate_model())
.map(|model| BpeModelBase { inner: BpeModelInner::Unicode(model) })
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}
}
pub struct VocabularyInner<C, I>(OrderMap<Word<C>, I>);
impl<C: std::hash::Hash + Eq, I: IdxLike> VocabularyInner<C, I> {
/// Build a reverse map from token bytes to token id.
pub fn new(vocab: &BTreeMap<I, Word<C>>) -> Self {
Self(vocab.iter().map(|(i, c)| (c.clone(), i.clone())).collect())
}
}
trait VocabularyImpl {
fn len(&self) -> usize;
fn get(&self, word: &str) -> Option<i64>;
fn items(&self) -> Vec<(Vec<u8>, i64)>;
}
impl<C: CanStrToWord + CharSplit + std::hash::Hash + Eq, I: IdxLike> VocabularyImpl for VocabularyInner<C, I> {
fn len(&self) -> usize {
self.0.len()
}
fn get(&self, word: &str) -> Option<i64> {
self.0.get(&word.to_word()).map(|i| i.to_u64() as i64)
}
fn items(&self) -> Vec<(Vec<u8>, i64)> {
self.0.iter().map(|(w, i)| (CharSplit::to_vec_u8(w), i.to_u64() as i64)).collect()
}
}
#[pyclass]
pub struct Vocabulary {
inner: Box<dyn VocabularyImpl + Send + Sync>,
}
#[pymethods]
impl Vocabulary {
#[getter]
/// Number of entries in the vocabulary.
pub fn len(&self) -> usize {
self.inner.len()
}
/// Look up a word/token and return its id if present.
pub fn get(&self, word: &str) -> Option<i64> {
self.inner.get(word)
}
/// Return all `(token_bytes, id)` entries.
pub fn items(&self) -> Vec<(Vec<u8>, i64)> {
self.inner.items()
}
}
#[pymodule_export]
pub use crate::pretokenizer::PreTokenizer;
#[pymodule_export]
pub use crate::counter::{BigramCounter, WordCounter};
#[pymethods]
impl PreTokenizer {
#[new]
/// Create a Python `PreTokenizer`.
///
/// - `special_tokens`: special tokens to treat as indivisible.
/// - `eot_token`: end-of-text token used for chunk boundary alignment.
/// - `pat_str`: optional regex pattern; defaults to the crate's default.
#[pyo3(signature = (special_tokens, eot_token=None, pat_str=None, unicode_bigrams=None, unicode_bigram_mixed_boundary="keep"))]
pub fn new_py(
special_tokens: Vec<String>,
eot_token: Option<String>,
pat_str: Option<String>,
unicode_bigrams: Option<Vec<String>>,
unicode_bigram_mixed_boundary: &str,
) -> PyResult<Self> {
let mut pretokenizer = Self::try_new(&special_tokens, eot_token.as_deref(), pat_str.as_deref())
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
let unicode_bigram_mixed_boundary = UnicodeBigramMixedBoundary::parse(unicode_bigram_mixed_boundary)
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
if let Some(unicode_bigrams) = unicode_bigrams {
pretokenizer = pretokenizer.with_unicode_bigrams(
parse_unicode_bigrams(&unicode_bigrams)
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?,
);
}
pretokenizer = pretokenizer.with_unicode_bigram_mixed_boundary(unicode_bigram_mixed_boundary);
Ok(pretokenizer)
}
#[pyo3(name = "find_chunk_boundaries", signature = (path, *, chunk_size=1048576, boundary="auto"))]
/// Python wrapper for [`PreTokenizer::find_chunk_boundaries`].
pub fn py_find_chunk_boundaries(
&self, py: Python, path: PathBuf, chunk_size: u64, boundary: &str,
) -> PyResult<Vec<(u64, usize)>> {
let options = chunk_options(chunk_size, boundary)?;
py.detach(||
self.find_chunk_boundaries_with_options(path, options)
).map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
}
#[pyo3(name = "get_words_from_segment")]
/// Python wrapper for [`PreTokenizer::get_words_from_segment`].
pub fn py_get_words_from_segment(
&self, py: Python, path: PathBuf, offset: u64, length: usize,
) -> PyResult<BTreeMap<String, i64>> {
py.detach(||
self.get_words_from_segment(path, offset, length)
).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}
#[pyo3(name = "get_words")]
/// Pretokenize text and return a word-frequency mapping.
pub fn py_get_words(&self, py: Python, text: &str) -> PyResult<BTreeMap<String, i64>> {
py.detach(|| self.get_words_owned(text))
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}
/// Create an empty mergeable Unicode bigram counter.
pub fn bigram_counter(&self) -> BigramCounter {
BigramCounter::new(self.clone())
}
/// Create an empty mergeable word counter.
pub fn word_counter(&self) -> WordCounter {
WordCounter::new(self.clone())
}
pub fn load_word_counter(&self, py: Python, path: PathBuf) -> PyResult<WordCounter> {
let pre_tokenizer = self.clone();
py.detach(|| {
let file = std::fs::File::open(path)?;
WordCounter::load(pre_tokenizer, std::io::BufReader::new(file))
}).map_err(|error| pyo3::exceptions::PyIOError::new_err(error.to_string()))
}
#[pyo3(name = "get_words_from_file", signature = (path, *, chunk_size=1048576, boundary="auto"))]
/// Python wrapper for [`PreTokenizer::get_words_from_file`].
pub fn py_get_words_from_file(
&self, py: Python, path: PathBuf, chunk_size: u64, boundary: &str,
) -> PyResult<BTreeMap<String, i64>> {
let options = chunk_options(chunk_size, boundary)?;
py.detach(||
self.get_words_from_file_with_options(path, options)
).map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
}
#[pyo3(name = "build_unicode_bigrams_from_file", signature = (path, *, chunk_size=1048576, boundary="auto", top_k=100000, min_freq=16))]
pub fn py_build_unicode_bigrams_from_file(
&self, py: Python, path: PathBuf, chunk_size: u64, boundary: &str, top_k: usize, min_freq: i64,
) -> PyResult<Vec<String>> {
Ok(
self
.py_select_unicode_bigrams_from_file(
py,
path,
chunk_size,
boundary,
top_k,
min_freq,
)?
.bigrams,
)
}
#[pyo3(name = "select_unicode_bigrams_from_file", signature = (path, *, chunk_size=1048576, boundary="auto", top_k=100000, min_freq=16))]
pub fn py_select_unicode_bigrams_from_file(
&self, py: Python, path: PathBuf, chunk_size: u64, boundary: &str, top_k: usize, min_freq: i64,
) -> PyResult<UnicodeBigramSelection> {
let options = chunk_options(chunk_size, boundary)?;
let selection = py.detach(||
self.build_unicode_bigram_selection_from_file_with_options(
path,
options,
top_k,
min_freq,
)
).map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))?;
Ok(selection.into())
}
}
fn source_options(max_records: usize, max_bytes: usize) -> PyResult<SourceBatchOptions> {
SourceBatchOptions { max_records, max_bytes }
.validate()
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))
}
fn should_flush_batch(batch_len: usize, batch_bytes: usize, next_bytes: usize, options: SourceBatchOptions) -> bool {
batch_len > 0 && (batch_len >= options.max_records || batch_bytes.saturating_add(next_bytes) > options.max_bytes)
}
trait SourceCounter: Send {
fn add_source_batch(&mut self, texts: &[String]) -> MyResult<()>;
}
impl SourceCounter for BigramCounter {
fn add_source_batch(&mut self, texts: &[String]) -> MyResult<()> {
self.add_batch(texts)
}
}
impl SourceCounter for WordCounter {
fn add_source_batch(&mut self, texts: &[String]) -> MyResult<()> {
self.add_batch(texts)
}
}
fn source_counter_error(error: MyError) -> PyErr {
pyo3::exceptions::PyRuntimeError::new_err(error.to_string())
}
fn validate_prefetch(prefetch: i64) -> PyResult<usize> {
if !(0..=1).contains(&prefetch) {
return Err(pyo3::exceptions::PyValueError::new_err("prefetch must be 0 or 1"));
}
Ok(prefetch as usize)
}
fn for_each_source_batch(
iterator: Bound<PyIterator>,
options: SourceBatchOptions,
mut consume: impl FnMut(Vec<String>) -> PyResult<()>,
) -> PyResult<()> {
let mut batch = Vec::new();
let mut batch_bytes = 0usize;
for item in iterator {
let text = item?.extract::<String>()?;
if should_flush_batch(batch.len(), batch_bytes, text.len(), options) {
consume(std::mem::take(&mut batch))?;
batch_bytes = 0;
}
batch_bytes = batch_bytes.checked_add(text.len())
.ok_or_else(|| pyo3::exceptions::PyOverflowError::new_err("batch byte size overflow"))?;
batch.push(text);
}
if !batch.is_empty() {
consume(batch)?;
}
Ok(())
}
fn add_source_sync<C: SourceCounter>(
counter: &mut C,
py: Python,
source: &Bound<PyAny>,
options: SourceBatchOptions,
) -> PyResult<()> {
for_each_source_batch(source.try_iter()?, options, |batch| {
py.detach(|| counter.add_source_batch(&batch)).map_err(source_counter_error)
})
}
fn add_source_prefetched<C: SourceCounter>(
counter: &mut C,
py: Python,
source: &Bound<PyAny>,
options: SourceBatchOptions,
) -> PyResult<()> {
let iterator = source.try_iter()?;
std::thread::scope(|scope| {
let (sender, receiver) = sync_channel::<Vec<String>>(0);
let worker = scope.spawn(move || -> MyResult<()> {
for batch in receiver {
counter.add_source_batch(&batch)?;
}
Ok(())
});
let producer_result = for_each_source_batch(iterator, options, |batch| {
py.detach(|| sender.send(batch))
.map_err(|_| pyo3::exceptions::PyRuntimeError::new_err("source counter worker stopped"))
});
drop(sender);
match py.detach(|| worker.join()) {
Ok(Ok(())) => producer_result,
Ok(Err(error)) => Err(source_counter_error(error)),
Err(_) => Err(pyo3::exceptions::PyRuntimeError::new_err("source counter worker panicked")),
}
})
}
#[cfg(test)]
fn run_prefetched_batches<C: SourceCounter, E>(
counter: &mut C,
produce: impl FnOnce(&std::sync::mpsc::SyncSender<Vec<String>>) -> Result<(), E>,
) -> (Result<(), E>, std::thread::Result<MyResult<()>>) {
std::thread::scope(|scope| {
let (sender, receiver) = sync_channel::<Vec<String>>(0);
let worker = scope.spawn(move || -> MyResult<()> {
for batch in receiver {
counter.add_source_batch(&batch)?;
}
Ok(())
});
let producer_result = produce(&sender);
drop(sender);
(producer_result, worker.join())
})
}
fn add_python_source<C: SourceCounter>(
counter: &mut C,
py: Python,
source: &Bound<PyAny>,
max_records: usize,
max_bytes: usize,
prefetch: i64,
) -> PyResult<()> {
let options = source_options(max_records, max_bytes)?;
let prefetch = validate_prefetch(prefetch)?;
match prefetch {
0 => add_source_sync(counter, py, source, options),
1 => add_source_prefetched(counter, py, source, options),
_ => unreachable!(),
}
}
#[cfg(test)]
mod source_prefetch_tests {
use std::{sync::mpsc::sync_channel, time::Duration};
use super::*;
struct BlockingCounter {
processing_started: std::sync::mpsc::SyncSender<()>,
next_batch_started: std::sync::mpsc::Receiver<()>,
calls: usize,
}
impl SourceCounter for BlockingCounter {
fn add_source_batch(&mut self, _texts: &[String]) -> MyResult<()> {
if self.calls == 0 {
self.processing_started.send(()).unwrap();
self.next_batch_started.recv_timeout(Duration::from_secs(1)).unwrap();
}
self.calls += 1;
Ok(())
}
}
#[test]
fn prefetched_batches_overlap_production_and_processing() {
let (processing_started_tx, processing_started_rx) = sync_channel(0);
let (next_batch_started_tx, next_batch_started_rx) = sync_channel(0);
let mut counter = BlockingCounter {
processing_started: processing_started_tx,
next_batch_started: next_batch_started_rx,
calls: 0,
};
let (producer, worker) = run_prefetched_batches(&mut counter, |sender| {
sender.send(vec!["first".to_string()]).unwrap();
processing_started_rx.recv_timeout(Duration::from_secs(1)).unwrap();
next_batch_started_tx.send(()).unwrap();
sender.send(vec!["second".to_string()]).unwrap();
Ok::<_, ()>(())
});
assert_eq!(producer, Ok(()));
assert!(worker.unwrap().is_ok());
assert_eq!(counter.calls, 2);
}
}
#[pymethods]
impl BigramCounter {
#[new]
pub fn new_py(pre_tokenizer: PreTokenizer) -> Self {
Self::new(pre_tokenizer)
}
#[pyo3(name = "add_text")]
pub fn py_add_text(&mut self, py: Python, text: &str) -> PyResult<()> {
py.detach(|| self.add_text(text))
.map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))
}
#[pyo3(name = "add_batch")]
pub fn py_add_batch(&mut self, py: Python, texts: Vec<String>) -> PyResult<()> {
py.detach(|| self.add_batch(&texts))
.map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))
}
#[pyo3(name = "add_source", signature = (source, *, max_records=4096, max_bytes=67108864, prefetch=1))]
pub fn py_add_source(
&mut self, py: Python, source: &Bound<PyAny>, max_records: usize, max_bytes: usize, prefetch: i64,
) -> PyResult<()> {
add_python_source(self, py, source, max_records, max_bytes, prefetch)
}
#[pyo3(name = "merge")]
pub fn py_merge(&mut self, py: Python, other: BigramCounter) -> PyResult<()> {
py.detach(|| self.merge(other))
.map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))
}
#[pyo3(name = "selected")]
pub fn py_selected(&self, top_k: usize, min_freq: i64) -> Vec<String> {
self.py_select(top_k, min_freq).bigrams
}
#[pyo3(name = "select")]
pub fn py_select(&self, top_k: usize, min_freq: i64) -> UnicodeBigramSelection {
self.selection(top_k, min_freq).into()
}
pub fn items(&self) -> Vec<(String, i64)> {
let mut items = self.counts().iter()
.map(|(bigram, frequency)| (unicode_bigram_to_string(*bigram), *frequency))
.collect::<Vec<_>>();
items.sort_unstable_by(|left, right| left.0.cmp(&right.0));
items
}
}
#[pymethods]
impl WordCounter {
#[new]
pub fn new_py(pre_tokenizer: PreTokenizer) -> Self {
Self::new(pre_tokenizer)
}
#[pyo3(name = "add_text")]
pub fn py_add_text(&mut self, py: Python, text: &str) -> PyResult<()> {
py.detach(|| self.add_text(text))
.map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))
}
#[pyo3(name = "add_batch")]
pub fn py_add_batch(&mut self, py: Python, texts: Vec<String>) -> PyResult<()> {
py.detach(|| self.add_batch(&texts))
.map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))
}
#[pyo3(name = "add_source", signature = (source, *, max_records=4096, max_bytes=67108864, prefetch=1))]
pub fn py_add_source(
&mut self, py: Python, source: &Bound<PyAny>, max_records: usize, max_bytes: usize, prefetch: i64,
) -> PyResult<()> {
add_python_source(self, py, source, max_records, max_bytes, prefetch)
}
#[pyo3(name = "merge")]
pub fn py_merge(&mut self, py: Python, other: WordCounter) -> PyResult<()> {
py.detach(|| self.merge(other))
.map_err(|error| pyo3::exceptions::PyRuntimeError::new_err(error.to_string()))
}
#[pyo3(name = "words")]
pub fn py_words(&self) -> BTreeMap<String, i64> {
WordCounter::words(self)
}
#[getter]
#[pyo3(name = "len")]
pub fn py_len(&self) -> usize {
WordCounter::len(self)
}
#[pyo3(name = "clear")]
pub fn py_clear(&mut self) {
WordCounter::clear(self)
}
#[pyo3(name = "save")]
pub fn py_save(&self, py: Python, path: PathBuf) -> PyResult<()> {
py.detach(|| {
let file = std::fs::File::create(path)?;
self.save(std::io::BufWriter::new(file))
}).map_err(|error| pyo3::exceptions::PyIOError::new_err(error.to_string()))
}
}
#[pyclass]
pub struct BpeEncoderBase(Arc<dyn Encoder<Idx> + Send + Sync>);
fn _arc_to_vec<I: Copy>(i: Arc<[I]>) -> Vec<I> {
i.iter().copied().collect()
}
fn new_bpe<C: Clone>(
vocab: Option<BTreeMap<Vec<u8>, Idx>>,
merges: Option<Vec<(Vec<u8>, Vec<u8>)>>,
vocab_file: Option<PathBuf>,
merges_file: Option<PathBuf>,
special_tokens: Option<Vec<String>>,
pat_str: Option<String>,
unicode_bigrams: Option<Vec<String>>,
unicode_bigram_mixed_boundary: &str,
split_on_vocab_bigrams: bool,
spec: &dyn Spec<C, Idx>,
) -> MyResult<BpeEncoderBase>
where
BpeEncoder<C>: CanEncode<C, Idx>
{
let mut builder = BpeBuilder::new();
if let Some(filename) = vocab_file {
builder = builder.load_vocab_file(filename, spec)?;
} else if let Some(vocab) = vocab {
builder = builder.set_vocab(vocab.into_iter().map(|(k, v)| (v, k)).collect());
} else {
return Err(MyError::BpeBuilder("Either vocab_file or vocab must be provided".to_string()));
}
if let Some(filename) = merges_file {
builder = builder.load_merges_file(filename, spec)?;
} else if let Some(merges) = merges {
builder = builder.set_merges_raw(merges);
} else {
return Err(MyError::BpeBuilder("Either merges_file or merges must be provided".to_string()));
}
builder = builder.set_special_tokens(special_tokens);
builder = builder.set_pat_str(pat_str);
builder = builder.set_split_on_vocab_bigrams(split_on_vocab_bigrams);
let bpe = configure_encoder(
builder.build(spec)?,
unicode_bigrams.as_deref(),
unicode_bigram_mixed_boundary,
)?;
Ok(BpeEncoderBase(Arc::new(bpe)))
}
fn configure_encoder<C>(
mut encoder: BpeEncoder<C>,
unicode_bigrams: Option<&[String]>,
unicode_bigram_mixed_boundary: &str,
) -> MyResult<BpeEncoder<C>>
where
BpeEncoder<C>: CanEncode<C, Idx>,
C: Clone,
{
if let Some(bigrams) = unicode_bigrams {
encoder.pre_tokenizer = encoder.pre_tokenizer.with_unicode_bigrams(
parse_unicode_bigrams(bigrams)?,
);
}
encoder.pre_tokenizer = encoder.pre_tokenizer.with_unicode_bigram_mixed_boundary(
UnicodeBigramMixedBoundary::parse(unicode_bigram_mixed_boundary)?,
);
Ok(encoder)
}
#[pymethods]
impl BpeEncoderBase {
#[new]
#[pyo3(signature = (format, unit, vocab, merges, vocab_file, merges_file, special_tokens, pat_str=None, unicode_bigrams=None, unicode_bigram_mixed_boundary="keep", split_on_vocab_bigrams=true))]
/// Create a Python BPE encoder.
///
/// The encoder can be created from in-memory `vocab`/`merges` or from file paths.
/// `format` and `unit` must be compatible.
/// Set `split_on_vocab_bigrams` to `false` to keep PAT words intact during BPE.
pub fn new_py(
py: Python,
format: &str, unit: &str,
vocab: Option<BTreeMap<Vec<u8>, Idx>>,
merges: Option<Vec<(Vec<u8>, Vec<u8>)>>,
vocab_file: Option<PathBuf>,
merges_file: Option<PathBuf>,
special_tokens: Option<Vec<String>>,
pat_str: Option<String>,
unicode_bigrams: Option<Vec<String>>,
unicode_bigram_mixed_boundary: &str,
split_on_vocab_bigrams: bool,
) -> PyResult<Self> {
py.detach(||
match (format, unit) {
("gpt2", "byte") => new_bpe::<u8>(vocab, merges, vocab_file, merges_file, special_tokens, pat_str, unicode_bigrams, unicode_bigram_mixed_boundary, split_on_vocab_bigrams, &Gpt2Spec),
("unitoken", "byte") => new_bpe::<u8>(vocab, merges, vocab_file, merges_file, special_tokens, pat_str, unicode_bigrams, unicode_bigram_mixed_boundary, split_on_vocab_bigrams, &UnitokenSpec),
("unitoken", "unicode") => new_bpe::<Character>(vocab, merges, vocab_file, merges_file, special_tokens, pat_str, unicode_bigrams, unicode_bigram_mixed_boundary, split_on_vocab_bigrams, &UnitokenSpec),
_ => Err(MyError::SpecError(format!("format {format} is not compatible with unit {unit}"))),
}
).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}
#[pyo3(name = "pre_tokenizer")]
/// Return the underlying pre-tokenizer.
pub fn py_pre_tokenizer(&self) -> PreTokenizer {
self.0.pre_tokenizer().clone()
}
#[pyo3(name = "encode_word")]
/// Encode one already-pretokenized word without PAT or special-token handling.
pub fn py_encode_word(&self, py: Python, word: &str) -> PyResult<Vec<Idx>> {
py.detach(||
self.0.encode_word(word).map(_arc_to_vec)
).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}
#[pyo3(name = "encode_words")]
/// Encode multiple already-pretokenized words.
pub fn py_encode_words(&self, py: Python, words: Vec<String>) -> PyResult<Vec<Vec<Idx>>> {
py.detach(|| {
let words = words.iter().map(|i| i.as_str()).collect::<Vec<_>>();
let result = self.0.encode_words(&words)?;
Ok(result.into_iter().map(_arc_to_vec).collect())
}).map_err(|e: MyError| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}
#[pyo3(name = "encode_to_numpy")]
/// Encode an arbitrary string into a NumPy `uint32` array of token ids.
pub fn py_encode_to_numpy<'py>(&self, py: Python<'py>, text: &str) -> PyResult<Bound<'py, PyArray1<Idx>>> {
let result = py.detach(|| {
self.0.encode_string(text)
});
match result {
Ok(v) => Ok(v.into_pyarray(py)),
Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(e.to_string())),
}
}
#[pyo3(name = "encode")]
/// Encode an arbitrary string into a Python list of token ids.
pub fn py_encode(&self, py: Python, text: &str) -> PyResult<Vec<Idx>> {
py.detach(|| {
self.0.encode_string(text)
}).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}
#[pyo3(name = "encode_file")]
/// Encode a file into a NumPy `uint32` array of token ids.
pub fn py_encode_file<'py>(&self, py: Python<'py>, path: PathBuf, num_chunks: usize) -> PyResult<Bound<'py, PyArray1<Idx>>> {
let result = py.detach(||
self.0.encode_file(&path, num_chunks)
);
match result {
Ok(v) => Ok(v.into_pyarray(py)),
Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(e.to_string())),
}
}
#[pyo3(name = "decode")]
/// Decode token ids back into a UTF-8 string.
///
/// Accepts either a Python sequence of ints or a NumPy `uint32` array.
pub fn py_decode(&self, py: Python, idxs: &Bound<PyAny>) -> PyResult<String> {
let vec: Vec<Idx> = if let Ok(v) = idxs.extract::<Vec<Idx>>() {
v
} else if let Ok(arr) = idxs.extract::<PyReadonlyArray1<Idx>>() {
arr.as_array().iter().copied().collect()
} else {
return Err(pyo3::exceptions::PyTypeError::new_err(
"idxs must be a sequence[int] or a numpy.ndarray[uint32]",
));
};
py.detach(|| self.0.decode(&vec))
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}
}
// #[pymodule(gil_used = false)]
// #[pyo3(name="_lib")]
// fn _tiktoken(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
// m.add_class::<BpeTrainerBase>()?;
// m.add_class::<BpeTrainer_u8_Idx>()?;
// m.add_class::<BpeTrainer_Character_CharIdx>()?;
// Ok(())
// }
}
#[test]
#[ignore = "manual"]
fn generate_py_stubs() {
println!("test");
let package_dir = std::path::Path::new("./python/ffbpe");
let mut libraries = std::fs::read_dir(package_dir)
.expect("Python package directory to exist")
.filter_map(Result::ok)
.map(|entry| entry.path())
.filter(|path| {
path.file_name()
.is_some_and(|name| name.to_string_lossy().starts_with("_lib."))
&& path.extension().is_some_and(|extension| {
matches!(extension.to_str(), Some("so" | "pyd" | "dylib" | "dll"))
})
})
.collect::<Vec<_>>();
libraries.sort();
let library = libraries.into_iter().next()
.expect("built Python extension to exist");
let module = pyo3_introspection::introspect_cdylib(
library,
"_lib",
)
.expect("introspection to succeed");
let result = pyo3_introspection::module_stub_files(&module);
println!("{result:?}");
let value = result.get(&std::path::PathBuf::from("__init__.pyi")).unwrap();
std::fs::write("./python/ffbpe/_lib.pyi", value).unwrap();
}