From 2597812b14273710f365d35f50a710cbbbc5ce28 Mon Sep 17 00:00:00 2001 From: Jiri Podivin Date: Sat, 8 Jun 2024 14:22:12 +0200 Subject: [PATCH] New JSON serializer Signed-off-by: Jiri Podivin --- src/evaluate/saving.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/evaluate/saving.py b/src/evaluate/saving.py index 4eea1a6a0..fc92fddf7 100644 --- a/src/evaluate/saving.py +++ b/src/evaluate/saving.py @@ -4,12 +4,26 @@ import sys from datetime import datetime from pathlib import Path +import numpy as np from datasets.utils.filelock import FileLock from . import __version__ +class NpEncoder(json.JSONEncoder): + """Numpy aware JSON encoder. + """ + def default(self, o): + if isinstance(o, np.floating): + return float(o) + if isinstance(o, np.integer): + return int(o) + if isinstance(o, np.ndarray): + return o.tolist() + return super().default(o) + + def save(path_or_file, **data): """ Saves results to a JSON file. Also saves system information such as current time, current commit @@ -40,7 +54,7 @@ def save(path_or_file, **data): with FileLock(str(file_path) + ".lock"): with open(file_path, "w") as f: - json.dump(data, f) + json.dump(data, f, cls=NpEncoder) # cleanup lock file try: