From 0e202b47ec57884c407ea45a84ffdec7d4ec8a7b Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Tue, 20 Oct 2026 09:41:00 +0200 Subject: [PATCH] Add encode_batch_ids and make Tokenizer a PythonCall wrapper encode_batch_ids calls encode_batch_fast and writes the ids of all texts into one Julia vector through the buffer protocol, returning a view per text. Tokenizer now defines ispy and Py, so it can be passed to Python. Fixes #13 Co-Authored-By: Claude Opus 5.5 --- CondaPkg.toml | 2 +- README.md | 19 ++++++++++++++++++ src/HuggingFaceTokenizers.jl | 29 +++++++++++++++++++++++++++- src/Tokenizer.jl | 37 ++++++++++++++++++++++++++++++++++++ test/runtests.jl | 31 +++++++++++++++++++++++++++++- 5 files changed, 115 insertions(+), 3 deletions(-) diff --git a/CondaPkg.toml b/CondaPkg.toml index 6cb6a34..8dadd5b 100644 --- a/CondaPkg.toml +++ b/CondaPkg.toml @@ -1,2 +1,2 @@ [deps] -tokenizers = "" +tokenizers = ">=0.20" diff --git a/README.md b/README.md index 45ec48e..86ea740 100644 --- a/README.md +++ b/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 diff --git a/src/HuggingFaceTokenizers.jl b/src/HuggingFaceTokenizers.jl index 0944777..df6fd47 100644 --- a/src/HuggingFaceTokenizers.jl +++ b/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 diff --git a/src/Tokenizer.jl b/src/Tokenizer.jl index 3032e82..4f8e358 100644 --- a/src/Tokenizer.jl +++ b/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} diff --git a/test/runtests.jl b/test/runtests.jl index 609ac45..5ae666c 100644 --- a/test/runtests.jl +++ b/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")