llm-arch-research/experiments/hf_integration/simple_hf_training.py

#!/usr/bin/env python3
"""
Experiment: simple_hf_training.py
Description: Упрощенное обучение GPT модели с использованием hf-proxy.
Использует ручное обучение вместо сложного HuggingFace Trainer.
"""

import torch
import torch.nn as nn
import os
import sys
import json

# Добавляем путь к shared модулям
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from llm.models.gpt import GPT
from llm.tokenizers import BPETokenizer
from hf_proxy import HFAdapter, HFTokenizerAdapter

from shared.configs import (
    TRAIN_TEXTS,
    BASE_GPT_CONFIG,
    BPE_CONFIG,
    TRAINING_CONFIG,
    PATHS,
    TEST_PROMPTS,
)


def create_dataset(hf_tokenizer, texts, max_length=128):
    """
    Создает простой датасет для обучения.

    Args:
        hf_tokenizer: Адаптированный токенизатор
        texts: Список текстов
        max_length: Максимальная длина последовательности

    Returns:
        list: Список тензоров input_ids
    """
    dataset = []

    for text in texts:
        # Токенизируем текст
        inputs = hf_tokenizer(
            text,
            max_length=max_length,
            truncation=True,
            padding=False,
            return_tensors="pt",
        )

        input_ids = inputs["input_ids"][0]

        # Создаем метки для языкового моделирования
        labels = input_ids.clone()

        dataset.append({"input_ids": input_ids, "labels": labels})

    return dataset


def manual_training_loop(hf_model, hf_tokenizer, train_texts, val_texts, config):
    """
    Ручной цикл обучения без использования Trainer.

    Args:
        hf_model: Адаптированная модель
        hf_tokenizer: Адаптированный токенизатор
        train_texts: Тексты для обучения
        val_texts: Тексты для валидации
        config: Конфигурация обучения

    Returns:
        dict: Результаты обучения
    """
    print("🎯 Запуск ручного обучения...")

    # Создаем датасеты
    train_dataset = create_dataset(hf_tokenizer, train_texts)
    val_dataset = create_dataset(hf_tokenizer, val_texts)

    print(f"📊 Данные: {len(train_dataset)} train, {len(val_dataset)} validation")

    # Оптимизатор
    optimizer = torch.optim.AdamW(hf_model.parameters(), lr=config["learning_rate"])

    # Функция потерь
    loss_fn = nn.CrossEntropyLoss()

    # Обучение
    hf_model.train()
    train_losses = []
    val_losses = []

    for epoch in range(config["num_epochs"]):
        print(f"\n📅 Эпоха {epoch + 1}/{config['num_epochs']}")

        # Обучение
        epoch_train_loss = 0
        for i, batch in enumerate(train_dataset):
            optimizer.zero_grad()

            input_ids = batch["input_ids"].unsqueeze(0)  # [1, seq_len]
            labels = batch["labels"].unsqueeze(0)  # [1, seq_len]

            # Forward pass
            outputs = hf_model(input_ids=input_ids, labels=labels)
            loss = outputs.loss

            # Backward pass
            loss.backward()
            optimizer.step()

            epoch_train_loss += loss.item()

            if i % 5 == 0:
                print(f"   Batch {i}/{len(train_dataset)}: loss = {loss.item():.4f}")

        avg_train_loss = epoch_train_loss / len(train_dataset)
        train_losses.append(avg_train_loss)
        print(f"   📊 Средняя train loss: {avg_train_loss:.4f}")

        # Валидация
        hf_model.eval()
        epoch_val_loss = 0
        with torch.no_grad():
            for batch in val_dataset:
                input_ids = batch["input_ids"].unsqueeze(0)
                labels = batch["labels"].unsqueeze(0)

                outputs = hf_model(input_ids=input_ids, labels=labels)
                epoch_val_loss += outputs.loss.item()

        avg_val_loss = epoch_val_loss / len(val_dataset)
        val_losses.append(avg_val_loss)
        print(f"   📊 Средняя val loss: {avg_val_loss:.4f}")

        hf_model.train()

    return {
        "train_losses": train_losses,
        "val_losses": val_losses,
        "final_train_loss": train_losses[-1],
        "final_val_loss": val_losses[-1],
    }


