llm-arch-research/experiments/shared/data.py

"""
Общие утилиты для работы с данными в экспериментах.
"""

import os
from typing import List, Tuple
from .configs import TRAIN_TEXTS, PATHS


def load_training_data(split_ratio: float = 0.8) -> Tuple[List[str], List[str]]:
    """
    Загружает данные для обучения и разделяет на train/validation.
    
    Args:
        split_ratio: Доля данных для обучения
        
    Returns:
        Tuple: (train_texts, val_texts)
    """
    train_size = int(len(TRAIN_TEXTS) * split_ratio)
    train_data = TRAIN_TEXTS[:train_size]
    val_data = TRAIN_TEXTS[train_size:]
    
    return train_data, val_data


def ensure_directories():
    """Создает необходимые директории если они не существуют."""
    directories = [
        "checkpoints",
        "checkpoints/gpt-bpe", 
        "checkpoints/hf-bpe-tokenizer",
        "checkpoints/hf-trained",
        "checkpoints/hf-trained-proxy",
        "logs"
    ]
    
    for directory in directories:
        os.makedirs(directory, exist_ok=True)


def get_model_paths(experiment_type: str = "llm_only") -> dict:
    """
    Возвращает пути для конкретного типа эксперимента.
    
    Args:
        experiment_type: Тип эксперимента ('llm_only' или 'hf_integration')
        
    Returns:
        dict: Словарь с путями
    """
    base_paths = PATHS.copy()
    
    if experiment_type == "hf_integration":
        base_paths.update({
            "model": base_paths["hf_model"],
            "tokenizer": base_paths["hf_tokenizer"]
        })
    else:  # llm_only
        base_paths.update({
            "model": base_paths["gpt_bpe_model"],
            "tokenizer": base_paths["bpe_tokenizer"]
        })
    
    return base_paths


def print_experiment_info(experiment_name: str, config: dict):
    """
    Выводит информацию о запускаемом эксперименте.
    
    Args:
        experiment_name: Название эксперимента
        config: Конфигурация эксперимента
    """
    print("=" * 70)
    print(f"🚀 Эксперимент: {experiment_name}")
    print("=" * 70)
    print("📊 Конфигурация:")
    for key, value in config.items():
        print(f"   {key}: {value}")
    print()


def save_experiment_results(results: dict, filepath: str):
    """
    Сохраняет результаты эксперимента в файл.
    
    Args:
        results: Словарь с результатами
        filepath: Путь для сохранения
    """
    import json
    
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    
    print(f"✅ Результаты эксперимента сохранены: {filepath}")


def load_experiment_results(filepath: str) -> dict:
    """
    Загружает результаты эксперимента из файла.
    
    Args:
        filepath: Путь к файлу с результатами
        
    Returns:
        dict: Загруженные результаты
    """
    import json
    
    if not os.path.exists(filepath):
        return {}
    
    with open(filepath, 'r', encoding='utf-8') as f:
        return json.load(f)


class ExperimentLogger:
    """
    Логгер для экспериментов.
    """
    
    def __init__(self, experiment_name: str):
        self.experiment_name = experiment_name
        self.metrics = {}
    
    def log_metric(self, name: str, value: float):
        """Логирует метрику."""
        if name not in self.metrics:
            self.metrics[name] = []
        self.metrics[name].append(value)
        print(f"📈 {name}: {value:.4f}")
    
    def log_step(self, step: int, loss: float, **kwargs):
        """Логирует шаг обучения."""
        print(f"📊 Step {step}: loss={loss:.4f}", end="")
        for key, value in kwargs.items():
            print(f", {key}={value:.4f}", end="")
        print()
    
    def log_epoch(self, epoch: int, train_loss: float, val_loss: float = None):
        """Логирует завершение эпохи."""
        print(f"🎯 Epoch {epoch}: train_loss={train_loss:.4f}", end="")
        if val_loss is not None:
            print(f", val_loss={val_loss:.4f}", end="")
        print()
    
    def save_logs(self, filepath: str):
        """Сохраняет логи эксперимента."""
        import json
        
        logs = {
            "experiment_name": self.experiment_name,
            "metrics": self.metrics
        }
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(logs, f, ensure_ascii=False, indent=2)
        
        print(f"✅ Логи эксперимента сохранены: {filepath}")