Repository navigation
[ML] Add byte_level_bpe ML tokenization with merges - #147217
Conversation
Adds ByteLevelBpeTokenization, ByteLevelBpeTokenizer, XContent and named-writeable registration, NlpTokenizer routing, and tests. Special token strings are configurable with RoBERTa-style defaults. Made-with: Cursor
NlpConfigUpdate.tokenizationFromMap lists known tokenizer names including byte_level_bpe; update the expected exception substring. Made-with: Cursor
There was a problem hiding this comment.
Pull request overview
This PR introduces a new byte_level_bpe tokenization option for ML NLP inference in Elasticsearch Java, including merge-table support, serialization/parsing plumbing, and tokenizer wiring alongside existing BPE-based tokenizers (e.g., RoBERTa).
Changes:
- Add
ByteLevelBpeTokenizationandByteLevelBpeTokenizationUpdatewith XContent + named writeable registration. - Implement
ByteLevelBpeTokenizerand wire it intoNlpTokenizer.build(...). - Add unit tests for tokenizer behavior and config/update serialization scaffolding.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/nlp/tokenizers/ByteLevelBpeTokenizerTests.java | Adds tokenizer parity tests vs RoBERTa-style defaults and multi-seq behavior. |
| x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/nlp/tokenizers/NlpTokenizer.java | Wires ByteLevelBpeTokenization to build a ByteLevelBpeTokenizer. |
| x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/nlp/tokenizers/ByteLevelBpeTokenizer.java | Implements the byte-level BPE tokenizer using existing BPE analyzer machinery and configurable special tokens. |
| x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/TextEmbeddingConfigTests.java | Includes ByteLevelBpeTokenization in random test generation for embedding configs. |
| x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/NlpConfigUpdateTests.java | Updates expected “unknown tokenization type” message to include byte_level_bpe. |
| x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/InferenceConfigTestScaffolding.java | Extends test scaffolding to clone/mutate byte-level BPE tokenization + updates. |
| x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenizationUpdateTests.java | Adds apply + wire serialization test coverage for tokenization updates. |
| x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenizationTests.java | Adds XContent + wire serialization tests for the new tokenization config. |
| x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/NlpConfigUpdate.java | Registers byte_level_bpe tokenization update parsing in config updates. |
| x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenizationUpdate.java | Implements update application logic for ByteLevelBpeTokenization. |
| x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenization.java | Adds the new tokenization config, parsing, serialization, and merges-required validation. |
| x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/MlInferenceNamedXContentProvider.java | Registers named XContent parsers and named writeables for the new tokenization + update. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| Tokenization.validateSpanAndTruncate(getTruncate(), getSpan()); | ||
|
|
||
| if (getTruncate() != null && getTruncate().isInCompatibleWithSpan() == false) { | ||
| return new ByteLevelBpeTokenization( | ||
| byteLevelBpeTokenization.doLowerCase(), | ||
| byteLevelBpeTokenization.withSpecialTokens(), | ||
| byteLevelBpeTokenization.maxSequenceLength(), | ||
| getTruncate(), | ||
| null, | ||
| byteLevelBpeTokenization.isAddPrefixSpace(), | ||
| byteLevelBpeTokenization.getUnkToken(), | ||
| byteLevelBpeTokenization.getPadToken(), | ||
| byteLevelBpeTokenization.getBosToken(), | ||
| byteLevelBpeTokenization.getEosToken(), | ||
| byteLevelBpeTokenization.getMaskToken() | ||
| ); | ||
| } | ||
|
|
||
| return new ByteLevelBpeTokenization( | ||
| byteLevelBpeTokenization.doLowerCase(), | ||
| byteLevelBpeTokenization.withSpecialTokens(), | ||
| byteLevelBpeTokenization.maxSequenceLength(), | ||
| Optional.ofNullable(this.getTruncate()).orElse(originalConfig.getTruncate()), | ||
| Optional.ofNullable(this.getSpan()).orElse(originalConfig.getSpan()), |
There was a problem hiding this comment.
The span-clearing logic here is inverted. Truncate.isInCompatibleWithSpan() returns true for FIRST/SECOND/BALANCED and false for NONE; the current == false branch triggers only for NONE and clears the span. This causes updates that set truncate to FIRST/SECOND/BALANCED to incorrectly keep an existing span (potentially throwing in the Tokenization constructor), and updates that set truncate to NONE to unnecessarily drop windowing. Flip the condition so span is cleared when the new truncate is incompatible with span, and ensure the returned config is always valid.
| Tokenization.validateSpanAndTruncate(getTruncate(), getSpan()); | |
| if (getTruncate() != null && getTruncate().isInCompatibleWithSpan() == false) { | |
| return new ByteLevelBpeTokenization( | |
| byteLevelBpeTokenization.doLowerCase(), | |
| byteLevelBpeTokenization.withSpecialTokens(), | |
| byteLevelBpeTokenization.maxSequenceLength(), | |
| getTruncate(), | |
| null, | |
| byteLevelBpeTokenization.isAddPrefixSpace(), | |
| byteLevelBpeTokenization.getUnkToken(), | |
| byteLevelBpeTokenization.getPadToken(), | |
| byteLevelBpeTokenization.getBosToken(), | |
| byteLevelBpeTokenization.getEosToken(), | |
| byteLevelBpeTokenization.getMaskToken() | |
| ); | |
| } | |
| return new ByteLevelBpeTokenization( | |
| byteLevelBpeTokenization.doLowerCase(), | |
| byteLevelBpeTokenization.withSpecialTokens(), | |
| byteLevelBpeTokenization.maxSequenceLength(), | |
| Optional.ofNullable(this.getTruncate()).orElse(originalConfig.getTruncate()), | |
| Optional.ofNullable(this.getSpan()).orElse(originalConfig.getSpan()), | |
| Tokenization.Truncate truncate = Optional.ofNullable(this.getTruncate()).orElse(originalConfig.getTruncate()); | |
| Integer span = Optional.ofNullable(this.getSpan()).orElse(originalConfig.getSpan()); | |
| if (getTruncate() != null && getTruncate().isInCompatibleWithSpan()) { | |
| span = null; | |
| } | |
| Tokenization.validateSpanAndTruncate(truncate, span); | |
| return new ByteLevelBpeTokenization( | |
| byteLevelBpeTokenization.doLowerCase(), | |
| byteLevelBpeTokenization.withSpecialTokens(), | |
| byteLevelBpeTokenization.maxSequenceLength(), | |
| truncate, | |
| span, |
There was a problem hiding this comment.
I'm deliberately leaving the logic as-is for parity with existing tokenizers and the tests (RobertaTokenizationUpdateTests / ByteLevelBpeTokenizationUpdateTests).
ByteLevelBpeTokenizationUpdate#apply intentionally matches RobertaTokenizationUpdate#apply (and the other *TokenizationUpdate types): validate the update’s own truncate/span, then take the special path when getTruncate() != null && getTruncate().isInCompatibleWithSpan() == false - i.e. when the update explicitly sets Truncate.NONE, which is the only enum whose isInCompatibleWithSpan() is false. That branch forces span to unset so a partial update that sets truncate: none doesn’t inherit a previous sliding-window span. Flipping the predicate would drop that NONE behaviour and instead clear span only for first/second/balanced, which would diverge from Roberta unless we changed all tokenizers together - and it would still be a behaviour change (e.g. updates that currently throw when merged truncate/span is invalid would start succeeding with span coerced away).
The Roberta comment above that if is misleading (it says “incompatible” but the condition is the NONE / span-compatible case); we can fix that separately rather than changing the predicate only on this class.
| var updatedTruncate = new ByteLevelBpeTokenizationUpdate(Tokenization.Truncate.FIRST, null).apply( | ||
| new ByteLevelBpeTokenization(true, true, 512, Tokenization.Truncate.SECOND, null, true, null, null, null, null, null) | ||
| ); | ||
| assertEquals( | ||
| new ByteLevelBpeTokenization(true, true, 512, Tokenization.Truncate.FIRST, null, true, null, null, null, null, null), | ||
| updatedTruncate | ||
| ); | ||
|
|
||
| var updatedNone = new ByteLevelBpeTokenizationUpdate(Tokenization.Truncate.NONE, null).apply( | ||
| new ByteLevelBpeTokenization(true, true, 512, Tokenization.Truncate.SECOND, null, false, null, null, null, null, null) | ||
| ); | ||
| assertEquals( | ||
| new ByteLevelBpeTokenization(true, true, 512, Tokenization.Truncate.NONE, null, false, null, null, null, null, null), | ||
| updatedNone |
There was a problem hiding this comment.
testApply doesn’t currently cover the important transition from a windowing config (truncate=none + span set) to an incompatible truncate (FIRST/SECOND/BALANCED), where apply() should clear the span rather than throwing. Adding an assertion for that case would have caught the span-handling bug in ByteLevelBpeTokenizationUpdate.apply() and helps prevent regressions.
There was a problem hiding this comment.
Today, with the code as it stands, apply() doesn’t strip the old span when you go from none + sliding window to first/second/balanced and leave span out of the update - it merges in the previous span and Tokenization’s ctor blows up with IllegalArgumentException. That’s eactly the same story as Roberta (and friends), not something ByteLevel BPE is doing differently.
So a test that expects span to be auto-cleared is really arguing for a new behaviour we’d want everywhere, not just BPE.
I've added tests for what we have today. An expectThrows for that transition has been added for all TokenizationUpdateTests.
Split the expected error substring so test sources stay within the 140-column limit. Made-with: Cursor
Guard with truncate != null so null truncate does not clear span. That restores span-only random instances for BWC serialization tests across all tokenizer update test classes. Made-with: Cursor
|
There is noticeable duplication across the various Consolidating that (for example a shared helper or a small per-tokenizer strategy for the |
Add expectThrows for NONE+sliding span merged with FIRST-only updates across all TokenizationUpdateTests, matching Tokenization validation. Made-with: Cursor
edsavage
changed the title
elasticsearchmachine
added
the
needs:triage
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 17 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| public ByteLevelBpeTokenization( | ||
| @Nullable Boolean doLowerCase, | ||
| @Nullable Boolean withSpecialTokens, | ||
| @Nullable Integer maxSequenceLength, | ||
| @Nullable Truncate truncate, | ||
| @Nullable Integer span, | ||
| @Nullable Boolean addPrefixSpace, | ||
| @Nullable String unkToken, | ||
| @Nullable String padToken, | ||
| @Nullable String bosToken, | ||
| @Nullable String eosToken, | ||
| @Nullable String maskToken | ||
| ) { | ||
| super(doLowerCase, withSpecialTokens, maxSequenceLength, truncate, span); | ||
| this.addPrefixSpace = Optional.ofNullable(addPrefixSpace).orElse(DEFAULT_ADD_PREFIX_SPACE); | ||
| this.unkToken = resolveSpecialToken(unkToken, DEFAULT_UNK_TOKEN); | ||
| this.padToken = resolveSpecialToken(padToken, DEFAULT_PAD_TOKEN); | ||
| this.bosToken = resolveSpecialToken(bosToken, DEFAULT_BOS_TOKEN); | ||
| this.eosToken = resolveSpecialToken(eosToken, DEFAULT_EOS_TOKEN); | ||
| this.maskToken = resolveSpecialToken(maskToken, DEFAULT_MASK_TOKEN); | ||
| } |
There was a problem hiding this comment.
do_lower_case is accepted here and passed to super(...), but the runtime tokenizer (ByteLevelBpeTokenizer/BpeAnalyzer) does not apply lowercasing (unlike WordPiece tokenizers). This makes do_lower_case=true a silent no-op that can change expected token IDs and is inconsistent with RobertaTokenization/XLMRobertaTokenization, which explicitly reject do_lower_case=true. Consider either (a) rejecting do_lower_case=true with an IllegalArgumentException and forcing super(false, ...), or (b) actually applying lowercasing in the analyzer/tokenizer pipeline so the setting has a real effect.
| if (withSpecialTokens) { | ||
| Set<String> missingSpecialTokens = Sets.difference(Set.of(eosToken, bosToken), vocab.keySet()); | ||
| if (missingSpecialTokens.isEmpty() == false) { | ||
| throw ExceptionsHelper.conflictStatusException("stored vocabulary is missing required {} token(s)", missingSpecialTokens); | ||
| } |
There was a problem hiding this comment.
Set.of(eosToken, bosToken) will throw an IllegalArgumentException if a user configures bos_token and eos_token to the same string (duplicate element). Since these token strings are configurable, it’s better to either validate they’re distinct with a clear error message, or build the required-token set in a way that tolerates duplicates (e.g., via a mutable set) before calling Sets.difference(...).
elasticsearchmachine
added
Team:ML
|
Hi @edsavage, I've created a changelog YAML for you. |
|
Pinging @elastic/ml-core (Team:ML) |
🔍 Preview links for changed docs⏳ Building and deploying preview... View progress This comment will be updated with preview links when the build is complete. |
ℹ️ Important: Docs version tagging👋 Thanks for updating the docs! Just a friendly reminder that our docs are now cumulative. This means all 9.x versions are documented on the same page and published off of the main branch, instead of creating separate pages for each minor version. We use applies_to tags to mark version-specific features and changes. Expand for a quick overviewWhen to use applies_to tags:✅ At the page level to indicate which products/deployments the content applies to (mandatory) What NOT to do:❌ Don't remove or replace information that applies to an older version 🤔 Need help?
|
Match Roberta by rejecting do_lower_case=true via delegating ctor and wire checks. Use HashSet for required BOS/EOS when validating vocab so identical configured strings do not trip Set.of. Add tests; fix BPE fixtures that passed true in the do_lower_case slot. Made-with: Cursor
Document RobertaTokenizationResult reuse and mask vocab behavior; pass false for do_lower_case in tokenizer builder; align NlpConfigUpdate map key with ByteLevelBpeTokenizationUpdate.NAME; add NlpConfigUpdateTests and ml yamlRestTest for byte_level_bpe text_embedding config. Made-with: Cursor
Replace x-pack plugin YAML suite coverage with ByteLevelBpeInferenceConfigRestIT under ml single-node-tests so classpath suite resolution is unchanged. Made-with: Cursor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 19 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| namedWriteables.add( | ||
| new NamedWriteableRegistry.Entry(Tokenization.class, ByteLevelBpeTokenization.NAME, ByteLevelBpeTokenization::new) | ||
| ); |
| ByteLevelBpeTokenizationUpdate.NAME.getPreferredName(), | ||
| ByteLevelBpeTokenizationUpdate::new | ||
| ) | ||
| ); | ||
| namedWriteables.add( | ||
| new NamedWriteableRegistry.Entry( | ||
| TokenizationUpdate.class, |
Add ml_byte_level_bpe_tokenization transport version, reject serialization to older nodes from ByteLevelBpeTokenization and ByteLevelBpeTokenizationUpdate, and raise TextEmbedding config minimal supported version when byte-level BPE is used. Update BWC tests and TextEmbeddingConfigTests bwcVersions filter.
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
WalkthroughThis PR adds comprehensive support for Byte-Level BPE tokenization to Elasticsearch's ML inference system, including new tokenization configuration classes, a tokenizer implementation, registry integration, and extensive test coverage. ChangesByte-Level BPE Tokenization Support
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenization.java (2)
104-132: 💤 Low valueValidate
doLowerCasebefore delegating, to surface the right error first.The public constructor delegates to the private ctor (which calls
super(false, ...)and runsvalidateSpanAndTruncate) before checkingBoolean.TRUE.equals(doLowerCase). If a caller passes bothdo_lower_case=trueand an invalid span/truncate combination, the user gets a span/truncate error rather than the more diagnostic "do_lower_case=true not allowed" error. Consider extracting the validation into astatichelper invoked as the first argument expression so it runs prior to delegation.♻️ Suggested refactor
public ByteLevelBpeTokenization( `@Nullable` Boolean doLowerCase, `@Nullable` Boolean withSpecialTokens, ... `@Nullable` String maskToken ) { this( - withSpecialTokens, + validateDoLowerCase(doLowerCase) ? withSpecialTokens : withSpecialTokens, maxSequenceLength, ... ); - if (Boolean.TRUE.equals(doLowerCase)) { - throw new IllegalArgumentException("unable to set [do_lower_case] to [true] for byte_level_bpe tokenizer"); - } + } + + private static boolean validateDoLowerCase(`@Nullable` Boolean doLowerCase) { + if (Boolean.TRUE.equals(doLowerCase)) { + throw new IllegalArgumentException("unable to set [do_lower_case] to [true] for byte_level_bpe tokenizer"); + } + return true; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenization.java` around lines 104 - 132, Move the doLowerCase check so it runs before the constructor delegation: add a small static helper method (e.g., validateDoLowerCaseOrThrow(Boolean doLowerCase)) that throws the IllegalArgumentException when Boolean.TRUE.equals(doLowerCase), and call that helper as the first argument expression in the public ByteLevelBpeTokenization(...) constructor invocation so the validation executes before delegating to the private constructor (the private ctor and validateSpanAndTruncate remain unchanged).
155-166: 💤 Low valueDefensive check against stream corruption or version mismatch is appropriate, but error message should reflect this.
The check in the StreamInput constructor (lines 157-159) is unreachable under normal operation:
writeToenforcesML_BYTE_LEVEL_BPE_TOKENIZATION_ADDEDversion compatibility, the public constructor rejectsdoLowerCase=true, and all internal construction paths explicitly passfalse. However, the check correctly handles corrupted or mismatched stream data. Consider clarifying the error message to indicate this is a version mismatch or data corruption scenario rather than a normal validation error—something like:"Unexpected [do_lower_case=true] in stream; indicates version mismatch or data corruption".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenization.java` around lines 155 - 166, The StreamInput constructor ByteLevelBpeTokenization(StreamInput in) currently throws an IllegalArgumentException with a generic message when doLowerCase is true; update that exception message to clearly indicate this is unexpected in a stream and likely due to a version mismatch or stream corruption (e.g., "Unexpected [do_lower_case=true] in stream; indicates version mismatch or data corruption") so callers can distinguish stream corruption/version issues from normal validation; locate the check in the ByteLevelBpeTokenization(StreamInput in) constructor and replace the existing message accordingly.x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/nlp/tokenizers/ByteLevelBpeTokenizer.java (1)
67-99: ⚡ Quick winMask token is never null; reorder validations before BpeAnalyzer allocation to improve resource handling.
The null mask-token concern is unfounded.
ByteLevelBpeTokenizationalways applies default values viaresolveSpecialTokenin its constructor, guaranteeing all token strings (includingmaskToken) are non-null when passed to this constructor. Therefore,Set.of(maskToken)on line 70 and subsequent uses cannot throwNullPointerException.The resource ordering concern is valid but minor.
BpeAnalyzeris allocated at line 67 before the unk/pad/bos/eos validations at lines 77–98 run. If any validation throws, the analyzer instance is created but never assigned to the field, andclose()is never called. WhileBpeAnalyzerdoes not hold critical resources (no file handles or network connections), moving the validations before instantiation would be cleaner:// Validate tokens against vocab before allocating BpeAnalyzer if (vocab.containsKey(unkToken) == false) { throw ExceptionsHelper.conflictStatusException("stored vocabulary is missing required [{}] token", unkToken); } // ... other validations ... // Then allocate analyzer this.bpeAnalyzer = new BpeAnalyzer(...);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/nlp/tokenizers/ByteLevelBpeTokenizer.java` around lines 67 - 99, The constructor currently instantiates BpeAnalyzer (this.bpeAnalyzer) before validating required tokens (unkToken, padToken, bosToken, eosToken), which can leave an owned analyzer created when a validation later throws; move all vocabulary/token existence checks (vocab.containsKey(unkToken), vocab.containsKey(padToken), and when withSpecialTokens check for eosToken/bosToken) to run before allocating new BpeAnalyzer so the analyzer is only created after validations pass; refer to the ByteLevelBpeTokenizer constructor, maskToken usage in the Set.of(...) call, the bpeAnalyzer field, and the token id fields (padTokenId, sepTokenId, clsTokenId) when making this change and ensure analyzer allocation and assignment happen after those checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenization.java`:
- Around line 26-37: The Javadoc incorrectly implies BOS/EOS vocabulary presence
is validated in ByteLevelBpeTokenization.validateVocabulary(); update the
docstring to state that the requirement is enforced when building/instantiating
the tokenizer (ByteLevelBpeTokenizer constructor) to match actual behavior (and
be consistent with RobertaTokenization and XLMRobertaTokenization), or
alternatively move the HashSet-based presence check into
ByteLevelBpeTokenization.validateVocabulary() if you want vocabulary-upload-time
validation; reference ByteLevelBpeTokenization.validateVocabulary() and
ByteLevelBpeTokenizer (the constructor where the current conflict exception is
thrown) when making the change.
In
`@x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/InferenceConfigTestScaffolding.java`:
- Around line 47-60: The clone of ByteLevelBpeTokenization incorrectly hardcodes
the first constructor argument to false, losing the original doLowerCase()
value; update the cloning block in InferenceConfigTestScaffolding so that the
first argument uses byteLevelBpeTokenization.doLowerCase() (preserving
null/true/false) instead of false while keeping all other fields the same
(withSpecialTokens(), maxSequenceLength(), truncate, getSpan(),
isAddPrefixSpace(), getUnkToken(), getPadToken(), getBosToken(), getEosToken(),
getMaskToken()) so only truncation is changed.
In
`@x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/TextEmbeddingConfigTests.java`:
- Around line 25-35: The current override of bwcVersions() filters out
pre-ML_BYTE_LEVEL_BPE_TOKENIZATION_ADDED versions which removes full BWC
coverage for non-byte-level-BPE TextEmbeddingConfig variants; revert
bwcVersions() to return getAllBWCVersions() (i.e., remove the .filter(...) call)
so general TextEmbeddingConfigTests keep full serialization matrix, and instead
add a dedicated test or branch that exercises createRandom() instances
containing ByteLevelBpeTokenization guarded by checking
ByteLevelBpeTokenization.ML_BYTE_LEVEL_BPE_TOKENIZATION_ADDED (e.g., a new test
method that filters versions and explicitly serializes/deserializes instances
produced by createRandom() that include ByteLevelBpeTokenization) so only that
case is version-gated while others remain unfiltered.
---
Nitpick comments:
In
`@x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenization.java`:
- Around line 104-132: Move the doLowerCase check so it runs before the
constructor delegation: add a small static helper method (e.g.,
validateDoLowerCaseOrThrow(Boolean doLowerCase)) that throws the
IllegalArgumentException when Boolean.TRUE.equals(doLowerCase), and call that
helper as the first argument expression in the public
ByteLevelBpeTokenization(...) constructor invocation so the validation executes
before delegating to the private constructor (the private ctor and
validateSpanAndTruncate remain unchanged).
- Around line 155-166: The StreamInput constructor
ByteLevelBpeTokenization(StreamInput in) currently throws an
IllegalArgumentException with a generic message when doLowerCase is true; update
that exception message to clearly indicate this is unexpected in a stream and
likely due to a version mismatch or stream corruption (e.g., "Unexpected
[do_lower_case=true] in stream; indicates version mismatch or data corruption")
so callers can distinguish stream corruption/version issues from normal
validation; locate the check in the ByteLevelBpeTokenization(StreamInput in)
constructor and replace the existing message accordingly.
In
`@x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/nlp/tokenizers/ByteLevelBpeTokenizer.java`:
- Around line 67-99: The constructor currently instantiates BpeAnalyzer
(this.bpeAnalyzer) before validating required tokens (unkToken, padToken,
bosToken, eosToken), which can leave an owned analyzer created when a validation
later throws; move all vocabulary/token existence checks
(vocab.containsKey(unkToken), vocab.containsKey(padToken), and when
withSpecialTokens check for eosToken/bosToken) to run before allocating new
BpeAnalyzer so the analyzer is only created after validations pass; refer to the
ByteLevelBpeTokenizer constructor, maskToken usage in the Set.of(...) call, the
bpeAnalyzer field, and the token id fields (padTokenId, sepTokenId, clsTokenId)
when making this change and ensure analyzer allocation and assignment happen
after those checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Enterprise
Run ID: 1c9c1b2b-d65e-465e-a887-d4eb028ac17a
⛔ Files ignored due to path filters (2)
server/src/main/resources/transport/definitions/referable/ml_byte_level_bpe_tokenization.csvis excluded by!**/*.csvserver/src/main/resources/transport/upper_bounds/9.5.csvis excluded by!**/*.csv
📒 Files selected for processing (21)
docs/changelog/147217.yamlx-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/MlInferenceNamedXContentProvider.javax-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenization.javax-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenizationUpdate.javax-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/NlpConfigUpdate.javax-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/TextEmbeddingConfig.javax-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/TextEmbeddingConfigUpdate.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/BertJapaneseTokenizationUpdateTests.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/BertTokenizationUpdateTests.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenizationTests.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenizationUpdateTests.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/InferenceConfigTestScaffolding.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/MPNetTokenizationUpdateTests.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/NlpConfigUpdateTests.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/RobertaTokenizationUpdateTests.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/TextEmbeddingConfigTests.javax-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/XLMRobertaTokenizationUpdateTests.javax-pack/plugin/ml/qa/single-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/ByteLevelBpeInferenceConfigRestIT.javax-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/nlp/tokenizers/ByteLevelBpeTokenizer.javax-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/nlp/tokenizers/NlpTokenizer.javax-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/nlp/tokenizers/ByteLevelBpeTokenizerTests.java
| /** | ||
| * Tokenization settings for byte-level BPE models (for example GPT-2 style vocabularies) that ship | ||
| * a merge table alongside the vocabulary. The in-process tokenizer applies UTF-8 byte mapping and | ||
| * BPE merges using the same engine as the RoBERTa tokenizer, but special token strings are | ||
| * configurable so deployments can align with Hugging Face tokenizer metadata. | ||
| * {@code do_lower_case} is not applied by the byte-level BPE analyzer and must not be set to {@code true}; | ||
| * requests that set it to {@code true} are rejected like {@link RobertaTokenization}. | ||
| * When {@code with_special_tokens} is {@code true}, the configured BOS and EOS strings must be present | ||
| * in the vocabulary; the configured mask string is optional for encoding-only inference, and the mask | ||
| * token id may be unset if that token is missing from the vocabulary (masking-oriented tasks should | ||
| * ensure it exists). | ||
| */ |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find where BOS/EOS validation is performed for ByteLevelBpe / Roberta-style tokenization.
fd -e java ByteLevelBpeTokenizer
rg -nP --type=java -C5 '\b(bosToken|eosToken|getBosToken|getEosToken)\b' -g '!**/build/**'
rg -nP --type=java -C3 'HashSet' -g '*ByteLevelBpe*'
rg -nP --type=java -C5 'validateVocabulary' -g '*Tokenization*.java'Repository: elastic/elasticsearch
Length of output: 44862
BOS/EOS-in-vocabulary validation is present; validation timing differs from documented location.
The class Javadoc documents that "When with_special_tokens is true, the configured BOS and EOS strings must be present in the vocabulary." This constraint is enforced, but in ByteLevelBpeTokenizer (lines 84–92 in x-pack/plugin/ml/src/main/java/...), not in ByteLevelBpeTokenization.validateVocabulary(). The validation uses a HashSet to check that configured BOS and EOS tokens exist in the vocabulary, throwing a conflict exception if missing. This validation occurs at tokenizer instantiation time (when the model is built/loaded), not during vocabulary upload validation. The pattern matches other tokenizations (RobertaTokenization, XLMRobertaTokenization), where validateVocabulary() checks only format-specific requirements (non-empty merges/scores) and semantic constraints are enforced in the tokenizer constructor. Consider clarifying the Javadoc to indicate that the constraint is enforced at tokenizer build time rather than at vocabulary upload, or add an explicit check in validateVocabulary() for earlier failure detection if desired.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/ByteLevelBpeTokenization.java`
around lines 26 - 37, The Javadoc incorrectly implies BOS/EOS vocabulary
presence is validated in ByteLevelBpeTokenization.validateVocabulary(); update
the docstring to state that the requirement is enforced when
building/instantiating the tokenizer (ByteLevelBpeTokenizer constructor) to
match actual behavior (and be consistent with RobertaTokenization and
XLMRobertaTokenization), or alternatively move the HashSet-based presence check
into ByteLevelBpeTokenization.validateVocabulary() if you want
vocabulary-upload-time validation; reference
ByteLevelBpeTokenization.validateVocabulary() and ByteLevelBpeTokenizer (the
constructor where the current conflict exception is thrown) when making the
change.
| } else if (tokenization instanceof ByteLevelBpeTokenization byteLevelBpeTokenization) { | ||
| return new ByteLevelBpeTokenization( | ||
| false, | ||
| byteLevelBpeTokenization.withSpecialTokens(), | ||
| byteLevelBpeTokenization.maxSequenceLength(), | ||
| truncate, | ||
| byteLevelBpeTokenization.getSpan(), | ||
| byteLevelBpeTokenization.isAddPrefixSpace(), | ||
| byteLevelBpeTokenization.getUnkToken(), | ||
| byteLevelBpeTokenization.getPadToken(), | ||
| byteLevelBpeTokenization.getBosToken(), | ||
| byteLevelBpeTokenization.getEosToken(), | ||
| byteLevelBpeTokenization.getMaskToken() | ||
| ); |
There was a problem hiding this comment.
Preserve doLowerCase() when cloning byte-level BPE tokenizations.
This helper is supposed to change truncation only, but it hard-codes the first constructor arg to false. ByteLevelBpeTokenization still allows do_lower_case to be unset, so this can produce the wrong expected config in update/version tests whenever the original instance had null there.
♻️ Suggested fix
} else if (tokenization instanceof ByteLevelBpeTokenization byteLevelBpeTokenization) {
return new ByteLevelBpeTokenization(
- false,
+ byteLevelBpeTokenization.doLowerCase(),
byteLevelBpeTokenization.withSpecialTokens(),
byteLevelBpeTokenization.maxSequenceLength(),
truncate,
byteLevelBpeTokenization.getSpan(),
byteLevelBpeTokenization.isAddPrefixSpace(),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| } else if (tokenization instanceof ByteLevelBpeTokenization byteLevelBpeTokenization) { | |
| return new ByteLevelBpeTokenization( | |
| false, | |
| byteLevelBpeTokenization.withSpecialTokens(), | |
| byteLevelBpeTokenization.maxSequenceLength(), | |
| truncate, | |
| byteLevelBpeTokenization.getSpan(), | |
| byteLevelBpeTokenization.isAddPrefixSpace(), | |
| byteLevelBpeTokenization.getUnkToken(), | |
| byteLevelBpeTokenization.getPadToken(), | |
| byteLevelBpeTokenization.getBosToken(), | |
| byteLevelBpeTokenization.getEosToken(), | |
| byteLevelBpeTokenization.getMaskToken() | |
| ); | |
| } else if (tokenization instanceof ByteLevelBpeTokenization byteLevelBpeTokenization) { | |
| return new ByteLevelBpeTokenization( | |
| byteLevelBpeTokenization.doLowerCase(), | |
| byteLevelBpeTokenization.withSpecialTokens(), | |
| byteLevelBpeTokenization.maxSequenceLength(), | |
| truncate, | |
| byteLevelBpeTokenization.getSpan(), | |
| byteLevelBpeTokenization.isAddPrefixSpace(), | |
| byteLevelBpeTokenization.getUnkToken(), | |
| byteLevelBpeTokenization.getPadToken(), | |
| byteLevelBpeTokenization.getBosToken(), | |
| byteLevelBpeTokenization.getEosToken(), | |
| byteLevelBpeTokenization.getMaskToken() | |
| ); |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/InferenceConfigTestScaffolding.java`
around lines 47 - 60, The clone of ByteLevelBpeTokenization incorrectly
hardcodes the first constructor argument to false, losing the original
doLowerCase() value; update the cloning block in InferenceConfigTestScaffolding
so that the first argument uses byteLevelBpeTokenization.doLowerCase()
(preserving null/true/false) instead of false while keeping all other fields the
same (withSpecialTokens(), maxSequenceLength(), truncate, getSpan(),
isAddPrefixSpace(), getUnkToken(), getPadToken(), getBosToken(), getEosToken(),
getMaskToken()) so only truncation is changed.
| /** | ||
| * {@link #createRandom()} may use {@link ByteLevelBpeTokenization}, which requires | ||
| * {@link ByteLevelBpeTokenization#ML_BYTE_LEVEL_BPE_TOKENIZATION_ADDED}; the base class filters BWC versions using only one | ||
| * {@link #createTestInstance()} sample, so we must assume byte-level BPE may appear in every serialization run. | ||
| */ | ||
| @Override | ||
| protected List<TransportVersion> bwcVersions() { | ||
| return getAllBWCVersions().stream() | ||
| .filter(v -> v.supports(ByteLevelBpeTokenization.ML_BYTE_LEVEL_BPE_TOKENIZATION_ADDED)) | ||
| .collect(Collectors.toList()); | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win
Keep full BWC coverage for non-byte-level-BPE configs.
This override drops every pre-ML_BYTE_LEVEL_BPE_TOKENIZATION_ADDED version for the entire TextEmbeddingConfigTests class. Once createRandom() started emitting ByteLevelBpeTokenization, the older serialization matrix for existing tokenizations stopped being exercised here too. I’d keep the full BWC version list for the general TextEmbeddingConfig path and isolate byte-level BPE into a dedicated version-gated case instead of narrowing coverage for all variants.
Also applies to: 93-103
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/trainedmodel/TextEmbeddingConfigTests.java`
around lines 25 - 35, The current override of bwcVersions() filters out
pre-ML_BYTE_LEVEL_BPE_TOKENIZATION_ADDED versions which removes full BWC
coverage for non-byte-level-BPE TextEmbeddingConfig variants; revert
bwcVersions() to return getAllBWCVersions() (i.e., remove the .filter(...) call)
so general TextEmbeddingConfigTests keep full serialization matrix, and instead
add a dedicated test or branch that exercises createRandom() instances
containing ByteLevelBpeTokenization guarded by checking
ByteLevelBpeTokenization.ML_BYTE_LEVEL_BPE_TOKENIZATION_ADDED (e.g., a new test
method that filters versions and explicitly serializes/deserializes instances
produced by createRandom() that include ByteLevelBpeTokenization) so only that
case is version-gated while others remain unfiltered.
valeriy42
left a comment
valeriy42
left a comment
There was a problem hiding this comment.
LGTM once you resolve the BWC issues flagged by the code rabbit.
Clarify Javadoc for BOS/EOS validation timing; preserve doLowerCase in InferenceConfigTestScaffolding clone; restore full TextEmbeddingConfig BWC matrix by moving byte-level BPE serialization to TextEmbeddingConfigByteLevelBpeBwcTests.
Regenerate ml_byte_level_bpe_tokenization after merge with main so it no longer shares id 9373000 with voyage_ai_embedding_type_non_nullable (IllegalStateException in TransportVersion holder).
edsavage
changed the title
Adds ByteLevelBpeTokenization, ByteLevelBpeTokenizer, XContent and named-writeable registration, NlpTokenizer routing, and tests. Special token strings are configurable with RoBERTa-style defaults.
Adds ByteLevelBpeTokenization, ByteLevelBpeTokenizer, XContent and named-writeable registration, NlpTokenizer routing, and tests. Special token strings are configurable with RoBERTa-style defaults.
Adds ByteLevelBpeTokenization, ByteLevelBpeTokenizer, XContent and named-writeable registration, NlpTokenizer routing, and tests. Special token strings are configurable with RoBERTa-style defaults.
Summary
ByteLevelBpeTokenization/ByteLevelBpeTokenizationUpdatewith XContent and named writeablesByteLevelBpeTokenizerwired fromNlpTokenizer.buildusing existing BPE byte-level + merge machineryMlInferenceNamedXContentProviderandNlpConfigUpdateBefore marking ready
Relates elastic/eland#818