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

chapter7/elo-leaderboard/optimized_elo.py

@@ -0,0 +1,303 @@
+"""
+Optimized Elo rating system using NumPy vectorization and Numba JIT
+"""
+import numpy as np
+import pandas as pd
+from typing import Dict, Tuple, List
+from numba import jit
+from tqdm import tqdm
+
+
+@jit(nopython=True)
+def expected_score_fast(rating_a: float, rating_b: float) -> float:
+    """
+    Fast expected score calculation using Numba JIT.
+    
+    Args:
+        rating_a: Rating of model A
+        rating_b: Rating of model B
+        
+    Returns:
+        Expected probability that A wins
+    """
+    return 1.0 / (1.0 + 10.0 ** ((rating_b - rating_a) / 400.0))
+
+
+@jit(nopython=True)
+def process_elo_updates_vectorized(ratings: np.ndarray,
+                                   model_a_indices: np.ndarray,
+                                   model_b_indices: np.ndarray,
+                                   outcomes: np.ndarray,
+                                   k_factor: float,
+                                   match_counts: np.ndarray,
+                                   win_counts: np.ndarray) -> np.ndarray:
+    """
+    Process Elo updates using vectorized NumPy operations with Numba JIT.
+    
+    This is the core hot loop optimized with Numba for maximum performance.
+    
+    Args:
+        ratings: Array of current ratings for all models
+        model_a_indices: Indices of model A for each match
+        model_b_indices: Indices of model B for each match
+        outcomes: Match outcomes (1.0 = A wins, 0.0 = B wins, 0.5 = tie)
+        k_factor: Elo K-factor
+        match_counts: Array to track match counts per model
+        win_counts: Array to track win counts per model
+        
+    Returns:
+        Updated ratings array
+    """
+    n_matches = len(model_a_indices)
+    
+    for i in range(n_matches):
+        idx_a = model_a_indices[i]
+        idx_b = model_b_indices[i]
+        outcome = outcomes[i]
+        
+        # Get current ratings
+        rating_a = ratings[idx_a]
+        rating_b = ratings[idx_b]
+        
+        # Calculate expected scores
+        expected_a = 1.0 / (1.0 + 10.0 ** ((rating_b - rating_a) / 400.0))
+        expected_b = 1.0 - expected_a
+        
+        # Update ratings
+        ratings[idx_a] += k_factor * (outcome - expected_a)
+        ratings[idx_b] += k_factor * ((1.0 - outcome) - expected_b)
+        
+        # Update counts
+        match_counts[idx_a] += 1
+        match_counts[idx_b] += 1
+        win_counts[idx_a] += outcome
+        win_counts[idx_b] += (1.0 - outcome)
+    
+    return ratings
+
+
+@jit(nopython=True)
+def calculate_expected_scores_vectorized(ratings_a: np.ndarray, 
+                                        ratings_b: np.ndarray) -> np.ndarray:
+    """
+    Vectorized calculation of expected scores for multiple matches.
+    
+    Args:
+        ratings_a: Array of ratings for model A
+        ratings_b: Array of ratings for model B
+        
+    Returns:
+        Array of expected scores for model A
+    """
+    return 1.0 / (1.0 + np.power(10.0, (ratings_b - ratings_a) / 400.0))
+
+
+class NumpyEloRatingSystem:
+    """
+    Highly optimized Elo rating system using NumPy arrays and Numba JIT.
+    
+    Optimizations:
+    - NumPy arrays for O(1) indexing instead of dictionary lookups
+    - Numba JIT compilation of hot loops
+    - Pre-allocated arrays to avoid memory reallocation
+    - Integer indexing for models instead of string lookups
+    """
+    
+    def __init__(self, initial_rating: float = 1000.0, k_factor: float = 4.0):
+        """Initialize NumPy-based Elo system."""
+        self.initial_rating = initial_rating
+        self.k_factor = k_factor
+        
+        # Model name to index mapping
+        self.model_to_idx: Dict[str, int] = {}
+        self.idx_to_model: Dict[int, str] = {}
+        
+        # NumPy arrays for fast access
+        self.ratings: np.ndarray = None
+        self.match_counts: np.ndarray = None
+        self.win_counts: np.ndarray = None
+        
+        self.n_models = 0
+    
+    def _prepare_data(self, df: pd.DataFrame):
+        """
+        Prepare NumPy arrays from DataFrame for fast processing.
+        
+        Args:
+            df: DataFrame with columns 'model_a', 'model_b', 'winner'
+        """
+        print("Preparing data structures...")
+        
+        # Get all unique models
+        all_models = sorted(set(df['model_a'].unique()) | set(df['model_b'].unique()))
+        self.n_models = len(all_models)
+        
+        print(f"Found {self.n_models} unique models")
+        
+        # Create model mappings
+        for idx, model in enumerate(all_models):
+            self.model_to_idx[model] = idx
+            self.idx_to_model[idx] = model
+        
+        # Initialize arrays
+        self.ratings = np.full(self.n_models, self.initial_rating, dtype=np.float64)
+        self.match_counts = np.zeros(self.n_models, dtype=np.int32)
+        self.win_counts = np.zeros(self.n_models, dtype=np.float64)
+        
+        # Convert DataFrame columns to NumPy arrays with integer indices
+        print("Converting model names to indices...")
+        model_a_indices = df['model_a'].map(self.model_to_idx).values.astype(np.int32)
+        model_b_indices = df['model_b'].map(self.model_to_idx).values.astype(np.int32)
+        
+        # Convert outcomes to numeric (1.0 for A wins, 0.0 for B wins, 0.5 for tie)
+        print("Converting outcomes to numeric...")
+        # 'tie (bothbad)' is a real Arena outcome; unmapped values become NaN and
+        # would silently poison every rating they touch, so fall back to a tie
+        # (same as EloRatingSystem.update_ratings).
+        outcome_map = {'model_a': 1.0, 'model_b': 0.0,
+                       'tie': 0.5, 'tie (bothbad)': 0.5}
+        outcomes = df['winner'].map(outcome_map).fillna(0.5).values.astype(np.float64)
+        
+        return model_a_indices, model_b_indices, outcomes
+    
+    def process_matches_vectorized(self, df: pd.DataFrame, show_progress: bool = True):
+        """
+        Process all matches using vectorized NumPy operations and Numba JIT.
+        
+        This is the fastest way to compute Elo ratings for large datasets.
+        
+        Args:
+            df: DataFrame with columns 'model_a', 'model_b', 'winner'
+            show_progress: Whether to show progress bar
+        """
+        # Prepare data
+        model_a_indices, model_b_indices, outcomes = self._prepare_data(df)
+        
+        print(f"\nProcessing {len(df)} matches with NumPy + Numba JIT...")
+        
+        # Process all matches using JIT-compiled function
+        # This is where the magic happens - Numba compiles this to machine code
+        if show_progress:
+            # Process in chunks to show progress
+            chunk_size = 50000
+            n_chunks = (len(model_a_indices) + chunk_size - 1) // chunk_size
+            
+            for i in tqdm(range(n_chunks), desc="Processing matches"):
+                start_idx = i * chunk_size
+                end_idx = min((i + 1) * chunk_size, len(model_a_indices))
+                
+                self.ratings = process_elo_updates_vectorized(
+                    self.ratings,
+                    model_a_indices[start_idx:end_idx],
+                    model_b_indices[start_idx:end_idx],
+                    outcomes[start_idx:end_idx],
+                    self.k_factor,
+                    self.match_counts,
+                    self.win_counts
+                )
+        else:
+            self.ratings = process_elo_updates_vectorized(
+                self.ratings,
+                model_a_indices,
+                model_b_indices,
+                outcomes,
+                self.k_factor,
+                self.match_counts,
+                self.win_counts
+            )
+        
+        print("✓ Processing complete!")
+    
+    def get_leaderboard(self) -> List[Tuple]:
+        """
+        Get sorted leaderboard using NumPy's fast sorting.
+        
+        Returns:
+            List of tuples (model, rating, matches, wins)
+        """
+        # Use NumPy's argsort for fast sorting
+        sorted_indices = np.argsort(-self.ratings)  # Negative for descending order
+        
+        leaderboard = []
+        for idx in sorted_indices:
+            model = self.idx_to_model[idx]
+            rating = float(self.ratings[idx])
+            matches = int(self.match_counts[idx])
+            wins = float(self.win_counts[idx])
+            leaderboard.append((model, rating, matches, wins))
+        
+        return leaderboard
+    
+    def calculate_win_probability(self, model_a: str, model_b: str) -> float:
+        """
+        Calculate win probability using fast NumPy operations.
+        
+        Args:
+            model_a: First model identifier
+            model_b: Second model identifier
+            
+        Returns:
+            Probability that model_a wins
+        """
+        if model_a not in self.model_to_idx or model_b not in self.model_to_idx:
+            return 0.5
+        
+        idx_a = self.model_to_idx[model_a]
+        idx_b = self.model_to_idx[model_b]
+        
+        rating_a = self.ratings[idx_a]
+        rating_b = self.ratings[idx_b]
+        
+        return expected_score_fast(rating_a, rating_b)
+    
+    def get_win_rate_matrix(self) -> Dict[Tuple[str, str], float]:
+        """
+        Calculate pairwise win probability matrix using vectorized operations.
+        
+        Returns:
+            Dictionary mapping (model_a, model_b) to win probability
+        """
+        matrix = {}
+        
+        # Vectorized calculation for all pairs
+        for i in range(self.n_models):
+            model_a = self.idx_to_model[i]
+            
+            # Calculate win probabilities against all other models at once
+            ratings_a = np.full(self.n_models, self.ratings[i])
+            win_probs = calculate_expected_scores_vectorized(ratings_a, self.ratings)
+            
+            for j in range(self.n_models):
+                if i != j:
+                    model_b = self.idx_to_model[j]
+                    matrix[(model_a, model_b)] = float(win_probs[j])
+        
+        return matrix
+
+
+def build_leaderboard_optimized(df: pd.DataFrame,
+                                initial_rating: float = 1000.0,
+                                k_factor: float = 4.0,
+                                show_progress: bool = True) -> NumpyEloRatingSystem:
+    """
+    Build Elo leaderboard using highly optimized NumPy + Numba algorithm.
+    
+    This implementation is significantly faster than the basic version:
+    - Uses NumPy arrays for O(1) indexing
+    - Numba JIT compilation for hot loops
+    - Pre-allocated arrays to avoid memory overhead
+    - Integer-based model indexing instead of string lookups
+    
+    Args:
+        df: DataFrame with match data (columns: model_a, model_b, winner)
+        initial_rating: Starting rating for all models
+        k_factor: Elo learning rate (K-factor)
+        show_progress: Whether to display progress bar
+        
+    Returns:
+        NumpyEloRatingSystem with final ratings
+    """
+    elo = NumpyEloRatingSystem(initial_rating=initial_rating, k_factor=k_factor)
+    elo.process_matches_vectorized(df, show_progress=show_progress)
+    return elo
+