def train_step(self, data):
X, y = data
training = tf.constant(True)
source_token_ids, source2target_ids = X
target_token_ids, target2source_ids = y
target_vocab_size = self.vocab.get_vocab_size("target")
unk_index = self.vocab.get_token_id(self.vocab._unk_token, "source")
start_index = self.vocab.get_token_id(self.vocab._start_token,
"source")
end_index = self.vocab.get_token_id(self.vocab._end_token, "source")
y_true = prep_y_true(source_token_ids, target_token_ids,
target2source_ids, target_vocab_size, unk_index,
start_index, end_index)
with tf.GradientTape() as tape:
output_dict = self(X, y, training)
# shape: ()
loss = self.compiled_loss(y_true, output_dict['ypred'])
gradients = tape.gradient(loss, self.trainable_variables)
self.optimizer.apply_gradients(
zip(gradients, self.trainable_variables))
self.compiled_metrics.update_state(y_true, output_dict['ypred'])
return {m.name: m.result() for m in self.metrics}
def test_step(self, data):
X, y = data
training = tf.constant(True)
source_token_ids, source2target_ids = X
target_token_ids, target2source_ids = y
target_vocab_size = self.vocab.get_vocab_size("target")
unk_index = self.vocab.get_token_id(self.vocab._unk_token, "source")
start_index = self.vocab.get_token_id(self.vocab._start_token,
"source")
end_index = self.vocab.get_token_id(self.vocab._end_token, "source")
y_true = prep_y_true(source_token_ids, target_token_ids,
target2source_ids, target_vocab_size, unk_index,
start_index, end_index)
output_dict = self(X, y, training)
self.compiled_loss(y_true, output_dict['ypred'])
self.compiled_metrics.update_state(y_true, output_dict['ypred'])
return {m.name: m.result() for m in self.metrics}
def train_step(self, data):
X, y = data
training = tf.constant(True)
cis, cit, ans, ner, pos, qpre = X
qit, qis = y
target_vocab_size = self.vocab.get_vocab_size("target")
unk_index = self.vocab.get_token_id(self.vocab._unk_token, "source")
start_index = self.vocab.get_token_id(self.vocab._start_token,
"source")
end_index = self.vocab.get_token_id(self.vocab._end_token, "source")
y_true = prep_y_true(cis, qit, qis, target_vocab_size, unk_index,
start_index, end_index)
with tf.GradientTape() as tape:
output_dict = self(X, y, training)
# shape: ()
loss = self.compiled_loss(y_true, output_dict['ypred'])
if (self._train_counter % 10 == 0
and (logging.get_absl_logger().getEffectiveLevel()
== logging.converter.STANDARD_DEBUG)):
samples = 3
tf.py_function(self.debug, [
cis[:samples], qit[:samples], ans[:samples],
qpre[:samples], output_dict["attentive_weights"][:samples],
output_dict["selective_weights"][:samples]
], [],
name="Debug")
gradients = tape.gradient(loss, self.trainable_variables)
self.optimizer.apply_gradients(
zip(gradients, self.trainable_variables))
self.compiled_metrics.update_state(y_true, output_dict['ypred'])
self._train_counter.assign_add(1)
return {m.name: m.result() for m in self.metrics}