Update README.md

README.md

@@ -79,17 +79,28 @@ The repository includes complete training and inference code. Key components:
 ## Example Code
 ```python
 import os
-import re
-import json
-import random
-from tqdm import tqdm
-import numpy as np
+import warnings
 from pathlib import Path
+from typing import List, Dict, Optional, Tuple
 
 import torch
 from torch import nn
-from torch.utils.data import DataLoader, Dataset
-from transformers import GPT2TokenizerFast
+from torch.utils.data import Dataset, DataLoader
+from transformers import (
+    GPT2TokenizerFast,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PretrainedConfig,
+    GenerationMixin,
+    pipeline
+)
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+from huggingface_hub import HfApi, login
+from datasets import load_dataset
+from tqdm import tqdm
 from adam_atan2_pytorch import AdoptAtan2
 
 from titans_pytorch import (
@@ -98,24 +109,17 @@ from titans_pytorch import (
     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
+# Отключаем предупреждения
+warnings.filterwarnings("ignore", category=UserWarning)
+torch._dynamo.config.suppress_errors = True
+torch._dynamo.config.cache_size_limit = 100000
+torch._dynamo.config.disable = True
 
-# Добавляем настройки для управления памятью CUDA
-import os
+# Настройки CUDA
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:32'
 
-
 # Константы
+repo_id = 'Grpp/memory-transformer-ru'
 NUM_BATCHES = int(1e5)
 BATCH_SIZE = 4
 GRADIENT_ACCUMULATE_EVERY = 4
@@ -146,173 +150,59 @@ 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 MemoryTransformerConfig(PretrainedConfig):
+    model_type = "memory_transformer"
+    
+    def __init__(
+        self,
+        vocab_size=50257,
+        dim=384,
+        depth=8,
+        segment_len=32,
+        num_persist_mem=4,
+        num_longterm_mem=4,
+        neural_mem_layers=(2, 4, 6),
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.depth = depth
+        self.segment_len = segment_len
+        self.num_persist_mem = num_persist_mem
+        self.num_longterm_mem = num_longterm_mem
+        self.neural_mem_layers = neural_mem_layers
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs
+        )
 
 
-class WikiTextDataset(Dataset):
-    def __init__(self, chunks: List[List[int]], seq_len: int = 512):
-        self.chunks = chunks
-        self.seq_len = seq_len
+class MemoryTransformerForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MemoryTransformerConfig
+    supports_gradient_checkpointing = True
     
-    def __len__(self):
-        return len(self.chunks)
-    
-    def __getitem__(self, idx):
-        chunk = self.chunks[idx]
+    def __init__(self, config):
+        super().__init__(config)
         
-        # Если чанк короче необходимой длины, дополняем его паддингом
-        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
+        neural_memory_model = (
+            MemoryAttention(dim=64) if USE_MEM_ATTENTION_MODEL 
+            else MemoryMLP(dim=64, depth=NEURAL_MEMORY_DEPTH)
+        )
             
-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,
+        self.transformer = MemoryAsContextTransformer(
+            num_tokens=config.vocab_size,
+            dim=config.dim,
+            depth=config.depth,
+            segment_len=config.segment_len,
+            num_persist_mem_tokens=config.num_persist_mem,
+            num_longterm_mem_tokens=config.num_longterm_mem,
+            neural_memory_layers=config.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,
@@ -331,139 +221,107 @@ def create_model():
                 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"
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.transformer(input_ids)
         
-        return model
+        if labels is not None:
+            loss = self.transformer(input_ids, return_loss=True)
+            return CausalLMOutputWithCrossAttentions(
+                loss=loss,
+                logits=outputs,
+                past_key_values=None,
+                hidden_states=None,
+                attentions=None,
+                cross_attentions=None
+            )
+        
+        return CausalLMOutputWithCrossAttentions(
+            loss=None,
+            logits=outputs,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+            cross_attentions=None
+        )
     
-    except Exception as e:
-        print(f"Error creating model: {e}")
-        raise e
+    def prepare_inputs_for_generation(
+        self, 
+        input_ids, 
+        past=None, 
+        attention_mask=None, 
+        **kwargs
+    ):
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past,
+            "attention_mask": attention_mask,
+        }
+        
+    @property
+    def device(self):
+        return next(self.parameters()).device
 
-def train_model(model, train_loader, val_loader, num_batches=int(1e4)):
-    optim = AdoptAtan2(model.parameters(), lr=2e-4)
-    
-    # Включаем автоматическую очистку кэша CUDA
-    torch.cuda.empty_cache()
+
+def setup_custom_model():
+    """Регистрация кастомной модели"""
+    AutoConfig.register("memory_transformer", MemoryTransformerConfig)
+    AutoModelForCausalLM.register(MemoryTransformerConfig, MemoryTransformerForCausalLM)
+
+
+def generate_example(model, tokenizer, text, max_length=100):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    model.eval()
     
-    pbar = tqdm(range(num_batches), desc='Training')
-    running_loss = 0.0
+    input_ids = tokenizer.encode(text, return_tensors="pt").to(device)
+    attention_mask = torch.ones_like(input_ids, device=device)
     
-    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
+    print(f"Model device: {next(model.parameters()).device}")
+    print(f"Input device: {input_ids.device}")
+
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_length=max_length,
+            num_return_sequences=1,
+            no_repeat_ngram_size=2,
+            do_sample=True,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.7,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
         )
-        
-        # Создание бесконечных итераторов
-        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
+
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
 
 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()
+
+    setup_custom_model()
+    config = AutoConfig.from_pretrained(repo_id)
+    model = AutoModelForCausalLM.from_pretrained(repo_id)
+    tokenizer = AutoTokenizer.from_pretrained(repo_id)
+
+    test_text = "Московский кремль является"
+    generated_text = generate_example(model, tokenizer, test_text)
+    print(generated_text)
 ```
 
 # Training