Merge branch 'main' into fix/arguments

foundation-model-stack · Mar 8, 2024 · 5a2097c · 5a2097c
2 parents e718d04 + 3ac941d
commit 5a2097c
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/build/launch_training.py b/build/launch_training.py
@@ -39,8 +39,12 @@
 def txt_to_obj(txt):
     base64_bytes = txt.encode("ascii")
     message_bytes = base64.b64decode(base64_bytes)
-    obj = pickle.loads(message_bytes)
-    return obj
+    try:
+        # If the bytes represent JSON string
+        return json.loads(message_bytes)
+    except UnicodeDecodeError:
+        # Otherwise the bytes are a pickled python dictionary
+        return pickle.loads(message_bytes)
 
 
 def get_highest_checkpoint(dir_path):