generate_initial_leaderboard.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Скрипт для генерации первоначального лидерборда DeathMath и загрузки данных в HuggingFace.
Использует результаты из директории results и загружает их в репозиторий Vikhrmodels/DeathMath-leaderboard-data.
"""

import os
import json
import logging
import pandas as pd
import argparse
from pathlib import Path
from huggingface_hub import HfApi, create_repo
from datetime import datetime

# Настройка логирования
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.FileHandler("leaderboard_generation.log"), logging.StreamHandler()],
)
logger = logging.getLogger(__name__)

# Константы
REPO_ID = "Vikhrmodels/DeathMath-leaderboard-data"
METAINFO_REPO_ID = "Vikhrmodels/DeathMath-leaderboard-metainfo"


def setup_repositories(token):
    """
    Создает необходимые репозитории на HuggingFace Hub, если они еще не существуют.

    Args:
        token (str): Токен для доступа к HuggingFace Hub
    """
    api = HfApi(token=token)

    try:
        # Проверка и создание репозитория для данных лидерборда
        try:
            api.repo_info(repo_id=REPO_ID, repo_type="dataset")
            logger.info(f"Репозиторий {REPO_ID} уже существует")
        except Exception:
            logger.info(f"Создание репозитория для данных лидерборда: {REPO_ID}")
            create_repo(repo_id=REPO_ID, repo_type="dataset", token=token)

        # Проверка и создание репозитория для метаданных лидерборда
        try:
            api.repo_info(repo_id=METAINFO_REPO_ID, repo_type="dataset")
            logger.info(f"Репозиторий {METAINFO_REPO_ID} уже существует")
        except Exception:
            logger.info(f"Создание репозитория для метаданных лидерборда: {METAINFO_REPO_ID}")
            create_repo(repo_id=METAINFO_REPO_ID, repo_type="dataset", token=token)

        return api
    except Exception as e:
        logger.error(f"Ошибка при создании репозиториев: {e}")
        raise


def load_results(results_file):
    """
    Загружает результаты из JSON файла и удаляет дубликаты.

    Args:
        results_file (str): Путь к файлу с результатами

    Returns:
        list: Список записей для лидерборда без дубликатов
    """
    try:
        with open(results_file, "r", encoding="utf-8") as f:
            data = json.load(f)

        leaderboard_entries = []
        seen_models = set()  # Множество для отслеживания уже обработанных моделей

        for key, value in data.items():
            if "_Combined_" in key:  # берем только комбинированные результаты
                model_name = value["model_name"]

                # Пропускаем модель, если она уже была добавлена
                if model_name in seen_models:
                    logger.info(f"Пропускаем дублирующуюся модель: {model_name}")
                    continue

                # Добавляем модель во множество обработанных
                seen_models.add(model_name)

                leaderboard_entry = {
                    "model_name": model_name,
                    "score": value["score"],
                    "math_score": value["math_score"],
                    "physics_score": value["physics_score"],
                    "total_tokens": value["total_tokens"],
                    "evaluation_time": value["evaluation_time"],
                    "system_prompt": value.get(
                        "system_prompt", "Вы - полезный помощник по математике и физике. Ответьте на русском языке."
                    ),
                }
                leaderboard_entries.append(leaderboard_entry)

        # Сортировка по общему баллу
        leaderboard_entries.sort(key=lambda x: x["score"], reverse=True)
        logger.info(f"Загружено {len(leaderboard_entries)} уникальных моделей после удаления дубликатов")
        return leaderboard_entries

    except Exception as e:
        logger.error(f"Ошибка при загрузке результатов: {e}")
        raise


def prepare_directory_structure():
    """
    Создает необходимую структуру директорий для внешних моделей.

    Returns:
        str: Путь к временной директории с подготовленной структурой
    """
    temp_dir = Path("./temp_leaderboard")
    model_data_dir = temp_dir / "model_data" / "external"

    # Очистка и создание директорий
    if temp_dir.exists():
        import shutil

        shutil.rmtree(temp_dir)

    model_data_dir.mkdir(parents=True, exist_ok=True)

    return str(temp_dir)


def upload_model_files(api, leaderboard_entries, temp_dir):
    """
    Загружает файлы моделей в репозиторий данных лидерборда.

    Args:
        api (HfApi): Экземпляр API для взаимодействия с HuggingFace
        leaderboard_entries (list): Список записей для лидерборда
        temp_dir (str): Путь к временной директории
    """
    model_data_dir = os.path.join(temp_dir, "model_data", "external")

    for entry in leaderboard_entries:
        model_name = entry["model_name"]
        safe_filename = model_name.replace("/", "_").replace(" ", "_")
        file_path = os.path.join(model_data_dir, f"{safe_filename}.json")

        with open(file_path, "w", encoding="utf-8") as f:
            json.dump(entry, f, ensure_ascii=False, indent=2)

        # Загрузка файла модели в репозиторий
        api.upload_file(
            path_or_fileobj=file_path,
            path_in_repo=f"model_data/external/{safe_filename}.json",
            repo_id=REPO_ID,
            repo_type="dataset",
        )
        logger.info(f"Загружен файл модели: {safe_filename}.json")


def generate_leaderboard_json(leaderboard_entries):
    """
    Создает JSON файл с данными лидерборда.

    Args:
        leaderboard_entries (list): Список записей для лидерборда

    Returns:
        str: Путь к созданному JSON файлу
    """
    leaderboard_file = "leaderboard.json"

    with open(leaderboard_file, "w", encoding="utf-8") as f:
        json.dump(leaderboard_entries, f, ensure_ascii=False, indent=2)

    return leaderboard_file


def generate_readme(leaderboard_entries):
    """
    Генерирует README.md с информацией о лидерборде.

    Args:
        leaderboard_entries (list): Список записей для лидерборда

    Returns:
        str: Путь к созданному README файлу
    """
    readme_file = "README.md"

    # Создаем DataFrame для удобного форматирования таблицы
    df = pd.DataFrame(leaderboard_entries)

    # Форматируем числовые колонки
    for col in ["score", "math_score", "physics_score"]:
        if col in df.columns:
            df[col] = df[col].apply(lambda x: f"{x:.3f}")

    if "total_tokens" in df.columns:
        df["total_tokens"] = df["total_tokens"].apply(lambda x: f"{int(x):,}")

    if "evaluation_time" in df.columns:
        df["evaluation_time"] = df["evaluation_time"].apply(lambda x: f"{x:.1f}s")

    # Создаем содержимое README
    current_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    readme_content = f"""# DeathMath Leaderboard

