基于ELMo嵌入与Bi-LSTM的句子级Token分类模型报错排查

问题背景 #

正在开展Kaggle平台的**术语检测（Jargon Detection）**项目，已完成数据清洗与序列填充操作。使用环境为Python 3.9.x、TensorFlow 2.12.x（含tensorflow-hub等依赖库）。构建了基于ELMo嵌入与Bi-LSTM的句子级Token分类模型，但运行训练代码时触发TypeError。

错误详情 #

运行模型时触发以下错误：

TypeError: Exception encountered when calling layer "elmo_embeddings" (type ElmoEmbeddings).

in user code:

    File "/var/folders/yq/qh4ny6ms3kz0yqqj4q2gqkn00000gn/T/ipykernel_4686/880792685.py", line 11, in call  *
        input_tensor = tf.constant(sequence)

    TypeError: Expected any non-tensor type, but got a tensor instead.


Call arguments received by layer "elmo_embeddings" (type ElmoEmbeddings):
  • inputs=tf.Tensor(shape=(None, 25), dtype=string)

输入参数：

• X_tr为包含18916条数据的列表，单条数据长度为25；
• y_tr为包含18916条数据的numpy数组，单条数据长度为25。

错误原因 #

1. 在Keras模型的ElmoEmbeddings层中，inputs是批量Tensor（shape=(None,25)），循环遍历inputs时，每个sequence已经是Tensor类型，但tf.constant()要求传入非Tensor类型（如Python列表、numpy数组），因此触发报错。
2. 独立测试时使用的是单条非批量数据（非Tensor形式），因此逻辑可行，但在模型上下文的批量Tensor输入下不兼容。
3. 手动循环Tensor元素的方式不符合TensorFlow的计算图机制，无法进行自动求导和优化，属于非向量化操作，性能低下且易出错。

解决方案 #

修改ElmoEmbeddings层的call方法，直接将批量Tensor输入传给ELMo模型——ELMo的default签名原生支持批量字符串输入，无需手动循环处理每个样本。

修改后的完整代码 #

from tensorflow.keras.layers import Input, LSTM, Dense, TimeDistributed, Dropout, Bidirectional, Layer
from tensorflow.keras.models import Model
import tensorflow as tf
import tensorflow_hub as hub

elmo_model = hub.load("https://tfhub.dev/google/elmo/3").signatures["default"]

class ElmoEmbeddings(Layer):
    def __init__(self, elmo_model, max_len, **kwargs):
        super(ElmoEmbeddings, self).__init__(**kwargs)
        self.elmo_model = elmo_model
        self.max_len = max_len
    
    def call(self, inputs):
        # 直接将批量输入传给ELMo模型，无需循环
        embeddings = self.elmo_model(inputs)['elmo']
        # ELMo返回的embeddings shape为(batch_size, max_len, 1024)，刚好符合后续LSTM输入要求
        return embeddings
    
    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.max_len, 1024)

# 定义输入层
input_layer = Input(shape=(max_len,), dtype=tf.string, name='input_layer')
elmo_embeddings = ElmoEmbeddings(elmo_model, max_len, name='elmo_embeddings')(input_layer)
lstm_output = Bidirectional(LSTM(units=512, return_sequences=True, recurrent_dropout=0.2, dropout=0.2, name='lstm'))(elmo_embeddings)
output = TimeDistributed(Dense(n_tags, activation='softmax', name='dense'))(lstm_output)

# 定义模型
tech_jargon_model = Model(inputs=input_layer, outputs=output)
tech_jargon_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 训练模型
history = tech_jargon_model.fit(
    tf.constant(X_tr), y_tr,
    validation_data=(tf.constant(X_val), y_val),
    batch_size=batch_size,
    epochs=3,
    verbose=1
)

关键修改说明 #

• 移除了call方法中的循环逻辑，直接将整个批量Tensorinputs传入ELMo模型；
• ELMo的default签名接受shape=(batch_size, sequence_length)的字符串Tensor，返回的elmo嵌入张量shape为(batch_size, sequence_length, 1024)，完全匹配后续Bi-LSTM的输入要求；
• 保留了compute_output_shape方法以明确输出形状，确保模型构建时的形状推断正常。

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