Support addSpecialTokens parameter

Author: budgetdevv

Native/src/lib.rs

@@ -253,31 +253,46 @@ pub unsafe extern "C" fn free_tokenizer(tokenizer_handle: *mut Tokenizer)
 #[no_mangle]
 pub unsafe extern "C" fn tokenizer_encode(
     tokenizer_ptr: *mut Tokenizer,
-    text_buffer: ReadOnlyBuffer<u8>) -> TokenizeOutput
+    text_buffer: ReadOnlyBuffer<u8>,
+    add_special_tokens: bool)
+    -> TokenizeOutput
 {
-    return tokenizer_encode_core(tokenizer_ptr, text_buffer, true);
+    return tokenizer_encode_core(
+        tokenizer_ptr,
+        text_buffer,
+        true,
+        add_special_tokens
+    );
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn tokenizer_encode_non_truncating(
     tokenizer_ptr: *mut Tokenizer,
-    text_buffer: ReadOnlyBuffer<u8>) -> TokenizeOutput
+    text_buffer: ReadOnlyBuffer<u8>,
+    add_special_tokens: bool)
+    -> TokenizeOutput
 {
-    return tokenizer_encode_core(tokenizer_ptr, text_buffer, false);
+    return tokenizer_encode_core(
+        tokenizer_ptr,
+        text_buffer,
+        false,
+        add_special_tokens
+    );
 }
 
 #[inline(always)]
 pub unsafe extern "C" fn tokenizer_encode_core(
     tokenizer_ptr: *mut Tokenizer,
     text_buffer: ReadOnlyBuffer<u8>,
-    truncate: bool)
+    truncate: bool,
+    add_special_tokens: bool)
     -> TokenizeOutput
 {
     let tokenizer = &*tokenizer_ptr;
 
     let text = std::str::from_utf8_unchecked(text_buffer.as_slice());
 
-    let encoded_result = tokenizer.encode_fast(text, true);
+    let encoded_result = tokenizer.encode_fast(text, add_special_tokens);
 
     let encoded_tokens = match encoded_result
     {
@@ -292,26 +307,41 @@ pub unsafe extern "C" fn tokenizer_encode_core(
 pub unsafe extern "C" fn tokenizer_encode_batch(
     tokenizer_ptr: *mut Tokenizer,
     text_buffers: ReadOnlyBuffer<ReadOnlyBuffer<u8>>,
-    output_buffer: Buffer<TokenizeOutput>)
+    output_buffer: Buffer<TokenizeOutput>,
+    add_special_tokens: bool)
 {
-    tokenizer_encode_batch_core(tokenizer_ptr, text_buffers, output_buffer, true);
+    tokenizer_encode_batch_core(
+        tokenizer_ptr,
+        text_buffers,
+        output_buffer,
+        true,
+        add_special_tokens
+    );
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn tokenizer_encode_batch_non_truncating(
     tokenizer_ptr: *mut Tokenizer,
     text_buffers: ReadOnlyBuffer<ReadOnlyBuffer<u8>>,
-    output_buffer: Buffer<TokenizeOutput>)
+    output_buffer: Buffer<TokenizeOutput>,
+    add_special_tokens: bool)
 {
-    tokenizer_encode_batch_core(tokenizer_ptr, text_buffers, output_buffer, false);
+    tokenizer_encode_batch_core(
+        tokenizer_ptr,
+        text_buffers,
+        output_buffer,
+        false,
+        add_special_tokens
+    );
 }
 
 #[inline(always)]
 pub unsafe extern "C" fn tokenizer_encode_batch_core(
     tokenizer_ptr: *mut Tokenizer,
     text_buffers: ReadOnlyBuffer<ReadOnlyBuffer<u8>>,
     output_buffer: Buffer<TokenizeOutput>,
-    truncate: bool)
+    truncate: bool,
+    add_special_tokens: bool)
 {
     let tokenizer = &*tokenizer_ptr;
 
@@ -321,7 +351,7 @@ pub unsafe extern "C" fn tokenizer_encode_batch_core(
         .map(|text_buffer| std::str::from_utf8_unchecked(text_buffer.as_slice()))
         .collect::<Vec<&str>>();
 
-    let encoded_result = tokenizer.encode_batch_fast(texts, true);
+    let encoded_result = tokenizer.encode_batch_fast(texts, add_special_tokens);
 
     let encoded_tokens = match encoded_result
     {

Tokenizers.NET/Tokenizer.cs

@@ -196,7 +196,7 @@ public Tokenizer()
         }
 
         [SkipLocalsInit]
-        public TokenizeOutput Tokenize(string input)
+        public TokenizeOutput Tokenize(string input, bool addSpecialTokens = true)
         {
             Span<byte> allocation;
 
@@ -236,7 +236,12 @@ public TokenizeOutput Tokenize(string input)
 
             var u8String = new ReadOnlyNativeBuffer<byte>(ref MemoryMarshal.GetReference(allocation), (nuint) bytesWritten);
 
-            var result = TokenizerNativeMethods.TokenizerEncode(TokenizerHandle, u8String, TRUNCATE);
+            var result = TokenizerNativeMethods.TokenizerEncode(
+                TokenizerHandle, 
+                u8String,
+                addSpecialTokens,
+                TRUNCATE
+            );
 
             if (useNativeMemory)
             {
@@ -246,35 +251,38 @@ public TokenizeOutput Tokenize(string input)
             return result;
         }
 
-        public void TokenizeBatch(ReadOnlySpan<string> inputs, Span<TokenizeOutput> outputs)
+        public void TokenizeBatch(ReadOnlySpan<string> inputs, Span<TokenizeOutput> outputs, bool addSpecialTokens = true)
         {
             TokenizeBatchInternal(
                 inputs, 
                 outputs, 
                 outputsPrePinned: false,
-                skipLengthCheck: false
+                skipLengthCheck: false,
+                addSpecialTokens: addSpecialTokens
             );
         }
 
-        public void TokenizeBatch(ReadOnlySpan<string> inputs, NativeMemory<TokenizeOutput> outputs)
+        public void TokenizeBatch(ReadOnlySpan<string> inputs, NativeMemory<TokenizeOutput> outputs, bool addSpecialTokens = true)
         {
             TokenizeBatchInternal(
                 inputs, 
                 outputs.Buffer.AsSpan(), 
                 outputsPrePinned: true,
-                skipLengthCheck: false
+                skipLengthCheck: false,
+                addSpecialTokens: addSpecialTokens
             );
         }
 
-        public NativeMemory<TokenizeOutput> TokenizeBatch(ReadOnlySpan<string> inputs)
+        public NativeMemory<TokenizeOutput> TokenizeBatch(ReadOnlySpan<string> inputs, bool addSpecialTokens = true)
         {
             var outputs = new NativeMemory<TokenizeOutput>((nuint) inputs.Length);
 
             TokenizeBatchInternal(
                 inputs, 
                 outputs.Buffer.AsSpan(), 
                 outputsPrePinned: true,
-                skipLengthCheck: true
+                skipLengthCheck: true,
+                addSpecialTokens: addSpecialTokens
             );
 
             return outputs;
@@ -286,7 +294,8 @@ private void TokenizeBatchInternal(
             ReadOnlySpan<string> inputs,
             Span<TokenizeOutput> outputs,
             bool outputsPrePinned,
-            bool skipLengthCheck)
+            bool skipLengthCheck,
+            bool addSpecialTokens)
         {
             var numInputs = inputs.Length;
 
@@ -358,6 +367,7 @@ stackalloc ReadOnlyNativeBuffer<byte>[numInputs]
                     tokenizerPtr: tokenizerHandle,
                     textNativeBuffers: readonlyU8Strings,
                     outputNativeBuffer: new(ref outputStart, outputLengthNative),
+                    addSpecialTokens,
                     truncate
                 );
             }
@@ -370,6 +380,7 @@ stackalloc ReadOnlyNativeBuffer<byte>[numInputs]
                         tokenizerPtr: tokenizerHandle,
                         textNativeBuffers: readonlyU8Strings,
                         outputNativeBuffer: new(outputsPtr, outputLengthNative),
+                        addSpecialTokens,
                         truncate
                     );
                 }

Tokenizers.NET/TokenizerNativeMethods.cs

@@ -18,55 +18,67 @@ internal static unsafe class TokenizerNativeMethods
         public static TokenizeOutput TokenizerEncode(
             nint tokenizerPtr,
             ReadOnlyNativeBuffer<byte> textNativeBuffer,
+            bool addSpecialTokens,
             bool truncate)
         {
             if (truncate)
             {
-                return TokenizerEncode(tokenizerPtr, textNativeBuffer);
+                return TokenizerEncode(tokenizerPtr, textNativeBuffer, addSpecialTokens);
             }
 
             else
             {
-                return TokenizerEncodeNonTruncating(tokenizerPtr, textNativeBuffer);
+                return TokenizerEncodeNonTruncating(tokenizerPtr, textNativeBuffer, addSpecialTokens);
             }
         }
 
         [DllImport(DLL_NAME, EntryPoint = "tokenizer_encode", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)]
-        private static extern TokenizeOutput TokenizerEncode(nint tokenizerPtr, ReadOnlyNativeBuffer<byte> textNativeBuffer);
+        private static extern TokenizeOutput TokenizerEncode(
+            nint tokenizerPtr,
+            ReadOnlyNativeBuffer<byte> textNativeBuffer,
+            [MarshalAs(UnmanagedType.U1)] bool addSpecialTokens
+        );
 
         [DllImport(DLL_NAME, EntryPoint = "tokenizer_encode_non_truncating", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)]
-        private static extern TokenizeOutput TokenizerEncodeNonTruncating(nint tokenizerPtr, ReadOnlyNativeBuffer<byte> textNativeBuffer);
+        private static extern TokenizeOutput TokenizerEncodeNonTruncating(
+            nint tokenizerPtr,
+            ReadOnlyNativeBuffer<byte> textNativeBuffer,
+            [MarshalAs(UnmanagedType.U1)] bool addSpecialTokens
+        );
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static void TokenizerEncodeBatch(
             nint tokenizerPtr, 
             ReadOnlyNativeBuffer<ReadOnlyNativeBuffer<byte>> textNativeBuffers, 
             NativeBuffer<TokenizeOutput> outputNativeBuffer,
+            bool addSpecialTokens,
             bool truncate)
         {
             if (truncate)
             {
-                TokenizerEncodeBatch(tokenizerPtr, textNativeBuffers, outputNativeBuffer);
+                TokenizerEncodeBatch(tokenizerPtr, textNativeBuffers, outputNativeBuffer, addSpecialTokens);
             }
 
             else
             {
-                TokenizerEncodeBatchNonTruncating(tokenizerPtr, textNativeBuffers, outputNativeBuffer);
+                TokenizerEncodeBatchNonTruncating(tokenizerPtr, textNativeBuffers, outputNativeBuffer, addSpecialTokens);
             }
         }
 
         [DllImport(DLL_NAME, EntryPoint = "tokenizer_encode_batch", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)]
         private static extern void TokenizerEncodeBatch(
             nint tokenizerPtr, 
             ReadOnlyNativeBuffer<ReadOnlyNativeBuffer<byte>> textNativeBuffers, 
-            NativeBuffer<TokenizeOutput> outputNativeBuffer
+            NativeBuffer<TokenizeOutput> outputNativeBuffer,
+            [MarshalAs(UnmanagedType.U1)] bool addSpecialTokens
         );
 
         [DllImport(DLL_NAME, EntryPoint = "tokenizer_encode_batch_non_truncating", CallingConvention = CallingConvention.Cdecl, ExactSpelling = true)]
         private static extern void TokenizerEncodeBatchNonTruncating(
             nint tokenizerPtr, 
             ReadOnlyNativeBuffer<ReadOnlyNativeBuffer<byte>> textNativeBuffers, 
-            NativeBuffer<TokenizeOutput> outputNativeBuffer
+            NativeBuffer<TokenizeOutput> outputNativeBuffer,
+            [MarshalAs(UnmanagedType.U1)] bool addSpecialTokens
         );
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]