Create train_model.py

train_model.py

@@ -0,0 +1,45 @@
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments
+from datasets import load_dataset
+
+# Загрузка датасета ImageNet
+dataset = load_dataset("imagenet-1k")
+
+# Инициализация модели и токенизатора
+model_name = "gpt2"
+model = GPT2LMHeadModel.from_pretrained(model_name)
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+
+# Предобработка данных
+def preprocess_data(examples):
+    inputs = examples["image"]
+    targets = examples["caption"]
+    inputs = tokenizer(inputs, padding=True, truncation=True, max_length=512, return_tensors="pt")
+    targets = tokenizer(targets, padding=True, truncation=True, max_length=512, return_tensors="pt")
+    inputs["labels"] = targets["input_ids"]
+    return inputs
+
+# Применение предобработки к датасету
+dataset = dataset.map(preprocess_data, batched=True)
+
+# Определение аргументов обучения
+training_args = TrainingArguments(
+    output_dir="./model",
+    num_train_epochs=5,
+    per_device_train_batch_size=4,
+    per_device_eval_batch_size=4,
+    warmup_steps=500,
+    weight_decay=0.01,
+    logging_dir="./logs",
+    logging_steps=100,
+    evaluation_strategy="epoch",
+)
+
+# Создание трейнера и обучение модели
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["validation"],
+    data_collator=None,
+)
+trainer.train()