"""ELO rating system with HuggingFace dataset persistence.""" from __future__ import annotations import json import random import threading import logging from pathlib import Path from itertools import combinations from huggingface_hub import CommitScheduler, hf_hub_download from src.config import DEFAULT_ELO, ELO_K_FACTOR, HF_LOG_REPO_ID, HF_TOKEN, HF_LOG_EVERY_MINUTES from src.galaxy_profiles import GALAXY_IDS logger = logging.getLogger(__name__) STATE_DIR = Path("state") STATE_FILE = STATE_DIR / "elo_state.json" _lock = threading.Lock() _elo_ratings: dict[str, float] = {} # CommitScheduler for pushing state to HF _state_scheduler = None def _init_scheduler(): """Initialize the CommitScheduler for ELO state persistence.""" global _state_scheduler if not HF_LOG_REPO_ID: return STATE_DIR.mkdir(parents=True, exist_ok=True) _state_scheduler = CommitScheduler( repo_id=HF_LOG_REPO_ID, repo_type="dataset", folder_path=STATE_DIR, path_in_repo="state", every=HF_LOG_EVERY_MINUTES, token=HF_TOKEN if HF_TOKEN else None, ) logger.info("ELO state scheduler initialized (repo=%s)", HF_LOG_REPO_ID) def load_elo_state(): """Load ELO state from HF dataset, falling back to all-default.""" global _elo_ratings # Try downloading from HF if HF_LOG_REPO_ID: try: local_path = hf_hub_download( repo_id=HF_LOG_REPO_ID, repo_type="dataset", filename="state/elo_state.json", token=HF_TOKEN if HF_TOKEN else None, ) with open(local_path) as f: raw = json.load(f) # Only keep keys that match current galaxy IDs _elo_ratings = {k: v for k, v in raw.items() if k in GALAXY_IDS} logger.info("Loaded ELO state from HF dataset (%d galaxies)", len(_elo_ratings)) except Exception as e: logger.warning("Could not load ELO state from HF: %s. Starting fresh.", e) _elo_ratings = {} # Ensure every galaxy has a rating for gid in GALAXY_IDS: if gid not in _elo_ratings: _elo_ratings[gid] = DEFAULT_ELO # Initialize scheduler after loading state _init_scheduler() # Write initial state file so scheduler has something to push _save_state() def _save_state(): """Write current ELO state to local JSON file.""" STATE_DIR.mkdir(parents=True, exist_ok=True) with open(STATE_FILE, "w") as f: json.dump(_elo_ratings, f, indent=2) def _expected_score(rating_a: float, rating_b: float) -> float: """Calculate expected score for player A.""" return 1.0 / (1.0 + 10.0 ** ((rating_b - rating_a) / 400.0)) def record_comparison(winner_id: str, loser_id: str) -> dict: """Record a comparison result and update ELO ratings. Returns dict with before/after ratings for both galaxies. """ with _lock: elo_w_before = _elo_ratings.get(winner_id, DEFAULT_ELO) elo_l_before = _elo_ratings.get(loser_id, DEFAULT_ELO) expected_w = _expected_score(elo_w_before, elo_l_before) expected_l = _expected_score(elo_l_before, elo_w_before) elo_w_after = elo_w_before + ELO_K_FACTOR * (1.0 - expected_w) elo_l_after = elo_l_before + ELO_K_FACTOR * (0.0 - expected_l) _elo_ratings[winner_id] = elo_w_after _elo_ratings[loser_id] = elo_l_after _save_state() return { "winner_elo_before": elo_w_before, "winner_elo_after": elo_w_after, "loser_elo_before": elo_l_before, "loser_elo_after": elo_l_after, } def get_rating(galaxy_id: str) -> float: """Get current ELO rating for a galaxy.""" with _lock: return _elo_ratings.get(galaxy_id, DEFAULT_ELO) def get_leaderboard() -> list[dict]: """Get all galaxies sorted by ELO descending. Returns list of {id, elo} dicts. """ with _lock: snapshot = dict(_elo_ratings) return sorted( [{"id": gid, "elo": elo} for gid, elo in snapshot.items()], key=lambda x: x["elo"], reverse=True, ) def select_pair(seen_pairs: set[tuple[str, str]], champion_id: str | None = None) -> tuple[str, str] | None: """Select a pair of galaxies for comparison. If champion_id is provided, returns (champion_id, challenger_id). Otherwise, prefers galaxies with close ELO ratings (70%) or random (30%). Skips pairs already seen in this session. Returns None if all pairs exhausted. """ if champion_id is not None: # King of the hill mode: find challenger for champion challenger_candidates = [gid for gid in GALAXY_IDS if gid != champion_id] # Filter out seen challengers available_challengers = [ gid for gid in challenger_candidates if (champion_id, gid) not in seen_pairs and (gid, champion_id) not in seen_pairs ] if not available_challengers: return None if random.random() < 0.3: # Pure random challenger challenger = random.choice(available_challengers) else: # Prefer challenger with close ELO to champion with _lock: champion_elo = _elo_ratings.get(champion_id, DEFAULT_ELO) rated_challengers = [ (gid, abs(_elo_ratings.get(gid, DEFAULT_ELO) - champion_elo)) for gid in available_challengers ] rated_challengers.sort(key=lambda x: x[1]) # Pick from top 20% closest top_n = max(1, len(rated_challengers) // 5) challenger = random.choice(rated_challengers[:top_n])[0] return (champion_id, challenger) # Original random pair selection logic all_pairs = list(combinations(GALAXY_IDS, 2)) # Normalize pair ordering for consistent comparison available = [ p for p in all_pairs if (p[0], p[1]) not in seen_pairs and (p[1], p[0]) not in seen_pairs ] if not available: return None if random.random() < 0.3: # Pure random pair = random.choice(available) else: # Prefer close ELO ratings with _lock: rated = [(p, abs(_elo_ratings.get(p[0], DEFAULT_ELO) - _elo_ratings.get(p[1], DEFAULT_ELO))) for p in available] rated.sort(key=lambda x: x[1]) # Pick from top 20% closest top_n = max(1, len(rated) // 5) pair = random.choice(rated[:top_n])[0] # Randomize left/right for initial random pairs if random.random() < 0.5: return (pair[1], pair[0]) return pair