almanach/benchmark-in-a-haystack
4
1import pandas as pd2import json3import matplotlib.pyplot as plt4import seaborn as sns5import os6import argparse7from pathlib import Path8from datetime import datetime9 10from rich.console import Console11 12console = Console()13 14# Set style for beautiful plots15plt.rcParams['font.family'] = 'sans-serif'16plt.rcParams['font.sans-serif'] = ['Arial', 'Helvetica', 'DejaVu Sans']17plt.rcParams['font.size'] = 1118plt.rcParams['axes.labelsize'] = 1319plt.rcParams['axes.titlesize'] = 1620plt.rcParams['xtick.labelsize'] = 1121plt.rcParams['ytick.labelsize'] = 1122plt.rcParams['legend.fontsize'] = 1123plt.rcParams['figure.titlesize'] = 1824 25def analyze_and_plot(results, documents, benchmark_positions, output_base_dir="results", inject_inside=True, prefilter_hq=False, num_docs=100000, dataset_name="fineweb"):26 """Output benchmark sample ranks across classifiers and create visualizations."""27 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")28 results_dir = os.path.join(output_base_dir, timestamp)29 os.makedirs(results_dir, exist_ok=True)30 31 mode_suffix = "injected" if inject_inside else "separate"32 prefilter_suffix = "_prefiltered" if prefilter_hq else ""33 file_suffix = f"_{mode_suffix}{prefilter_suffix}_{num_docs}docs"34 35 all_benchmark_ranks = []36 plot_data = []37 bench_ranks_dict = {}38 39 console.rule("[bold blue]Analyzing classifier results...[/bold blue]")40 41 for clf_name, scores in results.items():42 console.log(f"[yellow]Analyzing results for {clf_name}...[/yellow]")43 scores_df = pd.DataFrame(scores)44 scores_df = scores_df.dropna(subset=["score"])45 scores_df = scores_df.sort_values("score", ascending=False)46 scores_df["rank"] = range(1, len(scores_df) + 1)47 48 bench_df = scores_df[scores_df["contains_benchmark"] == True].copy()49 bench_df["classifier"] = clf_name50 bench_df["percentile"] = (len(scores_df) - bench_df["rank"]) / len(scores_df) * 10051 52 for _, row in bench_df.iterrows():53 key = (row["id"], row["benchmark_type"], row["benchmark_index"])54 if key not in bench_ranks_dict:55 bench_ranks_dict[key] = {56 "id": row["id"],57 "benchmark_type": row["benchmark_type"],58 "benchmark_index": row["benchmark_index"],59 }60 bench_ranks_dict[key][clf_name] = {61 "rank": int(row["rank"]),62 "percentile": float(row["percentile"]),63 "score": float(row["score"])64 }65 66 all_benchmark_ranks.append(bench_df)67 plot_data.append(bench_df[["classifier", "benchmark_type", "rank", "percentile"]])68 69 bench_ranks_json = os.path.join(results_dir, f"benchmark_ranks_all_classifiers{file_suffix}.json")70 with open(bench_ranks_json, "w") as f:71 json.dump(list(bench_ranks_dict.values()), f, indent=2)72 console.log(f"[green]Saved all benchmark ranks to {bench_ranks_json}[/green]")73 74 plot_rows = []75 for bench in bench_ranks_dict.values():76 for clf_name in results.keys():77 if clf_name in bench:78 plot_rows.append({79 "benchmark_id": bench["id"],80 "benchmark_type": bench["benchmark_type"],81 "classifier": clf_name,82 "rank": bench[clf_name]["rank"],83 "percentile": bench[clf_name]["percentile"],84 "score": bench[clf_name]["score"]85 })86 plot_df = pd.DataFrame(plot_rows)87 88 console.log("[yellow]Plotting benchmark sample ranks by classifier and benchmark type...[/yellow]")89 num_classifiers = len(results)90 fig_width = max(16, num_classifiers * 2.5) # More width for better spacing91 92 # Create figure with white background93 fig, ax = plt.subplots(figsize=(fig_width, 11), facecolor='white')94 ax.set_facecolor('#f8f9fa')95 96 # Use standard, easily distinguishable colors97 # Using tab10 and Set1 for better distinction98 standard_colors = [99 '#1f77b4', # blue100 '#ff7f0e', # orange101 '#2ca02c', # green102 '#d62728', # red103 '#9467bd', # purple104 '#8c564b', # brown105 '#e377c2', # pink106 '#7f7f7f', # gray107 '#bcbd22', # olive108 '#17becf', # cyan109 ]110 111 ax = sns.stripplot(112 data=plot_df,113 x="classifier",114 y="rank",115 hue="benchmark_type",116 dodge=True,117 jitter=0.3,118 size=13,119 alpha=0.75,120 linewidth=1.5,121 edgecolor="white",122 palette=standard_colors,123 ax=ax124 )125 126 # Title and labels127 plt.title(128 f"Benchmark Sample Ranks by Classifier\n{num_docs:,} Documents from {dataset_name} • {mode_suffix.capitalize()} Mode", 129 fontsize=18, 130 fontweight='bold', 131 pad=25,132 color='#2c3e50'133 )134 plt.xlabel("Classifier", fontsize=16, fontweight='bold', color='#34495e', labelpad=12)135 plt.ylabel("Rank (0 = best)", fontsize=15, fontweight='semibold', color='#34495e', labelpad=10)136 137 # Make x-axis labels bigger and more readable138 plt.xticks(rotation=45, ha='right', fontsize=14, fontweight='bold')139 plt.yticks(fontsize=12)140 141 # Invert y-axis so 0 is at the top (best rank)142 ax.invert_yaxis()143 144 # Enhanced legend145 plt.legend(146 title="Benchmark Type", 147 title_fontsize=13,148 bbox_to_anchor=(1.01, 1), 149 loc='upper left', 150 frameon=True, 151 shadow=True, 152 fontsize=12,153 fancybox=True,154 edgecolor='#bdc3c7'155 )156 157 # Grid styling158 plt.grid(axis='y', alpha=0.4, linestyle='--', linewidth=0.8, color='#95a5a6')159 160 # Add vertical lines between classifiers for better separation161 for i in range(len(plot_df['classifier'].unique()) - 1):162 plt.axvline(x=i + 0.5, color='#bdc3c7', linestyle='-', linewidth=1.2, alpha=0.5)163 164 # Add subtle border165 for spine in ax.spines.values():166 spine.set_edgecolor('#bdc3c7')167 spine.set_linewidth(1.5)168 169 # Adjust layout to accommodate larger labels170 plt.tight_layout()171 plt.subplots_adjust(bottom=0.15)172 173 plot_path = os.path.join(results_dir, f"benchmark_ranks_by_classifier{file_suffix}.png")174 plt.savefig(plot_path, dpi=300, bbox_inches='tight', facecolor='white', edgecolor='none')175 plt.close()176 console.log(f"[bold green]Saved plot to {plot_path}[/bold green]")177 178 # Create figure with white background for percentiles179 fig, ax = plt.subplots(figsize=(fig_width, 11), facecolor='white')180 ax.set_facecolor('#f8f9fa')181 182 # Use the same standard colors for consistency183 ax = sns.stripplot(184 data=plot_df,185 x="classifier",186 y="percentile",187 hue="benchmark_type",188 dodge=True,189 jitter=0.3,190 size=13,191 alpha=0.75,192 linewidth=1.5,193 edgecolor="white",194 palette=standard_colors,195 ax=ax196 )197 198 # Title and labels199 plt.title(200 f"Benchmark Sample Percentiles by Classifier\n{num_docs:,} Documents from {dataset_name} • {mode_suffix.capitalize()} Mode", 201 fontsize=18, 202 fontweight='bold', 203 pad=25,204 color='#2c3e50'205 )206 plt.xlabel("Classifier", fontsize=16, fontweight='bold', color='#34495e', labelpad=12)207 plt.ylabel("Percentile (higher is better)", fontsize=15, fontweight='semibold', color='#34495e', labelpad=10)208 209 # Make x-axis labels bigger and more readable210 plt.xticks(rotation=45, ha='right', fontsize=14, fontweight='bold')211 plt.yticks(fontsize=12)212 213 # Enhanced legend214 plt.legend(215 title="Benchmark Type", 216 title_fontsize=13,217 bbox_to_anchor=(1.01, 1), 218 loc='upper left', 219 frameon=True, 220 shadow=True, 221 fontsize=12,222 fancybox=True,223 edgecolor='#bdc3c7'224 )225 226 # Grid styling227 plt.grid(axis='y', alpha=0.4, linestyle='--', linewidth=0.8, color='#95a5a6')228 229 # Add vertical lines between classifiers for better separation230 for i in range(len(plot_df['classifier'].unique()) - 1):231 plt.axvline(x=i + 0.5, color='#bdc3c7', linestyle='-', linewidth=1.2, alpha=0.5)232 233 # Add subtle border234 for spine in ax.spines.values():235 spine.set_edgecolor('#bdc3c7')236 spine.set_linewidth(1.5)237 238 # Adjust layout to accommodate larger labels239 plt.tight_layout()240 plt.subplots_adjust(bottom=0.15)241 242 plot_path_pct = os.path.join(results_dir, f"benchmark_percentiles_by_classifier{file_suffix}.png")243 plt.savefig(plot_path_pct, dpi=300, bbox_inches='tight', facecolor='white', edgecolor='none')244 plt.close()245 console.log(f"[bold green]Saved plot to {plot_path_pct}[/bold green]")246 247def load_cache_data(cache_dir: str, dataset_name: str = None):248 """Load cached classifier results from JSON files.249 250 Args:251 cache_dir: Base cache directory (e.g., 'cache')252 dataset_name: Name of dataset subfolder (e.g., 'fineweb'). If None, auto-detect.253 254 Returns:255 results: Dictionary mapping classifier names to list of score dictionaries256 num_docs: Total number of documents257 inject_inside: Whether benchmarks were injected (inferred from data)258 """259 cache_path = Path(cache_dir)260 261 # Auto-detect dataset subfolder if not specified262 if dataset_name is None:263 subdirs = [d for d in cache_path.iterdir() if d.is_dir() and d.name != 'old']264 if not subdirs:265 raise ValueError(f"No dataset subdirectories found in {cache_dir}")266 if len(subdirs) > 1:267 console.log(f"[yellow]Multiple datasets found: {[d.name for d in subdirs]}[/yellow]")268 console.log(f"[yellow]Using: {subdirs[0].name}[/yellow]")269 dataset_path = subdirs[0]270 dataset_name = dataset_path.name271 else:272 dataset_path = cache_path / dataset_name273 if not dataset_path.exists():274 raise ValueError(f"Dataset directory not found: {dataset_path}")275 276 console.log(f"[cyan]Loading cache from: {dataset_path}[/cyan]")277 278 # Find all classifier JSON files279 json_files = list(dataset_path.glob("*Classifier.json"))280 if not json_files:281 raise ValueError(f"No classifier JSON files found in {dataset_path}")282 283 console.log(f"[green]Found {len(json_files)} classifier cache files[/green]")284 285 results = {}286 num_docs = 0287 288 for json_file in sorted(json_files):289 classifier_name = json_file.stem # e.g., "DCLMClassifier"290 console.log(f"[yellow]Loading {classifier_name}...[/yellow]")291 292 with open(json_file, 'r') as f:293 cache_data = json.load(f)294 295 # Convert cache format to results format296 scores_list = []297 for doc_hash, doc_data in cache_data.items():298 scores_list.append({299 'doc_hash': doc_hash,300 'id': doc_data['id'],301 'source': doc_data['source'],302 'contains_benchmark': doc_data['contains_benchmark'],303 'benchmark_type': doc_data.get('benchmark_type'),304 'benchmark_index': doc_data.get('benchmark_index'),305 'score': doc_data['score']306 })307 308 results[classifier_name] = scores_list309 num_docs = max(num_docs, len(scores_list))310 console.log(f"[green] → Loaded {len(scores_list)} documents[/green]")311 312 # Infer inject_inside from data (check if any fineweb docs contain benchmarks)313 inject_inside = False314 for scores in results.values():315 for doc in scores:316 if doc['source'] == 'fineweb' and doc['contains_benchmark']:317 inject_inside = True318 break319 if inject_inside:320 break321 322 console.log(f"[cyan]Total documents: {num_docs}[/cyan]")323 console.log(f"[cyan]Mode: {'injected' if inject_inside else 'separate'}[/cyan]")324 console.log(f"[cyan]Dataset: {dataset_name}[/cyan]")325 326 return results, num_docs, inject_inside, dataset_name327 328def main():329 """Run analysis standalone from cached data."""330 parser = argparse.ArgumentParser(331 description="Generate analysis plots from cached classifier results"332 )333 parser.add_argument(334 '--cache-dir',335 type=str,336 default='cache',337 help='Base cache directory (default: cache)'338 )339 parser.add_argument(340 '--dataset',341 type=str,342 default=None,343 help='Dataset subfolder name (e.g., fineweb). Auto-detect if not specified.'344 )345 parser.add_argument(346 '--output-dir',347 type=str,348 default='results',349 help='Output directory for plots (default: results)'350 )351 parser.add_argument(352 '--config',353 type=str,354 default='config.yaml',355 help='Config file for additional settings (default: config.yaml)'356 )357 358 args = parser.parse_args()359 360 console.rule("[bold blue]Standalone Analysis Mode[/bold blue]")361 362 # Load cached data363 try:364 results, num_docs, inject_inside, dataset_name = load_cache_data(args.cache_dir, args.dataset)365 except Exception as e:366 console.log(f"[bold red]Error loading cache: {e}[/bold red]")367 return 1368 369 # Try to load config for prefilter_hq setting370 prefilter_hq = False371 if os.path.exists(args.config):372 try:373 import yaml374 with open(args.config, 'r') as f:375 config = yaml.safe_load(f)376 prefilter_hq = config.get('dataset', {}).get('prefilter_hq', False)377 except Exception as e:378 console.log(f"[yellow]Could not load config: {e}. Using defaults.[/yellow]")379 380 # Generate plots (benchmark_positions not needed for plotting)381 analyze_and_plot(382 results=results,383 documents=None, # Not needed for plotting from cache384 benchmark_positions={}, # Not needed for plotting from cache385 output_base_dir=args.output_dir,386 inject_inside=inject_inside,387 prefilter_hq=prefilter_hq,388 num_docs=num_docs,389 dataset_name=dataset_name390 )391 392 console.rule("[bold green]Analysis completed successfully![/bold green]")393 return 0394 395if __name__ == "__main__":396 exit(main())397