Skip to content

[RNE Rewrite] feat: add tokenizer pipeline (#1248)#1274

Merged
msluszniak merged 3 commits into
rne-rewritefrom
@ms/issue1248-tokenizer
Jun 23, 2026
Merged

[RNE Rewrite] feat: add tokenizer pipeline (#1248)#1274
msluszniak merged 3 commits into
rne-rewritefrom
@ms/issue1248-tokenizer

Conversation

@msluszniak

@msluszniak msluszniak commented Jun 22, 2026

Copy link
Copy Markdown
Member

Description

Adds the tokenizer pipeline (issue #1248) using the new worklet-based architecture, with functional parity to the current TokenizerModule.

A new nlp extension exposes a loadTokenizer JSI primitive (top-level on __rnexecutorch_jsi__, like loadModel) returning a Tokenizer host object backed by tokenizers::HFTokenizer. On top of it sits a createTokenizer(config, runtime?) async factory (async + *Worklet variants + dispose) and a useTokenizer hook. Methods: encode, decode, getVocabSize, idToToken, tokenToId — same semantics as today (special tokens follow the tokenizer.json post_processor). The *Worklet variants let an upcoming text-embeddings task tokenize → build tensors → run forward within a single worklet.

  • C++: cpp/extensions/nlp/{tokenizer,install}.{h,cpp}, wired into RnExecutorch.cpp.
  • TS: src/extensions/nlp/{ops,tasks}/tokenizer.ts, src/hooks/useTokenizer.ts, exports in index.ts, example models.tokenizer.ALL_MINILM_L6_V2.
  • Build: tokenizer header-search paths added to android/CMakeLists.txt and the podspec — pytorch/tokenizers/include plus the bundled libs its public headers pull in (nlohmann/json, re2, and re2's abseil dep). Symbols link from the prebuilt libexecutorch. Documented in third-party/README.md.
  • Demo: a dedicated apps/nlp example app with a Tokenizer screen that drives the full pipeline on device.

Introduces a breaking change?

  • Yes
  • No

Type of change

  • Bug fix (change which fixes an issue)
  • New feature (change which adds functionality)
  • Documentation update (improves or adds clarity to existing documentation)
  • Other (chores, tests, code style improvements etc.)

Tested on

  • iOS
  • Android

Testing instructions

CI is TypeScript-only here (native isn't compiled in CI); yarn typecheck, root yarn lint, and yarn prepare (bob build) all pass.

To exercise the native tokenizer end-to-end via the demo app:

  1. Provision the ExecuTorch third-party artifacts into packages/react-native-executorch/third-party/ (see third-party/README.md). The existing PoC bundle already ships the llm/tokenizers extension (headers + symbols), so no rebuild is needed.
  2. yarn && cd apps/nlp && yarn ios — the iOS simulator works since the tokenizer is pure CPU (no GPU/Metal).
  3. Open the Tokenizer screen. It loads all-MiniLM-L6-v2 and auto-runs encode / decode / getVocabSize / idToToken / tokenToId, asserting: encode("Hello world") = [7592, 2088], decode round-trips to "hello world", getVocabSize() = 30522, and tokenToId(idToToken(id)) is the identity. All three assertions should read PASS (also logged as [TokenizerTest]).

The screen genuinely drives the new code — useTokenizercreateTokenizer → the loadTokenizer JSI primitive → the native TokenizerHostObject / HFTokenizer — so a green run validates the whole pipeline, not just the types. Verified locally on iOS (iPhone 16 Pro Max, Xcode 26.5): all assertions pass.

The apps/nlp app is intentionally minimal and exists only to prove this pipeline; it can be dropped after approval if you'd prefer not to keep a demo app in-tree.

Screenshots

Related issues

#1248, part of #1208

Checklist

  • I have performed a self-review of my code
  • I have commented my code, particularly in hard-to-understand areas
  • I have updated the documentation accordingly
  • My changes generate no new warnings

Additional notes

The C++ mirrors the current TokenizerModule, which is backed by pytorch/tokenizers (tokenizers::HFTokenizer) from the ExecuTorch llm/tokenizers extension. This PR consumes the same headers (the third-party bundle under extension/llm/tokenizers) and prebuilt symbols (from libexecutorch); it does not use the tokenizers-cpp submodule. Tokenizer download currently uses the temporary react-native-fs-based useResourceDownload introduced in #1264 (to be replaced by the ResourceFetcher in #1253).

@msluszniak msluszniak marked this pull request as draft June 22, 2026 12:49
@msluszniak msluszniak self-assigned this Jun 22, 2026
@msluszniak msluszniak added the feature PRs that implement a new feature label Jun 22, 2026
@msluszniak msluszniak linked an issue Jun 22, 2026 that may be closed by this pull request
@msluszniak msluszniak force-pushed the @ms/issue1248-tokenizer branch 4 times, most recently from c5817d8 to f426882 Compare June 22, 2026 13:30
@msluszniak msluszniak force-pushed the @ms/issue1248-tokenizer branch from f426882 to 66dfb9d Compare June 22, 2026 17:10
@msluszniak msluszniak force-pushed the @ms/issue1248-tokenizer branch from 66dfb9d to d394a7e Compare June 22, 2026 18:47
@msluszniak msluszniak force-pushed the @ms/issue1248-tokenizer branch from d394a7e to 81d7ab1 Compare June 22, 2026 18:56
@msluszniak msluszniak marked this pull request as ready for review June 22, 2026 19:14
@msluszniak msluszniak requested a review from barhanc June 22, 2026 19:14

@barhanc barhanc left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I only went over the library implementation. Tomorrow I will take a look at the example app and test it.

Comment thread packages/react-native-executorch/android/CMakeLists.txt Outdated
Comment thread packages/react-native-executorch/cpp/extensions/nlp/tokenizer.cpp Outdated
Comment thread packages/react-native-executorch/cpp/extensions/nlp/tokenizer.cpp Outdated
Comment thread packages/react-native-executorch/cpp/extensions/nlp/tokenizer.cpp Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/ops/tokenizer.ts Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/tasks/tokenizer.ts Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/ops/tokenizer.ts Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/tasks/tokenizer.ts Outdated
Comment thread packages/react-native-executorch/cpp/extensions/nlp/tokenizer.cpp Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/tasks/tokenizer.ts Outdated
Comment thread packages/react-native-executorch/cpp/extensions/nlp/install.cpp Outdated
Comment thread packages/react-native-executorch/src/extensions/nlp/ops/tokenizer.ts Outdated
@msluszniak msluszniak force-pushed the @ms/issue1248-tokenizer branch 2 times, most recently from 5fddd90 to cec7c37 Compare June 23, 2026 11:08
@msluszniak msluszniak requested a review from barhanc June 23, 2026 11:28

@barhanc barhanc left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tested and works good! One more thing I noticed (that isn't blocking on this PR because it will probably become clearer when we add some downstream tasks that use tokenizer) is the serialization cost of number[] returned by the async encode. I worry it might be quite large for long texts. I guess it's not a problem for tasks like embeddings or privacy-filter since there we would use the tokenizer.ts directly and tokens array would never pass the thread boundary since the tokens would then get consumed by some other model and the final returned object might be much lighter. Worth keeping in mind though. For now since the hook is mostly to test the tokenizer implementation it's fine.

Comment thread packages/react-native-executorch/src/extensions/nlp/tasks/tokenization.ts Outdated
Comment thread apps/nlp/app/tokenizer/index.tsx Outdated
- install loadTokenizer under the nlp JSI submodule (rnexecutorchJsi.nlp)
- move ops/tokenizer.ts to nlp/tokenizer.ts; standard-template JSDoc, drop
  low/high-level coupling and Model/Tensor mention
- default skipSpecialTokens and empty-input handling in C++ decode
- minimal tokenization task (no worklet wrappers; sync lightweight lookups);
  inline single-field config to a path string
- readable tokenizer error strings; drop unnecessary build-config comments
@msluszniak msluszniak force-pushed the @ms/issue1248-tokenizer branch from cec7c37 to b0548f1 Compare June 23, 2026 13:41
@msluszniak

Copy link
Copy Markdown
Member Author

Tested and works good! One more thing I noticed (that isn't blocking on this PR because it will probably become clearer when we add some downstream tasks that use tokenizer) is the serialization cost of number[] returned by the async encode. I worry it might be quite large for long texts. I guess it's not a problem for tasks like embeddings or privacy-filter since there we would use the tokenizer.ts directly and tokens array would never pass the thread boundary since the tokens would then get consumed by some other model and the final returned object might be much lighter. Worth keeping in mind though. For now since the hook is mostly to test the tokenizer implementation it's fine.

I'll add this point to the TODO.md. Maybe you have an idea how this might be fixed so we can quickly prototype this soon?

@barhanc

barhanc commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

I guess returning TypedArray instead of number[] is an easy fix, since with patch the transfer of ArrayBuffer between threads is fast.

@msluszniak

Copy link
Copy Markdown
Member Author

@barhanc last question, do you want to keep this demo app? I build it mainly for testing purpose. It is not interactive at all, so maybe we shouldn't keep it at all?

@barhanc

barhanc commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

For now it's fine. When we have some other screens like embeddings we can remove it.

@msluszniak msluszniak merged commit 1c95bb7 into rne-rewrite Jun 23, 2026
2 checks passed
@msluszniak msluszniak deleted the @ms/issue1248-tokenizer branch June 23, 2026 14:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature PRs that implement a new feature refactoring

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[RNE Rewrite] Add tokenizer pipeline implementation

2 participants