fix bleu score calculation (#1426)

Signed-off-by: naymaraq <dkaramyan@nvidia.com>
Co-authored-by: naymaraq <dkaramyan@nvidia.com>

Author: naymaraq

nemo_skills/dataset/covost2/__init__.py

@@ -16,7 +16,3 @@
 METRICS_TYPE = "audio"
 EVAL_ARGS = "++eval_type=audio ++eval_config.normalization_mode=multilingual"
 GENERATION_ARGS = "++prompt_format=openai ++enable_audio=true"
-JUDGE_PIPELINE_ARGS = {
-    "source_key": "extra_fields.src_text",
-    "reference_key": "extra_fields.tgt_text",
-}

nemo_skills/dataset/covost2/prepare.py

@@ -173,9 +173,11 @@ def _build_record(
     subset_for_metrics: str,
     task_type: str,
     extra_fields: dict,
+    source: str | None = None,
+    reference: str | None = None,
 ) -> dict:
     audio_metadata = {"path": container_audio_path, "duration": duration}
-    return {
+    record = {
         "expected_answer": expected_answer,
         "audio_path": container_audio_path,
         "duration": duration,
@@ -187,6 +189,11 @@ def _build_record(
         "task_type": f"Multilingual-{task_type.upper()}",
         "extra_fields": extra_fields,
     }
+    if source is not None:
+        record["source"] = source
+    if reference is not None:
+        record["reference"] = reference
+    return record
 
 
 def prepare_covost2(
@@ -269,6 +276,8 @@ def prepare_covost2(
                             "src_lang": src_lang,
                             "tgt_lang": tgt_lang,
                         },
+                        source=item["sentence"],
+                        reference=item["translation"],
                     )
                     out.write(json.dumps(record, ensure_ascii=False) + "\n")
     print(f"CoVoST2 {task_type} dataset prepared: {output_jsonl}")

nemo_skills/dataset/fleurs/__init__.py

@@ -16,7 +16,3 @@
 METRICS_TYPE = "audio"
 EVAL_ARGS = "++eval_type=audio ++eval_config.normalization_mode=multilingual"
 GENERATION_ARGS = "++prompt_format=openai ++enable_audio=true"
-JUDGE_PIPELINE_ARGS = {
-    "source_key": "extra_fields.src_raw_text",
-    "reference_key": "extra_fields.tgt_raw_text",
-}

nemo_skills/dataset/fleurs/prepare.py

@@ -151,9 +151,11 @@ def _build_record(
     subset_for_metrics: str,
     task_type: str,
     extra_fields: dict,
+    source: str | None = None,
+    reference: str | None = None,
 ) -> dict:
     audio_metadata = {"path": container_audio_path, "duration": duration}
-    return {
+    record = {
         "expected_answer": expected_answer,
         "audio_path": container_audio_path,
         "duration": duration,
@@ -165,6 +167,11 @@ def _build_record(
         "task_type": f"Multilingual-{task_type.upper()}",
         "extra_fields": extra_fields,
     }
+    if source is not None:
+        record["source"] = source
+    if reference is not None:
+        record["reference"] = reference
+    return record
 
 
 def prepare_fleurs(data_dir: Path, split: str, languages: list[str], no_audio: bool, task_type: str) -> None:
@@ -244,8 +251,12 @@ def prepare_fleurs(data_dir: Path, split: str, languages: list[str], no_audio: b
                             "tgt_lang_group": FLEURS_LANG_TO_GROUP[tgt_locale],
                         }
                     )
+                    source_text = source_row["raw_transcription"]
+                    reference_text = target_row["raw_transcription"]
                 else:
                     expected_answer = source_row[gt_key]
+                    source_text = None
+                    reference_text = None
 
                 record = _build_record(
                     expected_answer=expected_answer,
@@ -255,6 +266,8 @@ def prepare_fleurs(data_dir: Path, split: str, languages: list[str], no_audio: b
                     subset_for_metrics=subset_for_metrics,
                     task_type=task_type,
                     extra_fields=extra_fields,
+                    source=source_text,
+                    reference=reference_text,
                 )
                 out.write(json.dumps(record, ensure_ascii=False) + "\n")
     print(f"Fleurs {task_type} dataset prepared: {output_jsonl}")

nemo_skills/dataset/flores200/__init__.py

@@ -16,7 +16,3 @@
 METRICS_TYPE = "translation"
 GENERATION_ARGS = "++prompt_config=multilingual/segment-translation"
 EVAL_SPLIT = "devtest"
-JUDGE_PIPELINE_ARGS = {
-    "source_key": "text",
-    "reference_key": "translation",
-}

nemo_skills/dataset/flores200/prepare.py

@@ -25,10 +25,10 @@ def write_data_to_file(output_file, datasets, src_languages, tgt_languages):
         for src_lang in src_languages:
             for tgt_lang in tgt_languages:
                 if src_lang != tgt_lang:
-                    for src, tgt in zip(datasets[src_lang], datasets[tgt_lang], strict=True):
+                    for src, ref in zip(datasets[src_lang], datasets[tgt_lang], strict=True):
                         json_dict = {
-                            "text": src,
-                            "translation": tgt,
+                            "source": src,
+                            "reference": ref,
                             "source_language": src_lang,
                             "target_language": tgt_lang,
                             "source_lang_name": Language(src_lang).display_name(),

nemo_skills/dataset/wmt24pp/__init__.py

@@ -15,7 +15,3 @@
 
 METRICS_TYPE = "translation"
 GENERATION_ARGS = "++prompt_config=multilingual/segment-translation"
-JUDGE_PIPELINE_ARGS = {
-    "source_key": "text",
-    "reference_key": "translation",
-}

nemo_skills/dataset/wmt24pp/prepare.py

@@ -23,10 +23,10 @@
 def write_data_to_file(output_file, datasets, tgt_languages):
     with open(output_file, "wt", encoding="utf-8") as fout:
         for tgt_lang in tgt_languages:
-            for src, tgt in zip(datasets[tgt_lang]["source"], datasets[tgt_lang]["target"], strict=True):
+            for src, ref in zip(datasets[tgt_lang]["source"], datasets[tgt_lang]["target"], strict=True):
                 json_dict = {
-                    "text": src,
-                    "translation": tgt,
+                    "source": src,
+                    "reference": ref,
                     "source_language": "en",
                     "target_language": tgt_lang,
                     "source_lang_name": "English",

nemo_skills/evaluation/evaluator/audio.py

@@ -558,28 +558,33 @@ def evaluate_asr(
     return result
 
 
+_BLEU_TOKENIZE_BY_LANG = {
+    "ja": "ja-mecab",
+    "zh": "zh",
+    "cmn": "zh",
+    "yue": "zh",
+    "ko": "ko-mecab",
+}
+
+
+def resolve_bleu_tokenize(tgt_lang: str | None) -> str:
+    """Resolve sacrebleu tokenize from a target language code."""
+    if not isinstance(tgt_lang, str):
+        return "13a"
+    lang_code = tgt_lang.split("_")[0]
+    return _BLEU_TOKENIZE_BY_LANG.get(lang_code, "13a")
+
+
 def evaluate_translation(
     reference: str,
     hypothesis: str,
     tgt_lang: str | None = None,
 ) -> dict[str, Any]:
     """Evaluate translation: computes sentence-level BLEU score."""
+    tokenize = resolve_bleu_tokenize(tgt_lang)
     try:
         import sacrebleu
 
-        tokenize = "13a"
-        if isinstance(tgt_lang, str):
-            lang_code = tgt_lang.split("_")[0]
-            if lang_code in ["cmn", "yue"]:
-                lang_code = "zh"
-
-            if lang_code == "ja":
-                tokenize = "ja-mecab"
-            elif lang_code == "zh":
-                tokenize = "zh"
-            elif lang_code == "ko":
-                tokenize = "ko-mecab"
-
         text = reference.strip()
         pred_text = hypothesis.strip()
         bleu = sacrebleu.sentence_bleu(pred_text, [text], tokenize=tokenize)
@@ -590,6 +595,7 @@ def evaluate_translation(
             "is_correct": bleu_score > 0.3,
             "text": text,
             "pred_text": pred_text,
+            "bleu_tokenize": tokenize,
         }
     except Exception as e:
         return {
@@ -598,6 +604,7 @@ def evaluate_translation(
             "error": str(e),
             "text": reference.strip(),
             "pred_text": hypothesis.strip(),
+            "bleu_tokenize": tokenize,
         }
 
 

nemo_skills/evaluation/evaluator/comet.py

@@ -46,22 +46,11 @@ def load_comet_model(model_path: str):
     return model
 
 
-def _get_nested(sample: dict, key: str):
-    if "." in key:
-        value = sample
-        for part in key.split("."):
-            value = value[part]
-        return value
-    return sample[key]
-
-
 def process_file(
     input_file: Path,
     output_file: Path,
     comet_model,
     batch_size: int = 16,
-    source_key: str = "source",
-    reference_key: str = "reference",
 ):
     """Copy input file to output location and run xCOMET-XXL evaluation."""
     LOG.info(f"Processing {input_file} -> {output_file}")
@@ -87,9 +76,9 @@ def process_file(
         try:
             comet_list.append(
                 {
-                    "src": _get_nested(sample, source_key),
-                    "mt": _get_nested(sample, "generation"),
-                    "ref": _get_nested(sample, reference_key),
+                    "src": sample["source"],
+                    "mt": sample["generation"],
+                    "ref": sample["reference"],
                 }
             )
         except KeyError as e:
@@ -150,18 +139,6 @@ def main():
         default=1,
         help="Number of random seeds (for multiple seeds mode)",
     )
-    parser.add_argument(
-        "--source-key",
-        type=str,
-        default="source",
-        help="Sample field (supports dotted paths) holding the source text passed as COMET 'src'",
-    )
-    parser.add_argument(
-        "--reference-key",
-        type=str,
-        default="reference",
-        help="Sample field (supports dotted paths) holding the reference translation passed as COMET 'ref'",
-    )
     args = parser.parse_args()
 
     output_dir = Path(args.output_dir)
@@ -193,8 +170,6 @@ def main():
             output_file,
             comet_model,
             args.batch_size,
-            source_key=args.source_key,
-            reference_key=args.reference_key,
         )
 
     LOG.info("All files processed successfully")