unitoken 0.1.4

#[pyo3::pymodule(gil_used = false)]
mod _lib {
use std::{collections::BTreeMap, path::PathBuf, sync::Arc};

use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1};
use ordermap::OrderMap;
use pyo3::{prelude::*, pymethods, types::PyAny};

use crate::{MyError, MyResult, bpe::{BpeEncoder, BpeTrainer, BpeTrainerConfig, CharIdx, CharSplit, Character, Idx, IdxLike, InitialAlphabet, TieBreak, Word, encoder::BpeBuilder, utils::ToWord}, pretokenizer::{BoundaryMode, ChunkHint, ChunkOptions, UnicodeBigramMixedBoundary, parse_unicode_bigrams, unicode_bigram_to_string}, spec::{Spec, gpt2::Gpt2Spec, uni::UniSpec}, traits::{CanEncode, CanStrToWord, Encoder, Train as _}};

#[pyclass(subclass)]
pub struct BpeTrainerBase;

#[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 init_training(&mut self, py: Python);
  fn train_until(&mut self, py: Python, vocab_size: usize) -> PyResult<i64>;
  fn step(&mut self, py: Python) -> PyResult<i64>;
  fn get_vocabs(&self) -> Vocabs;
  fn save_vocab(&self, py: Python, path: PathBuf, spec: &str) -> PyResult<()>;
  fn save_merges_txt(&self, py: Python, path: PathBuf, spec: &str) -> PyResult<()>;
}

// #[pyclass(eq, eq_int)]
// #[derive(PartialEq)]
// pub enum SpecEnum {
//   #[pyo3(name = "gpt2")]
//   Gpt2,
//   #[pyo3(name = "uni")]
//   Uni,
// }

// #[pyclass(eq, eq_int)]
// #[derive(PartialEq)]
// pub enum CharLevel {
//   #[pyo3(name = "u8")]
//   U8,
//   #[pyo3(name = "char")]
//   Char,
// }

