TensorFlow TextVectorization层自定义standardize函数适配求助

一、正确的输入输出格式 #

• 输入：批量字符串张量（TensorFlow的tf.Tensor，形状为(batch_size,)，每个元素对应单条文本字符串）
• 输出：同形状的字符串张量，每个元素为处理后的单条文本字符串

注意：函数必须兼容TensorFlow张量操作，不能直接处理原生Python字符串（除非用tf.py_function做包装适配）

二、现有函数的改造方案 #

你的原函数基于Python字符串处理，需要调整适配TensorFlow张量输入，提供两种可行方案：

方案1：用tf.py_function快速包装（适合小批量场景） #

直接将现有Python函数包装为TensorFlow可调用形式，处理张量输入输出：

import tensorflow as tf
import nltk
from nltk.stem.wordnet import WordNetLemmatizer
import re
nltk.download('wordnet')

# 修正原函数的语法错误（添加冒号）
def my_preprocessing(text):
    text = re.split("\s+", text)
    wnl = WordNetLemmatizer()
    stemmed_words = [wnl.lemmatize(word) for word in text]
    testo = [w for w in stemmed_words if not w.isnumeric()]
    return ' '.join(testo)

# 包装成适配TextVectorization的函数
def tf_standardize_fn(input_tensor):
    # 将张量转为Python字符串处理，再转回张量
    processed_text = tf.py_function(
        func=lambda x: my_preprocessing(x.numpy().decode('utf-8')),
        inp=[input_tensor],
        Tout=tf.string
    )
    # 确保输出形状与输入一致
    processed_text.set_shape(input_tensor.shape)
    return processed_text

使用方式：

vectorizer = tf.keras.layers.TextVectorization(
    standardize=tf_standardize_fn,
    # 补充其他参数：max_tokens、output_mode等
)

方案2：改用TensorFlow原生操作重写（推荐，性能更高） #

用TF原生API替代Python逻辑，避免跨语言开销，同时支持模型序列化：

import tensorflow as tf

def tf_standardize_fn(input_tensor):
    # 1. 按空格分割字符串
    words = tf.strings.split(input_tensor, sep=' ')
    # 2. 简易词形还原（模拟NLTK效果，可扩展更复杂规则）
    def lemmatize_word(word):
        word = tf.strings.regex_replace(word, r's$', '')
        word = tf.strings.regex_replace(word, r'es$', '')
        return word
    stemmed_words = tf.map_fn(lemmatize_word, words, dtype=tf.string)
    # 3. 过滤数字文本
    non_numeric_mask = tf.logical_not(tf.strings.regex_match(stemmed_words, r'^\d+$'))
    filtered_words = tf.boolean_mask(stemmed_words, non_numeric_mask)
    # 4. 拼接回完整字符串
    return tf.strings.join(filtered_words, separator=' ')

三、关键注意事项 #

• 原函数存在语法错误：def my_preprocessing(text) 末尾缺少冒号，必须修正后才能运行
• 方案1中的tf.py_function会导致模型无法序列化（无法保存/加载），若需部署模型优先选方案2
• 函数会自动对张量中的每个元素（单条文本）执行处理逻辑，无需手动循环批量数据

内容的提问来源于stack exchange，提问作者Luca Toffanetti