Added AddedToken class

Author: ankane

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## 0.5.3 (unreleased)
 
+- Added `AddedToken` class
 - Added precompiled gem for Windows
 
 ## 0.5.2 (2024-08-26)

ext/tokenizers/src/lib.rs

@@ -15,7 +15,7 @@ mod utils;
 
 use encoding::RbEncoding;
 use error::RbError;
-use tokenizer::RbTokenizer;
+use tokenizer::{RbAddedToken, RbTokenizer};
 use utils::RbRegex;
 
 use magnus::{function, method, prelude::*, value::Lazy, Error, RModule, Ruby};
@@ -109,6 +109,10 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     let class = module.define_class("Regex", ruby.class_object())?;
     class.define_singleton_method("new", function!(RbRegex::new, 1))?;
 
+    let class = module.define_class("AddedToken", ruby.class_object())?;
+    class.define_singleton_method("_new", function!(RbAddedToken::new, 2))?;
+    class.define_method("content", method!(RbAddedToken::get_content, 0))?;
+
     let models = module.define_module("Models")?;
     let pre_tokenizers = module.define_module("PreTokenizers")?;
     let decoders = module.define_module("Decoders")?;

ext/tokenizers/src/tokenizer.rs

@@ -22,20 +22,21 @@ use super::processors::RbPostProcessor;
 use super::trainers::RbTrainer;
 use super::{RbError, RbResult};
 
+#[magnus::wrap(class = "Tokenizers::AddedToken")]
 pub struct RbAddedToken {
     pub content: String,
-    pub is_special_token: bool,
+    pub special: bool,
     pub single_word: Option<bool>,
     pub lstrip: Option<bool>,
     pub rstrip: Option<bool>,
     pub normalized: Option<bool>,
 }
 
 impl RbAddedToken {
-    pub fn from<S: Into<String>>(content: S, is_special_token: Option<bool>) -> Self {
+    pub fn from<S: Into<String>>(content: S, special: Option<bool>) -> Self {
         Self {
             content: content.into(),
-            is_special_token: is_special_token.unwrap_or(false),
+            special: special.unwrap_or(false),
             single_word: None,
             lstrip: None,
             rstrip: None,
@@ -44,7 +45,7 @@ impl RbAddedToken {
     }
 
     pub fn get_token(&self) -> tk::tokenizer::AddedToken {
-        let mut token = tk::AddedToken::from(&self.content, self.is_special_token);
+        let mut token = tk::AddedToken::from(&self.content, self.special);
 
         if let Some(sw) = self.single_word {
             token = token.single_word(sw);
@@ -71,11 +72,53 @@ impl From<tk::AddedToken> for RbAddedToken {
             lstrip: Some(token.lstrip),
             rstrip: Some(token.rstrip),
             normalized: Some(token.normalized),
-            is_special_token: !token.normalized,
+            special: !token.normalized,
         }
     }
 }
 
+impl RbAddedToken {
+    pub fn new(content: Option<String>, kwargs: RHash) -> RbResult<Self> {
+        let mut token = RbAddedToken::from(content.unwrap_or("".to_string()), None);
+
+        let value: Value = kwargs.delete(Symbol::new("single_word"))?;
+        if !value.is_nil() {
+            token.single_word = TryConvert::try_convert(value)?;
+        }
+
+        let value: Value = kwargs.delete(Symbol::new("lstrip"))?;
+        if !value.is_nil() {
+            token.lstrip = TryConvert::try_convert(value)?;
+        }
+
+        let value: Value = kwargs.delete(Symbol::new("rstrip"))?;
+        if !value.is_nil() {
+            token.rstrip = TryConvert::try_convert(value)?;
+        }
+
+        let value: Value = kwargs.delete(Symbol::new("normalized"))?;
+        if !value.is_nil() {
+            token.normalized = TryConvert::try_convert(value)?;
+        }
+
+        let value: Value = kwargs.delete(Symbol::new("special"))?;
+        if !value.is_nil() {
+            token.special = TryConvert::try_convert(value)?;
+        }
+
+        if !kwargs.is_empty() {
+            // TODO improve message
+            return Err(Error::new(exception::arg_error(), "unknown keyword"));
+        }
+
+        Ok(token)
+    }
+
+    pub fn get_content(&self) -> String {
+        self.content.to_string()
+    }
+}
+
 struct TextInputSequence<'s>(tk::InputSequence<'s>);
 
 impl<'s> TryConvert for TextInputSequence<'s> {

lib/tokenizers.rb

@@ -42,6 +42,7 @@
 require_relative "tokenizers/trainers/word_piece_trainer"
 
 # other
+require_relative "tokenizers/added_token"
 require_relative "tokenizers/char_bpe_tokenizer"
 require_relative "tokenizers/encoding"
 require_relative "tokenizers/from_pretrained"

lib/tokenizers/added_token.rb

@@ -0,0 +1,7 @@
+module Tokenizers
+  class AddedToken
+    def self.new(content, **kwargs)
+      _new(content, kwargs)
+    end
+  end
+end

test/added_token_test.rb

@@ -0,0 +1,8 @@
+require_relative "test_helper"
+
+class AddedTokenTest < Minitest::Test
+  def test_content
+    token = Tokenizers::AddedToken.new("test")
+    assert_equal "test", token.content
+  end
+end