Fix prediction metric

src/run_ed_recipe_nlg.py

@@ -832,14 +832,14 @@ def main():
         pred_metrics = jax.tree_map(jnp.mean, pred_metrics)
 
         # compute ROUGE metrics
-        rouge_desc = ""
+        mix_desc = ""
         if data_args.predict_with_generate:
-            rouge_metrics = compute_metrics(pred_generations, pred_labels)
-            pred_metrics.update(rouge_metrics)
-            rouge_desc = " ".join([f"Predict {key}: {value} |" for key, value in rouge_metrics.items()])
+            mix_metrics = compute_metrics(pred_generations, pred_labels)
+            pred_metrics.update(mix_metrics)
+            mix_desc = " ".join([f"Predict {key}: {value} |" for key, value in mix_metrics.items()])
 
         # Print metrics
-        desc = f"Predict Loss: {pred_metrics['loss']} | {rouge_desc})"
+        desc = f"Predict Loss: {pred_metrics['loss']} | {mix_desc})"
         logger.info(desc)
 
         # save checkpoint after each epoch and push checkpoint to the hub