Spaces:
Running
Running
| import matplotlib.pyplot as plt | |
| from matplotlib.ticker import ScalarFormatter | |
| from collections import defaultdict | |
| import numpy as np | |
| import json | |
| import os | |
| def get_aggregated_stats(data_list): | |
| if not data_list: | |
| return None, None | |
| max_len = 0 | |
| cleaned_list = [] | |
| for r, c in data_list: | |
| valid_len = min(len(r), len(c)) | |
| if valid_len > 0: | |
| max_len = max(max_len, valid_len) | |
| cleaned_list.append((r[:valid_len], c[:valid_len])) | |
| if max_len == 0: | |
| return None, None | |
| numerator = np.zeros(max_len) | |
| denominator = np.zeros(max_len) | |
| for rate, count in cleaned_list: | |
| l = len(rate) | |
| numerator[:l] += rate * count | |
| denominator[:l] += count | |
| with np.errstate(divide="ignore", invalid="ignore"): | |
| avg_rate = numerator / denominator | |
| return avg_rate, denominator | |
| def process_files(file_paths): | |
| model_raw_data = defaultdict(list) | |
| dataset_names = set() | |
| for file_path in file_paths: | |
| with open(file_path, "r", encoding="utf-8") as file: | |
| data = json.load(file) | |
| default_name = os.path.basename(file_path).replace(".json", "") | |
| model_name = data.get("model_name_or_path", default_name).split("/")[-1].replace(".pth", "") | |
| data_path = data.get("data_path", "") | |
| if data_path: | |
| dataset_names.add(data_path.split("/")[-1]) | |
| rate_list = data.get("byte_wise_compression_rate", []) | |
| counts_list = data.get("byte_wise_counts", []) | |
| if rate_list and counts_list: | |
| model_raw_data[model_name].append((np.array(rate_list), np.array(counts_list))) | |
| return model_raw_data, dataset_names | |
| def clean_label_for_path(l): | |
| if l.endswith(".json"): | |
| l = l.replace("UncheatableEval-", "").replace(".json", "")[:-20] | |
| return l | |
| def draw_long_context_plot( | |
| viz_mode, | |
| data_source_map, | |
| baseline_key, | |
| cutoff_ratio, | |
| smooth_window, | |
| start_offset, | |
| y_range, | |
| ): | |
| if not data_source_map: | |
| return None | |
| final_curves = {} | |
| dataset_info = set() | |
| for label, file_paths in data_source_map.items(): | |
| print(label, file_paths) | |
| raw_data_list = [] | |
| for fp in file_paths: | |
| with open(fp, "r", encoding="utf-8") as f: | |
| d = json.load(f) | |
| rate = d.get("byte_wise_compression_rate", []) | |
| counts = d.get("byte_wise_counts", []) | |
| d_path = d.get("data_path", "") | |
| if d_path: | |
| dataset_info.add(d_path.split("/")[-1]) | |
| if rate and counts: | |
| raw_data_list.append((np.array(rate), np.array(counts))) | |
| avg_rate, total_counts = get_aggregated_stats(raw_data_list) | |
| if avg_rate is not None: | |
| if len(total_counts) > 0: | |
| initial_count = total_counts[0] | |
| threshold = initial_count * cutoff_ratio | |
| cut_idx = np.where(total_counts < threshold)[0] | |
| if len(cut_idx) > 0: | |
| avg_rate = avg_rate[: cut_idx[0]] | |
| final_curves[label] = avg_rate | |
| plot_items = [] | |
| is_relative = "Relative" in viz_mode | |
| if is_relative: | |
| if baseline_key not in final_curves: | |
| return None | |
| baseline_curve = final_curves[baseline_key] | |
| for label, curve in final_curves.items(): | |
| if label == baseline_key: | |
| continue | |
| min_len = min(len(baseline_curve), len(curve)) | |
| if min_len <= start_offset: | |
| continue | |
| delta = curve[:min_len] - baseline_curve[:min_len] | |
| plot_items.append({"label": clean_label_for_path(label), "values": delta[start_offset:], "start_x": start_offset + 1}) | |
| else: # Absolute | |
| for label, curve in final_curves.items(): | |
| if len(curve) <= start_offset: | |
| continue | |
| plot_items.append({"label": clean_label_for_path(label), "values": curve[start_offset:], "start_x": start_offset + 1}) | |
| if not plot_items: | |
| return None | |
| # fig, ax = plt.subplots(figsize=(12, 6), dpi=300) | |
| fig, ax = plt.subplots(figsize=(12, 6), dpi=300, layout="constrained") | |
| if is_relative: | |
| ax.axhline(0, color="black", linestyle="--", linewidth=1, alpha=0.5, label=f"Baseline: {clean_label_for_path(baseline_key)}") | |
| colors = plt.cm.tab10(np.linspace(0, 1, max(1, len(plot_items)))) | |
| for idx, item in enumerate(plot_items): | |
| y_val = item["values"] | |
| x_val = np.arange(item["start_x"], item["start_x"] + len(y_val)) | |
| color = colors[idx % 10] | |
| if smooth_window > 1 and len(y_val) > smooth_window: | |
| kernel = np.ones(smooth_window) / smooth_window | |
| y_smooth = np.convolve(y_val, kernel, mode="valid") | |
| x_smooth = x_val[smooth_window - 1 :] | |
| ax.plot(x_val, y_val, color=color, alpha=0.15, linewidth=0.5) | |
| ax.plot(x_smooth, y_smooth, color=color, linewidth=2, label=item["label"]) | |
| else: | |
| ax.plot(x_val, y_val, color=color, linewidth=1.5, label=item["label"]) | |
| ax.set_xscale("log") | |
| ax.xaxis.set_major_formatter(ScalarFormatter()) | |
| if y_range: | |
| ymin, ymax = y_range | |
| if ymin is not None and ymax is not None: | |
| ax.set_ylim(ymin, ymax) | |
| elif ymin is not None: | |
| current_ymax = ax.get_ylim()[1] | |
| ax.set_ylim(bottom=ymin, top=current_ymax) | |
| elif ymax is not None: | |
| current_ymin = ax.get_ylim()[0] | |
| ax.set_ylim(bottom=current_ymin, top=ymax) | |
| # ds_str = ", ".join(sorted(dataset_info)) if len(dataset_info) < 5 else f"{len(dataset_info)} Datasets" | |
| mode_title = "Compression Rate" if not is_relative else "Compression Rate Delta" | |
| # sub_title = " (Averaged by Model)" if "Model" in viz_mode else "" | |
| ax.set_title(f"{mode_title}", fontsize=12, pad=8) | |
| # ax.legend(loc="lower center", bbox_to_anchor=(0.5, 1.02), ncol=4, fontsize="small", frameon=False, borderaxespad=0.0) | |
| ax.legend(loc="upper left", bbox_to_anchor=(1.02, 1), ncol=1, fontsize="small", frameon=False, borderaxespad=0.0) | |
| ax.set_xlabel("Byte Position") | |
| ax.set_ylabel("Compression Rate (%)" if not is_relative else "Compression Rate Delta (%)") | |
| ax.grid(True, which="major", ls="-", alpha=0.8) | |
| ax.grid(True, which="minor", ls="--", alpha=0.3) | |
| # plt.tight_layout() | |
| return fig | |