GAIR/Preference-Dissection-Visualization

1import os2import json3import pandas as pd4 5 6def print_colored_text(text, color="yellow", end=None):7    colors = {8        "red": "\033[91m",9        "green": "\033[92m",10        "yellow": "\033[93m",11        "blue": "\033[94m",12        "purple": "\033[95m",13        "cyan": "\033[96m",14        "white": "\033[97m",15        "reset": "\033[0m",16    }17 18    color_code = colors.get(color.lower(), colors["reset"])19    print(f"{color_code}{text}{colors['reset']}", end=end)20 21 22def read_parquet(parquet_file_path):23    data = pd.read_parquet(parquet_file_path)24    data = data.to_dict("records")25    return data26 27 28def write_jsonl(data, jsonl_file_path, mode="w"):29    # data is a list, each of the item is json-serilizable30    assert isinstance(data, list)31    if not os.path.exists(os.path.dirname(jsonl_file_path)):32        os.makedirs(os.path.dirname(jsonl_file_path))33    with open(jsonl_file_path, mode) as f:34        for item in data:35            f.write(json.dumps(item) + "\n")36 37 38def write_json(data, json_file_path):39    if not os.path.exists(os.path.dirname(json_file_path)):40        os.makedirs(os.path.dirname(json_file_path))41    with open(json_file_path, "w") as f:42        json.dump(data, f)43 44 45def read_jsonl(jsonl_file_path):46    s = []47    if not os.path.exists(jsonl_file_path):48        print_colored_text("File not exists: " + jsonl_file_path, "red")49        return s50    with open(jsonl_file_path, "r") as f:51        lines = f.readlines()52    for line in lines:53        linex = line.strip()54        if linex == "":55            continue56        s.append(json.loads(linex))57    return s58 59 60def read_json(json_file_path):61    with open(json_file_path, "r") as f:62        data = json.load(f)63    return data64 65 66def read_all(file_path):67    if file_path.endswith(".jsonl"):68        return read_jsonl(file_path)69    elif file_path.endswith(".json"):70        return read_json(file_path)71    elif file_path.endswith(".parquet"):72        return read_parquet(file_path)73    elif file_path.endswith(".txt"):74        with open(file_path, "r") as f:75            return [line.strip() for line in f.readlines() if line.strip() != ""]76    else:77        raise ValueError(f"Unrecognized file type: {file_path}")78 79 80# Function to convert JSON to Markdown formatted string with bold keys81def json_to_markdown_bold_keys(json_obj, depth=0):82    markdown_str = ""83    indent = "\t "84    for key, value in json_obj.items():85        if isinstance(value, dict):86            markdown_str += f"**{key}** :\n\n{indent * (depth)}- {json_to_markdown_bold_keys(value, depth + 1)}\n\n"87        elif isinstance(value, list):88            if len(value) > 0:89                markdown_str += (90                    f"**{key}** :\n\n "91                    + f"\n\n{indent * (depth)}- "92                    + f"\n\n{indent * (depth)}- ".join(93                        [94                            (95                                json_to_markdown_bold_keys(item, depth + 1)96                                if isinstance(item, dict)97                                else f"{indent * (depth + 1)}{item}"98                            )99                            for item in value100                        ]101                    )102                    + "\n\n"103                )104            else:105                markdown_str += f"**{key}** : None\n\n"106        else:107            if depth == 0:108                markdown_str += f"**{key}** : {value}\n\n"109            else:110                markdown_str += f"{indent * (depth)}- **{key}** : {value}\n\n"111    return markdown_str112 113 114def custom_md_with_color(text, color):115    return f"""116<div style="background-color:#{color};padding:10px;border-radius:5px;">117    <p style="color:black;font-size:16px;">📑\n\n{text}</p>118</div>"""119