def encode_example()

in summarize_from_feedback/datasets/jsonl_encoding.py [0:0]

• 25 lines of code
• 11 McCabe index (conditional complexity)

def encode_example(example):
    """
    Make numpy arrays in the example serializable.
    We use base64 since it compresses well and is JSON-compatible.
    """
    encoded_example = example.copy()
    for k in example:
        obj = example[k]
        if isinstance(obj, np.ndarray):
            # base64 encode ndarrays instead of having JSON them output in plaintext
            del encoded_example[k]
            encoded_example["__np_" + k] = dict(
                shape=obj.shape,
                dtype=obj.dtype.name,
                data=base64.b64encode(obj.tobytes()).decode("utf-8"),
            )
        elif isinstance(obj, bytes):
            if k.endswith("/utf8"):
                encoded_example[k] = obj.decode("utf-8")
            else:
                del encoded_example[k]
                encoded_example["__bytes_" + k] = base64.b64encode(obj).decode("utf-8")
        elif isinstance(obj, np.integer):
            # JSON refuses to serialize np.integer, so convert it to plain integer
            encoded_example[k] = int(obj)
        elif isinstance(obj, np.floating):
            # JSON refuses to serialize np.floating, convert it to plain float
            encoded_example[k] = float(obj)
        elif isinstance(obj, list):
            # Same as above, JSON doesn't know how to serialize np.integer so we convert each element to int
            if all(isinstance(i, np.integer) for i in obj):
                encoded_example[k] = [int(i) for i in obj]
    return encoded_example