Update README.md

README.md

@@ -324,6 +324,202 @@ if __name__ == "__main__":
     print(generated_text)
 ```
 
+
+## Finetine Code
+
+```python
+import os
+import torch
+from pathlib import Path
+from torch.utils.data import DataLoader
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+from tqdm import tqdm
+from adam_atan2_pytorch import AdoptAtan2
+
+# Импортируем классы из кода обучения
+from run_train_pep8 import (
+    WikiDatasetPreprocessor,
+    WikiTextDataset,
+    create_dataloaders,
+    cycle
+)  # From Train Code
+
+from test_load import setup_custom_model  # From Example Code
+
+# Настройки CUDA
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:32'
+
+# Константы для файнтьюнинга
+BATCH_SIZE = 2
+GRADIENT_ACCUMULATE_EVERY = 2
+LEARNING_RATE = 1e-5
+NUM_EPOCHS = 3
+STEPS_PER_EPOCH = 1000  # Количество шагов на эпоху
+SEQ_LEN = 256
+PROCESSED_DATA_DIR = 'processed_data'
+CACHE_DIR = 'cache'
+REPO_ID = 'Grpp/memory-transformer-ru'
+
+def finetune_model(
+    model,
+    train_loader,
+    val_loader,
+    num_epochs,
+    device,
+    save_path='finetuned_model'
+):
+    """Файнтьюнинг модели."""
+    
+    model = model.to(device)
+    optimizer = AdoptAtan2(model.parameters(), lr=LEARNING_RATE)
+    
+    best_val_loss = float('inf')
+    
+    for epoch in range(num_epochs):
+        model.train()
+        total_train_loss = 0
+        train_steps = 0
+        
+        # Прогресс-бар для фиксированного количества шагов
+        train_pbar = tqdm(range(STEPS_PER_EPOCH), 
+                         desc=f'Epoch {epoch+1}/{num_epochs} [Train]')
+        
+        for step in train_pbar:
+            total_loss = 0
+            
+            # Градиентное накопление
+            for _ in range(GRADIENT_ACCUMULATE_EVERY):
+                batch = next(train_loader)
+                batch = batch.to(device)
+                
+                # Получаем входные данные и метки
+                inputs = batch[:, :-1]
+                labels = batch[:, 1:]
+                
+                # Прямой проход
+                outputs = model(input_ids=inputs, labels=labels)
+                loss = outputs.loss / GRADIENT_ACCUMULATE_EVERY
+                
+                # Обратное распространение
+                loss.backward()
+                total_loss += loss.item()
+            
+            # Обновление параметров
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            optimizer.zero_grad()
+            
+            total_train_loss += total_loss
+            train_steps += 1
+            
+            # Обновление прогресс-бара
+            train_pbar.set_postfix({
+                'loss': f'{total_loss:.4f}',
+                'avg_loss': f'{total_train_loss/train_steps:.4f}'
+            })
+            
+            # Валидация каждые 100 шагов
+            if step % 100 == 0:
+                model.eval()
+                val_loss = 0
+                val_steps = 0
+                
+                with torch.no_grad():
+                    for _ in range(10):  # Ограничиваем количество валидационных шагов
+                        val_batch = next(val_loader)
+                        val_batch = val_batch.to(device)
+                        
+                        val_inputs = val_batch[:, :-1]
+                        val_labels = val_batch[:, 1:]
+                        
+                        val_outputs = model(input_ids=val_inputs, labels=val_labels)
+                        val_loss += val_outputs.loss.item()
+                        val_steps += 1
+                
+                avg_val_loss = val_loss / val_steps
+                
+                print(f"\nValidation loss: {avg_val_loss:.4f}")
+                
+                # Сохраняем лучшую модель
+                if avg_val_loss < best_val_loss:
+                    best_val_loss = avg_val_loss
+                    torch.save({
+                        'epoch': epoch,
+                        'model_state_dict': model.state_dict(),
+                        'optimizer_state_dict': optimizer.state_dict(),
+                        'loss': best_val_loss,
+                    }, f'{save_path}_best.pt')
+                
+                model.train()
+        
+        # Сохраняем чекпойнт после каждой эпохи
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'loss': total_train_loss / train_steps,
+        }, f'{save_path}_epoch_{epoch}.pt')
+        
+        print(f"\nEpoch {epoch+1} completed. Average loss: {total_train_loss/train_steps:.4f}")
+    
+    return model
+
+def main():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
+    # Загружаем и подготавливаем данные
+    processed_data_path = Path(PROCESSED_DATA_DIR) / 'processed_wiki.pt'
+    
+    if not processed_data_path.exists():
+        print("Processing dataset...")
+        preprocessor = WikiDatasetPreprocessor(CACHE_DIR, PROCESSED_DATA_DIR)
+        preprocessor.process_and_save(max_articles=10000)
+    
+    print("Creating dataloaders...")
+    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)
+    
+    # Загружаем предобученную модель
+    print("Loading pretrained model...")
+    setup_custom_model()
+    config = AutoConfig.from_pretrained(REPO_ID)
+    model = AutoModelForCausalLM.from_pretrained(REPO_ID)
+    
+    print("Starting finetuning...")
+    # Файнтьюним модель
+    model = finetune_model(
+        model,
+        train_loader,
+        val_loader,
+        NUM_EPOCHS,
+        device
+    )
+    
+    # Сохраняем финальную версию модели
+    print("Saving final model...")
+    model.save_pretrained('final_finetuned_model')
+    
+    return model
+
+if __name__ == "__main__":
+    torch.manual_seed(42)
+    torch.cuda.manual_seed_all(42)
+    torch.backends.cudnn.benchmark = True
+    
+    try:
+        model = main()
+        print("Finetuning completed successfully!")
+    except Exception as e:
+        print(f"An error occurred: {str(e)}")
+```
+
 # Training
 
 The model was trained on a cleaned subset of Russian Wikipedia articles using the following parameters: