Skip to content

feat(python): add Arrow string batch tokenization - #2400

Closed
granthamtaylor wants to merge 1 commit into
huggingface:feat/train_encode_splitfrom
granthamtaylor:feat/python-arrow-batch-tokenization
Closed

granthamtaylor wants to merge 1 commit into
huggingface:feat/train_encode_splitfrom
granthamtaylor:feat/python-arrow-batch-tokenization

Conversation

@granthamtaylor

@granthamtaylor granthamtaylor commented Sep 10, 2026

Copy link
Copy Markdown

Arrow callers currently need to convert arrays to Python lists before calling encode_batch. This adds Tokenizer.encode_batch_arrow, which accepts Arrow string and large_string arrays through __arrow_c_array__ and returns the existing pipeline Encoding objects.

Depends on #2119 and targets feat/train_encode_split. The diff contains one commit confined to the Python bindings.

Behavior

  • Supports sliced and empty arrays, special tokens, and the existing padding configuration and per-call overrides.
  • Nulls raise an error by default. null_handling="empty" encodes them as empty strings; "skip" omits them while preserving retained row order.
  • Validates string buffers and capsule ownership, with release callbacks running after reattaching to Python.
  • Avoids materializing Python strings. The existing pipeline still copies text into Rust-owned inputs.

Includes generated type hints, parity and ownership tests, and a standalone benchmark. PyArrow is a development dependency only. Chunked arrays and non-string Arrow types are outside the method's scope.

Validation

On macOS ARM64 with Python 3.13:

  • Local Python suite: 211 passed, 1 skipped, with 1 network test deselected.
  • Rust formatting and Clippy, Ruff, and ty check python tests tools: passed.
  • Release-build benchmark smoke test: passed for both Arrow string widths.

The full example type check still reports errors in the unchanged upstream examples/train_parity_bpe.py, which uses the previous binding API.

Related: #1415.

@ArthurZucker

Copy link
Copy Markdown
Collaborator

Happy to have a look once rebased on #2119

Accept Arrow string and large_string arrays through the C Data Interface
without materializing Python strings. Use the existing pipeline encoding
and padding APIs, with configurable null handling and Literal typing.

Add FFI lifetime and parity tests plus a Python benchmark matrix for
conversion, encoding, throughput, and process peak memory.
@granthamtaylor
granthamtaylor force-pushed the feat/python-arrow-batch-tokenization branch from f47c8a2 to f489f05 Compare September 15, 2026 17:34
@hf-security-analysis

Copy link
Copy Markdown
Contributor

🚨 Workflow security review — 16 critical, 20 high, 35 medium, 17 low, 2 informational

Scanned 10 workflow file(s) with zizmor, pinact, OSV/GHSA, Claude, in full.

.github/workflows/docs-check.yml

  • 🚨 CRITICAL impostor-commit — .github/workflows/docs-check.yml:31 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

.github/workflows/node.yml

  • 🚨 CRITICAL impostor-commit — .github/workflows/node.yml:24 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/node.yml:73 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/node.yml:111 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • ⚠️ HIGH unpinned-action — .github/workflows/node.yml:33 (via pinact)
    Action is not pinned to an immutable commit SHA: uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4

    Fix: Replace with uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0

  • ⚠️ HIGH unpinned-action — .github/workflows/node.yml:46 (via pinact)
    Action is not pinned to an immutable commit SHA: uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4

    Fix: Replace with uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0

  • ⚠️ HIGH unpinned-action — .github/workflows/node.yml:55 (via pinact · also flagged by zizmor)
    Action is not pinned to an immutable commit SHA: uses: astral-sh/setup-uv@v6

    Fix: Replace with uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6.8.0

  • ⚠️ HIGH unpinned-action — .github/workflows/node.yml:90 (via pinact)
    Action is not pinned to an immutable commit SHA: uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4

    Fix: Replace with uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0

