在AWS SageMaker部署集成sentence-transformers的TensorFlow推荐模型遇阻

背景 #

我需要在AWS SageMaker上部署一个TensorFlow模型，核心逻辑是：接收字符串输入，用sentence-transformers预训练模型计算文本嵌入，再通过TensorFlow Recommenders（TFR）从预计算好的嵌入集合中返回K近邻推荐结果。要求将分词等预处理步骤全部封装在模型内部。

在本地笔记本中能实现预测功能，但模型保存和SageMaker部署环节遇到了问题：Hugging Face的Tokenizer需要纯字符串列表输入，导致保存失败；用tf.py_function绕开后，模型能保存并部署，但SageMaker调用时出现OpKernel错误。

尝试的方法 #

1. 直接调用SentenceTransformer的方法 #

def text_to_startup_model(startups_ids: list, startup_vectors):
    import tensorflow as tf
    import tensorflow_recommenders as tfrs
    import numpy as np
    
    exported_model = tfrs.layers.factorized_top_k.BruteForce(SentenceTransformer("all-mpnet-base-v2").encode)
    exported_model.index(np.array(startup_vectors), np.array(startups_ids))
    
    # 测试模型
    test = exported_model(['Test Text Query'])
    print(test)
    
    return exported_model

# 测试预测正常
text_to_startup_model(search_db_ids, search_db_embeddings)

# 保存模型时报错
tf.saved_model.save(text_to_startup_model(search_db_ids, search_db_embeddings), export_dir="/home/nicholas/test_model_save/1")

保存时触发TypeError，关键错误信息：

TypeError: 'NoneType' object is not subscriptable
...
Call arguments received:
  • queries=['None']
  • k=None

错误溯源到SentenceTransformer的tokenize方法，处理输入时遇到None值无法下标访问。

2. 使用tf.py_function的方法 #

为了让模型适配TensorFlow的输入类型，用tf.py_function封装Tokenizer逻辑：

def approach_2(startups_ids: list, startup_vectors):
    import tensorflow as tf
    import tensorflow_recommenders as tfrs
    import numpy as np
    from transformers import MPNetTokenizer, TFMPNetModel

    tokenizer = MPNetTokenizer.from_pretrained(
        "sentence-transformers/all-mpnet-base-v2"
    )
    model = TFMPNetModel.from_pretrained(
        "sentence-transformers/all-mpnet-base-v2", from_pt=True
    )

    class SBert(tf.keras.layers.Layer):
        def __init__(self, tokenizer, model):
            super(SBert, self).__init__()
            self.tokenizer = tokenizer
            self.model = model

        def tf_encode(self, inputs):
            def encode(inputs):
                inputs = [x[0].decode("utf-8") for x in inputs.numpy()]
                outputs = self.tokenizer(
                    inputs, padding=True, truncation=True, return_tensors="tf"
                )
                return outputs["input_ids"], outputs["attention_mask"]

            return tf.py_function(
                func=encode, inp=[inputs], Tout=[tf.int32, tf.int32]
            )

        def process(self, i, a):
            def __call(i, a):
                model_output = self.model(
                    {"input_ids": i.numpy(), "attention_mask": a.numpy()}
                )
                return model_output[0]

            return tf.py_function(func=__call, inp=[i, a], Tout=[tf.float32])

        def mean_pooling(self, model_output, attention_mask):
            token_embeddings = tf.squeeze(tf.stack(model_output), axis=0)
            input_mask_expanded = tf.cast(
                tf.broadcast_to(
                    tf.expand_dims(attention_mask, -1), tf.shape(token_embeddings)
                ),
                tf.float32,
            )
            a = tf.math.reduce_sum(token_embeddings * input_mask_expanded, axis=1)
            b = tf.clip_by_value(
                tf.math.reduce_sum(input_mask_expanded, axis=1),
                1e-9,
                tf.float32.max,
            )
            embeddings = a / b
            embeddings, _ = tf.linalg.normalize(embeddings, 2, axis=1)
            return embeddings

        def call(self, inputs):
            input_ids, attention_mask = self.tf_encode(inputs)
            model_output = self.process(input_ids, attention_mask)
            embeddings = self.mean_pooling(model_output, attention_mask)
            return embeddings

    sbert = SBert(tokenizer, model)
    inputs = tf.keras.layers.Input((1,), dtype=tf.string)
    outputs = sbert(inputs)
    model = tf.keras.Model(inputs, outputs)

    exported_model = tfrs.layers.factorized_top_k.BruteForce(model)
    exported_model.index(np.array(startup_vectors), np.array(startups_ids))

    # 测试模型
    print(exported_model(tf.constant(["'Test Text Query'"])))
    return exported_model

model_to_store_1 = approach_2(search_db_ids, search_db_embeddings)
# 模型保存成功
tf.saved_model.save(model_to_store_1, export_dir="/home/nicholas/test_model_save/2")

模型成功部署到SageMaker，但调用时触发ModelError：

ModelError: An error occurred (ModelError) when calling the InvokeEndpoint operation: Received client error (400) from model with message "{
    "error": "No OpKernel was registered to support Op 'EagerPyFunc' used by {{node StatefulPartitionedCall/brute_force/model/s_bert/EagerPyFunc}} with these attrs: [is_async=false, Tin=[DT_STRING], _output_shapes=[<unknown>, <unknown>], Tout=[DT_INT32, DT_INT32], token=\"pyfunc_4\"]\nRegistered devices: [CPU]\nRegistered kernels:\n  <no registered kernels>\n\n\t [[StatefulPartitionedCall/brute_force/model/s_bert/EagerPyFunc]]\n\t [[StatefulPartitionedCall]]"
}"

3. 类实现方法 #

参考TensorFlow版SentenceTransformer的实现逻辑，封装自定义Layer，但保存模型时遇到和方法1相同的问题：Tokenizer的输入类型不兼容，导致保存失败。

问题 #

1. 将预处理（分词）包含在模型内部是否为最佳实践？还是应该直接传入分词后的句子张量？
2. 如何解决当前的模型保存失败、以及SageMaker调用时的OpKernel错误问题？

内容的提问来源于stack exchange，提问作者Nicholas Sartor