def test_generation_after_training(hf_model, hf_tokenizer, test_prompts):
    """
    Тестирует генерацию после обучения.

    Args:
        hf_model: Обученная модель
        hf_tokenizer: Токенизатор
        test_prompts: Тестовые промпты
    """
    print("\n🧪 Тестирование генерации после обучения...")
    hf_model.eval()

    for prompt in test_prompts[:3]:
        print(f"\n🔤 Промпт: '{prompt}'")

        try:
            inputs = hf_tokenizer(prompt, return_tensors="pt")

            with torch.no_grad():
                generated = hf_model.generate(
                    input_ids=inputs["input_ids"],
                    max_new_tokens=20,
                    do_sample=True,
                    temperature=0.8,
                )

            generated_text = hf_tokenizer.decode(generated[0], skip_special_tokens=True)
            print(f"🎯 Результат: '{generated_text}'")

        except Exception as e:
            print(f"❌ Ошибка генерации: {e}")


def main():
    """Основная функция эксперимента."""
    print("=" * 60)
    print("🚀 УПРОЩЕННОЕ ОБУЧЕНИЕ GPT С HF-PROXY")
    print("=" * 60)

    try:
        # === Подготовка данных ===
        print("🔧 Подготовка данных...")
        train_texts = TRAIN_TEXTS[
            :10
        ]  # Используем меньше данных для быстрого тестирования
        val_texts = TRAIN_TEXTS[10:12]

        print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")

        # === Подготовка токенизатора ===
        print("🔧 Подготовка токенизатора...")
        llm_tokenizer = BPETokenizer()
        llm_tokenizer.train(
            texts=train_texts,
            vocab_size=BPE_CONFIG["vocab_size"],
            special_tokens=BPE_CONFIG["special_tokens"],
        )

        hf_tokenizer = HFTokenizerAdapter(llm_tokenizer)
        print(f"✅ Токенизатор создан (vocab_size={hf_tokenizer.vocab_size})")

        # === Подготовка модели ===
        print("🔧 Подготовка модели...")
        model_config = BASE_GPT_CONFIG.copy()
        model_config["vocab_size"] = hf_tokenizer.vocab_size

        llm_model = GPT(model_config)
        hf_model = HFAdapter.from_llm_model(llm_model)
        print(f"✅ Модель создана")

        # === Тестирование до обучения ===
        print("\n🧪 Тестирование до обучения...")
        test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS)

        # === Обучение ===
        print(f"\n🎯 Обучение модели...")
        training_config = {
            "learning_rate": TRAINING_CONFIG["learning_rate"],
            "num_epochs": 2,  # Меньше эпох для быстрого тестирования
            "batch_size": TRAINING_CONFIG["batch_size"],
        }

        results = manual_training_loop(
            hf_model, hf_tokenizer, train_texts, val_texts, training_config
        )

        print(f"\n📊 Результаты обучения:")
        print(f"   Final train loss: {results['final_train_loss']:.4f}")
        print(f"   Final val loss: {results['final_val_loss']:.4f}")

        # === Тестирование после обучения ===
        print("\n🧪 Тестирование после обучения...")
        test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS)

        # === Сохранение модели ===
        print(f"\n💾 Сохранение модели...")

        # Создаем директории
        os.makedirs("checkpoints/hf_simple_trained", exist_ok=True)
        os.makedirs("checkpoints/hf_simple_tokenizer", exist_ok=True)

        # Сохраняем токенизатор
        hf_tokenizer.save_pretrained("checkpoints/hf_simple_tokenizer")
        print("✅ Токенизатор сохранен")

        # Сохраняем модель
        HFAdapter.save_pretrained(
            hf_model, "checkpoints/hf_simple_trained", tokenizer=hf_tokenizer
        )
        print("✅ Модель сохранена")

        # Сохраняем результаты
        results_path = "checkpoints/simple_training_results.json"
        with open(results_path, "w", encoding="utf-8") as f:
            json.dump(
                {
                    "training_config": training_config,
                    "model_config": model_config,
                    "results": results,
                },
                f,
                indent=2,
                ensure_ascii=False,
            )
        print(f"✅ Результаты сохранены в {results_path}")

        print(f"\n🎉 Упрощенное обучение завершено успешно!")
        print(f"\n💡 Для использования обученной модели:")
        print(f"   uv run python experiments/hf_integration/generate_with_hf_tools.py")

    except Exception as e:
        print(f"❌ Ошибка в эксперименте: {e}")
        import traceback

        traceback.print_exc()


if __name__ == "__main__":
    main()