Allow SFT_TRAINER_CONFIG_JSON_ENV_VAR to be encoded json string (#82)

* allow SFT_TRAINER_CONFIG_JSON_ENV_VAR to be encoded json string, not just pickled python Signed-off-by: Kelly A <[email protected]> * linting fix Signed-off-by: Kelly A <[email protected]> * fix error catching Signed-off-by: Kelly A <[email protected]> --------- Signed-off-by: Kelly A <[email protected]>
foundation-model-stack · Mar 8, 2024 · 3ac941d · 3ac941d
1 parent 9c5a3bf
commit 3ac941d
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):