Add binary variants for accessing encoding data (#32)

Author: jonatanklosko

lib/tokenizers/encoding.ex

@@ -23,25 +23,52 @@ defmodule Tokenizers.Encoding do
   @spec get_ids(Encoding.t()) :: [integer()]
   def get_ids(encoding), do: encoding |> Native.get_ids() |> Shared.unwrap()
 
+  @doc """
+  Same as `get_ids/1`, but returns binary with u32 values.
+  """
+  @spec get_u32_ids(Encoding.t()) :: binary()
+  def get_u32_ids(encoding), do: encoding |> Native.get_u32_ids() |> Shared.unwrap()
+
   @doc """
   Get the attention mask from an encoding.
   """
   @spec get_attention_mask(Encoding.t()) :: [integer()]
   def get_attention_mask(encoding), do: encoding |> Native.get_attention_mask() |> Shared.unwrap()
 
+  @doc """
+  Same as `get_attention_mask/1`, but returns binary with u32 values.
+  """
+  @spec get_u32_attention_mask(Encoding.t()) :: binary()
+  def get_u32_attention_mask(encoding),
+    do: encoding |> Native.get_u32_attention_mask() |> Shared.unwrap()
+
   @doc """
   Get token type ids from an encoding.
   """
   @spec get_type_ids(Encoding.t()) :: [integer()]
   def get_type_ids(encoding), do: encoding |> Native.get_type_ids() |> Shared.unwrap()
 
+  @doc """
+  Same as `get_type_ids/1`, but returns binary with u32 values.
+  """
+  @spec get_u32_type_ids(Encoding.t()) :: binary()
+  def get_u32_type_ids(encoding),
+    do: encoding |> Native.get_u32_type_ids() |> Shared.unwrap()
+
   @doc """
   Get special tokens mask from an encoding.
   """
   @spec get_special_tokens_mask(Encoding.t()) :: [integer()]
   def get_special_tokens_mask(encoding),
     do: encoding |> Native.get_special_tokens_mask() |> Shared.unwrap()
 
+  @doc """
+  Same as `get_special_tokens_mask/1`, but returns binary with u32 values.
+  """
+  @spec get_u32_special_tokens_mask(Encoding.t()) :: binary()
+  def get_u32_special_tokens_mask(encoding),
+    do: encoding |> Native.get_u32_special_tokens_mask() |> Shared.unwrap()
+
   @doc """
   Get offsets from an encoding.
 

lib/tokenizers/native.ex

@@ -16,10 +16,14 @@ defmodule Tokenizers.Native do
   def encode_batch(_tokenizer, _input, _add_special_tokens), do: err()
   def from_file(_path, _additional_special_tokens), do: err()
   def get_attention_mask(_encoding), do: err()
+  def get_u32_attention_mask(_encoding), do: err()
   def get_type_ids(_encoding), do: err()
+  def get_u32_type_ids(_encoding), do: err()
   def get_ids(_encoding), do: err()
+  def get_u32_ids(_encoding), do: err()
   def get_tokens(_encoding), do: err()
   def get_special_tokens_mask(_encoding), do: err()
+  def get_u32_special_tokens_mask(_encoding), do: err()
   def get_offsets(_encoding), do: err()
   def get_vocab(_tokenizer, _with_added_tokens), do: err()
   def get_vocab_size(_tokenizer, _with_added_tokens), do: err()

native/ex_tokenizers/src/encoding.rs

@@ -1,4 +1,6 @@
 use crate::error::ExTokenizersError;
+use rustler::resource::ResourceArc;
+use rustler::{Binary, Env};
 use tokenizers::utils::padding::PaddingDirection;
 use tokenizers::utils::truncation::TruncationDirection;
 use tokenizers::Encoding;
@@ -8,7 +10,7 @@ pub struct ExTokenizersEncodingRef(pub Encoding);
 #[derive(rustler::NifStruct)]
 #[module = "Tokenizers.Encoding"]
 pub struct ExTokenizersEncoding {
-    pub resource: rustler::resource::ResourceArc<ExTokenizersEncodingRef>,
+    pub resource: ResourceArc<ExTokenizersEncodingRef>,
 }
 
 impl ExTokenizersEncodingRef {
@@ -20,7 +22,7 @@ impl ExTokenizersEncodingRef {
 impl ExTokenizersEncoding {
     pub fn new(data: Encoding) -> Self {
         Self {
-            resource: rustler::resource::ResourceArc::new(ExTokenizersEncodingRef::new(data)),
+            resource: ResourceArc::new(ExTokenizersEncodingRef::new(data)),
         }
     }
 }
@@ -35,23 +37,60 @@ pub fn get_ids(encoding: ExTokenizersEncoding) -> Result<Vec<u32>, ExTokenizersE
     Ok(encoding.resource.0.get_ids().to_vec())
 }
 
+#[rustler::nif]
+pub fn get_u32_ids(env: Env, encoding: ExTokenizersEncoding) -> Result<Binary, ExTokenizersError> {
+    Ok(encoding
+        .resource
+        .make_binary(env, |r| slice_u32_to_u8(r.0.get_ids())))
+}
+
 #[rustler::nif]
 pub fn get_attention_mask(encoding: ExTokenizersEncoding) -> Result<Vec<u32>, ExTokenizersError> {
     Ok(encoding.resource.0.get_attention_mask().to_vec())
 }
 
+#[rustler::nif]
+pub fn get_u32_attention_mask(
+    env: Env,
+    encoding: ExTokenizersEncoding,
+) -> Result<Binary, ExTokenizersError> {
+    Ok(encoding
+        .resource
+        .make_binary(env, |r| slice_u32_to_u8(r.0.get_attention_mask())))
+}
+
 #[rustler::nif]
 pub fn get_type_ids(encoding: ExTokenizersEncoding) -> Result<Vec<u32>, ExTokenizersError> {
     Ok(encoding.resource.0.get_type_ids().to_vec())
 }
 
+#[rustler::nif]
+pub fn get_u32_type_ids(
+    env: Env,
+    encoding: ExTokenizersEncoding,
+) -> Result<Binary, ExTokenizersError> {
+    Ok(encoding
+        .resource
+        .make_binary(env, |r| slice_u32_to_u8(r.0.get_type_ids())))
+}
+
 #[rustler::nif]
 pub fn get_special_tokens_mask(
     encoding: ExTokenizersEncoding,
 ) -> Result<Vec<u32>, ExTokenizersError> {
     Ok(encoding.resource.0.get_special_tokens_mask().to_vec())
 }
 
+#[rustler::nif]
+pub fn get_u32_special_tokens_mask(
+    env: Env,
+    encoding: ExTokenizersEncoding,
+) -> Result<Binary, ExTokenizersError> {
+    Ok(encoding
+        .resource
+        .make_binary(env, |r| slice_u32_to_u8(r.0.get_special_tokens_mask())))
+}
+
 #[rustler::nif]
 pub fn get_offsets(
     encoding: ExTokenizersEncoding,
@@ -99,3 +138,7 @@ pub fn pad(
     new_encoding.pad(target_length, pad_id, pad_type_id, pad_token, direction);
     Ok(ExTokenizersEncoding::new(new_encoding))
 }
+
+fn slice_u32_to_u8(slice: &[u32]) -> &[u8] {
+    unsafe { std::slice::from_raw_parts(slice.as_ptr() as *const u8, slice.len() * 4) }
+}

native/ex_tokenizers/src/error.rs

@@ -23,6 +23,6 @@ pub enum ExTokenizersError {
 
 impl Encoder for ExTokenizersError {
     fn encode<'b>(&self, env: Env<'b>) -> Term<'b> {
-        format!("{:?}", self).encode(env)
+        format!("{self:?}").encode(env)
     }
 }

native/ex_tokenizers/src/lib.rs

@@ -26,9 +26,13 @@ rustler::init!(
         encode_batch,
         from_file,
         get_attention_mask,
+        get_u32_attention_mask,
         get_type_ids,
+        get_u32_type_ids,
         get_ids,
+        get_u32_ids,
         get_special_tokens_mask,
+        get_u32_special_tokens_mask,
         get_offsets,
         get_model,
         get_model_details,

test/tokenizers/tokenizer_test.exs

@@ -196,15 +196,43 @@ defmodule Tokenizers.TokenizerTest do
       ids = Enum.map(encodings, &Encoding.get_ids/1)
       {:ok, decoded} = Tokenizer.decode(tokenizer, ids)
       assert decoded == text
+
+      assert Enum.map(ids, &list_to_u32/1) == Enum.map(encodings, &Encoding.get_u32_ids/1)
     end
   end
 
   describe "encode metadata" do
+    test "can return attention mask", %{tokenizer: tokenizer} do
+      text = ["Hello world", "Yes sir hello indeed"]
+      {:ok, encodings} = Tokenizer.encode(tokenizer, text)
+
+      attention_mask = Enum.map(encodings, &Encoding.get_attention_mask/1)
+      assert [[1, 1, 1, 1], [1, 1, 1, 1, 1, 1]] == attention_mask
+
+      assert Enum.map(attention_mask, &list_to_u32/1) ==
+               Enum.map(encodings, &Encoding.get_u32_attention_mask/1)
+    end
+
+    test "can return type ids", %{tokenizer: tokenizer} do
+      text = [{"Hello", "world"}, {"Yes sir", "hello indeed"}]
+      {:ok, encodings} = Tokenizer.encode(tokenizer, text)
+
+      type_ids = Enum.map(encodings, &Encoding.get_type_ids/1)
+      assert [[0, 0, 0, 1, 1], [0, 0, 0, 0, 1, 1, 1]] == type_ids
+
+      assert Enum.map(type_ids, &list_to_u32/1) ==
+               Enum.map(encodings, &Encoding.get_u32_type_ids/1)
+    end
+
     test "can return special tokens mask", %{tokenizer: tokenizer} do
       text = ["This is a test", "And so is this"]
       {:ok, encodings} = Tokenizer.encode(tokenizer, text)
+
       special_tokens_mask = Enum.map(encodings, &Encoding.get_special_tokens_mask/1)
       assert [[1, 0, 0, 0, 0, 1], [1, 0, 0, 0, 0, 1]] == special_tokens_mask
+
+      assert Enum.map(special_tokens_mask, &list_to_u32/1) ==
+               Enum.map(encodings, &Encoding.get_u32_special_tokens_mask/1)
     end
 
     test "can return offsets", %{tokenizer: tokenizer} do
@@ -218,4 +246,8 @@ defmodule Tokenizers.TokenizerTest do
              ] == offsets
     end
   end
+
+  defp list_to_u32(list) do
+    for x <- list, into: <<>>, do: <<x::native-unsigned-32>>
+  end
 end