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()