ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter7/elo-leaderboard/data_loader.py

@@ -0,0 +1,223 @@
+"""
+Data loading and preprocessing for Chatbot Arena voting data
+"""
+import pandas as pd
+import requests
+import os
+from typing import Optional
+from tqdm import tqdm
+
+
+def download_arena_data(output_path: str = "arena_data.json", force_download: bool = False) -> str:
+    """
+    Download Chatbot Arena voting data via HTTPS.
+    
+    Args:
+        output_path: Path to save downloaded file
+        force_download: If True, re-download even if file exists
+        
+    Returns:
+        Path to downloaded file
+    """
+    if os.path.exists(output_path) and not force_download:
+        print(f"Data file already exists at {output_path}")
+        file_size = os.path.getsize(output_path) / (1024 * 1024)
+        print(f"File size: {file_size:.2f} MB")
+        return output_path
+    
+    print("Downloading Chatbot Arena voting data...")
+    url = "https://storage.googleapis.com/arena_external_data/public/clean_battle_20240814_public.json"
+    
+    try:
+        # Stream download with progress bar
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        
+        # Get total file size
+        total_size = int(response.headers.get('content-length', 0))
+        
+        # Download with progress bar
+        with open(output_path, 'wb') as f, tqdm(
+            desc=output_path,
+            total=total_size,
+            unit='B',
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as pbar:
+            for chunk in response.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+                    pbar.update(len(chunk))
+        
+        file_size = os.path.getsize(output_path) / (1024 * 1024)
+        print(f"\nDownloaded data to {output_path} ({file_size:.2f} MB)")
+        return output_path
+        
+    except Exception as e:
+        print(f"Error downloading data: {e}")
+        print("Please ensure you have internet connection and the URL is accessible.")
+        raise
+
+
+def load_arena_data(filepath: str) -> pd.DataFrame:
+    """
+    Load and preprocess Chatbot Arena voting data.
+    
+    Expected columns:
+    - model_a: Identifier for first model
+    - model_b: Identifier for second model
+    - winner: Which model won ('model_a', 'model_b', or 'tie')
+    - tstamp: Unix timestamp of the vote
+    - judge: User who made the vote
+    - turn: Conversation turn
+    - anony: Whether vote was anonymous (blind)
+    - language: Language of the conversation
+    
+    Args:
+        filepath: Path to data file
+        
+    Returns:
+        Preprocessed DataFrame sorted by timestamp
+    """
+    print(f"Loading data from {filepath}...")
+    print("Note: This is a large file (~2GB), loading may take 1-2 minutes...")
+    
+    # Try different file formats
+    if filepath.endswith('.json'):
+        df = pd.read_json(filepath)
+    elif filepath.endswith('.jsonl'):
+        df = pd.read_json(filepath, lines=True)
+    elif filepath.endswith('.csv'):
+        df = pd.read_csv(filepath)
+    else:
+        # Try JSON by default
+        try:
+            df = pd.read_json(filepath)
+        except (ValueError, KeyError):
+            df = pd.read_json(filepath, lines=True)
+    
+    print(f"Loaded {len(df)} records")
+    print(f"Columns: {df.columns.tolist()}")
+    
+    # Sort by timestamp
+    if 'tstamp' in df.columns:
+        df = df.sort_values('tstamp', ascending=True).reset_index(drop=True)
+        print(f"Data spans from {pd.to_datetime(df['tstamp'].min(), unit='s')} to {pd.to_datetime(df['tstamp'].max(), unit='s')}")
+    
+    # Basic statistics
+    if 'winner' in df.columns:
+        print(f"\nOutcome distribution:")
+        print(df['winner'].value_counts())
+    
+    if 'model_a' in df.columns and 'model_b' in df.columns:
+        all_models = set(df['model_a'].unique()) | set(df['model_b'].unique())
+        print(f"\nTotal unique models: {len(all_models)}")
+        print(f"Top 10 models by appearance:")
+        model_counts = pd.concat([df['model_a'], df['model_b']]).value_counts().head(10)
+        print(model_counts)
+    
+    return df
+
+
+def filter_data(df: pd.DataFrame, 
+                min_date: Optional[str] = None,
+                max_date: Optional[str] = None,
+                anony_only: bool = True,
+                language: Optional[str] = None,
+                min_turn: int = 1,
+                use_dedup: bool = True) -> pd.DataFrame:
+    """
+    Filter voting data based on various criteria (following official Chatbot Arena method).
+    
+    Args:
+        df: Input DataFrame
+        min_date: Minimum date (YYYY-MM-DD format)
+        max_date: Maximum date (YYYY-MM-DD format)
+        anony_only: If True, only include anonymous (blind) votes
+        language: If specified, filter by language
+        min_turn: Minimum conversation turn
+        use_dedup: If True, apply deduplication filter (official Arena method)
+        
+    Returns:
+        Filtered DataFrame
+    """
+    filtered = df.copy()
+    
+    print(f"Before filtering: {len(filtered)} records")
+    
+    # Filter by anonymous votes only (official method)
+    if anony_only and 'anony' in filtered.columns:
+        filtered = filtered[filtered['anony'] == True]
+        print(f"  After anony filter: {len(filtered)} records")
+    
+    # Apply deduplication (official method removes top 0.1% redundant prompts)
+    if use_dedup and 'dedup_tag' in filtered.columns:
+        try:
+            filtered = filtered[filtered["dedup_tag"].apply(lambda x: x.get("sampled", False) if isinstance(x, dict) else False)]
+            print(f"  After dedup filter: {len(filtered)} records")
+        except Exception as e:
+            print(f"  Warning: Could not apply dedup filter: {e}")
+    
+    # Filter by date
+    if 'tstamp' in filtered.columns:
+        if min_date:
+            min_timestamp = pd.to_datetime(min_date).timestamp()
+            filtered = filtered[filtered['tstamp'] >= min_timestamp]
+        if max_date:
+            max_timestamp = pd.to_datetime(max_date).timestamp()
+            filtered = filtered[filtered['tstamp'] <= max_timestamp]
+    
+    # Filter by language
+    if language and 'language' in filtered.columns:
+        filtered = filtered[filtered['language'] == language]
+    
+    # Filter by turn
+    if 'turn' in filtered.columns:
+        filtered = filtered[filtered['turn'] >= min_turn]
+    
+    pct = (len(filtered) / len(df) * 100) if len(df) else 0.0
+    print(f"After filtering: {len(filtered)} records ({pct:.1f}% of original)")
+    return filtered.reset_index(drop=True)
+
+
+def get_time_slices(df: pd.DataFrame, interval: str = 'W') -> list:
+    """
+    Split data into time slices for historical analysis.
+    
+    Args:
+        df: Input DataFrame with 'tstamp' column
+        interval: Pandas frequency string ('D' for daily, 'W' for weekly, 'M' for monthly)
+        
+    Returns:
+        List of (end_date, dataframe_slice) tuples
+    """
+    if 'tstamp' not in df.columns:
+        raise ValueError("DataFrame must have 'tstamp' column")
+
+    if len(df) == 0:
+        print(f"Created 0 time slices with interval '{interval}'")
+        return []
+    
+    df['datetime'] = pd.to_datetime(df['tstamp'], unit='s')
+    min_date = df['datetime'].min()
+    max_date = df['datetime'].max()
+    
+    # Pandas 2.2+ removed 'M' (month-end); keep the documented monthly alias.
+    freq = "ME" if interval == "M" else interval
+
+    # Generate date ranges
+    date_ranges = pd.date_range(start=min_date, end=max_date, freq=freq)
+    
+    slices = []
+    for end_date in date_ranges:
+        slice_df = df[df['datetime'] <= end_date].copy()
+        if len(slice_df) > 0:
+            slices.append((end_date, slice_df))
+    
+    # Empty date_ranges when span < interval; also cover trailing gap to max_date.
+    if len(date_ranges) == 0 or date_ranges[-1] < max_date:
+        slices.append((max_date, df.copy()))
+    
+    print(f"Created {len(slices)} time slices with interval '{interval}'")
+    return slices
+