n.exe e:/LLMS/Fine-tuning/arena/battlemode.py
2024-10-15 10:34:54,527 [INFO] Starting download from https://storage.googleapis.com/arena_external_data/public/clean_battle_20240814_public.json
local_battle_data.json: 100%|███████████████████████████████████| 1.99G/1.99G [02:19<00:00, 15.3MiB/s]
2024-10-15 10:37:14,804 [INFO] Data downloaded successfully and saved to local_battle_data.json
2024-10-15 10:37:14,819 [INFO] Loading data from local_battle_data.json
2024-10-15 10:37:16,020 [ERROR] Process terminated due to an error: module 'pandas.errors' has no attribute 'JSONDecodeError'
import os
import sys
import math
import json
import logging
from typing import Any

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from tqdm import tqdm
import requests

# Configure pandas display options
pd.options.display.float_format = "{:.2f}".format

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)

# Constants
DATA_URL: str = "https://storage.googleapis.com/arena_external_data/public/clean_battle_20240814_public.json"
LOCAL_FILE_NAME: str = "local_battle_data.json"
OUTPUT_FOLDER: str = "output_plots"


def download_data(url: str, local_path: str) -> None:
    """
    Downloads JSON data from the specified URL and saves it locally.

    Args:
        url (str): The URL to download data from.
        local_path (str): The local file path to save the downloaded data.

    Raises:
        requests.HTTPError: If the HTTP request to download data fails.
        Exception: If any other exception occurs during the download.
    """
    try:
        logging.info(f"Starting download from {url}")
        response = requests.get(url, stream=True)
        response.raise_for_status()
        total_size = int(response.headers.get('content-length', 0))
        with open(local_path, 'wb') as file, tqdm(
            desc=local_path,
            total=total_size,
            unit='iB',
            unit_scale=True,
            unit_divisor=1024,
        ) as bar:
            for data in response.iter_content(chunk_size=1024):
                size = file.write(data)
                bar.update(size)
        logging.info(f"Data downloaded successfully and saved to {local_path}")
    except requests.HTTPError as http_err:
        logging.error(f"HTTP error occurred while downloading data: {http_err}")
        raise
    except Exception as err:
        logging.error(f"An error occurred while downloading data: {err}")
        raise


def load_data(local_path: str) -> pd.DataFrame:
    """
    Loads JSON data from a local file into a pandas DataFrame and sorts it by timestamp.

    Args:
        local_path (str): The local file path to load data from.

    Returns:
        pd.DataFrame: The loaded and sorted battles DataFrame.

    Raises:
        FileNotFoundError: If the local file does not exist.
        pd.errors.JSONDecodeError: If the file content is not valid JSON.
        Exception: If any other exception occurs during data loading.
    """
    try:
        logging.info(f"Loading data from {local_path}")
        with open(local_path, 'r') as file:
            battles = pd.read_json(file).sort_values(by="tstamp", ascending=True)
        logging.info(f"Data loaded successfully with {len(battles)} records")
        return battles
    except FileNotFoundError as fnf_err:
        logging.error(f"File not found: {fnf_err}")
        raise
    except pd.errors.JSONDecodeError as json_err:
        logging.error(f"JSON decoding error: {json_err}")
        raise
    except Exception as err:
        logging.error(f"An error occurred while loading data: {err}")
        raise


def create_advanced_scatter_plot(battles: pd.DataFrame, output_folder: str) -> None:
    """
    Creates an advanced scatter plot with interactive controls and saves it as an HTML file.

    The scatter plot visualizes battles with:
    - X-axis: Model A
    - Y-axis: Model B
    - Marker size varies according to the battle rank
    - Interactive play and pause buttons to animate over time

    Args:
        battles (pd.DataFrame): The battles DataFrame containing battle records.
        output_folder (str): The directory where the plot will be saved.

    Raises:
        Exception: If any error occurs during plot creation or saving.
    """
    try:
        logging.info("Creating advanced scatter plot")

        # Ensure output directory exists
        os.makedirs(output_folder, exist_ok=True)

        # Example: Suppose we have a 'rank' column to determine scatter sizes
        # If not present, we'll create a dummy 'rank' based on 'turn' as an example
        if 'rank' not in battles.columns:
            battles['rank'] = battles['turn'].rank(method='dense').astype(int)

        # Aggregate data as needed, for example by model pairs and timestamp
        aggregated = (
            battles.groupby(['model_a', 'model_b', 'tstamp'])
            .size()
            .reset_index(name='battle_count')
        )

        # Normalize rank for marker sizes
        aggregated['size'] = aggregated['rank'].apply(lambda x: 10 + (x * 5))

        # Create scatter plot
        fig = px.scatter(
            aggregated,
            x='model_a',
            y='model_b',
            animation_frame='tstamp',
            size='size',
            hover_name='battle_count',
            title='Battle Outcomes Over Time',
            labels={
                'model_a': 'Model A',
                'model_b': 'Model B',
                'tstamp': 'Timestamp',
                'size': 'Rank'
            },
            size_max=60,
            color='rank',
            color_continuous_scale='Viridis'
        )

        # Update layout for better readability
        fig.update_layout(
            title={
                'text': "Advanced Battle Scatter Plot",
                'x': 0.5,
                'xanchor': 'center'
            },
            xaxis_title="Model A",
            yaxis_title="Model B",
            legend_title="Rank",
            height=700
        )

        # Add play and pause buttons
        fig.update_layout(
            updatemenus=[
                {
                    "buttons": [
                        {
                            "args": [None, {"frame": {"duration": 500, "redraw": True},
                                            "fromcurrent": True, "transition": {"duration": 300}}],
                            "label": "Play",
                            "method": "animate"
                        },
                        {
                            "args": [[None], {"frame": {"duration": 0, "redraw": False},
                                              "mode": "immediate",
                                              "transition": {"duration": 0}}],
                            "label": "Pause",
                            "method": "animate"
                        }
                    ],
                    "direction": "left",
                    "pad": {"r": 10, "t": 87},
                    "showactive": False,
                    "type": "buttons",
                    "x": 0.1,
                    "xanchor": "right",
                    "y": 0,
                    "yanchor": "top"
                }
            ]
        )

        # Save the plot to the output folder
        output_path = os.path.join(output_folder, "advanced_scatter_plot.html")
        fig.write_html(output_path)
        logging.info(f"Advanced scatter plot saved to {output_path}")

    except Exception as err:
        logging.error(f"An error occurred while creating the scatter plot: {err}")
        raise


def main() -> None:
    """
    Main function to orchestrate data download, loading, visualization, and saving.

    Raises:
        Exception: Propagates exceptions from called functions.
    """
    try:
        # Step 1: Download the latest data
        download_data(DATA_URL, LOCAL_FILE_NAME)

        # Step 2: Load the data into a DataFrame
        battles_df = load_data(LOCAL_FILE_NAME)

        # Step 3: Create advanced scatter plot and save it
        create_advanced_scatter_plot(battles_df, OUTPUT_FOLDER)

        logging.info("All tasks completed successfully.")

    except Exception as err:
        logging.error(f"Process terminated due to an error: {err}")
        sys.exit(1)


if __name__ == "__main__":
    main()