Fixing the public API of the Rust trainer.

Author: Narsil

bindings/node/Cargo.toml

@@ -2,7 +2,7 @@
 authors = ["Nicolas Patry <nicolas@huggingface.co>"]
 edition = "2021"
 name    = "node"
-version = "0.21.3-dev.0"
+version = "0.21.4-dev.0"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 

bindings/python/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tokenizers-python"
-version = "0.21.3-dev.0"
+version = "0.21.4-dev.0"
 authors = ["Anthony MOI <m.anthony.moi@gmail.com>"]
 edition = "2021"
 

tokenizers/Cargo.toml

@@ -2,7 +2,7 @@
 authors = ["Anthony MOI <m.anthony.moi@gmail.com>", "Nicolas Patry <patry.nicolas@protonmail.com>"]
 edition = "2018"
 name = "tokenizers"
-version = "0.21.3-dev.0"
+version = "0.21.4-dev.0"
 homepage = "https://github.com/huggingface/tokenizers"
 repository = "https://github.com/huggingface/tokenizers"
 documentation = "https://docs.rs/tokenizers/"

tokenizers/src/models/bpe/trainer.rs

@@ -9,6 +9,7 @@ use compact_str::CompactString;
 use dary_heap::OctonaryHeap;
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
+use std::collections::HashSet;
 
 #[derive(Debug, Eq)]
 struct Merge {
@@ -116,8 +117,10 @@ impl BpeTrainerBuilder {
 
     /// Set the initial alphabet
     #[must_use]
-    pub fn initial_alphabet(mut self, alphabet: AHashSet<char>) -> Self {
-        self.config.initial_alphabet = alphabet;
+    pub fn initial_alphabet(mut self, alphabet: HashSet<char>) -> Self {
+        let mut initial_alphabet = AHashSet::with_capacity(alphabet.len());
+        initial_alphabet.extend(alphabet);
+        self.config.initial_alphabet = initial_alphabet;
         self
     }
 

tokenizers/src/models/wordpiece/trainer.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use super::WordPiece;
 use crate::models::bpe::{BpeTrainer, BpeTrainerBuilder, BPE};
 use crate::tokenizer::{AddedToken, Result, Trainer};
@@ -61,7 +63,7 @@ impl WordPieceTrainerBuilder {
 
     /// Set the initial alphabet
     #[must_use]
-    pub fn initial_alphabet(mut self, alphabet: AHashSet<char>) -> Self {
+    pub fn initial_alphabet(mut self, alphabet: HashSet<char>) -> Self {
         self.bpe_trainer_builder = self.bpe_trainer_builder.initial_alphabet(alphabet);
         self
     }
@@ -138,8 +140,10 @@ impl WordPieceTrainer {
         &self.bpe_trainer.initial_alphabet
     }
 
-    pub fn set_initial_alphabet(&mut self, alphabet: AHashSet<char>) {
-        self.bpe_trainer.initial_alphabet = alphabet;
+    pub fn set_initial_alphabet(&mut self, alphabet: HashSet<char>) {
+        let mut initial_alphabet = AHashSet::with_capacity(alphabet.len());
+        initial_alphabet.extend(alphabet);
+        self.bpe_trainer.initial_alphabet = initial_alphabet;
     }
 
     pub fn continuing_subword_prefix(&self) -> &Option<String> {