Transformer语言模型适配可变输入及字符生成方案求助

核心问题定位 #

原模型的Embed层中，位置嵌入层position_embed硬编码使用了训练时的CONTEXT_LENGTH=200作为词汇表大小，导致输入长度超过200时会触发索引越界错误；同时输入层固定了shape=(200,)，限制了可变长度输入。

代码修改方案 #

1. 重构Embed层，支持可变长度输入 #

修改位置嵌入层的初始化逻辑，将最大上下文长度作为参数传入，同时添加长度校验避免越界：

class Embed(keras.layers.Layer):
    """词嵌入+位置嵌入 """
    def __init__(self, vocab_size, d_model, max_context_length):
        super().__init__()
        self.word_embed = keras.layers.Embedding(vocab_size, d_model)
        # 位置嵌入层使用最大上下文长度作为词汇表大小
        self.position_embed = keras.layers.Embedding(max_context_length, d_model)
        self.max_context_length = max_context_length

    def call(self, inputs):
        B, T = inputs.shape
        # 校验输入长度不超过最大上下文长度
        if T > self.max_context_length:
            raise ValueError(f"输入长度{T}超过模型支持的最大上下文长度{self.max_context_length}")
        tok_embed = self.word_embed(inputs)  # (B, T, d_model)
        pos_embed = self.position_embed(tf.range(T))  # (T, d_model)
        return tok_embed + pos_embed  # (B, T, d_model)
    
    def get_config(self):
        base_config = super().get_config()
        return {**base_config,
                "vocab_size": self.word_embed.input_dim,
                "d_model": self.word_embed.output_dim,
                "max_context_length": self.max_context_length}

2. 修改模型输入层，适配可变长度 #

将输入层的shape改为(None,)，允许任意长度的输入序列（不超过最大上下文长度）：

# 假设原常量定义保持不变：VOCAB_SIZE, d_model, d_k, d_v, h, N, d_in
CONTEXT_LENGTH = 200  # 训练时的上下文长度，同时作为模型支持的最大长度

inputs = keras.Input(shape=(None,))
# 初始化Embed层时传入参数
X = Embed(vocab_size=VOCAB_SIZE, d_model=d_model, max_context_length=CONTEXT_LENGTH)(inputs)
for _ in range(N):
    Z = MultiHeadAttention(mask=True)(X)
    X = keras.layers.LayerNormalization()(Z + X)

    Z = FeedForward()(X)
    X = keras.layers.LayerNormalization()(Z + X)
outputs = keras.layers.Dense(VOCAB_SIZE, activation="softmax")(X)

model = keras.Model(inputs=inputs, outputs=outputs, name="transformer")

3. 编写字符生成函数 #

实现循环预测逻辑，每次将生成的字符追加到输入序列后，持续生成直到达到指定长度：

def generate_text(model, start_tokens, max_gen_length=100, temperature=1.0):
    """
    Args:
        model: 训练好的Transformer模型
        start_tokens: 初始输入的token序列（整数列表，已完成字符转token）
        max_gen_length: 最大生成字符数
        temperature: 温度参数，控制生成随机性（值越大越随机）
    Returns:
        生成的完整token序列（包含初始序列）
    """
    current_seq = tf.convert_to_tensor([start_tokens], dtype=tf.int32)
    max_context_length = model.layers[1].max_context_length  # 从Embed层获取最大长度限制

    for _ in range(max_gen_length):
        seq_len = current_seq.shape[1]
        # 如果当前序列长度已达上限，停止生成
        if seq_len >= max_context_length:
            break
        
        # 预测下一个字符：取模型输出的最后一个位置的概率分布
        predictions = model(current_seq)[:, -1, :]  # shape: (1, VOCAB_SIZE)
        # 温度调整概率
        predictions = predictions / temperature
        # 采样选择下一个token
        next_token_idx = tf.random.categorical(predictions, num_samples=1)[0, 0].numpy()
        # 追加到当前序列
        current_seq = tf.concat([current_seq, tf.convert_to_tensor([[next_token_idx]], dtype=tf.int32)], axis=1)
    
    return current_seq[0].numpy()

注意事项 #

• 若需要支持超过训练时CONTEXT_LENGTH的输入，需重新训练模型，将max_context_length设为更大的值（否则位置嵌入层没有对应长度的权重，无法正常工作）。
• 生成时的temperature参数可灵活调整：值接近0时生成结果更保守，值大于1时生成更具随机性。

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