Update train.py

train.py

@@ -1,18 +1,21 @@
 from transformers import MT5Tokenizer, MT5ForConditionalGeneration, Trainer, TrainingArguments
 from transformers import ByT5Tokenizer, T5ForConditionalGeneration
-from transformers import T5ForConditionalGeneration
-from accelerate import init_empty_weights, infer_auto_device_map
 from datasets import load_dataset
 import os
 import wandb
 
+# 🔧 Название запуска (используется и как run_name, и как output_dir)
+run_name = "byt5-finetuning-run"
+
+# 🧠 Название модели для фантюнинга
+model_id = "google/byt5-small"
+
+# 📂 Куда сохранять результат обучения
+output_dir = f"./{run_name}"
+
 # Загружаем модель и токенизатор
-model = T5ForConditionalGeneration.from_pretrained(
-    "google/byt5-small",
-    device_map="auto", 
-    low_cpu_mem_usage=True
-)
-tokenizer = ByT5Tokenizer.from_pretrained("google/byt5-small")
+model = T5ForConditionalGeneration.from_pretrained(model_id)
+tokenizer = ByT5Tokenizer.from_pretrained(model_id)
 
 # Загружаем датасет
 data_files = {
@@ -30,15 +33,18 @@ def tokenize_function(examples):
 
 tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
+# Авторизация в Weights & Biases
 wandb.login(key="5f028bc0142fb7fa45bdacdde3c00dbbaf8bf98e")
 
+# Аргументы обучения
 training_args = TrainingArguments(
-    output_dir="./mt5-finetuned",
+    output_dir=output_dir,
     evaluation_strategy="steps",
     eval_steps=100,
     learning_rate=5e-5,
-    per_device_train_batch_size=140,
-    per_device_eval_batch_size=140,
+    per_device_train_batch_size=200,
+    per_device_eval_batch_size=200,
+    fp16=True,
     num_train_epochs=3,
     logging_steps=100,
     warmup_ratio=0.06,
@@ -47,13 +53,13 @@ training_args = TrainingArguments(
     logging_dir="./logs",
     save_total_limit=2,
     save_strategy="epoch",
-    report_to="wandb",                     
-    run_name="mt5-finetuning-run",   
+    report_to="wandb",
+    run_name=run_name,
     disable_tqdm=False,
-    max_grad_norm=1.0                       
+    max_grad_norm=1.0
 )
 
-
+# Инициализируем Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -63,9 +69,9 @@ trainer = Trainer(
 
 # Обучение
 trainer.train()
-#trainer.train(resume_from_checkpoint=True)
+# trainer.train(resume_from_checkpoint=True)
 
-# Сохраняем локально
-model.save_pretrained("./mt5-finetuned")
-tokenizer.save_pretrained("./mt5-finetuned")
-print("✅ Модель сохранена локально в ./mt5-finetuned")
+# Сохраняем модель
+model.save_pretrained(output_dir)
+tokenizer.save_pretrained(output_dir)
+print(f"✅ Модель сохранена локально в {output_dir}")