Create train.py

train.py

@@ -1 +1,71 @@
-print("✅ Training script is ready. Customize this file as needed.")
+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
+
+# Загружаем модель и токенизатор
+model = T5ForConditionalGeneration.from_pretrained(
+    "google/byt5-small",
+    device_map="auto", 
+    low_cpu_mem_usage=True
+)
+tokenizer = ByT5Tokenizer.from_pretrained("google/byt5-small")
+
+# Загружаем датасет
+data_files = {
+    "train": "mt5_training_data-1.jsonl",
+    "validation": "mt5_validation_data-1.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["labels"] = labels["input_ids"]
+    return model_inputs
+
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+
+wandb.login(key="5f028bc0142fb7fa45bdacdde3c00dbbaf8bf98e")
+
+training_args = TrainingArguments(
+    output_dir="./mt5-finetuned",
+    evaluation_strategy="steps",
+    eval_steps=100,
+    learning_rate=5e-5,
+    per_device_train_batch_size=140,
+    per_device_eval_batch_size=140,
+    num_train_epochs=3,
+    logging_steps=100,
+    warmup_ratio=0.06,
+    logging_first_step=True,
+    weight_decay=0.01,
+    logging_dir="./logs",
+    save_total_limit=2,
+    save_strategy="epoch",
+    report_to="wandb",                     
+    run_name="mt5-finetuning-run",   
+    disable_tqdm=False,
+    max_grad_norm=1.0                       
+)
+
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets["train"],
+    eval_dataset=tokenized_datasets["validation"]
+)
+
+# Обучение
+trainer.train()
+#trainer.train(resume_from_checkpoint=True)
+
+# Сохраняем локально
+model.save_pretrained("./mt5-finetuned")
+tokenizer.save_pretrained("./mt5-finetuned")
+print("✅ Модель сохранена локально в ./mt5-finetuned")