.github/workflows/python.yml

  • 🚨 CRITICAL impostor-commit — .github/workflows/python.yml:26 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/python.yml:60 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/python.yml:112 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/python.yml:161 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • ⚠️ HIGH unpinned-action — .github/workflows/python.yml:76 (via pinact · also flagged by zizmor)
    Action is not pinned to an immutable commit SHA: uses: astral-sh/setup-uv@v6

    Fix: Replace with uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6.8.0

  • ⚠️ HIGH unpinned-action — .github/workflows/python.yml:86 (via pinact)
    Action is not pinned to an immutable commit SHA: uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4

    Fix: Replace with uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0

  • ⚠️ HIGH unpinned-action — .github/workflows/python.yml:128 (via pinact · also flagged by zizmor)
    Action is not pinned to an immutable commit SHA: uses: astral-sh/setup-uv@v6

    Fix: Replace with uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6.8.0

.github/workflows/rust-release.yml

  • 🚨 CRITICAL impostor-commit — .github/workflows/rust-release.yml:19 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

.github/workflows/rust.yml

  • 🚨 CRITICAL impostor-commit — .github/workflows/rust.yml:26 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/rust.yml:115 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/rust.yml:138 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/rust.yml:155 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/rust.yml:186 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/rust.yml:210 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • 🚨 CRITICAL impostor-commit — .github/workflows/rust.yml:243 (via zizmor)
    commit with no history in referenced repository — uses a commit that doesn't belong to the specified org/repo

    Fix: The pinned SHA is not reachable from any branch or tag of the action's repository — it may come from a fork or a force-pushed ref. Re-pin from an upstream tag.

  • ⚠️ HIGH unpinned-action — .github/workflows/rust.yml:47 (via pinact)
    Action is not pinned to an immutable commit SHA: uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4

    Fix: Replace with uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0

  • ⚠️ HIGH unpinned-action — .github/workflows/rust.yml:56 (via pinact · also flagged by zizmor)
    Action is not pinned to an immutable commit SHA: uses: astral-sh/setup-uv@v6

    Fix: Replace with uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6.8.0

  • ⚠️ HIGH unpinned-action — .github/workflows/rust.yml:258 (via pinact)
    Action is not pinned to an immutable commit SHA: uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4

    Fix: Replace with uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0

  • ⚠️ HIGH unpinned-action — .github/workflows/rust.yml:267 (via pinact · also flagged by zizmor)
    Action is not pinned to an immutable commit SHA: uses: astral-sh/setup-uv@v6

    Fix: Replace with uses: astral-sh/setup-uv@d0cc045d04ccac9d8b7881df0226f9e82c39688e # v6.8.0

  • ⚠️ HIGH unpinnable-reference — .github/workflows/rust.yml:280 (via pinact · also flagged by zizmor)
    This reference cannot be pinned automatically — it is a branch or a non-semver ref, which moves with every upstream push.

    Fix: Replace the branch reference with a released tag, then pin that tag to its commit SHA. Set PINACT_BRANCH_TO_TAG to let the bot resolve these to the latest stable tag on its own.

.github/workflows/CI.yml

  • ⚠️ HIGH unpinned-action — .github/workflows/CI.yml:200 (via pinact)
    Action is not pinned to an immutable commit SHA: uses: actions/attest-build-provenance@ef244123eb79f2f7a7e75d99086184180e6d0018 # v1

    Fix: Replace with uses: actions/attest-build-provenance@ef244123eb79f2f7a7e75d99086184180e6d0018 # v1.4.4

