对话系统中的上下文生成问题

import tensorflow as tf

# 定义对话系统模型
class DialogModel(tf.keras.Model):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super(DialogModel, self).__init__()
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        self.gru = tf.keras.layers.GRU(hidden_dim, return_sequences=True, return_state=True)
        self.dense = tf.keras.layers.Dense(vocab_size)

    def call(self, inputs, hidden):
        embedded = self.embedding(inputs)
        output, state = self.gru(embedded, initial_state=hidden)
        logits = self.dense(output)
        return logits, state

# 定义损失函数
def loss_function(real, pred):
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
    loss_ = loss_object(real, pred)
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    mask = tf.cast(mask, dtype=loss_.dtype)
    loss_ *= mask

    return tf.reduce_mean(loss_)

# 定义训练过程
@tf.function
def train_step(inputs, targets, model, optimizer, hidden):
    with tf.GradientTape() as tape:
        predictions, hidden = model(inputs, hidden)
        loss = loss_function(targets, predictions)

    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    return loss, hidden

# 初始化模型和优化器
vocab_size = 10000
embedding_dim = 256
hidden_dim = 512
model = DialogModel(vocab_size, embedding_dim, hidden_dim)
optimizer = tf.keras.optimizers.Adam()

# 进行训练
EPOCHS = 10
for epoch in range(EPOCHS):
    hidden = model.reset_states()

    for inputs, targets in dataset:
        loss, hidden = train_step(inputs, targets, model, optimizer, hidden)

    print('Epoch {} Loss {:.4f}'.format(epoch + 1, loss.numpy()))