Update train.py

train.py

@@ -1,38 +1,48 @@
-from transformers import T5ForConditionalGeneration, ByT5Tokenizer, Trainer, TrainingArguments
+from transformers import T5ForConditionalGeneration, T5Tokenizer, Trainer, TrainingArguments
 from datasets import load_dataset
 import os
 import wandb
 import torch
 
 # 🔧 Название модели и путь
-model_name = "google/byt5-small"
-run_id = "byt5-autobatch"
+model_name = "google/flan-t5-large"
+run_id = "flan-t5-large-ru-autobatch"
 output_dir = f"./{run_id}"
-start_batch_size = 300
-step_batch_size = 5
+start_batch_size = 10   # ⚠️ Начинаем с небольшого batch, чтобы избежать OOM
+step_batch_size = 1
 
 # 📦 Загружаем модель и токенизатор
 model = T5ForConditionalGeneration.from_pretrained(model_name)
-tokenizer = ByT5Tokenizer.from_pretrained(model_name)
+tokenizer = T5Tokenizer.from_pretrained(model_name)
 
 # 📂 Загружаем датасет
 data_files = {
-    "train": "mt5_training_data-1.jsonl",
-    "validation": "mt5_validation_data-1.jsonl"
+    "train": "mt5_ru_gen_async.jsonl",
+    "validation": "mt5_ru_gen_eval.jsonl"
 }
 dataset = load_dataset("json", data_files=data_files)
 
 # 🔠 Токенизация
 def tokenize_function(examples):
-    model_inputs = tokenizer(examples["text"], max_length=256, truncation=True, padding="max_length")
-    labels = tokenizer(examples["target"], max_length=64, truncation=True, padding="max_length")
+    model_inputs = tokenizer(
+        examples["text"], max_length=256, truncation=True, padding="max_length"
+    )
+    with tokenizer.as_target_tokenizer():
+        labels = tokenizer(
+            examples["target"], max_length=256, truncation=True, padding="max_length"
+        )
+    # Заменяем PAD-токены на -100, чтобы не учитывать их в подсчёте loss
+    labels["input_ids"] = [
+        [(token if token != tokenizer.pad_token_id else -100) for token in label]
+        for label in labels["input_ids"]
+    ]
     model_inputs["labels"] = labels["input_ids"]
     return model_inputs
 
 tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
 # 🔑 Авторизация W&B
-wandb.login(key="5f028bc0142fb7fa45bdacdde3c00dbbaf8bf98e")
+wandb.login(key="ВАШ_WANDB_КЛЮЧ_ЗДЕСЬ")
 
 # 🚀 Функция автоподбора batch size
 def try_training_with_batch_size(batch_size_start):
@@ -44,11 +54,11 @@ def try_training_with_batch_size(batch_size_start):
                 output_dir=output_dir,
                 evaluation_strategy="steps",
                 eval_steps=100,
-                learning_rate=5e-5,
+                learning_rate=3e-5,
                 per_device_train_batch_size=batch_size,
                 per_device_eval_batch_size=batch_size,
-                #fp16=True,
-                num_train_epochs=3,
+                #fp16=True,  # Включайте при наличии подходящего GPU (A100 / V100 / T4)
+                num_train_epochs=10,
                 logging_steps=100,
                 warmup_ratio=0.06,
                 logging_first_step=True,