Add functions necessary for NER (#12)

Author: seanmor5

lib/tokenizers/encoding.ex

@@ -35,6 +35,19 @@ defmodule Tokenizers.Encoding do
   @spec get_type_ids(Encoding.t()) :: [integer()]
   def get_type_ids(encoding), do: encoding |> Native.get_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 """
+  Get offsets from an encoding.
+  """
+  @spec get_offsets(Encoding.t()) :: [{integer(), integer()}]
+  def get_offsets(encoding), do: encoding |> Native.get_offsets() |> Shared.unwrap()
+
   @doc """
   Truncate the encoding to the given length.
 

lib/tokenizers/native.ex

@@ -20,6 +20,8 @@ defmodule Tokenizers.Native do
   def get_type_ids(_encoding), do: err()
   def get_ids(_encoding), do: err()
   def get_tokens(_encoding), do: err()
+  def get_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()
   def id_to_token(_tokenizer, _id), do: err()

native/ex_tokenizers/src/encoding.rs

@@ -45,6 +45,20 @@ pub fn get_type_ids(encoding: ExTokenizersEncoding) -> Result<Vec<u32>, ExTokeni
     Ok(encoding.resource.0.get_type_ids().to_vec())
 }
 
+#[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_offsets(
+    encoding: ExTokenizersEncoding,
+) -> Result<Vec<(usize, usize)>, ExTokenizersError> {
+    Ok(encoding.resource.0.get_offsets().to_vec())
+}
+
 #[rustler::nif]
 pub fn n_tokens(encoding: ExTokenizersEncoding) -> Result<usize, ExTokenizersError> {
     Ok(encoding.resource.0.len())

native/ex_tokenizers/src/error.rs

@@ -21,7 +21,7 @@ pub enum ExTokenizersError {
     Unknown(#[from] anyhow::Error),
 }
 
-impl<'a> Encoder for ExTokenizersError {
+impl Encoder for ExTokenizersError {
     fn encode<'b>(&self, env: Env<'b>) -> Term<'b> {
         format!("{:?}", self).encode(env)
     }

native/ex_tokenizers/src/lib.rs

@@ -29,6 +29,8 @@ rustler::init!(
         get_attention_mask,
         get_type_ids,
         get_ids,
+        get_special_tokens_mask,
+        get_offsets,
         get_model,
         get_model_details,
         get_tokens,

test/tokenizers/tokenizer_test.exs

@@ -80,4 +80,24 @@ defmodule Tokenizers.TokenizerTest do
       assert decoded == text
     end
   end
+
+  describe "encode metadata" do
+    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
+    end
+
+    test "can return offsets", %{tokenizer: tokenizer} do
+      text = ["This is a test", "And so is this"]
+      {:ok, encodings} = Tokenizer.encode(tokenizer, text)
+      offsets = Enum.map(encodings, &Encoding.get_offsets/1)
+
+      assert [
+               [{0, 0}, {0, 4}, {5, 7}, {8, 9}, {10, 14}, {0, 0}],
+               [{0, 0}, {0, 3}, {4, 6}, {7, 9}, {10, 14}, {0, 0}]
+             ] == offsets
+    end
+  end
 end