Updating training metrics

run_summarization_flax.py

@@ -589,8 +589,10 @@ def main():
             desc="Running tokenizer on prediction dataset",
         )
 
-    # Metric
-    metric = load_metric("rouge")
+    # Metric    
+    rouge_metric = load_metric("rouge")
+    bleu_metric = load_metric("bleu")
+    meteor_metric = load_metric("meteor")
 
     def postprocess_text(preds, labels):
         preds = [pred.strip() for pred in preds]
@@ -609,14 +611,29 @@ def main():
         # Some simple post-processing
         decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)
 
-        result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
+        results = {}
+        rouge_scores = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer = True, \
+                                        rouge_types=['rougeL'])
         # Extract a few results from ROUGE
-        result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
-
-        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
-        result["gen_len"] = np.mean(prediction_lens)
-        result = {k: round(v, 4) for k, v in result.items()}
-        return result
+        rouge_scores = {key: value.mid.fmeasure * 100 for key, value in rouge_scores.items()}
+        rouge_scores = {k: round(v, 4) for k, v in rouge_scores.items()}
+        meteor_scores = meteor_metric.compute(predictions=decoded_preds, references=decoded_labels)
+        meteor_scores = {k: round(v, 4) for k, v in meteor_scores.items()}
+
+        # Compute bleu-1,2,3,4 scores
+        # Postprocess the predictions and references to compute bleu scores
+        tokenized_predictions = [decoded_preds[i].split() for i in range(len(decoded_preds))]
+        tokenized_labels = [[decoded_labels[i].split()] for i in range(len(decoded_labels))]
+        bleu_scores = {f'bleu-{i}' : \
+                bleu_metric.compute(predictions=tokenized_predictions, references=tokenized_labels, max_order=i)['bleu']\
+                for i in range(1,5)}
+        bleu_scores = {k: round(v, 4) for k, v in bleu_scores.items()}
+
+        results.update(bleu_scores)
+        results.update(rouge_scores)
+        results.update(meteor_scores)
+
+        return results
 
     # Enable tensorboard only on the master node
     has_tensorboard = is_tensorboard_available()