Create README.md

README.md

@@ -0,0 +1,853 @@
+---
+language:
+- ru
+license: mit
+datasets:
+- misterkirill/ru-wikipedia
+tags:
+- pytorch
+- neural-memory
+- titan
+- text-generation
+---
+
+# Neural Memory Model for Russian Text Generation
+
+This model implements a neural memory architecture for Russian text generation using PyTorch and the Titans library. The architecture is based on the implementation from [lucidrains/titans-pytorch](https://github.com/lucidrains/titans-pytorch).
+
+## Model Description
+
+The model uses a Transformer architecture enhanced with neural memory capabilities from the Titans library for improved context handling and long-range dependencies in text generation.
+
+### Architecture Source
+
+The core architecture is derived from the [Titans PyTorch implementation](https://github.com/lucidrains/titans-pytorch) by Phil Wang ([@lucidrains](https://github.com/lucidrains)). The original implementation provides the following key components that we utilize:
+- Memory-enhanced Transformer architecture
+- Flexible attention mechanisms
+- Neural memory layers
+
+### Key Features
+
+- Neural memory architecture with customizable depth and size
+- Sliding window attention mechanism
+- Gradient accumulation for stable training
+- CUDA-optimized implementation
+
+## Requirements
+
+### Environment
+
+- Python: 3.9.21
+- CUDA: 11.8
+- GPU with at least 16GB VRAM recommended
+
+### Key Dependencies
+```
+Python version: 3.9.21
+CUDA version: 11.8
+
+Requirements:
+adam-atan2-pytorch==0.1.18
+datasets==3.2.0
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparselt-cu12==0.6.2
+nvidia-nccl-cu12==2.21.5
+nvidia-nvtx-cu12==12.4.127
+titans-pytorch==0.3.25
+torchaudio==2.5.1
+torchvision==0.20.1
+transformers==4.48.3
+triton==3.1.0
+wandb==0.19.6
+```
+
+# Example
+The repository includes complete training and inference code. Key components:
+
+
+- Data preprocessing (WikiDatasetPreprocessor)
+- Custom dataset implementation (WikiTextDataset)
+- Training loop with gradient accumulation
+- Validation and checkpointing
+
+## Example Code
+```python
+import os
+import re
+import json
+import random
+from tqdm import tqdm
+import numpy as np
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+from transformers import GPT2TokenizerFast
+from adam_atan2_pytorch import AdoptAtan2
+
+from titans_pytorch import (
+    MemoryAsContextTransformer,
+    MemoryMLP,
+    MemoryAttention
+)
+
+import os
+import json
+import random
+from pathlib import Path
+from typing import List, Dict
+import numpy as np
+from tqdm import tqdm
+from datasets import load_dataset
+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import GPT2TokenizerFast
+
+# Добавляем настройки для управления памятью CUDA
+import os
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:32'
+
+
+# Константы
+NUM_BATCHES = int(1e5)
+BATCH_SIZE = 4
+GRADIENT_ACCUMULATE_EVERY = 4
+LEARNING_RATE = 2e-4
+VALIDATE_EVERY = 100
+GENERATE_EVERY = 500
+PRIME_LENGTH = 100
+GENERATE_LENGTH = 512
+SHOULD_GENERATE = True
+SEQ_LEN = 512
+
+# Константы для нейронной памяти
+NEURAL_MEMORY_DEPTH = 2
+NUM_PERSIST_MEM = 4
+NUM_LONGTERM_MEM = 4
+NEURAL_MEM_LAYERS = (2, 4, 6)
+NEURAL_MEM_GATE_ATTN_OUTPUT = False
+NEURAL_MEM_MOMENTUM = True
+NEURAL_MEM_MOMENTUM_ORDER = 1
+NEURAL_MEM_QK_NORM = True
+NEURAL_MEM_MAX_LR = 1e-1
+USE_MEM_ATTENTION_MODEL = False
+WINDOW_SIZE = 32
+NEURAL_MEM_SEGMENT_LEN = 4
+NEURAL_MEM_BATCH_SIZE = 128
+SLIDING_WINDOWS = True
+STORE_ATTN_POOL_CHUNKS = True
+MEMORY_MODEL_PER_LAYER_LEARNED_LR = True
+NEURAL_MEM_WEIGHT_RESIDUAL = True
+
+# Инициализация токенизатора
+tokenizer = GPT2TokenizerFast.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
+
+
+class WikiDatasetPreprocessor:
+    def __init__(self, cache_dir: str = 'cache', output_dir: str = 'processed_data'):
+        self.cache_dir = Path(cache_dir)
+        self.output_dir = Path(output_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Инициализация токенизатора
+        self.tokenizer = GPT2TokenizerFast.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
+        
+    def load_wiki_dataset(self):
+        """Загрузка датасета из Hugging Face"""
+        print("Loading Wikipedia dataset...")
+        dataset = load_dataset("misterkirill/ru-wikipedia", cache_dir=str(self.cache_dir))
+        print(f"Dataset loaded. Size: {len(dataset['train'])} articles")
+        return dataset
+
+    def clean_text(self, text: str) -> str:
+        """Базовая очистка текста"""
+        # Удаляем множественные пробелы и переносы строк
+        text = ' '.join(text.split())
+        return text
+
+    # В функции process_and_save уменьшаем размер чанков
+    def process_wiki_article(self, text: str) -> List[str]:
+        """Обработка одной статьи из википедии"""
+        processed_chunks = []
+
+        clean_text = self.clean_text(text)
+        tokens = self.tokenizer.encode(clean_text)
+
+        # Уменьшаем размер чанка
+        chunk_size = 256  # было 512
+        stride = 192      # было 384
+
+        for i in range(0, len(tokens), stride):
+            chunk = tokens[i:i + chunk_size]
+            if len(chunk) > 50:  # уменьшаем минимальную длину чанка
+                processed_chunks.append(chunk)
+
+        return processed_chunks
+
+    def process_and_save(self, batch_size: int = 1000, test_size: float = 0.1, max_articles: int = 10000):
+        """Обработка ограниченного количества статей из датасета и сохранение результатов"""
+        dataset = self.load_wiki_dataset()
+
+        # Ограничиваем размер датасета
+        total_articles = min(len(dataset['train']), max_articles)
+        print(f"Processing {total_articles} articles out of {len(dataset['train'])}")
+
+        # Сначала соберем все чанки
+        all_chunks = []
+
+        for i in tqdm(range(0, total_articles, batch_size), desc="Processing articles"):
+            batch = dataset['train'][i:i + batch_size]
+            for text in batch['text']:
+                chunks = self.process_wiki_article(text)
+                all_chunks.extend(chunks)
+
+                # Ограничиваем количество чанков для ускорения обучения
+                if len(all_chunks) > 50000:  # максимальное количество чанков
+                    break
+
+            if len(all_chunks) > 50000:
+                break
+
+        print(f"Total chunks created: {len(all_chunks)}")
+
+        # Перемешаем чанки
+        random.seed(42)
+        random.shuffle(all_chunks)
+
+        # Разделим на train и test
+        test_size = int(len(all_chunks) * test_size)
+        train_chunks = all_chunks[:-test_size]
+        test_chunks = all_chunks[-test_size:]
+
+        print(f"Saving {len(train_chunks)} training chunks and {len(test_chunks)} test chunks...")
+        torch.save({
+            'train': train_chunks,
+            'test': test_chunks
+        }, self.output_dir / 'processed_wiki.pt')
+
+
+class WikiTextDataset(Dataset):
+    def __init__(self, chunks: List[List[int]], seq_len: int = 512):
+        self.chunks = chunks
+        self.seq_len = seq_len
+    
+    def __len__(self):
+        return len(self.chunks)
+    
+    def __getitem__(self, idx):
+        chunk = self.chunks[idx]
+        
+        # Если чанк короче необходимой длины, дополняем его паддингом
+        if len(chunk) < self.seq_len + 1:
+            chunk = chunk + [50256] * (self.seq_len + 1 - len(chunk))
+        # Если длиннее - обрезаем
+        else:
+            chunk = chunk[:self.seq_len + 1]
+            
+        return torch.tensor(chunk, device='cuda').long()  # Добавляем device='cuda'
+
+def create_dataloaders(
+    processed_data_path: str,
+    batch_size: int = 4,
+    seq_len: int = 512,
+    train_test_split: float = 0.9
+) -> tuple:
+    """Создание загрузчиков данных для обучения и валидации"""
+    
+    print(f"Loading processed data from {processed_data_path}")
+    data = torch.load(processed_data_path)
+    train_chunks = data['train']
+    test_chunks = data['test']
+    
+    # Создание датасетов
+    train_dataset = WikiTextDataset(train_chunks, seq_len)
+    test_dataset = WikiTextDataset(test_chunks, seq_len)
+    
+    print(f"Created datasets with {len(train_dataset)} training and {len(test_dataset)} test samples")
+    
+    # Создание загрузчиков данных
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=0,  # Убираем многопоточность для отладки
+        pin_memory=False  # Отключаем pin_memory, так как данные уже на GPU
+    )
+    
+    val_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,  # Убираем многопоточность для отладки
+        pin_memory=False  # Отключаем pin_memory, так как данные уже на GPU
+    )
+    
+    return train_loader, val_loader
+
+def cycle(loader):
+    """Бесконечный итератор по загрузчику данных"""
+    while True:
+        for data in loader:
+            yield data
+            
+def create_model():
+    try:
+        if USE_MEM_ATTENTION_MODEL:
+            neural_memory_model = MemoryAttention(dim=64)
+        else:
+            neural_memory_model = MemoryMLP(dim=64, depth=NEURAL_MEMORY_DEPTH)
+
+        model = MemoryAsContextTransformer(
+            num_tokens=len(tokenizer),
+            dim=384,
+            depth=8,
+            segment_len=WINDOW_SIZE,
+            num_persist_mem_tokens=NUM_PERSIST_MEM,
+            num_longterm_mem_tokens=NUM_LONGTERM_MEM,
+            neural_memory_layers=NEURAL_MEM_LAYERS,
+            neural_memory_segment_len=NEURAL_MEM_SEGMENT_LEN,
+            neural_memory_batch_size=NEURAL_MEM_BATCH_SIZE,
+            neural_mem_gate_attn_output=NEURAL_MEM_GATE_ATTN_OUTPUT,
+            neural_mem_weight_residual=NEURAL_MEM_WEIGHT_RESIDUAL,
+            use_flex_attn=True,
+            sliding_window_attn=SLIDING_WINDOWS,
+            neural_memory_model=neural_memory_model,
+            neural_memory_kwargs=dict(
+                dim_head=64,
+                heads=4,
+                attn_pool_chunks=STORE_ATTN_POOL_CHUNKS,
+                qk_rmsnorm=NEURAL_MEM_QK_NORM,
+                momentum=NEURAL_MEM_MOMENTUM,
+                momentum_order=NEURAL_MEM_MOMENTUM_ORDER,
+                default_step_transform_max_lr=NEURAL_MEM_MAX_LR,
+                use_accelerated_scan=True,
+                per_parameter_lr_modulation=MEMORY_MODEL_PER_LAYER_LEARNED_LR
+            )
+        ).cuda()
+        
+        # Проверка, что модель на GPU
+        assert next(model.parameters()).is_cuda, "Model is not on CUDA"
+        
+        return model
+    
+    except Exception as e:
+        print(f"Error creating model: {e}")
+        raise e
+
+def train_model(model, train_loader, val_loader, num_batches=int(1e4)):
+    optim = AdoptAtan2(model.parameters(), lr=2e-4)
+    
+    # Включаем автоматическую очистку кэша CUDA
+    torch.cuda.empty_cache()
+    
+    pbar = tqdm(range(num_batches), desc='Training')
+    running_loss = 0.0
+    
+    try:
+        for i in pbar:
+            model.train()
+            
+            total_loss = 0
+            # Обучение с градиентным накоплением
+            for __ in range(4):
+                batch = next(train_loader)
+                loss = model(batch, return_loss=True)
+                loss = loss / 4  # нормализуем loss при градиентном накоплении
+                loss.backward()
+                total_loss += loss.item()
+            
+            # Клиппинг градиентов
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+            optim.step()
+            optim.zero_grad()
+            
+            # Очищаем кэш CUDA каждые 100 итераций
+            if i % 100 == 0:
+                torch.cuda.empty_cache()
+            
+            avg_loss = total_loss
+            running_loss = 0.9 * running_loss + 0.1 * avg_loss if i > 0 else avg_loss
+            
+            pbar.set_postfix({
+                'loss': f'{running_loss:.4f}',
+                'batch_loss': f'{avg_loss:.4f}'
+            })
+            
+            # Валидация
+            if i % 100 == 0:
+                model.eval()
+                with torch.no_grad():
+                    val_batch = next(val_loader)
+                    val_loss = model(val_batch, return_loss=True)
+                    pbar.set_postfix({
+                        'train_loss': f'{running_loss:.4f}',
+                        'val_loss': f'{val_loss.item():.4f}'
+                    })
+            
+            # Сохранение чекпойнта
+            if i % 1000 == 0 and i > 0:
+                torch.save({
+                    'epoch': i,
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optim.state_dict(),
+                    'loss': running_loss,
+                }, f'checkpoint_{i}.pt')
+                
+    except KeyboardInterrupt:
+        print("\nTraining interrupted by user")
+    except Exception as e:
+        print(f"\nTraining stopped due to error: {e}")
+        raise e
+    
+    return model
+
+def main():
+    try:
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA is not available. This code requires GPU.")
+        
+        print(f"Using CUDA device: {torch.cuda.get_device_name(0)}")
+        
+        # Параметры
+        BATCH_SIZE = 4
+        SEQ_LEN = 512
+        CACHE_DIR = 'cache'
+        PROCESSED_DATA_DIR = 'processed_data'
+        NUM_BATCHES = 10000  # уменьшаем количество итераций
+        
+        # Подготовка данных
+        preprocessor = WikiDatasetPreprocessor(CACHE_DIR, PROCESSED_DATA_DIR)
+        
+        processed_data_path = Path(PROCESSED_DATA_DIR) / 'processed_wiki.pt'
+        if not processed_data_path.exists():
+            print("Processing Wikipedia dataset...")
+            preprocessor.process_and_save(max_articles=10000)  # ограничиваем количество статей
+        
+        # Создание загрузчиков данных
+        train_loader, val_loader = create_dataloaders(
+            processed_data_path,
+            batch_size=BATCH_SIZE,
+            seq_len=SEQ_LEN
+        )
+        
+        # Создание бесконечных итераторов
+        train_loader = cycle(train_loader)
+        val_loader = cycle(val_loader)
+        
+        # Создание и обучение модели
+        model = create_model()
+        model = train_model(model, train_loader, val_loader, num_batches=NUM_BATCHES)
+        
+        # Сохранение финальной модели
+        torch.save(model.state_dict(), 'final_model.pt')
+        
+        return model, train_loader, val_loader
+        
+    except Exception as e:
+        print(f"Error in main: {e}")
+        raise e
+
+if __name__ == "__main__":
+    # Установка seed для воспроизводимости
+    torch.manual_seed(42)
+    torch.cuda.manual_seed_all(42)
+    
+    # Включение оптимизаций CUDA
+    torch.backends.cudnn.benchmark = True
+    
+    model, train_loader, val_loader = main()
+```
+
+# Training
+
+The model was trained on a cleaned subset of Russian Wikipedia articles using the following parameters:
+
+
+Batch size: 4
+Sequence length: 512
+Learning rate: 2e-4
+Gradient accumulation steps: 4
+Neural memory depth: 2
+Window size: 32
+
+## Train Code
+```python
+import json
+import os
+import random
+import re
+from pathlib import Path
+from typing import List, Dict
+
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+from transformers import GPT2TokenizerFast
+from tqdm import tqdm
+from datasets import load_dataset
+from adam_atan2_pytorch import AdoptAtan2
+from titans_pytorch import (
+    MemoryAsContextTransformer,
+    MemoryMLP,
+    MemoryAttention
+)
+
+# CUDA memory settings
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:32'
+
+# Training constants
+NUM_BATCHES = int(1e5)
+BATCH_SIZE = 4
+GRADIENT_ACCUMULATE_EVERY = 4
+LEARNING_RATE = 2e-4
+VALIDATE_EVERY = 100
+GENERATE_EVERY = 500
+PRIME_LENGTH = 100
+GENERATE_LENGTH = 512
+SHOULD_GENERATE = True
+SEQ_LEN = 512
+
+# Neural memory constants
+NEURAL_MEMORY_DEPTH = 2
+NUM_PERSIST_MEM = 4
+NUM_LONGTERM_MEM = 4
+NEURAL_MEM_LAYERS = (2, 4, 6)
+NEURAL_MEM_GATE_ATTN_OUTPUT = False
+NEURAL_MEM_MOMENTUM = True
+NEURAL_MEM_MOMENTUM_ORDER = 1
+NEURAL_MEM_QK_NORM = True
+NEURAL_MEM_MAX_LR = 1e-1
+USE_MEM_ATTENTION_MODEL = False
+WINDOW_SIZE = 32
+NEURAL_MEM_SEGMENT_LEN = 4
+NEURAL_MEM_BATCH_SIZE = 128
+SLIDING_WINDOWS = True
+STORE_ATTN_POOL_CHUNKS = True
+MEMORY_MODEL_PER_LAYER_LEARNED_LR = True
+NEURAL_MEM_WEIGHT_RESIDUAL = True
+
+# Initialize tokenizer
+tokenizer = GPT2TokenizerFast.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
+
+
+class WikiDatasetPreprocessor:
+    def __init__(self, cache_dir: str = 'cache', output_dir: str = 'processed_data'):
+        self.cache_dir = Path(cache_dir)
+        self.output_dir = Path(output_dir)
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.tokenizer = GPT2TokenizerFast.from_pretrained(
+            'sberbank-ai/rugpt3small_based_on_gpt2'
+        )
+
+    def load_wiki_dataset(self):
+        """Загрузка датасета из Hugging Face."""
+        print("Loading Wikipedia dataset...")
+        dataset = load_dataset(
+            "misterkirill/ru-wikipedia",
+            cache_dir=str(self.cache_dir)
+        )
+        print(f"Dataset loaded. Size: {len(dataset['train'])} articles")
+        return dataset
+
+    def clean_text(self, text: str) -> str:
+        """Базовая очистка текста."""
+        return ' '.join(text.split())
+
+    def process_wiki_article(self, text: str) -> List[str]:
+        """Обработка одной статьи из википедии."""
+        processed_chunks = []
+        clean_text = self.clean_text(text)
+        tokens = self.tokenizer.encode(clean_text)
+
+        chunk_size = 256
+        stride = 192
+
+        for i in range(0, len(tokens), stride):
+            chunk = tokens[i:i + chunk_size]
+            if len(chunk) > 50:
+                processed_chunks.append(chunk)
+
+        return processed_chunks
+
+    def process_and_save(
+        self,
+        batch_size: int = 1000,
+        test_size: float = 0.1,
+        max_articles: int = 10000
+    ):
+        """Обработка статей из датасета и сохранение результатов."""
+        dataset = self.load_wiki_dataset()
+        total_articles = min(len(dataset['train']), max_articles)
+        print(f"Processing {total_articles} articles out of {len(dataset['train'])}")
+
+        all_chunks = []
+        for i in tqdm(range(0, total_articles, batch_size), desc="Processing articles"):
+            batch = dataset['train'][i:i + batch_size]
+            for text in batch['text']:
+                chunks = self.process_wiki_article(text)
+                all_chunks.extend(chunks)
+
+                if len(all_chunks) > 50000:
+                    break
+
+            if len(all_chunks) > 50000:
+                break
+
+        print(f"Total chunks created: {len(all_chunks)}")
+
+        random.seed(42)
+        random.shuffle(all_chunks)
+
+        test_size = int(len(all_chunks) * test_size)
+        train_chunks = all_chunks[:-test_size]
+        test_chunks = all_chunks[-test_size:]
+
+        print(f"Saving {len(train_chunks)} training chunks and {len(test_chunks)} test chunks...")
+        torch.save(
+            {
+                'train': train_chunks,
+                'test': test_chunks
+            },
+            self.output_dir / 'processed_wiki.pt'
+        )
+
+
+class WikiTextDataset(Dataset):
+    def __init__(self, chunks: List[List[int]], seq_len: int = 512):
+        self.chunks = chunks
+        self.seq_len = seq_len
+
+    def __len__(self):
+        return len(self.chunks)
+
+    def __getitem__(self, idx):
+        chunk = self.chunks[idx]
+        if len(chunk) < self.seq_len + 1:
+            chunk = chunk + [50256] * (self.seq_len + 1 - len(chunk))
+        else:
+            chunk = chunk[:self.seq_len + 1]
+        return torch.tensor(chunk, device='cuda').long()
+
+
+def create_dataloaders(
+    processed_data_path: str,
+    batch_size: int = 4,
+    seq_len: int = 512,
+    train_test_split: float = 0.9
+) -> tuple:
+    """Создание загрузчиков данных для обучения и валидации."""
+    print(f"Loading processed data from {processed_data_path}")
+    data = torch.load(processed_data_path)
+    train_chunks = data['train']
+    test_chunks = data['test']
+
+    train_dataset = WikiTextDataset(train_chunks, seq_len)
+    test_dataset = WikiTextDataset(test_chunks, seq_len)
+
+    print(f"Created datasets with {len(train_dataset)} training and "
+          f"{len(test_dataset)} test samples")
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=0,
+        pin_memory=False
+    )
+
+    val_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=False
+    )
+
+    return train_loader, val_loader
+
+
+def cycle(loader):
+    """Бесконечный итератор по загрузчику данных."""
+    while True:
+        for data in loader:
+            yield data
+            
+            
+def create_model():
+    """Создание модели нейронной сети."""
+    try:
+        if USE_MEM_ATTENTION_MODEL:
+            neural_memory_model = MemoryAttention(dim=64)
+        else:
+            neural_memory_model = MemoryMLP(dim=64, depth=NEURAL_MEMORY_DEPTH)
+
+        model = MemoryAsContextTransformer(
+            num_tokens=len(tokenizer),
+            dim=384,
+            depth=8,
+            segment_len=WINDOW_SIZE,
+            num_persist_mem_tokens=NUM_PERSIST_MEM,
+            num_longterm_mem_tokens=NUM_LONGTERM_MEM,
+            neural_memory_layers=NEURAL_MEM_LAYERS,
+            neural_memory_segment_len=NEURAL_MEM_SEGMENT_LEN,
+            neural_memory_batch_size=NEURAL_MEM_BATCH_SIZE,
+            neural_mem_gate_attn_output=NEURAL_MEM_GATE_ATTN_OUTPUT,
+            neural_mem_weight_residual=NEURAL_MEM_WEIGHT_RESIDUAL,
+            use_flex_attn=True,
+            sliding_window_attn=SLIDING_WINDOWS,
+            neural_memory_model=neural_memory_model,
+            neural_memory_kwargs=dict(
+                dim_head=64,
+                heads=4,
+                attn_pool_chunks=STORE_ATTN_POOL_CHUNKS,
+                qk_rmsnorm=NEURAL_MEM_QK_NORM,
+                momentum=NEURAL_MEM_MOMENTUM,
+                momentum_order=NEURAL_MEM_MOMENTUM_ORDER,
+                default_step_transform_max_lr=NEURAL_MEM_MAX_LR,
+                use_accelerated_scan=True,
+                per_parameter_lr_modulation=MEMORY_MODEL_PER_LAYER_LEARNED_LR
+            )
+        ).cuda()
+
+        assert next(model.parameters()).is_cuda, "Model is not on CUDA"
+        return model
+
+    except Exception as e:
+        print(f"Error creating model: {e}")
+        raise e
+
+
+def train_model(model, train_loader, val_loader, num_batches=int(1e4)):
+    """Обучение модели."""
+    optim = AdoptAtan2(model.parameters(), lr=2e-4)
+    torch.cuda.empty_cache()
+    pbar = tqdm(range(num_batches), desc='Training')
+    running_loss = 0.0
+
+    try:
+        for i in pbar:
+            model.train()
+            total_loss = 0
+
+            for __ in range(4):
+                batch = next(train_loader)
+                loss = model(batch, return_loss=True)
+                loss = loss / 4
+                loss.backward()
+                total_loss += loss.item()
+
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+            optim.step()
+            optim.zero_grad()
+
+            if i % 100 == 0:
+                torch.cuda.empty_cache()
+
+            avg_loss = total_loss
+            running_loss = 0.9 * running_loss + 0.1 * avg_loss if i > 0 else avg_loss
+
+            pbar.set_postfix({
+                'loss': f'{running_loss:.4f}',
+                'batch_loss': f'{avg_loss:.4f}'
+            })
+
+            if i % 100 == 0:
+                model.eval()
+                with torch.no_grad():
+                    val_batch = next(val_loader)
+                    val_loss = model(val_batch, return_loss=True)
+                    pbar.set_postfix({
+                        'train_loss': f'{running_loss:.4f}',
+                        'val_loss': f'{val_loss.item():.4f}'
+                    })
+
+            if i % 1000 == 0 and i > 0:
+                torch.save({
+                    'epoch': i,
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optim.state_dict(),
+                    'loss': running_loss,
+                }, f'checkpoint_{i}.pt')
+
+    except KeyboardInterrupt:
+        print("\nTraining interrupted by user")
+    except Exception as e:
+        print(f"\nTraining stopped due to error: {e}")
+        raise e
+
+    return model
+
+
+def main():
+    """Основная функция программы."""
+    try:
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA is not available. This code requires GPU.")
+
+        print(f"Using CUDA device: {torch.cuda.get_device_name(0)}")
+
+        BATCH_SIZE = 4
+        SEQ_LEN = 512
+        CACHE_DIR = 'cache'
+        PROCESSED_DATA_DIR = 'processed_data'
+        NUM_BATCHES = 10000
+
+        preprocessor = WikiDatasetPreprocessor(CACHE_DIR, PROCESSED_DATA_DIR)
+        processed_data_path = Path(PROCESSED_DATA_DIR) / 'processed_wiki.pt'
+
+        if not processed_data_path.exists():
+            print("Processing Wikipedia dataset...")
+            preprocessor.process_and_save(max_articles=10000)
+
+        train_loader, val_loader = create_dataloaders(
+            processed_data_path,
+            batch_size=BATCH_SIZE,
+            seq_len=SEQ_LEN
+        )
+
+        train_loader = cycle(train_loader)
+        val_loader = cycle(val_loader)
+
+        model = create_model()
+        model = train_model(model, train_loader, val_loader, num_batches=NUM_BATCHES)
+
+        torch.save(model.state_dict(), 'final_model.pt')
+        return model, train_loader, val_loader
+
+    except Exception as e:
+        print(f"Error in main: {e}")
+        raise e
+
+
+if __name__ == "__main__":
+    torch.manual_seed(42)
+    torch.cuda.manual_seed_all(42)
+    torch.backends.cudnn.benchmark = True
+    model, train_loader, val_loader = main()
+```
+
+# License
+
+This project is licensed under the MIT License. See LICENSE file for details.
+
+
+# Citation
+
+If you use this model in your research, please cite:
+```bibtex
+@software{neural_memory_model,
+  title = {Neural Memory Model for Russian Text Generation},
+  year = {2024},
+  url = {https://huggingface.co/Grpp/memory-transformer-ru}
+}
+```