bigscience-data/process-pipeline-visualizer
6
1import os2import pprint as pp3from collections import OrderedDict, defaultdict4 5import json6import diff_viewer7import pandas as pd8import streamlit as st9from datasets import load_dataset, get_dataset_config_names10 11CHECK_DATASET_DIR_PATH_BEFORE_CLEAN_SELECT = st.secrets["CHECK_DATASET_DIR_PATH_BEFORE_CLEAN_SELECT"]12LOGS_DATASET_DIR_PATH_BEFORE_CLEAN_SELECT = st.secrets["LOGS_DATASET_DIR_PATH_BEFORE_CLEAN_SELECT"]13HF_API_TOKEN = st.secrets["HF_API_TOKEN"]14OPERATION_TYPES = [15 "Applied filter",16 "Applied deduplication function",17 "Applied map function",18]19MAX_LEN_DS_CHECKS = st.secrets["MAX_LEN_DS_CHECKS"]20 21 22def get_ds(config):23 ds = load_dataset(CHECK_DATASET_DIR_PATH_BEFORE_CLEAN_SELECT, config, use_auth_token=HF_API_TOKEN, trust_remote_code=True)24 return ds["train"]25 26 27def next_idx(idx: int):28 idx += 129 return idx % len(st.session_state["ds"])30 31 32def previous_idx(idx: int):33 idx -= 134 return idx % len(st.session_state["ds"])35 36 37def on_click_next():38 st.session_state["idx_1"] = next_idx(st.session_state["idx_1"])39 st.session_state["idx_2"] = next_idx(st.session_state["idx_2"])40 41 42def on_click_previous():43 st.session_state["idx_1"] = previous_idx(st.session_state["idx_1"])44 st.session_state["idx_2"] = previous_idx(st.session_state["idx_2"])45 46 47def on_ds_change(config):48 st.session_state["ds"] = get_ds(config)49 st.session_state["idx_1"] = 050 st.session_state["idx_2"] = 1 if len(st.session_state["ds"]) > 1 else 051 st.session_state["ds_check_config"] = config52 st.session_state["ds_max_docs"] = len(st.session_state["ds"])53 54 55def get_log_stats_df(raw_log):56 data = OrderedDict(57 {58 "Order": [],59 "Name": [],60 "Initial number of samples": [],61 "Final number of samples": [],62 "Initial size in bytes": [],63 "Final size in bytes": [],64 }65 )66 67 metric_dict = defaultdict(lambda: {})68 order = 069 for line in raw_log.split("\n"):70 for metric_name in list(data.keys()) + OPERATION_TYPES:71 72 if metric_name == "Name" or metric_name == "Order":73 continue74 75 if metric_name not in line:76 continue77 78 if (79 metric_name == "Removed percentage"80 and "Removed percentage in bytes" in line81 ):82 continue83 84 if (85 metric_name == "Deduplicated percentage"86 and "Deduplicated percentage in bytes" in line87 ):88 continue89 90 value = line.split(metric_name)[1].split(" ")[1]91 92 if metric_name in OPERATION_TYPES:93 operation_name = value94 metric_dict[operation_name]["Order"] = order95 order += 196 continue97 98 assert (99 metric_name not in metric_dict[operation_name]100 ), f"operation_name: {operation_name}\n\nvalue: {value}\n\nmetric_dict: {pp.pformat(metric_dict)} \n\nmetric_name: {metric_name} \n\nline: {line}"101 metric_dict[operation_name][metric_name] = value102 for name, data_dict in metric_dict.items():103 for metric_name in data.keys():104 if metric_name == "Name":105 data[metric_name].append(name)106 continue107 108 data[metric_name].append(data_dict[metric_name])109 df = pd.DataFrame(data)110 df.rename(111 {112 "Initial size in bytes": "Initial size (GB)",113 "Final size in bytes": "Final size (GB)",114 },115 axis=1,116 inplace=True,117 )118 df["% samples removed"] = (119 (120 df["Initial number of samples"].astype(float)121 - df["Final number of samples"].astype(float)122 )123 / df["Initial number of samples"].astype(float)124 * 100125 )126 df["Size (GB) % removed"] = (127 (df["Initial size (GB)"].astype(float) - df["Final size (GB)"].astype(float))128 / df["Initial size (GB)"].astype(float)129 * 100130 )131 return df132 133 134def get_logs_stats(raw_log):135 try:136 df = get_log_stats_df(raw_log)137 st.dataframe(df)138 except Exception as e:139 st.write(e)140 st.write("Subset of the logs:")141 subcontent = [142 line143 for line in raw_log.split("\n")144 if "INFO - __main__" in line145 and "Examples of" not in line146 and "Examples n°" not in line147 ]148 st.write(subcontent)149 150 151def meta_component(idx_key: str = "idx_1"):152 if "meta" not in st.session_state["ds"][st.session_state[idx_key]]:153 return154 155 with st.expander("See meta field of the example"):156 meta = st.session_state["ds"][st.session_state["idx_1"]]["meta"]157 st.write(meta)158 159 160def filter_page():161 index_example = st.number_input("Index of the chosen example", min_value=0, max_value=st.session_state["ds_max_docs"] -1, value=0, step=1)162 st.session_state["idx_1"] = index_example163 st.session_state["idx_2"] = next_idx(index_example) 164 idx_1 = st.session_state["idx_1"]165 idx_2 = st.session_state["idx_2"] 166 text_1 = st.session_state["ds"][idx_1]["text"]167 text_2 = st.session_state["ds"][idx_2]["text"]168 169 st.markdown(170 f"<h1 style='text-align: center'>Some examples of filtered out texts</h1>",171 unsafe_allow_html=True,172 )173 # col_button_previous, _, col_button_next = st.columns(3)174 175 176 # col_button_next.button(177 # "Go to next example",178 # key=None,179 # help=None,180 # on_click=on_click_next,181 # args=None,182 # kwargs=None,183 # )184 # col_button_previous.button(185 # "Go to previous example",186 # key=None,187 # help=None,188 # on_click=on_click_previous,189 # args=None,190 # kwargs=None,191 # )192 col_1, col_2 = st.columns(2)193 with col_1:194 st.subheader(f"Example n°{idx_1}")195 meta_component(idx_key="idx_1")196 text_1_show = text_1.replace("\n", "<br>")197 st.markdown(f"<div>{text_1_show}</div>", unsafe_allow_html=True)198 199 with col_2:200 st.subheader(f"Example n°{idx_2}")201 meta_component(idx_key="idx_2")202 text_2_show = text_2.replace("\n", "<br>")203 st.markdown(f"<div>{text_2_show}</div>", unsafe_allow_html=True)204 205 206def dedup_or_cleaning_page():207 index_example = st.number_input("Index of the chosen example", min_value=0, max_value=st.session_state["ds_max_docs"] -1, value=0, step=1)208 st.session_state["idx_1"] = index_example209 st.session_state["idx_2"] = next_idx(index_example) 210 211 # col_button_previous, col_title, col_button_next = st.columns(3)212 # col_title.markdown(213 # f"<h1 style='text-align: center'>Example n°{st.session_state['idx_1']}</h1>",214 # unsafe_allow_html=True,215 # )216 # col_button_next.button(217 # "Go to next example",218 # key=None,219 # help=None,220 # on_click=on_click_next,221 # args=None,222 # kwargs=None,223 # )224 # col_button_previous.button(225 # "Go to previous example",226 # key=None,227 # help=None,228 # on_click=on_click_previous,229 # args=None,230 # kwargs=None,231 # )232 233 text = st.session_state["ds"][st.session_state["idx_1"]]["text"]234 old_text = st.session_state["ds"][st.session_state["idx_1"]]["old_text"]235 st.markdown(236 f"<h2 style='text-align: center'>Changes applied</h1>", unsafe_allow_html=True237 )238 col_text_1, col_text_2 = st.columns(2)239 with col_text_1:240 st.subheader("Old text")241 with col_text_2:242 st.subheader("New text")243 diff_viewer.diff_viewer(old_text=old_text, new_text=text, lang="none")244 meta_component(idx_key="idx_1")245 246 with st.expander("See full old and new texts of the example"):247 text_show = text.replace("\n", "<br>")248 old_text_show = old_text.replace("\n", "<br>")249 250 col_1, col_2 = st.columns(2)251 with col_1:252 st.subheader("Old text")253 st.markdown(f"<div>{old_text_show}</div>", unsafe_allow_html=True)254 with col_2:255 st.subheader("New text")256 st.markdown(f"<div>{text_show}</div>", unsafe_allow_html=True)257 258 259# Streamlit page260st.set_page_config(page_title="Dataset explorer", page_icon=":hugging_face:", layout="wide")261st.write(262 "The purpose of this application is to sequentially view the changes made to a dataset."263)264 265 266# st.write(CHECK_DATASET_DIR_PATH_BEFORE_CLEAN_SELECT)267# ds_log = load_dataset(CHECK_DATASET_DIR_PATH_BEFORE_CLEAN_SELECT, 'clean_v1_dsname_lm_en_multi_un_2', use_auth_token=HF_API_TOKEN)268# st.write(ds_log)269 270 271 272col_option_clean, col_option_ds = st.columns(2)273 274with open("dataset_configs.json", "r") as f:275 CHECK_CONFIGS = json.load(f)276# CHECK_CONFIGS = get_dataset_config_names(CHECK_DATASET_DIR_PATH_BEFORE_CLEAN_SELECT, use_auth_token=HF_API_TOKEN)277 278CLEANING_VERSIONS = set()279dataset_names = defaultdict(set)280checks_names = defaultdict(lambda: defaultdict(set))281 282for check_config in CHECK_CONFIGS:283 cleaning_version, check_config = check_config.split("_dsname_")284 dataset_name, checks_name = check_config.split("_operation_")285 CLEANING_VERSIONS.add(cleaning_version)286 dataset_names[cleaning_version].add(dataset_name)287 checks_names[cleaning_version][dataset_name].add(checks_name)288 289# CLEANING_VERSIONS = sorted(list(os.listdir(DATASET_DIR_PATH_BEFORE_CLEAN_SELECT)), reverse=True)290option_clean = col_option_clean.selectbox(291 "Select the cleaning version", sorted(CLEANING_VERSIONS, reverse=True)292)293 294# DATASET_DIR_PATH = os.path.join(DATASET_DIR_PATH_BEFORE_CLEAN_SELECT, option_clean)295# dataset_names = sorted(list(os.listdir(DATASET_DIR_PATH)))296option_ds = col_option_ds.selectbox("Select the dataset", sorted(dataset_names[option_clean]))297 298# checks_path = os.path.join(DATASET_DIR_PATH, option_ds, "checks")299# checks_names = sorted(list(os.listdir(checks_path)))300 301# log_path = os.path.join(DATASET_DIR_PATH, option_ds, "logs.txt")302ds_log = load_dataset(LOGS_DATASET_DIR_PATH_BEFORE_CLEAN_SELECT, f"{option_clean}_dsname_{option_ds}", use_auth_token=HF_API_TOKEN, trust_remote_code=True)303log = ds_log["train"][0]["log"]304get_logs_stats(raw_log=log)305 306option_check = st.selectbox("Select the operation applied to inspect", sorted(checks_names[option_clean][option_ds]))307 308ds_check_config = f"{option_clean}_dsname_{option_ds}_operation_{option_check}"309 310if "ds" not in st.session_state or ds_check_config != st.session_state["ds_check_config"]:311 on_ds_change(ds_check_config)312 313if len(st.session_state["ds"]) == MAX_LEN_DS_CHECKS:314 st.warning(315 f"Note: only a subset of size {MAX_LEN_DS_CHECKS} of the modified / filtered examples can be shown in this application"316 )317with st.expander("See details of the available checks"):318 st.write(st.session_state["ds"])319 320 321_ = filter_page() if "_filter_" in option_check else dedup_or_cleaning_page()322