fn trainer_config(
  initial_alphabet: Option<&str>,
  tie_break: Option<&str>,
  parallel_merge_min_occurs_in: Option<usize>,
) -> 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}"))),
  };
  Ok(BpeTrainerConfig {
    initial_alphabet,
    tie_break,
    parallel_merge_min_occurs_in,
  })
}

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))]
  pub fn new_py(
    special_tokens: Vec<String>,
    initial_alphabet: Option<&str>,
    tie_break: Option<&str>,
    parallel_merge_min_occurs_in: Option<usize>,
  ) -> PyResult<(Self, BpeTrainerBase)> {
    let config = trainer_config(initial_alphabet, tie_break, parallel_merge_min_occurs_in)?;
    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)))
    )
  }

  /// Current vocabulary size.
  pub fn vocab_size(&self) -> usize {
    self.inner.vocab_size()
  }

  /// Initialize internal training state.
  pub fn init_training(&mut self, py: Python) {
    py.detach(|| self.inner.init_training())
  }

  /// Train until the vocabulary reaches `vocab_size`.
  ///
  /// 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_vocabs(&self) -> Vocabs {
    Vocabs {
      inner: Box::new(VocabsInner::new(&self.inner.vocab)),
    }
  }

  /// Save the vocabulary JSON to `path` using the requested spec (`"gpt2"` or `"uni"`).
  pub fn save_vocab(&self, py: Python, path: PathBuf, spec: &str) -> PyResult<()> {
    py.detach(|| {
      let mut file = std::fs::File::create(&path)?;
      let mut writer = std::io::BufWriter::new(&mut file);
      match spec {
        "gpt2" => self.inner.save_vocab_json(&Gpt2Spec, &mut writer),
        "uni" => self.inner.save_vocab_json(&UniSpec, &mut writer),
        _ => Err(MyError::SpecError(format!("Unknown spec: {}", spec))),
      }
    }).map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
  }

  /// Save merges to `path` using the requested spec (`"gpt2"` or `"uni"`).
  pub fn save_merges_txt(&self, py: Python, path: PathBuf, spec: &str) -> PyResult<()> {
    py.detach(|| {
      let mut file = std::fs::File::create(&path)?;
      let mut writer = std::io::BufWriter::new(&mut file);
      match spec {
        "gpt2" => self.inner.save_merges_txt(&Gpt2Spec, &mut writer),
        "uni" => self.inner.save_merges_txt(&UniSpec, &mut writer),
        _ => Err(MyError::SpecError(format!("Unknown spec: {}", spec))),
      }
    }).map_err(|e| pyo3::exceptions::PyIOError::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))]
  pub fn new_py(
    special_tokens: Vec<String>,
    initial_alphabet: Option<&str>,
    tie_break: Option<&str>,
    parallel_merge_min_occurs_in: Option<usize>,
  ) -> PyResult<(Self, BpeTrainerBase)> {
    let config = trainer_config(initial_alphabet, tie_break, parallel_merge_min_occurs_in)?;
    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)))
    )
  }

  /// Current vocabulary size.
  pub fn vocab_size(&self) -> usize {
    self.inner.vocab_size()
  }

  /// Initialize internal training state.
  pub fn init_training(&mut self, py: Python) {
    py.detach(|| self.inner.init_training())
  }

  /// Train until the vocabulary reaches `vocab_size`.
  ///
  /// 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_vocabs(&self) -> Vocabs {
    Vocabs {
      inner: Box::new(VocabsInner::new(&self.inner.vocab)),
    }
  }

  /// Save the vocabulary JSON to `path`.
  ///
  /// Note: `"gpt2"` is not supported for the character tokenizer.
  pub fn save_vocab(&self, py: Python, path: PathBuf, spec: &str) -> PyResult<()> {
    py.detach(|| {
      let mut file = std::fs::File::create(&path)?;
      let mut writer = std::io::BufWriter::new(&mut file);
      match spec {
        "gpt2" => Err(MyError::SpecError("gpt2 spec not supported for Character tokenizer".to_string())),
        "uni" => self.inner.save_vocab_json(&UniSpec, &mut writer),
        _ => Err(MyError::SpecError(format!("Unknown spec: {}", spec))),
      }
    }).map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
  }

  /// Save merges to `path`.
  ///
  /// Note: `"gpt2"` is not supported for the character tokenizer.
  pub fn save_merges_txt(&self, py: Python, path: PathBuf, spec: &str) -> PyResult<()> {
    py.detach(|| {
      let mut file = std::fs::File::create(&path)?;
      let mut writer = std::io::BufWriter::new(&mut file);
      match spec {
        "gpt2" => Err(MyError::SpecError("gpt2 spec not supported for Character tokenizer".to_string())),
        "uni" => self.inner.save_merges_txt(&UniSpec, &mut writer),
        _ => Err(MyError::SpecError(format!("Unknown spec: {}", spec))),
      }
    }).map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
  }
}

pub struct VocabsInner<C, I>(OrderMap<Word<C>, I>);

impl<C: std::hash::Hash + Eq, I: IdxLike> VocabsInner<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 VocabsImpl {
  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> VocabsImpl for VocabsInner<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 Vocabs {
  inner: Box<dyn VocabsImpl + Send + Sync>,
}

#[pymethods]
impl Vocabs {
  #[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;

#[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`: optional regex pattern; defaults to the crate's default.
  #[pyo3(signature = (special_tokens, eot_token=None, pat=None, unicode_bigrams=None, unicode_bigram_mixed_boundary="keep"))]
  pub fn new_py(
    special_tokens: Vec<String>,
    eot_token: Option<String>,
    pat: 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.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_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>> {
    let options = chunk_options(chunk_size, boundary)?;
    let bigrams = py.detach(||
      self.build_unicode_bigram_set_from_file_with_options(path, options, top_k, min_freq)
    ).map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))?;
    let mut bigrams = bigrams.into_iter().collect::<Vec<_>>();
    bigrams.sort();
    Ok(bigrams.into_iter().map(unicode_bigram_to_string).collect())
  }
}

