Merge pull request #52 from sbintuitions/fix/prediction_logging_to_jsonl

[Fix] 予測値をjsonlに保存

src/jmteb/__main__.py

@@ -37,6 +37,10 @@ def main(
             dataset_name=eval_name,
             task_name=evaluator.__class__.__name__.replace("Evaluator", ""),
         )
+        if getattr(evaluator, "log_predictions", False):
+            score_recorder.record_predictions(
+                metrics, eval_name, evaluator.__class__.__name__.replace("Evaluator", "")
+            )
 
         logger.info(f"Results for {eval_name}\n{json.dumps(metrics.as_dict(), indent=4, ensure_ascii=False)}")
 

src/jmteb/utils/score_recorder.py

@@ -3,6 +3,7 @@
 import json
 from abc import ABC, abstractmethod
 from collections import defaultdict
+from dataclasses import asdict
 from os import PathLike
 from pathlib import Path
 from typing import Any
@@ -30,6 +31,12 @@ def save_to_json(scores: EvaluationResults | dict[Any, Any], filename: str | Pat
         with open(filename, "w") as fout:
             json.dump(scores, fout, indent=4, ensure_ascii=False)
 
+    @staticmethod
+    def save_prediction_to_jsonl(predictions: list[Any], filename: str | PathLike[str]) -> None:
+        with open(filename, "w") as fout:
+            for prediction in predictions:
+                fout.write(json.dumps(asdict(prediction), ensure_ascii=False) + "\n")
+
     def record_task_scores(self, scores: EvaluationResults, dataset_name: str, task_name: str) -> None:
         if not self.save_dir:
             return
@@ -39,6 +46,13 @@ def record_task_scores(self, scores: EvaluationResults, dataset_name: str, task_
         self.scores[task_name][dataset_name] = scores
         self.save_to_json(self.scores[task_name][dataset_name].as_dict(), save_filename)
 
+    def record_predictions(self, results: EvaluationResults, dataset_name: str, task_name: str) -> None:
+        if not self.save_dir:
+            return
+        save_filename = Path(self.save_dir) / task_name / f"predictions_{dataset_name}.jsonl"
+        save_filename.parent.mkdir(parents=True, exist_ok=True)
+        self.save_prediction_to_jsonl(results.predictions, save_filename)
+
     def record_summary(self):
         if not self.save_dir:
             return