import pandas as pd import re from pathlib import Path from typing import List, Tuple, Optional import json import math import os from typing import Dict 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: 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 >= 12: return "14b" if p >= 9: 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 is not None: if len(model_groups) == 0: return pd.DataFrame() 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 = [ # 代码 (Code) "github cpp", "github python", "github javascript", # 科研 (Research) "arxiv physics", "arxiv cs", "arxiv math", # 写作 (Writing) "ao3 english", "github markdown", # 世界知识 (World Knowledge) "bbc news", "wikipedia english", ] 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 class LongContextDataManager: def __init__(self, data_dir: str): self.data_dir = data_dir # { period: { "Display Name": "full_path" } } self.period_file_map = {} # { period: { "Model Name": ["path1", "path2"] } } self.period_model_map = {} # { period: { "Dataset Name": set(["path1", "path2"]) } } self.period_dataset_map = {} # { period: { "Model Name": { "Dataset Name": ["path1", "path2"] } } } self.period_model_dataset_map = {} self._scan_directories() def _scan_directories(self): for root, dirs, files in os.walk(self.data_dir): json_files = [f for f in files if f.endswith(".json")] if json_files: period = os.path.basename(root) if period not in self.period_file_map: self.period_file_map[period] = {} self.period_model_map[period] = {} self.period_dataset_map[period] = {} self.period_model_dataset_map[period] = {} for jf in json_files: full_path = os.path.join(root, jf) with open(full_path, "r", encoding="utf-8") as f: meta = json.load(f) model_name = meta.get("model_name_or_path", jf.replace(".json", "")).split("/")[-1].replace(".pth", "") data_path = meta.get("data_path", "").replace("UncheatableEval-", "") dataset_name = data_path.split("/")[-1] if data_path else "Unknown" file_display_label = f"{model_name}-{dataset_name}" print(file_display_label) self.period_file_map[period][file_display_label] = full_path if model_name not in self.period_model_map[period]: self.period_model_map[period][model_name] = [] self.period_model_map[period][model_name].append(full_path) if dataset_name not in self.period_dataset_map[period]: self.period_dataset_map[period][dataset_name] = [] self.period_dataset_map[period][dataset_name].append(full_path) if model_name not in self.period_model_dataset_map[period]: self.period_model_dataset_map[period][model_name] = {} if dataset_name not in self.period_model_dataset_map[period][model_name]: self.period_model_dataset_map[period][model_name][dataset_name] = [] self.period_model_dataset_map[period][model_name][dataset_name].append(full_path) def get_available_periods(self): return sorted(list(self.period_file_map.keys())) def get_file_choices(self, period): """返回 [(Display Name, Full Path), ...]""" if period not in self.period_file_map: return [] return [(k, v) for k, v in self.period_file_map[period].items()] def get_model_choices(self, period): """返回 [(Model Name, Model Name), ...]""" if period not in self.period_model_map: return [] return [(k, k) for k in self.period_model_map[period].keys()] def get_paths_for_model(self, period, model_name): return self.period_model_map.get(period, {}).get(model_name, []) def get_dataset_choices(self, period): """返回某个period下的所有数据集列表 [(Dataset Name, Dataset Name), ...]""" if period not in self.period_dataset_map: return [] return [(k, k) for k in sorted(self.period_dataset_map[period].keys())] def get_paths_for_model_and_datasets(self, period, model_name, dataset_names): """根据模型名称和数据集名称列表，返回对应的文件路径列表""" if period not in self.period_model_dataset_map: return [] if model_name not in self.period_model_dataset_map[period]: return [] paths = [] for dataset_name in dataset_names: if dataset_name in self.period_model_dataset_map[period][model_name]: paths.extend(self.period_model_dataset_map[period][model_name][dataset_name]) return paths 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)) lcm = LongContextDataManager("longctx_data")