From d8be2fa0bbe0cc472bb4458b6412afec8d82ed3b Mon Sep 17 00:00:00 2001
From: Li Yin <li.yin.gravity@gmail.com>
Date: Tue, 4 Jun 2024 20:34:06 -0700
Subject: [PATCH] improve serialization code

---
 lightrag/utils/serialization.py               |  28 ++----------------
 .../classification/{ => evals}/zero_shot.json |   0
 use_cases/classification/train.py             |   2 +-
 use_cases/classification/zero_shot.pickle     | Bin 110 -> 0 bytes
 4 files changed, 3 insertions(+), 27 deletions(-)
 rename use_cases/classification/{ => evals}/zero_shot.json (100%)
 delete mode 100644 use_cases/classification/zero_shot.pickle

diff --git a/lightrag/utils/serialization.py b/lightrag/utils/serialization.py
index f1206ce2..46b1829a 100644
--- a/lightrag/utils/serialization.py
+++ b/lightrag/utils/serialization.py
@@ -41,31 +41,6 @@ def serialize(obj: Mapping[str, Any]) -> str:
     return json.dumps(obj, indent=4, default=default)
 
 
-# TODO: make this more clear
-def save(obj: Mapping[str, Any], f: str = "task") -> None:
-    __doc__ = """Save the object to a json file.
-
-    We save two versions of the object:
-    - task.json: the object itself with Parameter serialized to dict
-    - task.pickle: the object itself with Parameter as is
-    """
-
-    # save the object to a json file
-    json_f = f"{f}.json"
-
-    os.makedirs(os.path.dirname(json_f) or ".", exist_ok=True)
-    with open(json_f, "w") as file:
-        # serialize the object with the default function
-        # serialized_obj = serialize(obj)
-        # file.write(serialized_obj, indent=4)
-        json.dump(obj, file, indent=4, default=default)
-
-    # save the object to a pickle file
-    pickle_f = f"{f}.pickle"
-    with open(pickle_f, "wb") as file:
-        pickle.dump(obj, file)
-
-
 def save_json(obj: Mapping[str, Any], f: str = "task.json") -> None:
     __doc__ = """Save the object to a json file.
 
@@ -76,7 +51,8 @@ def save_json(obj: Mapping[str, Any], f: str = "task.json") -> None:
     os.makedirs(os.path.dirname(f) or ".", exist_ok=True)
     try:
         with open(f, "w") as file:
-            json.dump(obj, file, indent=4, default=default)
+            serialized_obj = serialize(obj)
+            file.write(serialized_obj)
     except IOError as e:
         raise IOError(f"Error saving object to JSON file {f}: {e}")
 
diff --git a/use_cases/classification/zero_shot.json b/use_cases/classification/evals/zero_shot.json
similarity index 100%
rename from use_cases/classification/zero_shot.json
rename to use_cases/classification/evals/zero_shot.json
diff --git a/use_cases/classification/train.py b/use_cases/classification/train.py
index 0a31fc38..26edf7b0 100644
--- a/use_cases/classification/train.py
+++ b/use_cases/classification/train.py
@@ -190,7 +190,7 @@ def batch_eval(self, batch: Dict[str, Any]) -> Tuple[float, float]:
         return accuracy, macro_f1_score
 
     def eval_zero_shot(self, save_path: str = None):
-        save_path = save_path or "use_cases/classification/evals/zero_shot"
+        save_path = save_path or "use_cases/classification/evals/zero_shot.json"
         json_obj: Dict[str, Any] = {}
         self.task.eval()  # not using any trained examples
         acc, macro_f1, best_weights_per_class = self.eval()  # zero shot, 0.542
diff --git a/use_cases/classification/zero_shot.pickle b/use_cases/classification/zero_shot.pickle
deleted file mode 100644
index ab6869699a6afe7c865178ac388c3f054cf50732..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 110
zcmZo*nVQT10ku<lIIB{N^5cs$@=K=FPSNOLNi9pv0kfDBlar^o+dr3@zpw1XmL87W
z#AKktG((VBtK5#j*R7>JEG4PMV0{^^?)J}2?Q7mn%*bE^GIjaFSR^KwmKN&)0OJ%Y
AkpKVy