#[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>(
  vocabs: Option<BTreeMap<Vec<u8>, Idx>>,
  merges: Option<Vec<(Vec<u8>, Vec<u8>)>>,
  vocab_filename: Option<PathBuf>,
  merges_filename: Option<PathBuf>,
  special_tokens: Option<Vec<String>>,
  pat_str: Option<String>,
  spec: &dyn Spec<C, Idx>,
) -> MyResult<BpeEncoderBase>
where
  BpeEncoder<C>: CanEncode<C, Idx>
{
  let mut builder = BpeBuilder::new();
  if let Some(filename) = vocab_filename {
    builder = builder.load_vocab_file(filename, spec)?;
  } else if let Some(vocabs) = vocabs {
    builder = builder.set_vocab(vocabs.into_iter().map(|(k, v)| (v, k)).collect());
  } else {
    return Err(MyError::BpeBuilder("Either vocab_filename or vocabs must be provided".to_string()));
  }
  if let Some(filename) = merges_filename {
    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_filename or merges must be provided".to_string()));
  }
  builder= builder.set_special_tokens(special_tokens);
  builder = builder.set_pat_str(pat_str);
  let bpe = builder.build(spec)?;
  Ok(BpeEncoderBase(Arc::new(bpe)))
}

#[pymethods]
impl BpeEncoderBase {
  #[new]
  #[pyo3(signature = (spec, char_level, vocabs, merges, vocab_filename, merges_filename, special_tokens, pat_str=None))]
  /// Create a Python BPE encoder.
  ///
  /// The encoder can be created from in-memory `vocabs`/`merges` or from file paths.
  /// `spec` and `char_level` must be compatible (e.g. `("gpt2", "u8")`, `("uni", "char")`).
  pub fn new_py(
    py: Python,
    spec: &str, char_level: &str,
    vocabs: Option<BTreeMap<Vec<u8>, Idx>>,
    merges: Option<Vec<(Vec<u8>, Vec<u8>)>>,
    vocab_filename: Option<PathBuf>,
    merges_filename: Option<PathBuf>,
    special_tokens: Option<Vec<String>>,
    pat_str: Option<String>,
  ) -> PyResult<Self> {
    py.detach(||
      match (spec, char_level) {
        ("gpt2", "u8") => new_bpe::<u8>(vocabs, merges, vocab_filename, merges_filename, special_tokens, pat_str, &Gpt2Spec),
        ("uni", "u8") => new_bpe::<u8>(vocabs, merges, vocab_filename, merges_filename, special_tokens, pat_str, &UniSpec),
        ("uni", "char") => new_bpe::<Character>(vocabs, merges, vocab_filename, merges_filename, special_tokens, pat_str, &UniSpec),
        _ => Err(MyError::SpecError(format!("spec {spec} not compatibale with {char_level}"))),
      }
    ).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 a single word into token ids.
  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 words into token ids.
  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_string")]
  /// Encode an arbitrary string into a NumPy `uint32` array of token ids.
  pub fn py_encode_string<'py>(&self, py: Python<'py>, s: &str) -> PyResult<Bound<'py, PyArray1<Idx>>> {
    let result = py.detach(|| {
      self.0.encode_string(s)
    });
    match result {
      Ok(v) => Ok(v.into_pyarray(py)),
      Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(e.to_string())),
    }
  }

  #[pyo3(name = "encode_string_to_list")]
  /// Encode an arbitrary string into a Python list of token ids.
  pub fn py_encode_string_to_list(&self, py: Python, s: &str) -> PyResult<Vec<Idx>> {
    py.detach(|| {
      self.0.encode_string(s)
    }).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 module = pyo3_introspection::introspect_cdylib(
      "./python/uni_tokenizer/_lib.cpython-313-darwin.so",
      "_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/uni_tokenizer/_lib.pyi", value).unwrap();
}