Limits of custom split functions in tf.keras TextVectorization

Direct answer

A custom split callable in tf.keras.layers.TextVectorization runs after standardize, so any normalization that removes whitespace or punctuation changes the token boundaries the split function sees. During adapt(), the layer counts token frequencies across the dataset, sorts by descending frequency, and truncates to max_tokens. A custom split that emits many short tokens can therefore push the vocabulary past the intended limit, and rare short tokens may be dropped instead.

Confirmed behavior

• split receives the standardized string, not the raw input.
• adapt() builds the vocabulary from frequency-ranked tokens and applies max_tokens as a hard ceiling.
• get_vocabulary() returns the post-adaptation token-to-index mapping as a list in index order.
• Token order across dataset elements follows input order, but internal reduction operations are not guaranteed to be deterministic across TensorFlow minor releases.

Likely explanation for TF 2.13 vs 2.14 differences

If the same split function produces different vocabularies under TensorFlow 2.13 and 2.14, the cause is most likely a change in internal segment reduction, hashing, or sorting behavior rather than the split callable itself. This is a version-sensitivity risk, not a documented contract. Locking the TensorFlow version is the only practical way to make the vocabulary reproducible across environments.

Steps to bound vocabulary size

1. Set standardize=None if your split logic depends on whitespace that standardization would remove.
2. Filter short or noise tokens inside the custom split before returning the token list.
3. Choose max_tokens to cover the expected head of the frequency distribution, accepting that long-tail tokens will be mapped to out-of-vocabulary indices.
4. Call adapt() on a fixed, unshuffled dataset, then inspect the result with get_vocabulary().

Verification

import tensorflow as tf
from tensorflow.keras.layers import TextVectorization

def my_split(text):
    return tf.strings.split(text)

layer = TextVectorization(
    standardize=None,
    split=my_split,
    max_tokens=1000,
    output_sequence_length=32
)
layer.adapt(["hello world", "hello tensorflow", "world tensorflow"])
print(layer.get_vocabulary()[:20])
print("vocab size:", len(layer.get_vocabulary()))

Run this under each TensorFlow version you support and compare the printed vocabulary. If the lists differ, pin the TensorFlow version in your requirements file.

Missing diagnostic detail

Does your dataset pass through shuffle() before adapt()? Shuffling changes token arrival order and can change which tokens survive frequency truncation when counts are tied.