import pickle
import json
import sys


def make_serializable(obj):
    if isinstance(obj, dict):
        return {str(k): make_serializable(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [make_serializable(item) for item in obj]
    elif isinstance(obj, (int, float, bool, str, type(None))):
        return obj
    else:
        return repr(obj)


input_path = "llama3_8b_fp8_without_fix.pickle"
output_path = "llama3_8b_fp8_without_fix.json"

print(f"Loading {input_path}...")
with open(input_path, "rb") as f:
    data = pickle.load(f)

print(f"Loaded object of type: {type(data)}")

print("Converting to JSON-serializable format...")
serializable = make_serializable(data)

print(f"Writing to {output_path}...")
with open(output_path, "w") as f:
    json.dump(serializable, f, indent=2)

print("Done.")
