Add encode_batch_ids and make Tokenizer a PythonCall wrapper

CondaPkg.toml

@@ -1,2 +1,2 @@
 [deps]
-tokenizers = ""
+tokenizers = ">=0.20"

README.md

@@ -67,6 +67,25 @@ ids_batch = [result.ids for result in batch_results]
 decoded_texts = decode_batch(tokenizer, ids_batch)
 ```
 
+#### Fast Batch Encoding of IDs
+
+To tokenize a large corpus, use `encode_batch_ids`, which returns only the token IDs and is much faster and lighter on memory than `encode_batch`:
+
+```julia
+ids_batch = encode_batch_ids(tokenizer, texts; T=UInt16)  # one view per text
+all_ids = parent(ids_batch[1])                            # the IDs of all texts, concatenated
+```
+
+### Python Interop
+
+A `Tokenizer` can be passed directly to Python functions called through PythonCall, and `Py(tokenizer)` returns the underlying Python object:
+
+```julia
+using PythonCall
+get_ids = pyeval("lambda tokenizer, text: tokenizer.encode(text).ids", Main)
+get_ids(tokenizer, "Hello, how are you?")
+```
+
 ### Saving a Tokenizer
 
 ```julia

src/HuggingFaceTokenizers.jl

@@ -10,8 +10,34 @@ module HuggingFaceTokenizers
 using PythonCall
 
 const tokenizers = PythonCall.pynew()
+const write_lengths = PythonCall.pynew()
+const write_ids = PythonCall.pynew()
 
-__init__() = PythonCall.pycopy!(tokenizers, pyimport("tokenizers"))
+# Writing into Julia arrays through the buffer protocol avoids creating
+# a Julia wrapper for each `Encoding` and its `ids` list.
+const HELPERS = """
+from array import array
+
+def write_lengths(encodings, out):
+    out = memoryview(out)
+    for i, e in enumerate(encodings):
+        out[i] = len(e)
+
+def write_ids(encodings, out):
+    out = memoryview(out)
+    i = 0
+    for e in encodings:
+        ids = array(out.format, e.ids)
+        out[i:i + len(ids)] = ids
+        i += len(ids)
+"""
+
+function __init__()
+    PythonCall.pycopy!(tokenizers, pyimport("tokenizers"))
+    helpers = pyexec(@NamedTuple{write_lengths::Py, write_ids::Py}, HELPERS, @__MODULE__)
+    PythonCall.pycopy!(write_lengths, helpers.write_lengths)
+    PythonCall.pycopy!(write_ids, helpers.write_ids)
+end
 
 include("Tokenizer.jl")
 export Tokenizer
@@ -21,6 +47,7 @@ export save
 export encode
 export decode
 export encode_batch
+export encode_batch_ids
 export decode_batch
 
 end

src/Tokenizer.jl

@@ -2,11 +2,17 @@
     Tokenizer
 
 A wrapper around a Python tokenizer.
+
+It can be passed directly to Python functions called through PythonCall,
+and `Py(tokenizer)` returns the underlying Python object.
 """
 struct Tokenizer
     py_tokenizer::Py
 end
 
+PythonCall.ispy(::Tokenizer) = true
+PythonCall.Py(tokenizer::Tokenizer) = tokenizer.py_tokenizer
+
 """
     from_file(::Type{Tokenizer}, path) -> Tokenizer
 
@@ -72,6 +78,37 @@ function encode_batch(tokenizer::Tokenizer, args...; kwargs...)
     end
 end
 
+"""
+    encode_batch_ids(tokenizer::Tokenizer, text_batch; T=UInt32, kwargs...) -> Vector{<:SubArray{T}}
+
+Encode multiple texts in batch, returning only the token IDs of each text.
+
+The IDs of all texts are stored contiguously in a single `Vector{T}`, available as
+`parent(ids_batch[1])`, and the result holds a view into it for each text.
+This is much faster and uses much less memory than `encode_batch`,
+which makes it suited to tokenizing large corpora.
+`T=UInt16` is enough for vocabularies of up to 65536 tokens, such as GPT-2's,
+and an error is thrown if an ID does not fit in `T`.
+
+Keyword arguments other than `T` are passed to the Python tokenizer's `encode_batch_fast`,
+e.g. `add_special_tokens=false`.
+
+```julia
+ids_batch = encode_batch_ids(tokenizer, ["Hello, how are you?", "I'm doing great!"]; T=UInt16)
+```
+"""
+function encode_batch_ids(tokenizer::Tokenizer, args...;
+        T::Type{<:Union{Int8,Int16,Int32,Int64,UInt8,UInt16,UInt32,UInt64}}=UInt32, kwargs...)
+    encodings = tokenizer.py_tokenizer.encode_batch_fast(args...; kwargs...)
+    lengths = Vector{Int}(undef, pylen(encodings))
+    write_lengths(encodings, lengths)
+    ids = Vector{T}(undef, sum(lengths))
+    write_ids(encodings, ids)
+    PythonCall.pydel!(encodings)
+    stops = cumsum(lengths)
+    return [view(ids, stop-len+1:stop) for (len, stop) in zip(lengths, stops)]
+end
+
 """
     decode_batch(tokenizer::Tokenizer, ids_batch::Vector{Vector{Int}}) -> Vector{String}
 

test/runtests.jl

@@ -46,7 +46,36 @@ using Test
         @test length(decoded_texts) == length(texts)
         @test all(!isempty, decoded_texts)
     end
-
+
+    @testset "encode_batch_ids" begin
+        tokenizer = from_pretrained(Tokenizer, "bert-base-uncased")
+        texts = ["Hello, how are you?", "", "I'm doing great!"]
+
+        ids_batch = encode_batch_ids(tokenizer, texts)
+        @test ids_batch isa Vector{<:SubArray{UInt32}}
+        @test ids_batch == [result.ids for result in encode_batch(tokenizer, texts)]
+        @test parent(ids_batch[1]) == reduce(vcat, ids_batch)
+        @test all(ids -> parent(ids) === parent(ids_batch[1]), ids_batch)
+
+        @test encode_batch_ids(tokenizer, texts; T=UInt16) isa Vector{<:SubArray{UInt16}}
+        @test encode_batch_ids(tokenizer, texts; T=Int) == ids_batch
+        @test encode_batch_ids(tokenizer, texts; add_special_tokens=false) ==
+            [result.ids for result in encode_batch(tokenizer, texts; add_special_tokens=false)]
+        @test isempty(encode_batch_ids(tokenizer, String[]))
+        @test_throws HuggingFaceTokenizers.PythonCall.PyException encode_batch_ids(tokenizer, texts; T=Int8)
+
+        @test decode_batch(tokenizer, ids_batch) == decode_batch(tokenizer, [result.ids for result in encode_batch(tokenizer, texts)])
+    end
+
+    @testset "Python interop" begin
+        PythonCall = HuggingFaceTokenizers.PythonCall
+        tokenizer = from_pretrained(Tokenizer, "bert-base-uncased")
+        @test PythonCall.ispy(tokenizer)
+        @test PythonCall.Py(tokenizer) === tokenizer.py_tokenizer
+        f = PythonCall.pyeval("lambda tokenizer, text: tokenizer.encode(text).ids", Main)
+        @test PythonCall.pyconvert(Vector{Int}, f(tokenizer, "Hello")) == encode(tokenizer, "Hello").ids
+    end
+
     @testset "File Operations" begin
         tokenizer = from_pretrained(Tokenizer, "bert-base-uncased")