import pandas as pd import re from pathlib import Path from typing import List, Tuple, Optional import json import math class DataManager: def __init__(self, data_dir: str): self.data_dir = Path(data_dir) self.master_df = self._load_all_data() def _load_old_format_folder(self, period_dir: Path) -> pd.DataFrame: all_xlsx_data = [] period = period_dir.name for file_path in period_dir.iterdir(): if file_path.suffix != ".xlsx": continue model_group = file_path.stem xls = pd.read_excel(file_path, sheet_name=None) for sheet_name, df in xls.items(): if df.empty: continue df = self._clean_dataframe(df) df["Period"] = period df["Metric"] = sheet_name df["Model Group"] = model_group all_xlsx_data.append(df) return all_xlsx_data def _load_new_format_folder(self, period_dir: Path) -> pd.DataFrame: """ 读取新格式 JSON 数据并转换为宽表。 """ raw_records = [] period = period_dir.name for file_path in period_dir.rglob("*.json"): with open(file_path, "r", encoding="utf-8") as f: data = json.load(f) full_path = data["model_name_or_path"] model_name = full_path.split("/")[-1].replace(".pth", "") params = data["parameters count"] data_path = data["data_path"] source_col = data_path.split("-")[-1] neg_log_prob = data["neg_log_prob_sum"] avg_char = data["avg character count"] avg_bytes = data["avg bytes"] metrics = {} metrics["cr"] = data["compression_rate"] metrics["bpc"] = (neg_log_prob / avg_char) * (1 / math.log(2)) metrics["bpb"] = (neg_log_prob / avg_bytes) * (1 / math.log(2)) for metric_type, value in metrics.items(): if value is not None: raw_records.append( { "Name": model_name, "Params (B)": params, "Period": period, "Metric": metric_type, "Model Group": "other", "Source": source_col, "Value": value, } ) if not raw_records: return [] df_long = pd.DataFrame(raw_records) df_wide = df_long.pivot_table( index=["Name", "Params (B)", "Period", "Metric", "Model Group"], columns="Source", values="Value", ).reset_index() df_wide.columns.name = None def assign_group(p): if p >= 13: return "14b" if p >= 8: return "9b" if p >= 6: return "7b" if p >= 2.5: return "3b" if p >= 1: return "1b5" return "other" df_wide["Model Group"] = df_wide["Params (B)"].apply(assign_group) metadata_cols = ["Name", "Params (B)", "Period", "Metric", "Model Group"] new_columns = {} for col in df_wide.columns: if col not in metadata_cols: new_columns[col] = col.replace("_", " ") df_wide = df_wide.rename(columns=new_columns) return [df_wide] def _load_all_data(self) -> pd.DataFrame: all_records = [] if not self.data_dir.exists(): print(f"Warning: Directory {self.data_dir} does not exist.") return pd.DataFrame() period_dirs = [d for d in self.data_dir.iterdir() if d.is_dir() and re.match(r"^\d{4}-\d{2}$", d.name)] for period_dir in period_dirs: if period_dir.name <= "2025-11": all_records.extend(self._load_old_format_folder(period_dir)) else: all_records.extend(self._load_new_format_folder(period_dir)) if not all_records: return pd.DataFrame() final_df = pd.concat(all_records, ignore_index=True) exclude_cols = ["Name", "Period", "Metric", "Model Group"] numeric_cols = [c for c in final_df.columns if c not in exclude_cols] for col in numeric_cols: final_df[col] = pd.to_numeric(final_df[col], errors="coerce") return final_df def _clean_dataframe(self, df: pd.DataFrame) -> pd.DataFrame: df = df.dropna(axis=1, how="all") new_columns = [] for col in df.columns: col_str = str(col) if "Parameters" in col_str: new_columns.append("Params (B)") elif col_str == "Average (The lower the better)": new_columns.append("Average (lower=better)") else: new_columns.append(col_str.rsplit("_", maxsplit=1)[0].replace("\u200b", "")) # 去除_202xxxxx后缀 df.columns = new_columns column_mapping = {col: col.replace("_", " ") for col in df.columns} df = df.rename(columns=column_mapping) return df def get_available_periods(self) -> List[str]: """返回所有可用的时间周期，已排序，从旧到新""" if self.master_df.empty: return [] return sorted(self.master_df["Period"].unique().tolist()) def get_available_columns(self, period: str) -> List[str]: """获取特定时间段内的数据列（排除元数据列和全部为NaN的列）""" if self.master_df.empty: return [] subset = self.master_df[self.master_df["Period"] == period] if subset.empty: return [] metadata_cols = ["Name", "Params (B)", "Period", "Metric", "Model Group", "Average (lower=better)"] return [c for c in subset.columns if c not in metadata_cols and not subset[c].isna().all()] def query( self, period: str, metric_code: str, param_range: Tuple[float, float], model_groups: Optional[List[str]] = None, visible_columns: Optional[List[str]] = None, ) -> pd.DataFrame: """ 统一查询接口。 Args: period: 时间周期 (e.g. "2025-12") metric_display_name: bpc, bpb, cr param_range: 参数量范围元组 (min, max) model_groups: (可选) 文件名列表，如 ['14b', '7b'] visible_columns: (可选) 需要参与计算平均值的列名列表 """ mask = ( (self.master_df["Period"] == period) & (self.master_df["Metric"] == metric_code) & (self.master_df["Params (B)"].between(param_range[0], param_range[1])) ) if model_groups: mask = mask & (self.master_df["Model Group"].isin(model_groups)) filtered_df = self.master_df.loc[mask].copy() if filtered_df.empty: return filtered_df # 始终排除的列 exclude_cols = ["Period", "Metric", "Model Group"] # 始终保留的列 metadata_cols = ["Name", "Params (B)", "Average (lower=better)"] if visible_columns is not None: valid_visible_cols = [c for c in visible_columns if c in filtered_df.columns] columns_to_keep = metadata_cols + valid_visible_cols cols_for_average = valid_visible_cols else: all_cols = [c for c in filtered_df.columns if c not in exclude_cols] columns_to_keep = all_cols cols_for_average = [c for c in all_cols if c not in metadata_cols] if not cols_for_average: return pd.DataFrame() filtered_df["Average (lower=better)"] = filtered_df[cols_for_average].mean(axis=1).round(3) columns_to_keep = [c for c in columns_to_keep if c not in exclude_cols] columns_to_keep = list(dict.fromkeys(columns_to_keep)) filtered_df = filtered_df[columns_to_keep] if "Name" in filtered_df.columns: filtered_df["Name"] = filtered_df["Name"].apply(lambda x: x.replace(".pth", "")) filtered_df = filtered_df.sort_values(by="Average (lower=better)", ascending=True, kind="mergesort", na_position="last").reset_index( drop=True ) # 定义列优先级顺序 fixed_cols = ["Name", "Params (B)", "Average (lower=better)"] column_priority = [ "ao3 english", "bbc news", "wikipedia english", "arxiv physics", "arxiv cs", "arxiv math", "github cpp", "github python", "github javascript", "github markdown", ] existing_cols = filtered_df.columns.tolist() ordered_cols = [] for col in fixed_cols: if col in existing_cols: ordered_cols.append(col) for col in column_priority: if col in existing_cols and col not in ordered_cols: ordered_cols.append(col) for col in existing_cols: if col not in ordered_cols: ordered_cols.append(col) filtered_df = filtered_df[ordered_cols] return filtered_df if __name__ == "__main__": dm = DataManager("data") periods = dm.get_available_periods() print(f"Total records loaded: {len(dm.master_df)}") print(f"Available periods: {periods}") print(f"Available columns: {dm.get_available_columns('2025-11')}") result = dm.query( period="2025-11", metric_code="cr", param_range=(0, 20), model_groups=["7b"], visible_columns=["wikipedia_english"], ) print(result.head(20))