Traceback (most recent call last):
  File "C:\Users\Dell\AI\eridAI-models\LEP_new\main.py", line 206, in <module>
    adapt_gpt2_for_russian(
  File "C:\Users\Dell\AI\eridAI-models\LEP_new\main.py", line 196, in adapt_gpt2_for_russian
    adapter.save_adapted_model(output_path)
  File "C:\Users\Dell\AI\eridAI-models\LEP_new\main.py", line 156, in save_adapted_model
    self.russian_tokenizer.save_model(tokenizer_path)
AttributeError: 'SentencePieceProcessor' object has no attribute 'save_model'import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from sentencepiece import SentencePieceTrainer, SentencePieceProcessor
from datasets import load_dataset
import numpy as np
from typing import Dict, List
import os
import tempfile
from tqdm import tqdm
import time


class RussianLEPAdapter:
    def __init__(self, base_model_path: str, dataset_name: str, device: str = None):
        """
        Инициализация LEP адаптера для русского языка

        Args:
            base_model_path: Путь к предобученной модели GPT-2
            dataset_name: Название датасета из HuggingFace
            device: Устройство для вычислений ('cuda' или 'cpu')
        """
        self.base_model_path = base_model_path
        self.dataset_name = dataset_name
        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"Using device: {self.device}")

        # Прогресс адаптации
        self.total_steps = 100  # Общее количество шагов процесса
        self.current_step = 0
        self.start_time = None
        self.pbar = None

        # Загрузка базовой модели и токенизатора
        self.base_model = GPT2LMHeadModel.from_pretrained(base_model_path).to(self.device)
        self.base_tokenizer = GPT2Tokenizer.from_pretrained(base_model_path)

        self.russian_tokenizer = None
        self.token_mapping = None

    def update_progress(self, increment: int, description: str = None):
        """
        Обновление общего прогресса адаптации

        Args:
            increment: На сколько шагов продвинулись
            description: Описание текущего этапа
        """
        self.current_step += increment
        if description:
            self.pbar.set_description(f"[{description}]")
        self.pbar.update(increment)

    def prepare_corpus(self, text_column: str = "text") -> str:
        """Подготовка корпуса из датасета HuggingFace"""
        dataset = load_dataset(self.dataset_name, trust_remote_code=True)

        self.update_progress(1, "Подготовка корпуса")
        with tempfile.NamedTemporaryFile(mode='w', delete=False, encoding='utf-8') as f:
            if isinstance(dataset, dict):
                total = sum(len(split) for split in dataset.values())
                items_per_step = max(total // 20, 1)  # Распределяем на 20 шагов
                processed = 0

                for split in dataset.keys():
                    for item in dataset[split]:
                        f.write(item[text_column] + '\n')
                        processed += 1
                        if processed % items_per_step == 0:
                            self.update_progress(1)
            else:
                items_per_step = max(len(dataset) // 20, 1)
                for i, item in enumerate(dataset):
                    f.write(item[text_column] + '\n')
                    if i % items_per_step == 0:
                        self.update_progress(1)

            return f.name

    def train_russian_tokenizer(self, vocab_size: int = 32000, model_type: str = "unigram"):
        """Обучение нового токенизатора для русского языка"""
        corpus_path = self.prepare_corpus()

        self.update_progress(1, "Обучение токенизатора")
        model_prefix = "russian_tokenizer"
        SentencePieceTrainer.train(
            input=corpus_path,
            model_prefix=model_prefix,
            vocab_size=vocab_size,
            model_type=model_type,
            character_coverage=0.9995,
            pad_id=0,
            unk_id=1,
            bos_id=2,
            eos_id=3,
            normalization_rule_name="nmt_nfkc"
        )

        self.russian_tokenizer = SentencePieceProcessor()
        self.russian_tokenizer.load(f"{model_prefix}.model")
        os.remove(corpus_path)
        self.update_progress(5)

    def create_token_mapping(self) -> Dict[int, List[int]]:
        """Создание маппинга между токенами"""
        mapping = {}
        vocab_size = self.russian_tokenizer.get_piece_size()
        tokens_per_step = max(vocab_size // 30, 1)  # Распределяем на 30 шагов

        self.update_progress(1, "Создание маппинга токенов")
        for i in range(vocab_size):
            new_token = self.russian_tokenizer.id_to_piece(i)
            old_tokens = self.base_tokenizer.encode(new_token, add_special_tokens=False)
            mapping[i] = old_tokens

            if i % tokens_per_step == 0:
                self.update_progress(1)

        self.token_mapping = mapping
        return mapping

    def propagate_embeddings(self):
        """Распространение эмбеддингов"""
        orig_embeddings = self.base_model.transformer.wte.weight.data
        new_vocab_size = self.russian_tokenizer.get_piece_size()
        embedding_dim = orig_embeddings.shape[1]
        new_embeddings = torch.zeros((new_vocab_size, embedding_dim), device=self.device)

        tokens_per_step = max(new_vocab_size // 20, 1)  # Распределяем на 20 шагов
        self.update_progress(1, "Распространение эмбеддингов")

        for i in range(new_vocab_size):
            old_ids = self.token_mapping[i]
            if len(old_ids) > 0:
                new_embeddings[i] = orig_embeddings[old_ids].mean(dim=0)
            else:
                new_embeddings[i] = torch.randn(embedding_dim, device=self.device)

            if i % tokens_per_step == 0:
                self.update_progress(1)

        self.base_model.transformer.wte = torch.nn.Embedding.from_pretrained(new_embeddings)
        self.base_model.lm_head = torch.nn.Linear(embedding_dim, new_vocab_size, bias=False)
        self.base_model.lm_head.weight = torch.nn.Parameter(new_embeddings.clone())
        self.update_progress(5)

    def save_adapted_model(self, output_path: str):
        """Сохранение адаптированной модели и токенизатора"""
        self.update_progress(1, "Сохранение модели")
        os.makedirs(output_path, exist_ok=True)

        self.base_model = self.base_model.cpu()
        self.base_model.save_pretrained(output_path)

        tokenizer_path = os.path.join(output_path, "tokenizer.model")
        self.russian_tokenizer.save_model(tokenizer_path)
        self.update_progress(5)


def adapt_gpt2_for_russian(
        base_model_path: str,
        dataset_name: str,
        output_path: str,
        device: str = None,
        trust_remote_code=True,
        vocab_size: int = 32000,
        text_column: str = "text"
):
    """Основная функция для адаптации GPT-2 к русскому языку"""
    # Оценка общего времени выполнения (в секундах)
    estimated_time = 100  # примерное время выполнения

    adapter = RussianLEPAdapter(base_model_path, dataset_name, device)

    # Инициализация общего прогресс бара
    adapter.start_time = time.time()
    adapter.pbar = tqdm(
        total=adapter.total_steps,
        desc="[Инициализация]",
        # Исправленный формат строки с добавлением пробела перед [
        bar_format="{desc}: {percentage:3.0f}%%[{bar}] {n_fmt}/{total_fmt} шагов "
                   "[Прошло: {elapsed}, Осталось: {remaining}, Всего: %ds]" % estimated_time
    )

    try:
        # Обучение русского токенизатора
        adapter.train_russian_tokenizer(vocab_size=vocab_size, model_type="unigram")

        # Создание маппинга токенов
        adapter.create_token_mapping()

        # Распространение эмбеддингов
        adapter.propagate_embeddings()

        # Сохранение адаптированной модели
        adapter.save_adapted_model(output_path)

        print(f"\nАдаптация успешно завершена за {time.time() - adapter.start_time:.1f} секунд!")

    finally:
        adapter.pbar.close()


# Пример использования:
if __name__ == "__main__":
    adapt_gpt2_for_russian(
        base_model_path="gpt2",
        dataset_name="IlyaGusev/gazeta",
        output_path="./gpt2-russian",
        device=None,
        trust_remote_code=True,
        vocab_size=32000,
        text_column="text"
    )