Supporting adding special tokens when creating a tokenizer (#26)

Author: travismorton

lib/tokenizers/native.ex

@@ -15,7 +15,7 @@ defmodule Tokenizers.Native do
   def decode_batch(_tokenizer, _ids, _skip_special_tokens), do: err()
   def encode(_tokenizer, _input, _add_special_tokens), do: err()
   def encode_batch(_tokenizer, _input, _add_special_tokens), do: err()
-  def from_file(_path), do: err()
+  def from_file(_path, _additional_special_tokens), do: err()
   def get_attention_mask(_encoding), do: err()
   def get_type_ids(_encoding), do: err()
   def get_ids(_encoding), do: err()

lib/tokenizers/tokenizer.ex

@@ -55,6 +55,8 @@ defmodule Tokenizers.Tokenizer do
       even if `:use_cache` is false. By default it uses `:filename.basedir/3` to get
       a cache dir based in the "tokenizers_elixir" application name.
 
+    * `:additional_special_tokens` - A list of special tokens to append to the tokenizer.
+      Defaults to `[]`.
   """
   @spec from_pretrained(String.t(), Keyword.t()) :: {:ok, Tokenizer.t()} | {:error, term()}
   def from_pretrained(identifier, opts \\ []) do
@@ -63,7 +65,8 @@ defmodule Tokenizers.Tokenizer do
         revision: "main",
         use_cache: true,
         cache_dir: :filename.basedir(:user_cache, "tokenizers_elixir"),
-        http_client: {Tokenizers.HTTPClient, []}
+        http_client: {Tokenizers.HTTPClient, []},
+        additional_special_tokens: []
       )
 
     {http_client, http_opts} = opts[:http_client]
@@ -87,19 +90,21 @@ defmodule Tokenizers.Tokenizer do
       Path.join(cache_dir, entry_filename(url, etag))
     end
 
+    tokenizer_opts = Keyword.take(opts, [:additional_special_tokens])
+
     if opts[:use_cache] do
       with {:ok, response} <- request(http_client, Keyword.put(http_opts, :method, :head)) do
         etag = fetch_etag(response.headers)
         file_path = file_path_fun.(etag)
 
         if File.exists?(file_path) do
-          from_file(file_path)
+          from_file(file_path, tokenizer_opts)
         else
           with {:ok, response} <- request(http_client, http_opts) do
             File.mkdir_p!(cache_dir)
             File.write!(file_path, response.body)
 
-            from_file(file_path)
+            from_file(file_path, tokenizer_opts)
           end
         end
       end
@@ -111,7 +116,7 @@ defmodule Tokenizers.Tokenizer do
         File.mkdir_p!(cache_dir)
         File.write!(file_path, response.body)
 
-        from_file(file_path)
+        from_file(file_path, tokenizer_opts)
       end
     end
   end
@@ -156,9 +161,17 @@ defmodule Tokenizers.Tokenizer do
 
   @doc """
   Instantiate a new tokenizer from the file at the given path.
+
+  ## Options
+
+    * `:additional_special_tokens` - A list of special tokens to append to the tokenizer.
+      Defaults to `[]`.
   """
-  @spec from_file(String.t()) :: {:ok, Tokenizer.t()} | {:error, term()}
-  def from_file(path), do: Native.from_file(path)
+  @spec from_file(String.t(), Keyword.t()) :: {:ok, Tokenizer.t()} | {:error, term()}
+  def from_file(path, opts \\ []) do
+    opts = Keyword.validate!(opts, additional_special_tokens: [])
+    Native.from_file(path, opts[:additional_special_tokens])
+  end
 
   @doc """
   Save the tokenizer to the provided path.

native/ex_tokenizers/src/tokenizer.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 use rustler::Term;
 
+use tokenizers::tokenizer::AddedToken;
 use tokenizers::{EncodeInput, Tokenizer};
 
 use crate::encoding::ExTokenizersEncoding;
@@ -31,8 +32,14 @@ impl ExTokenizersTokenizer {
 }
 
 #[rustler::nif(schedule = "DirtyIo")]
-pub fn from_file(path: &str) -> Result<ExTokenizersTokenizer, ExTokenizersError> {
-    let tokenizer = Tokenizer::from_file(path)?;
+pub fn from_file(
+    path: &str,
+    additional_special_tokens: Vec<String>,
+) -> Result<ExTokenizersTokenizer, ExTokenizersError> {
+    let mut tokenizer = Tokenizer::from_file(path)?;
+    for token in additional_special_tokens {
+        tokenizer.add_special_tokens(&[AddedToken::from(token, true)]);
+    }
     Ok(ExTokenizersTokenizer::new(tokenizer))
 }
 

test/tokenizers/tokenizer_test.exs

@@ -25,6 +25,38 @@ defmodule Tokenizers.TokenizerTest do
     end
   end
 
+  describe "modify tokenizer" do
+    test "can add special tokens" do
+      special_tokens = ["<|test|>"]
+
+      {:ok, tokenizer} =
+        Tokenizer.from_file("test/fixtures/bert-base-cased.json",
+          additional_special_tokens: special_tokens
+        )
+
+      assert Tokenizer.get_vocab_size(tokenizer) == 28997
+    end
+
+    test "can decode special tokens" do
+      text = ["This <|test|>is a test<|also|>", "<|test|>And so<|also|> is this<|test|>"]
+      special_tokens = ["<|test|>", "<|also|>"]
+
+      {:ok, tokenizer} =
+        Tokenizer.from_file("test/fixtures/bert-base-cased.json",
+          additional_special_tokens: special_tokens
+        )
+
+      {:ok, encodings} = Tokenizer.encode(tokenizer, text)
+
+      {:ok, decodings} =
+        Tokenizer.decode(tokenizer, Enum.map(encodings, &Encoding.get_ids/1),
+          skip_special_tokens: true
+        )
+
+      assert ["This is a test", "And so is this"] == decodings
+    end
+  end
+
   describe "from_pretrained/2" do
     defmodule SuccessHTTPClient do
       def request(opts) do