在Keras模型中集成sentence-transformers模型时的TensorFlow符号张量输入错误问题求助

Let's break down what's going wrong and walk through the fixes step by step:

The Core Problem #

Your error stems from two key issues:

• When using tf.py_function, the input passed to your tokenizer is a Tensor object, not a native Python list of strings. Hugging Face tokenizers can't process Tensor inputs directly—they expect regular strings or lists of strings.
• Your mean_pooling method is missing the self parameter (required for class methods in Python), and your tf_encode function returns a list instead of the dictionary of tensors the model expects.

The Corrected Code #

Here's the fully working version of your layer, with explanations embedded:

import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModel

MODEL_PATH = 'sentence-transformers/all-MiniLM-L6-v2'
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = TFAutoModel.from_pretrained(MODEL_PATH, 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):
            # Convert Tensor to numpy array, then to a list of strings
            # This bridges the gap between TensorFlow's tensor format and the tokenizer's requirements
            inputs_list = inputs.numpy().tolist()
            # Handle byte tensors (common in some TF contexts) by decoding to UTF-8
            inputs_list = [s.decode('utf-8') if isinstance(s, bytes) else s for s in inputs_list]
            return self.tokenizer(
                inputs_list, padding=True, truncation=True, return_tensors='tf'
            )
        
        # Return the tokenizer's output as a dictionary of tensors (matches what the model expects)
        encoded_result = tf.py_function(
            func=encode, 
            inp=[inputs], 
            Tout={
                'input_ids': tf.int64,
                'attention_mask': tf.int64
            }
        )
        # Set static shapes so TensorFlow can build a valid computation graph
        encoded_result['input_ids'].set_shape([None, None])
        encoded_result['attention_mask'].set_shape([None, None])
        return encoded_result

    def mean_pooling(self, model_output, attention_mask):
        # Added self parameter to make this a valid class instance method
        token_embeddings = model_output[0]
        input_mask_expanded = tf.cast(
            tf.broadcast_to(tf.expand_dims(attention_mask, -1), token_embeddings.shape),
            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):
        encoded_input = self.tf_encode(inputs)
        model_output = self.model(encoded_input)
        embeddings = self.mean_pooling(model_output, encoded_input['attention_mask'])
        return embeddings

# Test the layer
sbert = SBert(tokenizer, model)
test_output = sbert(tf.constant(['some text', 'more text']))
print(test_output.shape)  # Should output (2, 384) for all-MiniLM-L6-v2

Key Fixes Explained #

1. Tensor to String Conversion:
   In the encode function, we convert the input Tensor to a numpy array (.numpy()) then to a list of strings. We also handle byte tensors by decoding them to UTF-8, which ensures the tokenizer gets valid text input.

2. Returning a Dictionary from tf.py_function:
   Instead of returning a list, we define Tout as a dictionary matching the tokenizer's output keys (input_ids and attention_mask). This matches the input format the pretrained model expects.

3. Setting Static Shapes:
   After getting the tokenized result, we set static shapes for the tensors. This helps TensorFlow build a stable computation graph, which is critical for integrating this layer into a larger Keras model.

4. Fixing the mean_pooling Method:
   Added the self parameter to make it a valid instance method of the SBert class—this was causing a silent error that would break graph execution.

Using the Layer in a Larger Keras Model #

Once this layer works, you can use it like any other Keras layer. For example, building a text classification model:

# Build a complete classification model
input_layer = tf.keras.layers.Input(shape=(), dtype=tf.string)
sbert_embeddings = SBert(tokenizer, model)(input_layer)
dense_layer = tf.keras.layers.Dense(64, activation='relu')(sbert_embeddings)
output_layer = tf.keras.layers.Dense(2, activation='softmax')(dense_layer)

full_model = tf.keras.Model(inputs=input_layer, outputs=output_layer)
full_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

This will work seamlessly in both eager execution and graph mode.

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