DeathMath - это бенчмарк для оценки способности моделей решать сложные математические и физические задачи на русском языке.

## Текущий лидерборд

Последнее обновление: {current_date}

| Модель | Общий балл | Математика | Физика | Токены | Время оценки |
|--------|------------|------------|---------|---------|--------------|
"""

    # Добавляем строки таблицы
    for _, row in df.iterrows():
        readme_content += f"| {row['model_name']} | {row['score']} | {row['math_score']} | {row['physics_score']} | {row.get('total_tokens', 'N/A')} | {row.get('evaluation_time', 'N/A')} |\n"

    readme_content += """
## Как принять участие в бенчмарке

Для участия в бенчмарке DeathMath:

1. Клонируйте репозиторий и запустите тесты вашей модели
2. Загрузите результаты через [HuggingFace Space](https://huggingface.co/spaces/Vikhrmodels/DeathMath-leaderboard)
3. Дождитесь проверки и добавления результатов в лидерборд

## Формат результатов

Результаты должны быть в формате JSON со следующей структурой:
```json
{
  "score": 0.586,
  "math_score": 0.8,
  "physics_score": 0.373,
  "total_tokens": 1394299,
  "evaluation_time": 4533.2,
  "system_prompt": "Вы - полезный помощник по математике и физике. Ответьте на русском языке."
}
```

## Лицензия

Бенчмарк распространяется под лицензией Apache 2.0
"""

    with open(readme_file, "w", encoding="utf-8") as f:
        f.write(readme_content)

    return readme_file


def upload_leaderboard_files(api, leaderboard_file, readme_file):
    """
    Загружает файлы лидерборда в репозиторий метаданных.

    Args:
        api (HfApi): Экземпляр API для взаимодействия с HuggingFace
        leaderboard_file (str): Путь к JSON файлу лидерборда
        readme_file (str): Путь к README файлу
    """
    # Загрузка JSON лидерборда
    api.upload_file(
        path_or_fileobj=leaderboard_file, path_in_repo="leaderboard.json", repo_id=METAINFO_REPO_ID, repo_type="dataset"
    )
    logger.info(f"Загружен файл лидерборда: leaderboard.json в {METAINFO_REPO_ID}")

    # Загрузка README
    api.upload_file(
        path_or_fileobj=readme_file, path_in_repo="README.md", repo_id=METAINFO_REPO_ID, repo_type="dataset"
    )
    logger.info(f"Загружен README: README.md в {METAINFO_REPO_ID}")


def main():
    # Парсинг аргументов командной строки
    parser = argparse.ArgumentParser(description="Генерация первоначального лидерборда DeathMath")
    parser.add_argument(
        "--results",
        default="../results/leaderboard_results.json",
        help="Путь к файлу с результатами (по умолчанию: ../results/leaderboard_results.json)",
    )
    parser.add_argument("--token", required=True, help="Токен для доступа к HuggingFace Hub")

    args = parser.parse_args()

    try:
        logger.info("Начинаем генерацию лидерборда DeathMath")

        # Настраиваем репозитории
        api = setup_repositories(args.token)
        logger.info("Репозитории успешно настроены")

        # Загружаем результаты
        leaderboard_entries = load_results(args.results)
        logger.info(f"Загружено {len(leaderboard_entries)} записей для лидерборда")

        # Подготавливаем структуру директорий
        temp_dir = prepare_directory_structure()
        logger.info(f"Создана временная директория: {temp_dir}")

        # Загружаем файлы моделей
        upload_model_files(api, leaderboard_entries, temp_dir)
        logger.info("Файлы моделей успешно загружены")

        # Генерируем JSON лидерборда
        leaderboard_file = generate_leaderboard_json(leaderboard_entries)
        logger.info(f"Создан файл лидерборда: {leaderboard_file}")

        # Генерируем README
        readme_file = generate_readme(leaderboard_entries)
        logger.info(f"Создан README: {readme_file}")

        # Загружаем файлы лидерборда
        upload_leaderboard_files(api, leaderboard_file, readme_file)
        logger.info("Файлы лидерборда успешно загружены")

        logger.info("Генерация лидерборда успешно завершена!")

    except Exception as e:
        logger.error(f"Ошибка при генерации лидерборда: {e}")
        raise


if __name__ == "__main__":
    main()