feat: wire in the thread pool

[McPatate]

Wiring in the thread pool.

For now we only parallelise on special tokens, the intra-seq splitting will come in a subsequent PR. I also need to handle Pairs properly, but also for another PR.

Left to do in this PR:

• add some testing

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Copilot]

Pull request overview

This PR introduces a streaming/parallel encoding path for the experimental PipelineTokenizer, wiring it to the shared thread pool and planning work units by splitting inputs on special tokens to enable parallel execution.

Changes:

• Refactors PipelineTokenizer to be Clone via an Arc-backed inner struct, enabling sharing across worker threads.
• Adds a new streaming EncodeHandle mode backed by a Job work-queue with atomic coordination and per-unit outputs.
• Updates special-token segmentation to track input offsets (Segment::Text { text, input_offset }) for range-based unit encoding.

Suppressed comments (4)

tokenizers/tk-encode/src/tokenizer/pipeline.rs:1143

• Parallel path runs encode_generic at STAGE_POSTPROCESS for each unit. When inputs are split around special tokens, this applies the post-processor prefix/suffix per unit (and can also place prefix/suffix after leading special tokens or before trailing special tokens), producing different output than the serial encode.

            let Plan { units, outputs } = self.plan_work(&inputs);
            if units.len() < 2 {
                return EncodeHandle::blocking(self.encode_serial(inputs, add_special_tokens));
            }

tokenizers/tk-encode/src/tokenizer/pipeline.rs:1181

• encode_serial currently panics on Input::Pair via todo!("handle input pairs"). This is a public API surface (encode accepts tuples via Into<Inputs>), so it should return an Err rather than panic.

            Inputs::Single(seq) => {
                let Input::Single(seq) = seq else {
                    todo!("handle input pairs")
                };

tokenizers/tk-encode/src/tokenizer/pipeline.rs:1191

• encode_serial panics on Input::Pair entries inside a batch. This makes mixed batches a runtime panic instead of a per-item error.

                for seq in batch {
                    let Input::Single(seq) = seq else {
                        todo!("handle input pairs")
                    };

tokenizers/tk-encode/src/tokenizer/pipeline.rs:1087

• plan_work panics on Input::Pair via todo!("handle pair input"), so passing a pair can crash before falling back to serial. Consider planning a single unit that will deterministically return an Err in the worker, matching the serial behavior.

        for (seq_idx, input) in inputs.into_iter().enumerate() {
            let Input::Single(input) = input else {
                todo!("handle pair input")
            };

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[SBrandeis]

first pass of review 🥵

[SBrandeis]

Unused imports I think, don't forget to run clippy and cargo fmt!

[SBrandeis]

Pre-existing to this PR, but does this require to copy the string from Python to Rust?

[McPatate]

You can make it 0-copy from a python string all the way to the argument in rust, although somewhere in the process you need to materialise the string to utf8 from the python side which will necessarily copy the string. We simply make access from it's utf8 storage 0-copy via some unsafe, but we can hold a reference to the object making it safe.

[SBrandeis]

That can happen if you try to take the result before the work producing it is completed, right?

Should we return an Error here instead of panicking?

[McPatate]

It should not be possible for this value to be None. If we call take_result before the work has been completed, this is a bug. Imo it shouldn't be an error for the user to handle because it is unrecoverable anyways and hard panic is valid behaviour.

[SBrandeis]

What does this do?

[McPatate]

from the fn's doc:

core::hint

pub fn spin_loop()

Emits a machine instruction to signal the processor that it is running in
a busy-wait spin-loop ("spin lock").

Upon receiving the spin-loop signal the processor can optimize its behavior by,
for example, saving power or switching hyper-threads.

This function is different from thread::yield_now which directly
yields to the system's scheduler, whereas spin_loop does not interact
with the operating system.

A common use case for spin_loop is implementing bounded optimistic
spinning in a CAS loop in synchronization primitives. To avoid problems
like priority inversion, it is strongly recommended that the spin loop is
terminated after a finite amount of iterations and an appropriate blocking
syscall is made.

Note: On platforms that do not support receiving spin-loop hints this
function does not do anything at all.

[SBrandeis]

🙈

Suggested change

	seq: usize,
	sequence_idx: usize,

[McPatate]

I fear it makes the Unit::idx less clear afterwards actually, I think I'd rather leave it as seq!

[SBrandeis]

naming suggestion, feel free to disregard

Suggested change

	struct Unit {
	struct WorkUnit {

[McPatate]

Well, it's a unit to work on, not a unit of work technically 🤓