Internal change

PiperOrigin-RevId: 307918408

official/nlp/transformer/misc.py

@@ -261,18 +261,15 @@ def get_callbacks(steps_per_epoch):
   return callbacks
 
 
-def build_stats(history, callbacks):
-  """Normalizes and returns dictionary of stats.
+def update_stats(history, stats, callbacks):
+  """Normalizes and updates dictionary of stats.
 
   Args:
     history: Results of the training step.
+    stats: Dict with pre-existing training stats.
     callbacks: a list of callbacks which might include a time history callback
       used during keras.fit.
-
-  Returns:
-    Dictionary of normalized results.
   """
-  stats = {}
 
   if history and history.history:
     train_hist = history.history
@@ -280,7 +277,7 @@ def build_stats(history, callbacks):
     stats['loss'] = float(train_hist['loss'][-1])
 
   if not callbacks:
-    return stats
+    return
 
   # Look for the time history callback which was used during keras.fit
   for callback in callbacks:
@@ -293,4 +290,3 @@ def build_stats(history, callbacks):
             callback.batch_size * callback.log_steps *
             (len(callback.timestamp_log)-1) /
             (timestamp_log[-1].timestamp - timestamp_log[0].timestamp))
-  return stats

official/nlp/transformer/transformer_main.py

@@ -363,7 +363,8 @@ class TransformerTask(object):
 
     stats = ({
         "loss": train_loss
-    } if history is None else misc.build_stats(history, callbacks))
+    } if history is None else {})
+    misc.update_stats(history, stats, callbacks)
     if uncased_score and cased_score:
       stats["bleu_uncased"] = uncased_score
       stats["bleu_cased"] = cased_score