.github/workflows/bitsplit.yml

  • ⚠️ HIGH unpinned-action — .github/workflows/bitsplit.yml:18 (via pinact · also flagged by zizmor)
    Action is not pinned to an immutable commit SHA: - uses: actions/checkout@v4

    Fix: Replace with - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0

  • ⚠️ HIGH unpinnable-reference — .github/workflows/bitsplit.yml:19 (via pinact · also flagged by zizmor)
    This reference cannot be pinned automatically — it is a branch or a non-semver ref, which moves with every upstream push.

    Fix: Replace the branch reference with a released tag, then pin that tag to its commit SHA. Set PINACT_BRANCH_TO_TAG to let the bot resolve these to the latest stable tag on its own.

  • ⚠️ HIGH unpinned-action — .github/workflows/bitsplit.yml:32 (via pinact · also flagged by zizmor)
    Action is not pinned to an immutable commit SHA: - uses: actions/checkout@v4

    Fix: Replace with - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0

  • ⚠️ HIGH unpinnable-reference — .github/workflows/bitsplit.yml:33 (via pinact · also flagged by zizmor)
    This reference cannot be pinned automatically — it is a branch or a non-semver ref, which moves with every upstream push.

    Fix: Replace the branch reference with a released tag, then pin that tag to its commit SHA. Set PINACT_BRANCH_TO_TAG to let the bot resolve these to the latest stable tag on its own.

  • ⚠️ HIGH unpinned-action — .github/workflows/bitsplit.yml:53 (via pinact · also flagged by zizmor)
    Action is not pinned to an immutable commit SHA: - uses: actions/checkout@v4

    Fix: Replace with - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0

  • ⚠️ HIGH unpinnable-reference — .github/workflows/bitsplit.yml:54 (via pinact · also flagged by zizmor)
    This reference cannot be pinned automatically — it is a branch or a non-semver ref, which moves with every upstream push.

    Fix: Replace the branch reference with a released tag, then pin that tag to its commit SHA. Set PINACT_BRANCH_TO_TAG to let the bot resolve these to the latest stable tag on its own.

  • ⚠️ HIGH unpinned-action — .github/workflows/bitsplit.yml:56 (via pinact · also flagged by zizmor)
    Action is not pinned to an immutable commit SHA: - uses: bytecodealliance/actions/wasmtime/setup@v1

    Fix: Replace with - uses: bytecodealliance/actions/wasmtime/setup@9152e710e9f7182e4c29ad218e4f335a7b203613 # v1.1.3

  • 🟡 MEDIUM excessive-permissions — .github/workflows/bitsplit.yml:1 (via zizmor)
    overly broad permissions — default permissions used due to no permissions: block

    Fix: Declare an explicit permissions: block with the narrowest scopes the job needs, at job level rather than workflow level.

  • 🟡 MEDIUM excessive-permissions — .github/workflows/bitsplit.yml:15 (via zizmor)
    overly broad permissions — default permissions used due to no permissions: block

    Fix: Declare an explicit permissions: block with the narrowest scopes the job needs, at job level rather than workflow level.

  • 🟡 MEDIUM excessive-permissions — .github/workflows/bitsplit.yml:29 (via zizmor)
    overly broad permissions — default permissions used due to no permissions: block

    Fix: Declare an explicit permissions: block with the narrowest scopes the job needs, at job level rather than workflow level.

  • 🟡 MEDIUM excessive-permissions — .github/workflows/bitsplit.yml:50 (via zizmor)
    overly broad permissions — default permissions used due to no permissions: block

    Fix: Declare an explicit permissions: block with the narrowest scopes the job needs, at job level rather than workflow level.

…and 50 more finding(s) not shown.

ℹ️ This PR comes from a fork, so the bot cannot push a fix branch to it. Apply the suggestions below manually.

Pinning and static findings are deterministic. Findings marked (via Claude) are model judgements — check them before acting.

@granthamtaylor
granthamtaylor changed the base branch from main to feat/train_encode_split September 15, 2026 17:48
@granthamtaylor

Copy link
Copy Markdown
Author

@ArthurZucker thank you so much !

I have rebased.

@ArthurZucker

Copy link
Copy Markdown
Collaborator

Ty! this won't make it to rc0 but we'll se after that!

@ArthurZucker
ArthurZucker deleted the branch huggingface:feat/train_encode_split September 18, 2026 11:22
@granthamtaylor

granthamtaylor commented Sep 18, 2026

Copy link
Copy Markdown
Author

I understand!

Please feel free to tag me with any next steps.

I'd love to see this feature integrated for my current project in which I am trying to tokenize many, many small strings.

Thank you for your consideration here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants